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[]> Permutations<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);
        }
    }
    
    private static TValue[] GetCombination<TValue>(int[] innerIndices, TValue[] innerPool) => 
        innerIndices.Select(i => innerPool[i]).ToArray();
}