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

package algo

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.dotya.ml/wanderer/math-optim/report"
	"git.dotya.ml/wanderer/math-optim/stats"
	"git.dotya.ml/wanderer/math-optim/util"
	"gonum.org/v1/gonum/floats"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/plotutil"
	"gonum.org/v1/plot/vg"
)

const (
	preferredFont      = "Mono"
	titlePreferredFont = "Sans"
	yAxisLabel         = "f(x)"
	xAxisLabel         = "Generations"
	xAxisLabelAlt      = "FES"
)

// PlotMeanValsMulti creates plots for every member of 'stats.AlgoMeanVals' it
// is handed and saves them as 'result.Pic's results into a package-global
// slice.
func PlotMeanValsMulti(
	w *sync.WaitGroup,
	dimens, iterations int,
	bench, fPrefix, fExt string,
	algoMeanVals ...stats.AlgoMeanVals,
) {
	defer w.Done()

	// track time.
	start := time.Now()

	pWidth := 13 * vg.Centimeter
	pHeight := 13 * vg.Centimeter

	plotter.DefaultFont.Typeface = preferredFont
	plotter.DefaultLineStyle.Width = vg.Points(2.0)

	p := plot.New()
	pic := report.NewPic()

	p.Title.Text = fmt.Sprintf(
		"Comparison of Means (%dI) -\n%s (%dD)",
		iterations, bench, dimens,
	)

	// p.X.Label.Text = "Objective func. evaluations"
	p.X.Label.Text = xAxisLabelAlt
	p.X.Label.TextStyle.Font.Variant = preferredFont
	p.X.Label.TextStyle.Font.Weight = 1 // Medium
	p.X.Tick.Label.Font.Variant = preferredFont
	p.Y.Label.Text = yAxisLabel
	p.Y.Label.TextStyle.Font.Variant = preferredFont
	p.Y.Label.TextStyle.Font.Weight = 1 // Medium
	p.Y.Tick.Label.Font.Variant = preferredFont

	p.Title.TextStyle.Font.Size = 14.5
	p.Title.TextStyle.Font.Variant = titlePreferredFont
	p.Title.TextStyle.Font.Weight = 2 // SemiBold
	p.Title.Padding = 3 * vg.Millimeter

	p.Legend.TextStyle.Font.Variant = preferredFont
	p.Legend.TextStyle.Font.Size = 10
	p.Legend.Top = true
	p.Legend.Padding = 2 * vg.Millimeter

	lines := make([]interface{}, 0)

	for _, v := range algoMeanVals {
		// mark the end of the X axis with the greatest of len(v)
		if greatest := float64(len(v.MeanVals)); greatest > p.X.Max {
			p.X.Max = greatest
		}

		if min := floats.Min(v.MeanVals); min < p.Y.Min {
			p.Y.Min = min
		}

		if max := floats.Max(v.MeanVals); max > p.Y.Max {
			p.Y.Max = max
		}

		pts := make(plotter.XYs, len(v.MeanVals))

		// fill the plotter with datapoints.
		for k, res := range v.MeanVals {
			pts[k].X = float64(k)
			pts[k].Y = res
		}

		lines = append(lines, v.Title, pts)
	}

	err := plotutil.AddLines(
		p,
		lines...,
	)
	if err != nil {
		log.Panic(err)
	}

	filename := util.SanitiseFName(
		fmt.Sprintf("%s%scomparison-of-means-%dI-%s-%02dD",
			report.GetPicsDir(),
			fPrefix,
			iterations,
			bench,
			dimens,
		),
	)

	// set pic file path and caption.
	pic.FilePath = filename
	// pic.Caption = strings.ReplaceAll(filename, " ", "~")
	pic.Caption = strings.ReplaceAll(
		fmt.Sprintf("Comparison of Means (%dI) - %s (%dD)",
			iterations, bench, dimens,
		),
		" ", "~")
	pic.Bench = bench

	elapsed := time.Since(start)
	info := fmt.Sprintf("saving img to file: %s%s [generated in %s]",
		filename, fExt, elapsed,
	)
	log.Println(info)

	// Save the plot to a file using the above-constructed 'filename'.
	if err := p.Save(
		pWidth,
		pHeight,
		filename+fExt,
	); err != nil {
		panic(err)
	}

	mCoMPL.Lock()
	comparisonOfMeansPicList.Pics = append(comparisonOfMeansPicList.Pics, *pic)
	mCoMPL.Unlock()
}

func plotMeanVals(meanVals []float64, title, algo string, fes int) *plot.Plot {
	plotter.DefaultFont.Typeface = preferredFont
	plotter.DefaultLineStyle.Width = vg.Points(2.0)

	if fes != len(meanVals) {
		log.Fatalf("meanVals - FES mismatch: %d vs %d , bailing\n", fes, len(meanVals))
	}

	p := plot.New()
	// pic := report.NewPic()

	p.Title.Text = "Mean - " + title

	if strings.Contains(algo, "DE") || strings.Contains(algo, "SOMA") {
		p.X.Label.Text = xAxisLabelAlt
	} else {
		p.X.Label.Text = xAxisLabel
	}

	// p.X.Label.Padding = 8 * vg.Millimeter
	p.X.Label.TextStyle.Font.Variant = preferredFont
	p.X.Label.TextStyle.Font.Weight = 1 // Medium
	p.X.Tick.Label.Font.Variant = preferredFont
	// p.X.Padding = 2 * vg.Millimeter
	p.Y.Label.Text = yAxisLabel
	// p.Y.Label.Padding = 2 * vg.Millimeter
	p.Y.Label.TextStyle.Font.Variant = preferredFont
	p.Y.Label.TextStyle.Font.Weight = 1 // Medium
	p.Y.Tick.Label.Font.Variant = preferredFont
	// p.Y.Padding = 1 * vg.Millimeter

	p.Title.TextStyle.Font.Size = 14.5
	p.Title.TextStyle.Font.Variant = titlePreferredFont
	p.Title.TextStyle.Font.Weight = 2 // SemiBold
	// p.Title.Padding = 5 * vg.Millimeter
	p.Title.Padding = 3 * vg.Millimeter

	// mark the end of the X axis with len(meanVals).
	p.X.Max = float64(len(meanVals))
	p.Y.Min = floats.Min(meanVals)
	p.Y.Max = floats.Max(meanVals)

	pts := make(plotter.XYs, len(meanVals))

	// fill the plotter with datapoints.
	for k, res := range meanVals {
		pts[k].X = float64(k)
		pts[k].Y = res
	}

	lines := make([]interface{}, 0)
	lines = append(lines, pts)

	err := plotutil.AddLines(
		p,
		lines...,
	)
	if err != nil {
		log.Panic(err)
	}

	return p
}

// violating gocognit 30, TODO(me): still split this up. fPrefix is always "plot" according to unparam.
// nolint: gocognit,unparam
func plotAllDims(algoStats []stats.Stats, fPrefix, fExt string, ch chan report.PicList, chMean chan report.PicList) {
	start := time.Now()

	picsDir := report.GetPicsDir()

	// create picsDir (if not exists), fail on error, no point in going on
	// parsing algoStats and computing images if we cannot save them.
	if err := util.CreatePath(picsDir); err != nil {
		log.Fatalln("went to create picsDir, there was an issue: ", err)
	}

	pL := report.NewPicList()
	pics := make([]report.Pic, 0)
	pLMean := report.NewPicList()
	picsMean := make([]report.Pic, 0)

	// since the algoStats only contains results of a single algo, it's safe to
	// set the value like this.
	pL.Algo = algoStats[0].Algo
	pLMean.Algo = algoStats[0].Algo

	pWidth := 13 * vg.Centimeter
	pHeight := 13 * vg.Centimeter

	plotter.DefaultFont.Typeface = preferredFont
	plotter.DefaultLineStyle.Width = vg.Points(1.5)

	for _, s := range algoStats {
		p := plot.New()
		pic := report.NewPic()

		p.Title.Text = fmt.Sprintf("D: %d, G: %d, I: %d",
			s.Dimens,
			s.Generations,
			s.Iterations,
		)

		// For differential evolution, add params to title.
		if strings.Contains(s.Algo, "DE") {
			p.Title.Text += fmt.Sprintf(",\nNP: %d, F: ", s.NP) +
				strconv.FormatFloat(s.F, 'f', -1, 64) + ", CR: " +
				strconv.FormatFloat(s.CR, 'f', -1, 64)
		}

		// this is latex-rendered.
		pic.Caption = strings.ReplaceAll(p.Title.Text, " ", "~")

		// since a single stat slice of algoStats only contains results of a
		// single bench func, it's safe to set the value like this.
		pL.Bench = s.BenchFuncStats[0].BenchName
		pLMean.Bench = s.BenchFuncStats[0].BenchName

		if strings.Contains(s.Algo, "DE") {
			p.X.Label.Text = xAxisLabelAlt
		} else {
			p.X.Label.Text = xAxisLabel
		}

		p.X.Label.TextStyle.Font.Variant = preferredFont
		p.X.Label.TextStyle.Font.Weight = 1 // Medium
		p.X.Tick.Label.Font.Variant = preferredFont
		p.Y.Label.Text = yAxisLabel
		p.Y.Label.TextStyle.Font.Variant = preferredFont
		p.Y.Label.TextStyle.Font.Weight = 1 // Medium
		p.Y.Tick.Label.Font.Variant = preferredFont
		// p.Y.Padding = 1 * vg.Millimeter

		p.Title.TextStyle.Font.Size = 14.5
		p.Title.TextStyle.Font.Variant = titlePreferredFont
		p.Title.TextStyle.Font.Weight = 2 // SemiBold
		// p.Title.Padding = 5 * vg.Millimeter
		p.Title.Padding = 3 * vg.Millimeter

		// p.Legend.TextStyle.Font.Variant = preferredFontStyle
		// p.Legend.TextStyle.Font.Size = 8
		// p.Legend.Top = true
		// p.Legend.Padding = 0 * vg.Centimeter
		// p.Add(plotter.NewGrid())

		for _, dim := range s.BenchFuncStats {
			// infinite thanks to this SO comment for the interface "hack":
			//   https://stackoverflow.com/a/44872993
			lines := make([]interface{}, 0)

			for _, iter := range dim.BenchResults {
				// mark the end of the X axis with len(iter.Results).
				p.X.Max = float64(len(iter.Results))

				if floats.Min(iter.Results) < p.Y.Min {
					p.Y.Min = floats.Min(iter.Results)
				}

				if floats.Max(iter.Results) > p.Y.Max {
					p.Y.Max = floats.Max(iter.Results)
				}

				pts := make(plotter.XYs, len(iter.Results))

				// fill the plotter with datapoints.
				for k, res := range iter.Results {
					pts[k].X = float64(k)
					pts[k].Y = res
				}

				// lines = append(lines, "#"+fmt.Sprint(j), pts)
				lines = append(lines, pts)
			}

			err := plotutil.AddLines(
				p,
				lines...,
			)
			if err != nil {
				log.Printf("issue constructing a plot, bench: %s", dim.BenchName)

				// panic (don't panic, I know) instead of a hard exit.
				log.Panic(err)
			}

			// TODO(me): add Neighbourhood param
			// TODO(me): add search space percent param
			filename := fmt.Sprintf("%s%s-%s-%s-%dD-%dG-%dI",
				picsDir,
				fPrefix,
				util.SanitiseFName(s.Algo),
				util.SanitiseFName(dim.BenchName),
				s.Dimens,
				s.Generations,
				len(dim.BenchResults),
			)
			filenameMean := filename + util.SanitiseFName(" Mean")

			// NEVER EVER ATTEMPT TO INITIALISE THIS WITH `pic`!
			picMean := report.NewPic()
			meanTitle := fmt.Sprintf("D: %d, G: %d, I: %d",
				s.Dimens,
				s.Generations,
				s.Iterations,
			)

			// this is latex-rendered.
			picMean.Caption = "Mean - " + strings.ReplaceAll(meanTitle, " ", "~")

			// set pic file path (later used in tmpl generation)
			pic.FilePath = filename
			picMean.FilePath = filenameMean

			// get the *mean* plot.
			pMean := plotMeanVals(dim.MeanVals, meanTitle, s.Algo, s.Generations)

			elapsed := time.Since(start)
			info := fmt.Sprintf("saving img to file: %s(-Mean)%s [generated in %s]",
				filename, fExt, elapsed,
			)

			switch {
			case s.Algo == "Random Search":
				printRandomSearch(info)

			case strings.Contains(s.Algo, "Stochastic Hill Climbing"):
				printSHC(info)

			default:
				log.Println(info)
			}

			// Save the plot to a file using the above-constructed 'filename'.
			if err := p.Save(
				pWidth,
				pHeight,
				filename+fExt,
			); err != nil {
				panic(err)
			}
			// log.Println(filename + fExt)

			// save pic.
			pics = append(pics, *pic)

			// save plot of pMean.
			if err := pMean.Save(
				pWidth,
				pHeight,
				filenameMean+fExt,
			); err != nil {
				panic(err)
			}

			// save mean pic.
			picsMean = append(picsMean, *picMean)
		}
	}

	pL.Pics = pics
	pLMean.Pics = picsMean

	ch <- *pL
	chMean <- *pLMean
}