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

namespace Dacryptero.Data.Generator
{
    public class Markow<T>
        where T : notnull
    {
        private readonly ReadOnlyMemory<double> propability;

        private readonly ReadOnlyMemory<T> tokens;

        public Markow(IEnumerable<T> tokens)
        {
            var set = new HashSet<T>(tokens);
            var tokenTarget = new T[set.Count];
            set.CopyTo(tokenTarget);
            this.tokens = tokenTarget;

            var dict = new Dictionary<T, int>(set.Count);
            for (int i = 0; i < this.tokens.Length; ++i)
                dict.Add(this.tokens.Span[i], i);

            Memory<double> map = new double[set.Count * set.Count];
            Span<int> column = new int[set.Count];
            int previous = -1;
            bool first = true;
            foreach (var token in tokens)
            {
                if (first)
                {
                    previous = dict[token];
                    first = false;
                }
                else
                {
                    var next = dict[token];
                    map.Span[GetPropabilityIndex(previous, next)]++;
                    column[previous]++;
                    previous = next;
                }
            }

            for (int i = 0; i < this.tokens.Length; ++i)
                if (column[i] > 0)
                    for (int j = 0; j < this.tokens.Length; ++j)
                        map.Span[GetPropabilityIndex(i, j)] /= column[i];

            propability = map;
        }

        private int GetPropabilityIndex(int i, int j)
        {
            return i * tokens.Length + j;
        }
    
        public IEnumerable<T> Generate(Random? rng = null, int limit = -1)
        {
            rng ??= new Random();
            var index = rng.Next(tokens.Length);
            while (limit != 0)
            {
                if (limit > 0)
                    limit--;
                var rv = rng.NextDouble();
                var sum = 0.0;
                var emmitted = false;
                for (int i = 0; i < tokens.Length; ++i)
                {
                    sum += propability.Span[GetPropabilityIndex(index, i)];
                    if (sum > rv)
                    {
                        yield return tokens.Span[i];
                        index = i;
                        emmitted = true;
                        break;
                    }
                }
                if (!emmitted)
                {
                    yield return tokens.Span[index = rng.Next(tokens.Length)];
                }
            }
        }
    }
}