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)
        {
            selectedRow = null;

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

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

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

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

                totalWeight += Mathf.Max(1, row.Weight);
            }

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

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

                cumulativeWeight += Mathf.Max(1, row.Weight);
                if (randomWeight <= cumulativeWeight)
                {
                    selectedRow = row;
                    return true;
                }
            }

            selectedRow = _eligibleRowBuffer[_eligibleRowBuffer.Count - 1];
            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 (displayPhaseIndex < row.MinPhase || displayPhaseIndex > row.MaxPhase)
                {
                    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];
                float curveWeight = GetRarityCurveWeight(row.Rarity, phaseT);
                if (curveWeight <= 0f)
                {
                    continue;
                }

                float rowWeight = Mathf.Max(1, row.Weight) * curveWeight;
                if (_rarityWeightBuffer.TryGetValue(row.Rarity, out float existingWeight))
                {
                    _rarityWeightBuffer[row.Rarity] = existingWeight + rowWeight;
                }
                else
                {
                    _rarityWeightBuffer[row.Rarity] = rowWeight;
                }
            }

            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 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;
            }
        }
    }
}