package main

import (
	"fmt"
	"math"
	"math/rand/v2"
	"os"

	"github.com/gofiber/fiber/v2/log"
	"gonum.org/v1/gonum/optimize"
)

type Detector struct {
	Name           string
	End            float64    // samples
	Position       [3]float64 // m
	EventEpochs    []float64  // samples
	ClockSpeedMult float64
}

func main() {
	samplerate := 44100. // Hz
	start := 1253385.    // samples
	sos := 347.99        // m/s
	// Detector 1 is reference clock
	data := []Detector{{
		Name:     "Computer",
		End:      4036375,
		Position: [3]float64{-0.817, -0.701, 1.343},
		EventEpochs: []float64{
			2718755,
			2886568,
			3093112,
		},
	}, {
		Name:     "Xiaomi Phone",
		End:      4036348,
		Position: [3]float64{-0.964, 0.394, 0.031},
		EventEpochs: []float64{
			2718813,
			2886539,
			3093112,
		},
	}, {
		Name:     "Google Pixel Phone",
		End:      4036401,
		Position: [3]float64{1.216, -0.475, 0.051},
		EventEpochs: []float64{
			2718843,
			2886703,
			3093149,
		},
	}, {
		Name:     "Camera",
		End:      4036372,
		Position: [3]float64{1.027, 0.618, 0.767},
		EventEpochs: []float64{
			2718799,
			2886606,
			3093109,
		},
	},
	}

	// Compute clock speed multiplier
	for detector_index := range data {
		data[detector_index].ClockSpeedMult = (data[detector_index].End - start) / (data[0].End - start) // How much to divide the time since start by to get the same sample count as the reference clock
	}

	realEventPositions := [][3]float64{{-0.14957, -0.36392, 1.1292}, {-0.419, 0.2958, 0.9319}, {0.021, 0.0058, 0.7719}}

	for event_index := range data[0].EventEpochs {
		problem := optimize.Problem{
			Func: func(x []float64) float64 {
				guessedX := x[0]
				guessedY := x[1]
				guessedZ := x[2]
				guessedEpoch := x[3] // samples
				errorSum := 0.
				for detector_index := range data {
					distance := math.Sqrt((guessedX-data[detector_index].Position[0])*(guessedX-data[detector_index].Position[0]) + (guessedY-data[detector_index].Position[1])*(guessedY-data[detector_index].Position[1]) + (guessedZ-data[detector_index].Position[2])*(guessedZ-data[detector_index].Position[2]))
					expectedEpoch := guessedEpoch + (distance/sos)*samplerate
					offset := ((data[detector_index].EventEpochs[event_index]-start)/data[detector_index].ClockSpeedMult + start - expectedEpoch)
					errorSum += offset * offset
				}
				return errorSum
			},
		}
		result, err := optimize.Minimize(problem, []float64{0, 0, 0, data[0].EventEpochs[event_index]}, nil, &optimize.NelderMead{})
		if err != nil {
			log.Error(err)
		}
		fmt.Printf("[Event %d]: (%0.4f m, %0.4f m, %0.4f m, %0.4f s) | Error: %0.4fm\n", event_index, result.X[0], result.X[1], result.X[2], result.X[3]/samplerate, math.Sqrt((realEventPositions[event_index][0]-result.X[0])*(realEventPositions[event_index][0]-result.X[0])+(realEventPositions[event_index][1]-result.X[1])*(realEventPositions[event_index][1]-result.X[1])+(realEventPositions[event_index][2]-result.X[2])*(realEventPositions[event_index][2]-result.X[2])))
	}

	expected_error_samples_sd_per_detector := 2. // to be multiplied by *sqrt(3) when doing the uniform distrib
	montecarlosims := 1000
	stepsPerDirection := 10
	max_error_exposure := 1. // m (for point cloud vis)
	boundsX := []float64{-1, 1}
	boundsY := []float64{-1, 1}
	boundsZ := []float64{0, 2}
	resStr := fmt.Sprintf(`ply
format ascii 1.0
element vertex %d          
property float x
property float y
property float z
property uchar red
property uchar green
property uchar blue
end_header
`, stepsPerDirection*stepsPerDirection*stepsPerDirection)

	for i := range stepsPerDirection {
		for j := range stepsPerDirection {
			fmt.Printf("%.1f %%\n", 100*float64(i)/float64(stepsPerDirection)+100*float64(j)/float64(stepsPerDirection)/float64(stepsPerDirection))
			for k := range stepsPerDirection {
				S := 0.
				realPos := []float64{boundsX[0] + float64(i)/float64(stepsPerDirection-1)*(boundsX[1]-boundsX[0]), boundsY[0] + float64(j)/float64(stepsPerDirection-1)*(boundsY[1]-boundsY[0]), boundsZ[0] + float64(k)/float64(stepsPerDirection-1)*(boundsZ[1]-boundsZ[0])}
				epochs := []float64{}

				for detector_index := range data {
					distance := math.Sqrt((realPos[0]-data[detector_index].Position[0])*(realPos[0]-data[detector_index].Position[0]) + (realPos[1]-data[detector_index].Position[1])*(realPos[1]-data[detector_index].Position[1]) + (realPos[2]-data[detector_index].Position[2])*(realPos[2]-data[detector_index].Position[2]))
					expectedEpoch := (distance/sos)*samplerate + 2*(rand.Float64()-0.5)*expected_error_samples_sd_per_detector*math.Sqrt(3)
					epochs = append(epochs, expectedEpoch)
				}

				for range montecarlosims {
					problem := optimize.Problem{
						Func: func(x []float64) float64 {
							guessedX := x[0]
							guessedY := x[1]
							guessedZ := x[2]
							guessedEpoch := x[3] // samples
							errorSum := 0.
							for detector_index := range data {
								distance := math.Sqrt((guessedX-data[detector_index].Position[0])*(guessedX-data[detector_index].Position[0]) + (guessedY-data[detector_index].Position[1])*(guessedY-data[detector_index].Position[1]) + (guessedZ-data[detector_index].Position[2])*(guessedZ-data[detector_index].Position[2]))
								expectedEpoch := guessedEpoch + (distance/sos)*samplerate
								offset := (epochs[detector_index] - expectedEpoch)
								errorSum += offset * offset
							}
							return errorSum
						},
					}
					result, err := optimize.Minimize(problem, []float64{realPos[0], realPos[1], realPos[2], 0}, nil, &optimize.NelderMead{})
					if err != nil {
						log.Error(err)
					}
					S += math.Sqrt((realPos[0]-result.X[0])*(realPos[0]-result.X[0]) + (realPos[1]-result.X[1])*(realPos[1]-result.X[1]) + (realPos[2]-result.X[2])*(realPos[2]-result.X[2]))
				}
				avg_error := S / float64(montecarlosims)
				color := int(min(avg_error/max_error_exposure*255., 255))
				resStr += fmt.Sprintf("%.6f %.6f %.6f %d %d %d\n", realPos[0], realPos[1], realPos[2], color, color, color)
			}
		}
	}
	file, err := os.Create("error.ply") // Usage: https://blender.stackexchange.com/questions/310858/how-to-visualize-point-cloud-colors-in-blender-4-0-after-ply-data-import
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	file.WriteString(resStr)
}
