// Copyright 2023 wanderer <a_mirre at utb dot cz>
// SPDX-License-Identifier: GPL-3.0-or-later

package stats

import (
	"encoding/json"
	"io/ioutil"
	"log"

	"git.dotya.ml/wanderer/math-optim/util"
	"gonum.org/v1/gonum/stat"
)

// BenchRound holds the iteration couter value and Results of size 'maxFES'.
type BenchRound struct {
	Iteration int
	Results   []float64
}

type FuncStats struct {
	BenchName    string
	BenchResults []BenchRound
	// MeanVals holds the mean value for each step of the iteration.
	MeanVals []float64
}

// BenchMean structure holds mean vals and metadata of a bench-params
// combination.
type BenchMean struct {
	Bench  string
	Dimens int

	Iterations  int
	Generations int
	// Neighbours param is used in Hill Climbing type algos, set to -1 if not
	// applicable.
	Neighbours int
	// MeanVals holds the mean value for each step of the iteration.
	MeanVals []float64
}

// AlgoBenchMean holds BenchMean structs of different benchmarks but the same algorithm.
type AlgoBenchMean struct {
	Algo       string
	BenchMeans []BenchMean
}

// MeanStats aggregates AlgoMean structs of different algos.
type MeanStats struct {
	AlgoMeans []AlgoBenchMean
}

// AlgoMeanVals holds computed mean values of an Algo.
type AlgoMeanVals struct {
	Title string
	// MeanVals of a particular bench func to compare.
	MeanVals []float64
}

type Stats struct {
	Algo           string
	Dimens         int
	BenchFuncStats []FuncStats
	Iterations     int
	// this should perhaps be named FES as that is a smarter thing to be
	// comparing than algo generations, which can vary based on the type of
	// algo, number of neighbours, etc..
	Generations int
	// NP is the initial population size, disable with 0. Only applicable to
	// GAs (see algo/de for more details).
	NP int
	// F is the mutation factor, disable with 0. Only applicable to GAs.
	F float64
	// CR is the crossover probability, disable with 0. Only applicable to GAs.
	CR float64
}

// Len implements the sort.Interface.
func (a AlgoBenchMean) Len() int {
	return len(a.BenchMeans)
}

// Less implements the sort.Interface.
func (a AlgoBenchMean) Less(i, j int) bool {
	return a.BenchMeans[i].Bench < a.BenchMeans[j].Bench
}

// Swap implements the sort.Interface.
func (a AlgoBenchMean) Swap(i, j int) {
	a.BenchMeans[i], a.BenchMeans[j] = a.BenchMeans[j], a.BenchMeans[i]
}

// Len implements the sort.Interface.
func (m MeanStats) Len() int {
	return len(m.AlgoMeans)
}

// Less implements the sort.Interface.
// note: this in fact sorts in reverse (fits with our use case atm).
func (m MeanStats) Less(i, j int) bool {
	// return m.AlgoMeans[i].Algo < m.AlgoMeans[j].Algo
	return m.AlgoMeans[i].Algo > m.AlgoMeans[j].Algo
}

// Swap implements the sort.Interface.
func (m MeanStats) Swap(i, j int) {
	m.AlgoMeans[i], m.AlgoMeans[j] = m.AlgoMeans[j], m.AlgoMeans[i]
}

func GetFuncStats(funcName string, benchResults []BenchRound) FuncStats {
	f := FuncStats{
		BenchName:    funcName,
		BenchResults: benchResults,
	}

	return f
}

func GetStats(algo string, dimens int, benchFuncStats []FuncStats, iterations, generations int) Stats {
	s := Stats{
		Algo:           algo,
		Dimens:         dimens,
		BenchFuncStats: benchFuncStats,
		Iterations:     iterations,
		Generations:    generations,
	}

	return s
}

func SaveStats(stats []Stats, fName string) {
	prefix := "out/stats/"
	ext := ".json"

	// create required folders, bail early in case of an error.
	if err := util.CreatePath(prefix); err != nil {
		log.Fatalln("went to create stats dir, there was an issue: ", err)
	}

	if j, err := json.MarshalIndent(stats, "", "  "); err != nil {
		log.Fatal(err)
	} else {
		log.Println("saving json stats to:", prefix+fName+ext)

		if err = ioutil.WriteFile(prefix+fName+ext, j, 0o600); err != nil {
			log.Println("error saving stats to file:", err)
		}
	}
}

func GetMeanVals(values []BenchRound, fes int) []float64 {
	res := make([]float64, fes)

	for i := 0; i < fes; i++ {
		iterVals := make([]float64, len(values))

		for j, v := range values {
			iterVals[j] = v.Results[i]
		}

		res[i] = stat.Mean(iterVals, nil)
	}

	return res
}