geometry-tower-defense-base/src-ref/CustomComponent/InventoryGeneration/DropPoolRoller.cs

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

namespace GeometryTD.CustomComponent
{
    public sealed class DropPoolRoller
    {
        private const float RarityCurveScalePhase = 30f;
        private static readonly RarityType[] OrderedRarities =
        {
            RarityType.White,
            RarityType.Green,
            RarityType.Blue,
            RarityType.Purple,
            RarityType.Red
        };

        private readonly List<DROutGameDropPool> _eligibleRowBuffer = new();
        private readonly float[] _rarityWeightBuffer = new float[OrderedRarities.Length];
        private readonly IDataTable<DROutGameDropPool> _dropPoolTable;

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

        public bool TryRollRow(
            int displayPhaseIndex,
            LevelThemeType themeType,
            Random random,
            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, random);
            if (selectedRarity == RarityType.None)
            {
                return false;
            }

            int totalWeight = 0;
            DROutGameDropPool fallbackRow = null;
            foreach (var row in _eligibleRowBuffer)
            {
                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.Next(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();

            foreach (var row in allRows)
            {
                if (row == null)
                {
                    continue;
                }

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

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

                _eligibleRowBuffer.Add(row);
            }
        }

        private RarityType RollRarity(int displayPhaseIndex, Random random)
        {
            for (int i = 0; i < _rarityWeightBuffer.Length; i++)
            {
                _rarityWeightBuffer[i] = 0f;
            }

            float phaseT = Mathf.Clamp01((displayPhaseIndex - 1) / RarityCurveScalePhase);

            foreach (var row in _eligibleRowBuffer)
            {
                for (int rarityIndex = 0; rarityIndex < OrderedRarities.Length; rarityIndex++)
                {
                    RarityType rarity = OrderedRarities[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;
                    }

                    _rarityWeightBuffer[rarityIndex] += rowWeight * curveWeight;
                }
            }

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

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

            float randomWeight = (float)(random.NextDouble() * totalWeight);
            float cumulativeWeight = 0f;
            for (int rarityIndex = 0; rarityIndex < _rarityWeightBuffer.Length; rarityIndex++)
            {
                cumulativeWeight += Mathf.Max(0f, _rarityWeightBuffer[rarityIndex]);
                if (randomWeight <= cumulativeWeight)
                {
                    return OrderedRarities[rarityIndex];
                }
            }

            for (int rarityIndex = 0; rarityIndex < _rarityWeightBuffer.Length; rarityIndex++)
            {
                if (_rarityWeightBuffer[rarityIndex] > 0f)
                {
                    return OrderedRarities[rarityIndex];
                }
            }

            return RarityType.None;
        }

        private static bool IsEligibleAtPhase(DROutGameDropPool row, int displayPhaseIndex)
        {
            for (int rarityIndex = 0; rarityIndex < OrderedRarities.Length; rarityIndex++)
            {
                RarityType rarity = OrderedRarities[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;
            }
        }
    }
}