package main

import (
	"flag"
	"log"
	"os/exec"

	"git.dotya.ml/wanderer/ak9im/p2/stats"
	"git.dotya.ml/wanderer/ak9im/p3/lrls"
)

var (
	datafile = flag.String("datafile", "", "read data from this file")
	vis      = flag.Bool("vis", false, "run 'python visualise.py' to produce visualisations")
)

func run() error {
	flag.Parse()

	if *datafile != "" {
		data, err := readFile(datafile)
		if err != nil {
			return err
		}

		u := 0
		y := 1

		meanU := stats.Mean(data[u])
		meanY := stats.Mean(data[y])
		varianceU := stats.Variance(data[u])
		varianceY := stats.Variance(data[y])

		maxShift := 0.1
		autocorrelationU := stats.Autocorrelate(data[u], maxShift)
		autocorrelationY := stats.Autocorrelate(data[y], maxShift)

		mutCorrelationUY, err := stats.MutCorrelate(data[u], data[y], maxShift)
		if err != nil {
			return err
		}

		mutCorrelationYU, err := stats.MutCorrelate(data[y], data[u], maxShift)
		if err != nil {
			return err
		}

		cov := stats.Covariance(data[u], data[y])

		theta, errVals, err := lrls.Estimate(1, data[u], data[y])
		if err != nil {
			return err
		}

		thetaT := lrls.TransposeTheta(theta)

		// explicit method does not return theta and errs directly, instead
		// it's computed by a Matlab script and the resulting data is saved in
		// `./data`.
		_, _, err = lrls.Estimate(0, data[u], data[y])
		if err != nil {
			return err
		}

		log.Printf("len(data): %d", len(data[u]))

		log.Printf("means - u: %v, y: %v", meanU, meanY)

		log.Printf("variance - u: %v, y: %v", varianceU, varianceY)

		log.Printf("covariance U,Y: %v", cov)
		log.Printf("correlation U,Y: %v", stats.Correlation(data[u], data[y]))

		err = saveStuff(
			meanU, meanY,
			varianceU, varianceY,
			cov,
			autocorrelationU, autocorrelationY,
			mutCorrelationUY, mutCorrelationYU,
			errVals,
			theta, thetaT,
		)
		if err != nil {
			return err
		}
	}

	if *vis {
		err := visualise()
		if err != nil {
			return err
		}
	}

	return nil
}

func visualise() error {
	red := "\033[31m"
	cyan := "\033[36m"
	reset := "\033[0m"
	f := "visualise.py"

	log.Printf("running %s`python %s`%s", cyan, f, reset)

	out, err := exec.Command("python", f).CombinedOutput()
	if err != nil {
		log.Printf("visualise failed with:\n\n%s%s%s\n", red, out, reset)
		return err
	}

	return nil
}