package main

import (
	"fmt"
	"math"
	"os"

	coolbox "github.com/MeanderingAI/CoolBox/_deliverables/libraries/bindings/go_bindings"
)

var training = [][]float64{
	{0, 0}, {0, 1}, {1, 0}, {1, 1},
	{0.1, 0.1}, {0.1, 0.9}, {0.9, 0.1}, {0.9, 0.9},
}
var labels = []int{0, 1, 1, 0, 0, 1, 1, 0}
var validation = [][]float64{
	{0.25, 0.25}, {0.25, 0.75}, {0.75, 0.25}, {0.75, 0.75},
	{0.35, 0.35}, {0.35, 0.65}, {0.65, 0.35}, {0.65, 0.65},
}
var test = [][]float64{{0.05, 0.8}, {0.8, 0.05}, {0.05, 0.05}, {0.8, 0.8}}
var testLabels = []int{1, 1, 0, 0}

func decode(genes []float64) coolbox.BoostTreeParameters {
	return coolbox.BoostTreeParameters{
		NumEstimators: 20 + uint32(80*genes[0]+0.5),
		LearningRate:  0.03 + 0.17*genes[1],
		MaxDepth:      1 + uint32(3*genes[2]+0.5),
	}
}

func fit(genes []float64) (*coolbox.BoostTreeClassifier, error) {
	model, err := coolbox.NewBoostTreeClassifier(decode(genes))
	if err != nil {
		return nil, err
	}
	if err := model.Fit(training, labels); err != nil {
		model.Close()
		return nil, err
	}
	return model, nil
}

func loss(model *coolbox.BoostTreeClassifier) (float64, error) {
	classes, err := model.Classes()
	if err != nil {
		return 0, err
	}
	total := 0.0
	for i, row := range validation {
		probabilities, err := model.PredictProba(row)
		if err != nil {
			return 0, err
		}
		index := -1
		for j, label := range classes {
			if label == labels[i] {
				index = j
				break
			}
		}
		if index < 0 {
			return 0, fmt.Errorf("missing class %d", labels[i])
		}
		total -= math.Log(math.Max(1e-15, probabilities[index]))
	}
	return total / float64(len(validation)), nil
}

func run() error {
	initial := []float64{0.5, (0.1 - 0.03) / 0.17, 0}
	baseline, err := fit(initial)
	if err != nil {
		return err
	}
	defer baseline.Close()
	baselineLoss, err := loss(baseline)
	if err != nil {
		return err
	}
	config := coolbox.DefaultGeneticSearchConfig()
	config.Dimensions, config.PopulationSize, config.Generations = 3, 12, 8
	config.MutationRate, config.MutationSigma = 0.2, 0.15
	config.LowerBound, config.UpperBound, config.Seed = 0, 1, 42
	search, err := coolbox.NewSearchOptimizer(coolbox.GeneticSearchOptimizer, config)
	if err != nil {
		return err
	}
	defer search.Close()
	var objectiveErr error
	winner, err := search.Optimize(func(genes []float64) float64 {
		if objectiveErr != nil {
			return math.Inf(1)
		}
		model, fitErr := fit(genes)
		if fitErr != nil {
			objectiveErr = fitErr
			return math.Inf(1)
		}
		defer model.Close()
		score, scoreErr := loss(model)
		if scoreErr != nil {
			objectiveErr = scoreErr
			return math.Inf(1)
		}
		return score
	}, initial)
	if objectiveErr != nil {
		return fmt.Errorf("search objective: %w", objectiveErr)
	}
	if err != nil {
		return err
	}
	candidate, err := fit(winner)
	if err != nil {
		return err
	}
	defer candidate.Close()
	candidateLoss, err := loss(candidate)
	if err != nil {
		return err
	}
	selected, selectedLoss := baseline, baselineLoss
	if candidateLoss < baselineLoss {
		selected, selectedLoss = candidate, candidateLoss
	}
	fmt.Printf("Validation loss: %.6f -> %.6f\n", baselineLoss, selectedLoss)
	for _, entry := range []struct {
		name  string
		model *coolbox.BoostTreeClassifier
	}{{"Baseline", baseline}, {"Selected", selected}} {
		correct := 0
		for i, row := range test {
			predicted, err := entry.model.Predict(row)
			if err != nil {
				return err
			}
			if predicted == testLabels[i] {
				correct++
			}
		}
		fmt.Printf("%s test accuracy: %.0f%%\n", entry.name, 100*float64(correct)/float64(len(test)))
	}
	return nil
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}