geometry-tower-defense/Assets/GameMain/Scripts/CustomComponent/InventoryGeneration/DropPoolRoller.cs

using System.Collections.Generic;
using GameFramework.DataTable;
using GeometryTD.DataTable;
using GeometryTD.Definition;
using UnityEngine;

namespace GeometryTD.CustomComponent
{
    public sealed class DropPoolRoller
    {
        private const float RarityCurveScalePhase = 30f;

        private readonly List<DROutGameDropPool> _eligibleRowBuffer = new();
        private readonly Dictionary<RarityType, float> _rarityWeightBuffer = new();
        private readonly IDataTable<DROutGameDropPool> _dropPoolTable;

        public DropPoolRoller(IDataTable<DROutGameDropPool> dropPoolTable)
        {
            _dropPoolTable = dropPoolTable;
        }

        public bool TryRollRow(
            int displayPhaseIndex,
            LevelThemeType themeType,
            out DROutGameDropPool selectedRow,
            out RarityType selectedRarity)
        {
            selectedRow = null;
            selectedRarity = RarityType.None;

            DROutGameDropPool[] allRows = _dropPoolTable.GetAllDataRows();
            if (allRows == null || allRows.Length <= 0)
            {
                return false;
            }

            CollectEligibleRows(allRows, displayPhaseIndex, themeType);
            if (_eligibleRowBuffer.Count <= 0)
            {
                return false;
            }

            selectedRarity = RollRarity(displayPhaseIndex);
            if (selectedRarity == RarityType.None)
            {
                return false;
            }

            int totalWeight = 0;
            DROutGameDropPool fallbackRow = null;
            for (int i = 0; i < _eligibleRowBuffer.Count; i++)
            {
                DROutGameDropPool row = _eligibleRowBuffer[i];
                if (!IsEligibleAtPhase(row, selectedRarity, displayPhaseIndex))
                {
                    continue;
                }

                int rowWeight = row.GetWeight(selectedRarity);
                if (rowWeight <= 0)
                {
                    continue;
                }

                fallbackRow = row;
                totalWeight += rowWeight;
            }

            if (totalWeight <= 0)
            {
                return false;
            }

            int randomWeight = Random.Range(1, totalWeight + 1);
            int cumulativeWeight = 0;
            foreach (var row in _eligibleRowBuffer)
            {
                if (!IsEligibleAtPhase(row, selectedRarity, displayPhaseIndex))
                {
                    continue;
                }

                int rowWeight = row.GetWeight(selectedRarity);
                if (rowWeight <= 0)
                {
                    continue;
                }

                cumulativeWeight += rowWeight;
                if (randomWeight <= cumulativeWeight)
                {
                    selectedRow = row;
                    return true;
                }
            }

            selectedRow = fallbackRow;
            return selectedRow != null;
        }

        private void CollectEligibleRows(
            DROutGameDropPool[] allRows,
            int displayPhaseIndex,
            LevelThemeType themeType)
        {
            _eligibleRowBuffer.Clear();

            for (int i = 0; i < allRows.Length; i++)
            {
                DROutGameDropPool row = allRows[i];
                if (row == null)
                {
                    continue;
                }

                if (row.LevelThemeType != themeType)
                {
                    continue;
                }

                if (!IsEligibleAtPhase(row, displayPhaseIndex))
                {
                    continue;
                }

                _eligibleRowBuffer.Add(row);
            }
        }

        private RarityType RollRarity(int displayPhaseIndex)
        {
            _rarityWeightBuffer.Clear();
            float phaseT = Mathf.Clamp01((displayPhaseIndex - 1) / RarityCurveScalePhase);

            for (int i = 0; i < _eligibleRowBuffer.Count; i++)
            {
                DROutGameDropPool row = _eligibleRowBuffer[i];
                for (int rarityIndex = (int)RarityType.White; rarityIndex <= (int)RarityType.Red; rarityIndex++)
                {
                    RarityType rarity = (RarityType)rarityIndex;
                    if (!IsEligibleAtPhase(row, rarity, displayPhaseIndex))
                    {
                        continue;
                    }

                    int rowWeight = row.GetWeight(rarity);
                    if (rowWeight <= 0)
                    {
                        continue;
                    }

                    float curveWeight = GetRarityCurveWeight(rarity, phaseT);
                    if (curveWeight <= 0f)
                    {
                        continue;
                    }

                    float rarityWeight = rowWeight * curveWeight;
                    if (_rarityWeightBuffer.TryGetValue(rarity, out float existingWeight))
                    {
                        _rarityWeightBuffer[rarity] = existingWeight + rarityWeight;
                    }
                    else
                    {
                        _rarityWeightBuffer[rarity] = rarityWeight;
                    }
                }
            }

            float totalWeight = 0f;
            foreach (var pair in _rarityWeightBuffer)
            {
                totalWeight += Mathf.Max(0f, pair.Value);
            }

            if (totalWeight <= 0f)
            {
                return RarityType.None;
            }

            float randomWeight = Random.value * totalWeight;
            float cumulativeWeight = 0f;
            foreach (var pair in _rarityWeightBuffer)
            {
                cumulativeWeight += Mathf.Max(0f, pair.Value);
                if (randomWeight <= cumulativeWeight)
                {
                    return pair.Key;
                }
            }

            foreach (var pair in _rarityWeightBuffer)
            {
                return pair.Key;
            }

            return RarityType.None;
        }

        private static bool IsEligibleAtPhase(DROutGameDropPool row, int displayPhaseIndex)
        {
            for (int rarityIndex = (int)RarityType.White; rarityIndex <= (int)RarityType.Red; rarityIndex++)
            {
                RarityType rarity = (RarityType)rarityIndex;
                if (IsEligibleAtPhase(row, rarity, displayPhaseIndex))
                {
                    return true;
                }
            }

            return false;
        }

        private static bool IsEligibleAtPhase(DROutGameDropPool row, RarityType rarity, int displayPhaseIndex)
        {
            if (row.GetWeight(rarity) <= 0)
            {
                return false;
            }

            int minPhase = row.GetMinPhase(rarity);
            int maxPhase = row.GetMaxPhase(rarity);
            return displayPhaseIndex >= minPhase && displayPhaseIndex <= maxPhase;
        }

        private static float GetRarityCurveWeight(RarityType rarityType, float phaseT)
        {
            float hump = Mathf.Exp(-Mathf.Pow((phaseT - 0.35f) / 0.28f, 2f));
            switch (rarityType)
            {
                case RarityType.White:
                    return Mathf.Max(0.05f, 0.18f + 1.25f * hump);
                case RarityType.Green:
                    return Mathf.Max(0.05f, 0.35f + 0.55f * hump);
                case RarityType.Blue:
                    return 0.18f + 0.55f * phaseT;
                case RarityType.Purple:
                    return 0.05f + 0.22f * phaseT;
                case RarityType.Red:
                    return 0.01f + 0.08f * phaseT * phaseT;
                default:
                    return 0f;
            }
        }
    }
}