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

namespace AdventOfCode.Extensions;

public static class EnumerableExtensions
{
    public static IEnumerable<TValue[]> Combinations<TValue>(this IEnumerable<TValue> values, int count)
    {
        var pool = values.ToArray();
        var poolLength = pool.Length;
        if (count > poolLength)
            yield break;
        int[] indices = [..Enumerable.Range(0, count)];
        yield return GetCombination(indices, pool);
        while (true)
        {
            var idx = count - 1;
            for(;idx >= 0; --idx)
            {
                var maxPossibleIndex = idx + poolLength - count;
                var isIndexBelowMax = indices[idx] != maxPossibleIndex;
                if (isIndexBelowMax) 
                    break;
            }
            if(idx < 0)
                yield break;
            
            indices[idx] += 1;
            for (var j = idx + 1; j < indices.Length; j++)
            {
                indices[j] = indices[j - 1] + 1;
            }

            yield return GetCombination(indices, pool);
        }
    }
    
    public static IEnumerable<TValue[]> CombinationsWithRepeats<TValue>(this IEnumerable<TValue> values, int count)
    {
        var pool = values.ToArray();
        var poolLength = pool.Length;
        if (count > poolLength)
            yield break;
        int[] indices = [..Enumerable.Repeat(0, count)];
        yield return GetCombination(indices, pool);
        while (true)
        {
            var idx = count - 1;
            for(;idx >= 0; --idx)
            {
                var maxPossibleIndex = poolLength - 1;
                var isIndexBelowMax = indices[idx] != maxPossibleIndex;
                if (isIndexBelowMax) 
                    break;
            }
            if(idx < 0)
                yield break;
            
            var num = indices[idx];
            for (var j = idx; j < indices.Length; j++)
            {
                indices[j] = num + 1;
            }

            yield return GetCombination(indices, pool);
        }
    }
    
    public static IEnumerable<TValue[]> PermutationsWithRepeats<TValue>(this IEnumerable<TValue> values, int count)
    {
        var pool = values.ToArray();
        if (count > pool.Length)
            yield break;
        var indices = Enumerable.Repeat(0, count).ToArray();
        yield return GetCombination(indices, pool);
        while (true)
        {
            var idx = count - 1;
            for(;idx > 0; --idx)
            {
                var isIndexBelowMax = indices[idx] != idx + pool.Length - count;
                if (isIndexBelowMax) 
                    break;
            }
            if(indices.All(i => i >= pool.Length - 1))
                yield break;
            indices[idx] += 1;
            for (var j = idx + 1; j < count; j++)
            {
                indices[j] = 0;
            }

            yield return GetCombination(indices, pool);
        }
    }
    
    public static IEnumerable<TValue[]> Permutations<TValue>(this IEnumerable<TValue> values, int count)
    {
        var pool = values.ToArray();
        var poolLength = pool.Length;
        if (count > poolLength)
            yield break;
        var indices = Enumerable.Range(0, poolLength).ToArray();
        var cycles = Enumerable.Range(0, poolLength + 1)
            .Skip(poolLength - count + 1)
            .Reverse()
            .ToArray();
        yield return GetCombination(indices[..count], pool);
        var notDone = true;
        while (notDone)
        {
            notDone = false;
            for (var i = count - 1; i >= 0; i--)
            {
                cycles[i] -= 1;
                if (cycles[i] == 0)
                {
                    RotateIndices(indices, i);
                    cycles[i] = poolLength - i;
                }
                else
                {
                    var j = cycles[i];
                    (indices[i], indices[^j]) = (indices[^j], indices[i]);
                    yield return GetCombination(indices[..count], pool);
                    notDone = true;
                    break;
                }
            }
        }

        yield break;

        static void RotateIndices(int[] indices, int from)
        {
            var tmp = indices[from];
            for (var j = from; j < indices.Length - 1; j++)
            {
                indices[j] = indices[j + 1];
            }
            indices[^1] = tmp;
        }
    }
    
    private static TValue[] GetCombination<TValue>(int[] innerIndices, TValue[] innerPool) => 
        innerIndices.Select(i => innerPool[i]).ToArray();
}