using System;
using System.Collections.Generic;
using System.Linq;

namespace AibisDream.Framework
{
    public class RandomSelector<T> : ISelector<T>
    {
        private readonly List<T> _items;
        private readonly List<int> _prefixSums;
        private readonly int _totalWeight;
        private readonly Random _random;
        
        public int Count => _items.Count;
        
        public RandomSelector(IEnumerable<T> items, Func<T, int> weightSelector)
        {
            if (items == null) throw new ArgumentNullException(nameof(items));
            if (weightSelector == null) throw new ArgumentNullException(nameof(weightSelector));

            var itemList = items.ToList();
            _items = new List<T>(itemList.Count);
            _prefixSums = new List<int>(itemList.Count);
        
            int sum = 0;
            foreach (var item in itemList)
            {
                int weight = weightSelector(item);
                if (weight < 0)
                    throw new ArgumentException("Weights cannot be negative", nameof(items));

                if (weight > 0)
                {
                    sum += weight;
                    _prefixSums.Add(sum);
                    _items.Add(item);
                }
            }

            _totalWeight = sum;
            _random = new Random();
        }
        
        public T Select()
        {
            if (_totalWeight == 0)
                throw new InvalidOperationException("All weights are zero");

            int randomValue = _random.Next(0, _totalWeight);
            int index = BinarySearch(randomValue);
            return _items[index];
        }

        private int BinarySearch(int value)
        {
            int left = 0;
            int right = _prefixSums.Count - 1;

            while (left < right)
            {
                int mid = left + (right - left) / 2;
                if (_prefixSums[mid] <= value)
                {
                    left = mid + 1;
                }
                else
                {
                    right = mid;
                }
            }

            return left;
        }
    }

    public class BatchSelector<T> : ISelector<T>
    {
        private readonly List<T> _items;
        private readonly List<int> _batchWeight;
        private readonly int _totalWeight;
        private readonly Random _random;
        
        private int _currentWeight;
        private List<int> _curBatchWeight;

        public int Count => _totalWeight;
        
        public BatchSelector(IEnumerable<T> items, Func<T, int> weightSelector)
        { 
            if (items == null) throw new ArgumentNullException(nameof(items));
            if (weightSelector == null) throw new ArgumentNullException(nameof(weightSelector));
            
            var itemList = items.ToList();
            _items = new List<T>(itemList.Count);
            _batchWeight = new List<int>(itemList.Count);
        
            int sum = 0;
            foreach (var item in itemList)
            {
                int weight = weightSelector(item);
                if (weight < 0)
                    throw new ArgumentException("Weights cannot be negative", nameof(items));

                if (weight > 0)
                {
                    sum += weight;
                    _batchWeight.Add(weight);
                    _items.Add(item);
                }
            }

            _curBatchWeight = new List<int>(_batchWeight);
            
            _totalWeight = sum;
            _currentWeight = sum;
            _random = new Random();

            if (_totalWeight == 0)
            {
                // 出错了
                throw new InvalidOperationException("All weights are zero");
            }
        }

        public T Select()
        {
            // 如果已经抽完就重置
            if (_currentWeight == 0)
            {
                ResetSelector();
            }
            
            int randomValue = _random.Next(0, _currentWeight);
            int index = BinarySearch(randomValue);
            
            // 选择后减去权重
            _currentWeight -= 1;
            _curBatchWeight[index] -= 1;
            
            return _items[index];
        }

        private void ResetSelector()
        {
            _currentWeight = _totalWeight;
            _curBatchWeight = new List<int>(_batchWeight);
        }

        private int BinarySearch(int value)
        {
            var sum = 0;
            for (var i = 0; i < _curBatchWeight.Count; i++)
            {
                sum += _curBatchWeight[i];
                if (sum >= value)
                {
                    return i;
                }
            }

            return _curBatchWeight.Count - 1;
        }
    }

    public interface ISelector<out T>
    {
        T Select();
        int Count { get; }
    }
}