package main

import (
	"image/color"
	"log"
	"math"
	"time"

	"gonum.org/v1/gonum/mat"

	"github.com/alltom/oklab"
	"github.com/hajimehoshi/ebiten/v2"
	"github.com/hajimehoshi/ebiten/v2/inpututil"
	"github.com/hajimehoshi/ebiten/v2/vector"
)

const (
	N                       = 100
	L                       = 1e-9            // 1 nm box
	hbar                    = 1.054571817e-34 // J·s
	mass                    = 9.10938356e-31  // kg
	eV                      = 1.602176634e-19 // J
	V0Max                   = 30 * eV
	ProbabilityDensityScale = 0.3       // Percentage of the screen space
	TimeFactor              = 2 * 1e-25 // Time factor for animation
	Width                   = 1200
	Height                  = 1200
)

// Game implements ebiten.Game interface.
type Game struct {
	changedPotential bool

	potential              []float64
	eigenvalues            []float64
	eigenvectors           mat.Dense
	maxdisplayedeigenvalue int
	computedColors         []color.RGBA

	renormalizedEigenvectors mat.Dense

	mousePressed bool
	prevX        int
	prevY        float64

	touchIDs         []ebiten.TouchID
	releasedTouchIDs []ebiten.TouchID
}

// Update proceeds the game state.
// Update is called every tick (1/60 [s] by default).
func (g *Game) Update() error {
	// Write your game's logical update.
	if g.changedPotential {
		dx := L / float64(N)

		H := mat.NewSymDense(N, nil)

		// kinetic prefactor
		t := hbar * hbar / (2 * mass * dx * dx)

		for i := range N {
			V := g.potential[i] * V0Max

			H.SetSym(i, i, 2*t+V)

			if i > 0 {
				H.SetSym(i, i-1, -t)
			}
		}

		var eig mat.EigenSym
		ok := eig.Factorize(H, true)
		if !ok {
			panic("eigendecomposition failed")
		}
		g.eigenvalues = eig.Values(nil)
		eig.VectorsTo(&g.eigenvectors)

		g.maxdisplayedeigenvalue = -1
		g.computedColors = []color.RGBA{}
		for i := range g.eigenvalues {
			if g.eigenvalues[i] <= V0Max {
				g.maxdisplayedeigenvalue = i
				g.computedColors = append(g.computedColors, OKLCHToColor(1, 0.5, 2*math.Pi*g.eigenvalues[i]/V0Max))
				norm := 0.0
				for j := range N {
					norm += math.Pow(g.eigenvectors.At(j, i), 2)
				}
				norm = math.Sqrt(norm) / ProbabilityDensityScale
				for j := range N {
					y := float64(Height) * float64(g.eigenvectors.At(j, i)/norm)
					g.eigenvectors.Set(j, i, y)
				}
			}
		}

		g.changedPotential = false
	}

	g.touchIDs = ebiten.AppendTouchIDs(g.touchIDs[:0])

	if ebiten.IsMouseButtonPressed(ebiten.MouseButtonLeft) || len(g.touchIDs) != 0 {
		x, y := ebiten.CursorPosition()
		if len(g.touchIDs) != 0 {
			x, y = ebiten.TouchPosition(g.touchIDs[0])
		}
		newX := int(float64(x) / float64(Width) * (N - 1))
		newY := float64(y) / float64(Height)
		if g.mousePressed {
			diff := newX - g.prevX
			if diff > 0 {
				for i := range diff {
					j := i + g.prevX
					if j >= 0 && j < N {
						g.potential[j] = 1 - (newY*float64(i)/float64(diff) + g.prevY*(1-float64(i)/float64(diff)))
					}
				}
			} else {
				for i := range -(diff) {
					j := g.prevX - i - 1
					if j >= 0 && j < N {
						g.potential[j] = 1 - (newY*float64(i)/float64(-diff) + g.prevY*(1-float64(i)/float64(-diff)))
					}
				}
			}
		}
		g.prevX = newX
		g.prevY = newY
		g.mousePressed = true

	}

	g.releasedTouchIDs = inpututil.AppendJustReleasedTouchIDs(g.releasedTouchIDs[:0])

	if inpututil.IsMouseButtonJustReleased(ebiten.MouseButtonLeft) || len(g.releasedTouchIDs) != 0 {
		g.changedPotential = true
		g.mousePressed = false
	}

	return nil
}

func OKLCHToColor(L, C, H float64) color.RGBA {
	c := oklab.Oklch{
		L: L,
		C: C,
		H: H,
	}

	r, g, b, a := c.RGBA()

	return color.RGBA{
		R: uint8(r >> 8),
		G: uint8(g >> 8),
		B: uint8(b >> 8),
		A: uint8(a >> 8),
	}
}

// Draw draws the game screen.
// Draw is called every frame (typically 1/60[s] for 60Hz display).
func (g *Game) Draw(screen *ebiten.Image) {
	screen.Fill(color.RGBA{0, 0, 0, 255})
	for i := range N - 1 {
		vector.StrokeLine(screen,
			float32(Width)*(float32(i)/float32(N-1)),
			float32(Height)-float32(Height)*float32(g.potential[i]),
			float32(Width)*(float32(i+1)/float32(N-1)),
			float32(Height)-float32(Height)*float32(g.potential[i+1]),
			2,
			color.RGBA{255, 255, 255, 255},
			true)
	}
	for i := range g.maxdisplayedeigenvalue {
		if g.eigenvalues[i] <= V0Max {
			phase := math.Sin(float64(time.Now().UnixNano()*10^-9) * TimeFactor * eV / hbar)
			centerY := float32(g.eigenvalues[i] / V0Max)
			yPrev := float32(g.eigenvectors.At(0, i) * phase)
			for j := range N - 1 {
				y := float32(g.eigenvectors.At(j+1, i) * phase)

				vector.StrokeLine(screen,
					float32(Width)*(float32(j)/float32(N-1)),
					float32(Height)-(float32(Height)*(centerY)+yPrev),
					float32(Width)*(float32(j+1)/float32(N-1)),
					float32(Height)-(float32(Height)*(centerY)+y),
					2,
					g.computedColors[i],
					true)
				yPrev = y
			}
		}
	}

}

// Layout takes the outside size (e.g., the window size) and returns the (logical) screen size.
// If you don't have to adjust the screen size with the outside size, just return a fixed size.
func (g *Game) Layout(outsideWidth, outsideHeight int) (screenWidth, screenHeight int) {
	return Width, Height
}

func main() {
	potential := make([]float64, N)
	for i := range potential {
		if i < 3*N/4 && i >= N/4 {
			potential[i] = 0.1
		} else {
			potential[i] = 0.9
		}
	}
	game := &Game{
		potential:        potential,
		changedPotential: true,
	}
	// Specify the window size as you like. Here, a doubled size is specified.
	ebiten.SetWindowSize(Width, Height)
	ebiten.SetWindowTitle("Poisson solver")
	ebiten.SetTPS(240)
	// Call ebiten.RunGame to start your game loop.
	if err := ebiten.RunGame(game); err != nil {
		log.Fatal(err)
	}

}
