package main

import (
	"bufio"
	"errors"
	"fmt"
	"git.nunosempere.com/NunoSempere/fermi/sample"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"
)

const NORMAL90CONFIDENCE = 1.6448536269514727
const GENERAL_ERR_MSG = "Valid inputs: 2 || * 2 || / 2 || 2 20 || * 2 20 || / 2 20 || clean || =: var || op var || clean || help || debug || exit"
const N_SAMPLES = 10 // 1_000_000

// Distribution interface
// https://go.dev/tour/methods/9

type Dist interface {
	Samples() []float64
}

type Scalar struct {
	p float64
}

func (p Scalar) Samples() []float64 {
	xs := make([]float64, N_SAMPLES)
	for i := 0; i < N_SAMPLES; i++ {
		xs[i] = p.p
	}
	return xs
}

type Lognormal struct {
	low  float64
	high float64
}

func (ln Lognormal) Samples() []float64 {
	sampler := func(r sample.Src) float64 { return sample.Sample_to(ln.low, ln.high, r) }
	// return sample.Sample_parallel(sampler, N_SAMPLES)
	// Can't do parallel because then I'd have to await throughout the code
	return sample.Sample_serially(sampler, N_SAMPLES)
}

type Beta struct {
	a float64
	b float64
}

func (beta Beta) Samples() []float64 {
	sampler := func(r sample.Src) float64 { return sample.Sample_beta(beta.a, beta.b, r) }
	// return sample.Sample_parallel(sampler, N_SAMPLES)
	return sample.Sample_serially(sampler, N_SAMPLES)
}

type FilledSamples struct {
	xs []float64
}

func (fs FilledSamples) Samples() []float64 {
	return fs.xs
}

// Parse line into Distribution
func parseLineErr(err_msg string) (string, Dist, error) {
	fmt.Println(GENERAL_ERR_MSG)
	fmt.Println(err_msg)
	var errorDist Dist
	return "", errorDist, errors.New(err_msg)
}
func parseLine(line string, vars map[string]Dist) (string, Dist, error) {

	words := strings.Split(strings.TrimSpace(line), " ")
	op := ""
	var dist Dist

	switch words[0] {
	case "*":
		op = "*"
		words = words[1:]
	case "/":
		op = "/"
		words = words[1:]
	case "+":
		return parseLineErr("+ operation not implemented yet")
	case "-":
		return parseLineErr("- operation not implemented yet")
	default:
		op = "*" // later, change the below to
	}

	switch len(words) {
	case 0:
		return parseLineErr("Operator must have operand; can't operate on nothing")
	case 1:
		var_word, var_word_exists := vars[words[0]]
		single_float, err1 := strconv.ParseFloat(words[0], 64) // abstract this away to search for K/M/B/T/etc.
		switch {
		case var_word_exists:
			dist = var_word
		case err1 == nil:
			dist = Lognormal{low: single_float, high: single_float}
		case err1 != nil && !var_word_exists:
			return parseLineErr("Trying to operate on a scalar, but scalar is neither a float nor an assigned variable")
		}
	case 2:
		new_low, err1 := strconv.ParseFloat(words[0], 64)
		new_high, err2 := strconv.ParseFloat(words[1], 64)
		if err1 != nil || err2 != nil {
			return parseLineErr("Trying to operate by a distribution, but distribution is not specified as two floats")
		}
		dist = Lognormal{low: new_low, high: new_high}
	case 3:
		if words[0] == "beta" || words[0] == "b" {
			a, err1 := strconv.ParseFloat(words[1], 64)
			b, err2 := strconv.ParseFloat(words[2], 64)
			if err1 != nil || err2 != nil {
				return parseLineErr("Trying to specify a beta distribution? Try beta 1 2")
			}
			dist = Beta{a: a, b: b}
		} else {
			return parseLineErr("Input not understood or not implemented yet")
		}
	default:
		return parseLineErr("Input not understood or not implemented yet")
	}
	return op, dist, nil

}

// Join distributions
// Multiply lognormals

func multiplyLogDists(l1 Lognormal, l2 Lognormal) Lognormal {
	logmean1 := (math.Log(l1.high) + math.Log(l1.low)) / 2.0
	logstd1 := (math.Log(l1.high) - math.Log(l1.low)) / (2.0 * NORMAL90CONFIDENCE)

	logmean2 := (math.Log(l2.high) + math.Log(l2.low)) / 2.0
	logstd2 := (math.Log(l2.high) - math.Log(l2.low)) / (2.0 * NORMAL90CONFIDENCE)

	logmean_product := logmean1 + logmean2
	logstd_product := math.Sqrt(logstd1*logstd1 + logstd2*logstd2)

	h := logstd_product * NORMAL90CONFIDENCE
	loglow := logmean_product - h
	loghigh := logmean_product + h
	return Lognormal{low: math.Exp(loglow), high: math.Exp(loghigh)}

}

func multiplyBetaDists(beta1 Beta, beta2 Beta) Beta {
	return Beta{a: beta1.a + beta2.a, b: beta1.b + beta2.b}
}

func multiplyAsSamples(dist1 Dist, dist2 Dist) Dist {
	// dist2 = Beta{a: 1, b: 2}
	// fmt.Printf("dist1: %v\n", dist1)
	// fmt.Printf("dist2: %v\n", dist2)
	xs := dist1.Samples()
	ys := dist2.Samples()
	// fmt.Printf("xs: %v\n", xs)
	// fmt.Printf("ys: %v\n", ys)
	zs := make([]float64, N_SAMPLES)

	for i := 0; i < N_SAMPLES; i++ {
		zs[i] = xs[i] * ys[i]
	}

	// fmt.Printf("%v\n", zs)
	return FilledSamples{xs: zs}
}

func multiplyDists(old_dist Dist, new_dist Dist) (Dist, error) {

	switch o := old_dist.(type) {
	case Lognormal:
		{
			switch n := new_dist.(type) {
			case Lognormal:
				return multiplyLogDists(o, n), nil
			case Scalar:
				return multiplyLogDists(o, Lognormal{low: n.p, high: n.p}), nil
			default:
				return multiplyAsSamples(o, n), nil
			}
		}
	case Scalar:
		{
			if o.p == 1 {
				return new_dist, nil
			}
			switch n := new_dist.(type) {
			case Lognormal:
				return multiplyLogDists(Lognormal{low: o.p, high: o.p}, n), nil
			case Scalar:
				return Scalar{p: o.p * n.p}, nil
			default:
				return multiplyAsSamples(o, n), nil
			}
		}
	case Beta:
		switch n := new_dist.(type) {
		case Beta:
			return multiplyBetaDists(o, n), nil
		default:
			return multiplyAsSamples(o, n), nil
		}
	default:
		return multiplyAsSamples(old_dist, new_dist), nil
		// return nil, errors.New("Can't multiply dists")
	}
}

func joinDists(old_dist Dist, new_dist Dist, op string) (Dist, error) {

	switch op {
	case "*":
		return multiplyDists(old_dist, new_dist)
	default:
		return old_dist, errors.New("Can't combine distributions in this way")
	}
	/*
		switch {
		case old_dist.Type == "Lognormal" && new_dist.Type == "Lognormal" && op == "*":
			return Dist{Type: "Lognormal", Lognormal: multiplyLogDists(old_dist.Lognormal, new_dist.Lognormal), Samples: nil}, nil
		case old_dist.Type == "Lognormal" && new_dist.Type == "Lognormal" && op == "/":
			tmp_dist := Lognormal{low: 1.0 / new_dist.Lognormal.high, high: 1.0 / new_dist.Lognormal.low}
			return Dist{Type: "Lognormal", Lognormal: multiplyLogDists(old_dist.Lognormal, tmp_dist), Samples: nil}, nil
		default:
			fmt.Printf("For now, can't do anything besides multiplying lognormals\n")
		}
	*/
	// return old_dist, errors.New("Can't combine distributions in this way")
}

/* Pretty print distributions */
func prettyPrint2Floats(low float64, high float64) {
	// fmt.Printf("=> %.1f %.1f\n", low, high)
	switch {
	case math.Abs(low) >= 1_000_000_000_000:
		fmt.Printf("%.1fT", low/1_000_000_000_000)
	case math.Abs(low) >= 1_000_000_000:
		fmt.Printf("%.1fB", low/1_000_000_000)
	case math.Abs(low) >= 1_000_000:
		fmt.Printf("%.1fM", low/1_000_000)
	case math.Abs(low) >= 1_000:
		fmt.Printf("%.1fK", low/1_000)
	case math.Abs(low) >= 1_000:
		fmt.Printf("%.1fK", low/1_000)
	default:
		fmt.Printf("%.1f", low)
	}
	fmt.Printf(" ")
	switch {
	case math.Abs(high) >= 1_000_000_000_000:
		fmt.Printf("%.1fT", high/1_000_000_000_000)
	case math.Abs(high) >= 1_000_000_000:
		fmt.Printf("%.1fB", high/1_000_000_000)
	case math.Abs(high) >= 1_000_000:
		fmt.Printf("%.1fM", high/1_000_000)
	case math.Abs(high) >= 1_000:
		fmt.Printf("%.1fK", high/1_000)
	case math.Abs(high) >= 1_000:
		fmt.Printf("%.1fK", high/1_000)
	default:
		fmt.Printf("%.1f", high)
	}
	fmt.Printf("\n")
	// fmt.Printf("=> %.1f %.1f\n", low, high)
}

func prettyPrintDist(dist Dist) {
	switch v := dist.(type) {
	case Lognormal:
		fmt.Printf("=> ")
		prettyPrint2Floats(v.low, v.high)
	case FilledSamples:
		tmp_xs := make([]float64, N_SAMPLES)
		copy(tmp_xs, v.xs)
		sort.Slice(tmp_xs, func(i, j int) bool {
			return tmp_xs[i] < tmp_xs[j]
		})
		low_int := N_SAMPLES / 20
		low := tmp_xs[low_int]
		high_int := N_SAMPLES * 19 / 20
		high := tmp_xs[high_int]
		prettyPrint2Floats(low, high)
	case Beta:
		fmt.Printf("=> beta ")
		prettyPrint2Floats(v.a, v.b)
	default:
		fmt.Printf("%v", v)
	}
}

/* Main event loop */
func main() {

	reader := bufio.NewReader(os.Stdin)
	var init_dist Dist
	init_dist = Scalar{p: 1} // Lognormal{low: 1, high: 1}
	old_dist := init_dist
	vars := make(map[string]Dist)
	// Could eventually be a more complex struct with:
	// { Dist, VariableMaps, ConfigParams } or smth
EventForLoop:
	for {
		input, _ := reader.ReadString('\n')
		if strings.TrimSpace(input) == "" {
			continue EventForLoop
		}

		{
			words := strings.Split(strings.TrimSpace(input), " ")
			switch {
			case words[0] == "exit" || words[0] == "e":
				break EventForLoop
			case words[0] == "help" || words[0] == "h":
				fmt.Println(GENERAL_ERR_MSG)
				continue EventForLoop
			case words[0] == "debug" || words[0] == "d":
				fmt.Printf("Old dist: %v\n", old_dist)
				fmt.Printf("Vars: %v\n", vars)
				continue EventForLoop
			case words[0] == "=:" && len(words) == 2:
				vars[words[1]] = old_dist
				fmt.Printf("%s ", words[1])
				prettyPrintDist(old_dist)
				continue EventForLoop
			case words[0] == "." || words[0] == "clean" || words[0] == "c":
				old_dist = init_dist
				fmt.Println()
				continue EventForLoop
			case words[0] == "=." && len(words) == 2:
				vars[words[1]] = old_dist
				fmt.Printf("%s ", words[1])
				prettyPrintDist(old_dist)
				old_dist = init_dist
				fmt.Println()
				continue EventForLoop
				// Other possible cases:
				// Save to file
				// Sample n samples
				// Save stack to a variable?
				// clean stack
				// Define a function? No, too much of a nerdsnipea
			}
		}

		op, new_dist, err := parseLine(input, vars)
		if err != nil {
			continue EventForLoop
		}

		joint_dist, err := joinDists(old_dist, new_dist, op)
		if err != nil {
			fmt.Printf("%v\n", err)
			fmt.Printf("Dist on stack: ")
			prettyPrintDist(old_dist)
			continue EventForLoop
		}
		old_dist = joint_dist
		prettyPrintDist(old_dist)
	}
}