package main

import (
	"fmt"
	"image"
	"image/color"
	"log"
	"math"
	"math/rand"
	"os"

	"gioui.org/app"
	"gioui.org/font/gofont"
	"gioui.org/layout"
	"gioui.org/op"
	"gioui.org/op/clip"
	"gioui.org/op/paint"
	"gioui.org/text"
	"gioui.org/widget/material"
)

const pi Radians = 3.14159

type C = layout.Context
type D = layout.Dimensions

var darkPalette = material.Palette{
	Bg:         color.NRGBA{R: 0x28, G: 0x28, B: 0x28, A: 0xff},
	Fg:         color.NRGBA{R: 0xeb, G: 0xdb, B: 0xb2, A: 0xff},
	ContrastBg: color.NRGBA{R: 0x48, G: 0x85, B: 0x88, A: 0xff},
	ContrastFg: color.NRGBA{R: 0xeb, G: 0xdb, B: 0xb2, A: 0xff},
}

func main() {
	go func() {
		window := new(app.Window)
		window.Option(app.Title("Physarum"))

		err := run(window)
		if err != nil {
			log.Fatal(err)
		}
		os.Exit(0)
	}()

	app.Main()
}

func run(window *app.Window) error {
	theme := material.NewTheme()
	theme.Palette = darkPalette
	theme.Shaper =
		text.NewShaper(text.WithCollection(gofont.Collection()))

	seed := int64(0)
	s := Simulation{}
	s.random = rand.New(rand.NewSource(seed))
	s.trails = s._trails[:]
	s.headings = s._headings[:]
	s.halfFoV = Radians(3 * pi / 4)
	s.focus = 2
	s.offset++
	s.move = 4
	s.turn = Radians(pi / 16)
	s.deposit = 128

	initialPos := GridId(524288) // GridId(len(s.grid) >> 1)
	initialHeading := Radians(0)
	for i := range s.headings {
		s.headings[i] = initialHeading
		s.trails[i*TrailLenMax] = initialPos
	}

	var ops op.Ops

	for {
		switch e := window.Event().(type) {
		case app.DestroyEvent:
			return e.Err

		case app.FrameEvent:
			gtx := app.NewContext(&ops, e)

			if s.grid == nil {
				/*squareMax := min(
					gtx.Constraints.Max.X,
					gtx.Constraints.Max.Y,
				)*/
				// Turn this into a power of 2 for fast modulo
				/*nearestPow2 := log2(uint32(squareMax))
				squareMax = int(1 << nearestPow2)
				s.grid = make([]uint8, squareMax<<nearestPow2)*/
				s.dimension = 1024
				s.dimLog2 = 10
				s.grid = make([]uint8, s.dimension<<s.dimLog2)
			}
			paint.Fill(&ops, theme.Bg)
			s.Step()
			s.Layout(gtx, theme)

			e.Source.Execute(op.InvalidateCmd{})

			e.Frame(gtx.Ops)
		}
	}
}

type GridId uint32
type TrailId uint8
type Radians float32

const (
	AgentCount  int = 128
	TrailLenMax int = 64
)

type Pos struct{ x, y int32 }

func (p1 Pos) Add(p2 Pos) Pos {
	return Pos{p1.x + p2.x, p1.y + p2.y}
}

func (p1 Pos) Sub(p2 Pos) Pos {
	return Pos{p1.x - p2.x, p1.y - p2.y}
}

type Simulation struct {
	random *rand.Rand

	_trails   [TrailLenMax * AgentCount]GridId
	_headings [AgentCount]Radians
	trails    []GridId  // Grid ids for agent trails
	headings  []Radians // Agent headings

	grid      []uint8 // Screen grid
	dimension int32   // Square dimension of client space
	dimLog2   uint32  // Log base-2 of dimension

	halfFoV Radians // Agent sensor angle
	focus   float64 // Agent sensor distance
	move    float64 // Agent move distance
	turn    Radians // Agent turn angle
	offset  uint8   // Index for trail circular buffers
	deposit uint8   // Agent deposit
}

func (s *Simulation) Step() {
	dxx := math.Cos(float64(s.halfFoV))
	dyx := math.Sin(float64(s.halfFoV))
	dxy := -dyx
	dyy := dxx

	for i, heading := range s.headings {
		// This may roll over, so we compute mod TrailLenMax.
		prevOffset := (int(s.offset) + TrailLenMax - 1) & (TrailLenMax - 1)
		nextOffset := (int(s.offset) + TrailLenMax) & (TrailLenMax - 1)
		prev := i*TrailLenMax + prevOffset
		next := i*TrailLenMax + nextOffset

		gridId := s.trails[prev]
		p0 := gridToPos(gridId, s.dimension, s.dimLog2)

		// Sense
		cx := s.focus * math.Cos(float64(heading))
		cy := s.focus * math.Sin(float64(heading))
		lx, ly := int32(cx*dxx+cy*dxy), int32(cy*dyy+cx*dyx)
		rx, ry := int32(cx*dxx-cy*dxy), int32(cy*dyy-cx*dyx)

		centerId :=
			posToGrid(
				Pos{p0.x + int32(cx), p0.y - int32(cy)},
				s.dimLog2,
			)
		leftId := posToGrid(Pos{p0.x + lx, p0.y - ly}, s.dimLog2)
		rightId := posToGrid(Pos{p0.x + rx, p0.y - ry}, s.dimLog2)

		var left, center, right uint8
		if 0 <= leftId && leftId <= GridId(len(s.grid)-1) {
			left = s.grid[leftId]
		}
		if 0 <= centerId && centerId <= GridId(len(s.grid)-1) {
			center = s.grid[centerId]
		}
		if 0 <= rightId && rightId <= GridId(len(s.grid)-1) {
			right = s.grid[rightId]
		}

		// Rotate
		switch {
		case left < center:
			if center <= right {
				heading -= s.turn // Steer right
			}
			// Otherwise, stay the course.
		case left >= center:
			if center > right {
				heading += s.turn // Steer left
				break
			}
			fallthrough
		default:
			if s.random.Float32() < 0.5 {
				heading += s.turn // Steer left
			} else {
				heading -= s.turn // Steer right
			}
		}
		s.headings[i] = heading

		// Move
		cos := math.Cos(float64(heading))
		sin := math.Sin(float64(heading))
		mx := s.move * cos
		my := s.move * sin
		p1 := Pos{
			int32(float64(p0.x) + mx),
			int32(float64(p0.y) + my),
		}
		// If we exceed the bounds of the grid, stop at the
		// boundary and change heading according to the angle of
		// incidence.
		gridMin := Pos{0, 0}
		gridMax := Pos{s.dimension - 1, s.dimension - 1}
		fromWallMin := p1.Sub(gridMin)
		fromWallMax := gridMax.Sub(p1)
		cnt := 0
		for !inBounds(p1, s.dimension) && cnt < 4 {
			cnt++
			if !inBounds(p1, s.dimension) && cnt >= 2 {
				fmt.Fprintf(os.Stderr, "\nstart at %v; out of bounds at %v\n", p0, p1)
			}
			if p1.x < gridMin.x { // Left wall
				s.headings[i] = pi - heading
				// Cosine cannot be zero when distance to wall is zero
				distToImpact := math.Abs(float64(fromWallMin.x) / cos)
				p1.x = gridMin.x                  // 0
				p1.y += int32(sin * distToImpact) // -1
				if inBounds(p1, s.dimension) {
					break
				}
			}
			if gridMax.x < p1.x { // Right wall
				s.headings[i] = pi - heading
				distToImpact := math.Abs(float64(fromWallMax.x) / cos)
				p1.x = gridMax.x
				p1.y += int32(sin * distToImpact)
				if inBounds(p1, s.dimension) {
					break
				}
			}
			if p1.y < gridMin.y { // Bottom wall
				s.headings[i] = 2*pi - heading
				distToImpact := math.Abs(float64(fromWallMin.y) / sin)
				p1.x += int32(cos * distToImpact)
				p1.y = gridMin.y
				if inBounds(p1, s.dimension) {
					break
				}
			}
			if gridMax.y < p1.y { // Top wall
				s.headings[i] = 2*pi - heading
				distToImpact := math.Abs(float64(fromWallMax.y) / sin)
				p1.x += int32(cos * distToImpact)
				p1.y = gridMax.y
				if inBounds(p1, s.dimension) {
					break
				}
			}
		}
		if !inBounds(p1, s.dimension) {
			fmt.Fprintf(os.Stderr, "\nstart at %v; out of bounds at %v\n", p0, p1)
			panic("blarg")
		}
		gridId = posToGrid(p1, s.dimLog2)
		s.trails[next] = gridId

		// Deposit and diffuse
		s.grid[gridId] = s.deposit
		s.diffuse(gridId)

		// Decay trail
		for j := 1; j < TrailLenMax; j++ {
			tid := (int(s.offset) + j) & (TrailLenMax - 1)
			s.grid[s.trails[tid]]--
		}
	}
	s.offset++
	s.offset &= uint8(TrailLenMax - 1)
}

func (s *Simulation) diffuse(id GridId) {
	value := s.grid[id] >> 4
	s.grid[id] = s.grid[id] >> 1

	p := gridToPos(id, s.dimension, s.dimLog2)
	for i := -1; i < 2; i++ {
		for j := -1; j < 2; j++ {
			bp := Pos{p.x + int32(j), p.y + int32(i)}

			if !inBounds(bp, s.dimension) {
				// For now, just send deposits that fall outside the
				// grid back to the center of the ball.
				s.grid[id] += value
				continue
			}
			bid := posToGrid(bp, s.dimLog2)
			// Protect against overflow.
			s.grid[bid] = max(s.grid[bid], s.grid[bid]+value)
		}
	}
}

func (s *Simulation) gridMean(id GridId) (mean uint8) {
	p := gridToPos(id, s.dimension, s.dimLog2)
	n := 0
	for i := -1; i < 2; i++ {
		for j := -1; j < 2; j++ {
			bp := Pos{p.x + int32(j), p.y + int32(i)}

			if !inBounds(bp, s.dimension) {
				continue
			}
			mean += s.grid[posToGrid(bp, s.dimLog2)]
			n++
		}
	}
	return uint8(float32(mean) / float32(n))
}

func (s *Simulation) Layout(
	gtx layout.Context,
	th *material.Theme,
) D {
	circle := clip.Ellipse{Max: image.Pt(8, 8)}

	// Outer margin
	defer op.Offset(image.Pt(10, 10)).Push(gtx.Ops).Pop()

	halfW, halfH := circle.Max.X>>1, circle.Max.Y>>1

	trailColor := th.Fg

	for _, gridId := range s.trails {
		if gridId == 0 {
			continue
		}
		trailColor.A = s.gridMean(gridId)

		p := gridToPos(gridId, s.dimension, s.dimLog2)
		pos := image.Pt(int(p.x)-halfW, int(s.dimension-1-p.y)+halfH)

		macro := op.Record(gtx.Ops)
		stack := op.Offset(pos).Push(gtx.Ops)
		paint.FillShape(gtx.Ops, trailColor, circle.Op(gtx.Ops))
		stack.Pop()
		c := macro.Stop()
		op.Defer(gtx.Ops, c)
	}
	return layout.Dimensions{
		Size: image.Pt(int(s.dimension), int(s.dimension)),
	}
}

func inBounds(p Pos, dim int32) bool {
	return -1 < p.x && p.x < dim &&
		-1 < p.y && p.y < dim
}

func posToGrid(p Pos, logOfDim uint32) GridId {
	id := uint32(p.y) << logOfDim
	id |= uint32(p.x)
	return GridId(id)
}

func gridToPos(id GridId, dim int32, logOfDim uint32) Pos {
	x := int32(id) & (dim - 1)
	y := int32(id) >> logOfDim
	return Pos{x: x, y: y}
}

func log2(x uint32) uint32 {
	b := []uint32{0x2, 0xC, 0xF0, 0xFF00, 0xFFFF0000}
	s := []uint32{1, 2, 4, 8, 16}

	var r uint32 = 0

	for i := 4; i >= 0; i-- {
		if x&b[i] != 0 {
			x >>= s[i]
			r |= s[i]
		}
	}
	return r
}