package main

import (
	"math"
)

// Event type represents different events
type Event string

type Solver string

type SolverTimes map[Solver]map[Event]float64

type Assignment map[Solver][]Event

type PermWithTime struct {
	Perm []Solver
	Time float64
}

type AssignmentWithTime struct {
	Assignment Assignment
	Time       float64
}

type Config struct {
	SolverTimes  SolverTimes `yaml:"solvers"`
	Events       []Event     `yaml:"events"`
	Participants []Solver    `yaml:"participants"`
}

//
// Generic Functions
//

func Combinations[T comparable](lst []T, n int) func() ([]T, bool) {
	length := len(lst)
	current := make([]T, n)
	index := make([]int, n)

	// Initialize the first combination: ["lst[0]", "lst[0]", ..., "lst[0]"]
	for i := 0; i < n; i++ {
		current[i] = lst[0]
		index[i] = 0
	}

	return func() ([]T, bool) {
		if current == nil {
			return nil, false
		}

		// Copy current combination to return
		result := make([]T, n)
		copy(result, current)

		// Find the next combination
		j := n - 1
		for j >= 0 && index[j] == length-1 {
			j--
		}

		if j < 0 {
			current = nil
		} else {
			index[j]++
			current[j] = lst[index[j]]

			for k := j + 1; k < n; k++ {
				index[k] = 0
				current[k] = lst[0]
			}
		}

		return result, current != nil
	}
}

func CombinationsWithStart[T comparable](start []T, lst []T, n int) func() ([]T, bool) {
	hasNext := true
	iter := Combinations(lst, n-len(start))
	var end []T
	return func() ([]T, bool) {
		end, hasNext = iter()
		return append(start, end...), hasNext
	}
}

func countCombinations[T comparable](lst []T, n int) int {
	return int(math.Pow(float64(len(lst)), float64(n)))
}

//
// Relay Functions
//

func CalculateRelayTime(perm []Solver, events []Event, solverTimes SolverTimes) float64 {
	solversTotalTimes := make(map[Solver]float64, len(events))
	for i, solver := range perm {
		solversTotalTimes[solver] += solverTimes[solver][events[i]]
	}

	var relayTime float64
	for _, time := range solversTotalTimes {
		if time > relayTime {
			relayTime = time
		}
	}
	return relayTime
}

//
// Worker Functions
//

func GenerateSolverCombinationsWorker(solverCombs chan<- func() ([]Solver, bool), solvers []Solver, events []Event) {
	hasNext := true
	iter := Combinations(solvers, len(events))
	for hasNext {
		var start []Solver
		start, hasNext = iter()
		solverCombs <- CombinationsWithStart(start, solvers, len(events))
	}
}

func CalculateRelayTimesWorker(solverCombs <-chan func() ([]Solver, bool), assignments chan<- PermWithTime, events []Event, solverTimes SolverTimes) {
	for solverComb := range solverCombs {
		hasNext := true
		for hasNext {
			var comb []Solver
			comb, hasNext = solverComb()
			relayTime := CalculateRelayTime(comb, events, solverTimes)
			assignments <- PermWithTime{comb, relayTime}
		}
	}
}

func Assign(solvers []Solver, solverTimes SolverTimes, events []Event) AssignmentWithTime {
	solverCombs := make(chan func() ([]Solver, bool))
	permsWithTime := make(chan PermWithTime)

	go GenerateSolverCombinationsWorker(solverCombs, solvers, events)

	bestAssignment := AssignmentWithTime{
		Assignment: Assignment{},
		Time:       math.Inf(1),
	}
	// var bestAssignmentMu sync.Mutex

	for i := 0; i < 64; i++ {
		go CalculateRelayTimesWorker(solverCombs, permsWithTime, events, solverTimes)
	}

	totalPerms := countCombinations(solvers, len(events))

	for i := 0; i < totalPerms; i++ {
		// fmt.Println("checking assignment")
		p := <-permsWithTime
		if p.Time < bestAssignment.Time {
			assignment := make(Assignment, len(p.Perm))
			for i, solver := range p.Perm {
				assignment[solver] = append(assignment[solver], events[i])
			}
			bestAssignment = AssignmentWithTime{assignment, p.Time}
		}
	}

	return bestAssignment
}