package main

import (
	"aoc/utils"
	"fmt"
	"math"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"sync"

	"gonum.org/v1/gonum/mat"
	"gonum.org/v1/gonum/optimize/convex/lp"
)

type button []int8

type problem struct {
	want                int
	buttons             []button
	joltageRequirements []int8
}

func parseLights(lights string) (want int) {
	want = 0
	for idx, ch := range lights {
		if ch == '#' {
			want += (1 << idx)
		}
	}
	return want
}

func parseProblem(line string) problem {
	reg := regexp.MustCompile(`\[([.#]+)\] ((?:(?:\([^)]+\))\s*)+) {([^}]+)}`)
	matches := reg.FindStringSubmatch(line)

	lights := matches[1]
	buttonString := matches[2]
	joltage := matches[3]

	buttonString = string(slices.DeleteFunc(
		[]rune(buttonString),
		func(ch rune) bool {
			return ch == '(' || ch == ')'
		}))

	buttons := []button{}
	for switchString := range strings.SplitSeq(buttonString, " ") {
		switches := []int8{}
		for sw := range strings.SplitSeq(switchString, ",") {
			val, _ := strconv.Atoi(sw)
			switches = append(switches, int8(val))
		}
		buttons = append(buttons, switches)
	}

	joltageRequirements := []int8{}
	for req := range strings.SplitSeq(joltage, ",") {
		val, _ := strconv.Atoi(req)
		joltageRequirements = append(joltageRequirements, int8(val))
	}

	return problem{
		want:                parseLights(lights),
		buttons:             buttons,
		joltageRequirements: joltageRequirements,
	}
}

func getProblemChannel(input chan string) chan problem {
	ch := make(chan problem, cap(input))
	go func() {
		for line := range input {
			ch <- parseProblem(line)
		}
		close(ch)
	}()
	return ch
}

func (b button) value() (val int) {
	val = 0
	for _, sw := range b {
		val += 1 << sw
	}
	return val
}

func bfs(goal int, buttons []button) int {
	type bfsState struct {
		epoch int8
		value int
	}

	visited := []int{0}
	queue := []bfsState{}

	for _, bt := range buttons {
		value := bt.value()
		queue = append(queue, bfsState{epoch: 1, value: value})
		visited = append(visited, value)
	}

	for len(queue) > 0 {
		state := queue[0]
		queue = queue[1:]

		if state.value == goal {
			return int(state.epoch)
		}

		for _, bt := range buttons {
			nextValue := state.value ^ bt.value()
			if !slices.Contains(visited, nextValue) {
				visited = append(visited, nextValue)
				queue = append(queue, bfsState{
					epoch: state.epoch + 1,
					value: nextValue,
				})
			}
		}
	}

	return math.MaxInt
}

func (p problem) solve() int {
	depth := bfs(p.want, p.buttons)
	return int(depth)
}

func parallel(threadCount int, f func(thrId int)) {
	var wg sync.WaitGroup
	for id := range threadCount {
		wg.Go(func() {
			f(id)
		})
	}
	wg.Wait()
}

func stepOne(input chan string) (int, error) {
	sum := 0
	pbChan := getProblemChannel(input)
	for pb := range pbChan {
		depth := pb.solve()
		fmt.Println(depth)
		sum += depth
	}
	return sum, nil
}

type linearEquation struct {
	solution int
	coeffs   []int
}

func (p problem) toLinearEquations() linearProblem {
	equations := make([]linearEquation, len(p.joltageRequirements))
	varCount := len(p.buttons)

	for equIdx, solution := range p.joltageRequirements {
		equ := linearEquation{
			solution: int(solution),
			coeffs:   make([]int, varCount),
		}

		for btIdx, bt := range p.buttons {
			if slices.Contains(bt, int8(equIdx)) {
				equ.coeffs[btIdx] = 1
			}
		}
		equations[equIdx] = equ
	}
	return equations
}

type linearProblem []linearEquation

func (lp linearProblem) A() mat.Matrix {
	rows := len(lp)
	cols := len(lp[0].coeffs)
	data := make([]float64, rows*cols)

	for r := range rows {
		for c := range cols {
			data[r*cols+c] = float64(lp[r].coeffs[c])
		}
		fmt.Println(data[r*cols : (r+1)*cols])
	}

	return mat.NewDense(rows, cols, data)
}

func (lp linearProblem) c() []float64 {
	count := len(lp[0].coeffs)
	consts := make([]float64, count)
	for idx := range count {
		consts[idx] = 1.0
	}
	return consts
}

func (lp linearProblem) b() []float64 {
	bees := make([]float64, len(lp))
	for idx := range bees {
		bees[idx] = float64(lp[idx].solution)
	}
	return bees
}

func stepTwo(input chan string) (int, error) {
	sum := 0
	pbChan := getProblemChannel(input)
	for pb := range pbChan {
		lpb := pb.toLinearEquations()
		c := lpb.c()
		A := lpb.A()
		b := lpb.b()

		optF, _, err := lp.Simplex(c, A, b, 1.0, nil)
		if err != nil {
			fmt.Errorf("simplex error: %v\n", err)
		}
		fmt.Println(optF)
		sum += int(optF)
	}

	return sum, nil
}

func main() {
	input := utils.FilePath("day10.txt")
	utils.RunStep(utils.ONE, input, stepOne)
	utils.RunStep(utils.TWO, input, stepTwo)
}