using System.Data;
using AdvenOfCode.Contracts;
using AdventOfCode.Extensions;
using AdventOfCode.HelperClasses;
using Microsoft.Z3;

namespace AoC_2025;

public class Day10 : IPuzzleSolver<long>
{
    private record ButtonCombination(uint Mask, int[] Joltages, int Cost);
    private (bool[] lamps, int[][] buttons, int[] joltages)[] ParsePuzzleInput(string path)
    {
        var puzzleInput = File.ReadAllLines(path)
            .Where(line => !string.IsNullOrWhiteSpace(line))
            .Select(line => line.Split(' ', StringSplitOptions.RemoveEmptyEntries))
            .Select(instruction => instruction switch
            {
                [var lamps, ..var buttons, var joltages] => (lamps: lamps[1..^1], buttons, joltages: joltages[1..^1]),
                _ => throw new DataException("Misaligned data")
            })
            .OrderBy(inst => inst.lamps.Length)
            .ToArray();
        var lamps = puzzleInput
            .Select(input => input.lamps)
            .Select(lamps => lamps.Index())
            .Select(ixlamps =>
                ixlamps.Select(ixlamp => ixlamp.Item == '#').ToArray())
            .ToArray();
        var buttons = puzzleInput
            .Select(input => input.buttons)
            .Select(buttons => buttons
                .Select(button => button[1..^1]
                    .Split(',', StringSplitOptions.RemoveEmptyEntries)
                    .Select(int.Parse)
                    .ToArray())
                .ToArray())
            .ToArray();
        var joltages = puzzleInput
            .Select(input => input.joltages)
            .Select(joltage => joltage
                .Split(',', StringSplitOptions.RemoveEmptyEntries)
                .Select(int.Parse)
                .ToArray())
            .ToArray();
        var res = lamps
            .Zip(buttons, (lamp, button) => (lamp, button))
            .Zip(joltages, (bl, joltage) => (bl.lamp, bl.button, joltage))
            .ToArray();
        return res;
    }

    public long SolvePart1(string pathToPuzzleInput)
    {
        var instructions = ParsePuzzleInput(pathToPuzzleInput);
        var minCounts = instructions
            .Select(instruction => GetMinCountButtonPressesForLamps(instruction.lamps, instruction.buttons))
            .ToArray();
        return minCounts.Sum();
    }
    
    public long SolvePart2(string pathToPuzzleInput)
    {
        var instructions = ParsePuzzleInput(pathToPuzzleInput);
        var minCounts = instructions
            .Select(instruction => GetMinCountButtonPressesForJoltageLevels(instruction.joltages, instruction.buttons))
            .ToArray();
        return minCounts.Sum();
    }
    
    // not proud of you
    public long SolvePart2Z3(string pathToPuzzleInput)
    {
        var instructions = ParsePuzzleInput(pathToPuzzleInput);
        var minCounts = instructions
            .Select(ins => GetMinCountButtonPressesForJoltagesZ3(ins.joltages, ins.buttons))
            .ToArray();
        return minCounts.Sum();
    }

    private long GetMinCountButtonPressesForLamps(bool[] lamps, int[][] buttons)
    {
        var lampMask = AsMask(lamps, x => x);
        var pattern = GetButtonCombinations(lamps.Length, buttons)
            .First(pattern => pattern.Mask == lampMask);
        return pattern.Cost;
    }

    private long GetMinCountButtonPressesForJoltageLevels(int[] joltages, int[][] buttons)
    {
        var buttonCombinations = GetButtonCombinations(joltages.Length, buttons)
            .GroupBy(x => x.Mask)
            .ToDictionary(x => x.Key, x => x.ToArray());
        var res = GetMinCountInternal(joltages, buttonCombinations, new Dictionary<int[], long>(new IntArrayEqualityComparer()));
        return res;

        static long GetMinCountInternal(int[] targetJoltages, Dictionary<uint, ButtonCombination[]> buttonCombinations, Dictionary<int[], long> memory)
        {
            if (targetJoltages.Any(j => j < 0)) return 1_000_000_000L;
            if (targetJoltages.All(j => j == 0)) return 0;
            if (memory.TryGetValue(targetJoltages, out var knownResult)) return knownResult;
            
            var minCount = 1_000_000_000L;
            var joltageMask = AsMask(targetJoltages, level => level % 2 == 1);
            var combinationsToCheck = buttonCombinations.GetValueOrDefault(joltageMask, []);
            foreach (var (_, joltageLevels, cost) in combinationsToCheck)
            {
                var nextTarget = targetJoltages.Zip(joltageLevels, (jolts, level) => (jolts - level) / 2).ToArray();
                var nextCount = 2 * GetMinCountInternal(nextTarget, buttonCombinations, memory) + cost;
                minCount = Math.Min(minCount, nextCount);
            }

            memory[targetJoltages] = minCount;
            return minCount;
        }
    }
    
    private IEnumerable<ButtonCombination> GetButtonCombinations(int patternLength, int[][] buttons)
    {
        yield return new ButtonCombination(0, [..Enumerable.Repeat(0, patternLength)], 0);
        for (var i = 1; i <= buttons.Length; i++)
        {
            foreach (var buttonPattern in buttons.Combinations(i))
            {
                var joltageLevels = GetJoltageLevelsForButtonPattern(patternLength, buttonPattern);
                var mask = AsMask(joltageLevels, level => level % 2 == 1);
                yield return new ButtonCombination(mask, joltageLevels, buttonPattern.Length);
            }
        }
        
        yield break;
        
        static int[] GetJoltageLevelsForButtonPattern(int countLamps, int[][] buttonPattern)
        {
            var joltagesProduced = Enumerable.Repeat(0, countLamps).ToArray();
            foreach (var button in buttonPattern)
            {
                foreach (var index in button)
                {
                    joltagesProduced[index]++;
                }
            }

            return joltagesProduced;
        }
    }
    
    private static uint AsMask<TValue>(TValue[] toBeMasked, Func<TValue, bool> maskPredicate) => 
        toBeMasked.Length <= 32 
            ? toBeMasked.Aggregate((uint)0L, (acc, val) => (acc << 1) | (uint)(maskPredicate(val) ? 1 : 0))
            : throw new DataException("Mask can only be 32 bits long");

    private long GetMinCountButtonPressesForJoltagesZ3(int[] joltages, int[][] buttons)
    {
        using var context = new Context();
        var optimize = context.MkOptimize();
        var variables = Enumerable.Range(0, buttons.Length)
            .Select(ix => context.MkIntConst($"n{ix}"))
            .ToArray();
        foreach (var variable in variables)
        {
            optimize.Add(variable >= 0);
        }

        var zero = context.MkInt(0);
        foreach (var (i, joltage) in joltages.Index())
        {
            var equation = zero + zero;
            
            foreach (var (b, button) in buttons.Index())
            {
                if(button.Contains(i))
                    equation += variables[b];
            }

            optimize.Add(context.MkEq(equation, context.MkInt(joltage)));
        }
        var sumVariables = variables.Aggregate(zero + zero, (current, variable) => current + variable);

        optimize.MkMinimize(sumVariables);
        optimize.Check();
        var answerExpression = optimize.Model.Eval(sumVariables);
        var answer = answerExpression is IntNum number ? number.Int64 : 0;
        return answer;
    }
}