package experiments

import (
	"flag"
	"git.openprivacy.ca/sarah/microworlds/core"
	"git.openprivacy.ca/sarah/microworlds/graphics"
	"github.com/veandco/go-sdl2/sdl"
	"github.com/wcharczuk/go-chart"
	"image/color"
	"log"
	"math/rand"
	"os"
	"runtime/pprof"
	"sync"
	"time"
)

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
var width = flag.Int("width", 300, "width of environment")
var height = flag.Int("height", 300, "height of environment")
var pxsize = flag.Int("pxsize", 1, "pixels per tile edge")
var numTurtles = flag.Int("numTurtles", 5000, "number of turtles")

type Experiment struct {
	env                *core.Environment
	turtles            []*core.Turtle
	graphics           *graphics.Graphics
	initializedTurtles int
	pheromones         []string
	OnStep             func(*core.Environment, []*core.Turtle, int)
	plots              []*graphics.Plot
	running            bool
}

func (e *Experiment) InitializeExperiment() {
	// We don't need real randomness
	rand.Seed(time.Now().Unix())

	flag.Parse()
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}

	if err := sdl.Init(sdl.INIT_EVERYTHING); err != nil {
		panic(err)
	}

	e.OnStep = func(env *core.Environment, turtles []*core.Turtle, i int) {
		for _, name := range e.pheromones {
			env.EvaporateAndDiffuse(0.95, name)
		}
	}

	e.env = core.NewEnvironment(*width, *height)
	e.turtles = make([]*core.Turtle, 0)
	e.graphics = graphics.NewGraphics(int32(*width), int32(*height), int32(*pxsize))
}

func (e *Experiment) GetNumTurtles() int {
	return (*numTurtles)
}

func (e *Experiment) InitPheromone(name string, col color.Color) {
	e.env.InitPheromone(name)
	r, g, b, a := col.RGBA()
	e.pheromones = append(e.pheromones, name)
	e.graphics.ColorPheromone(name, [4]uint8{uint8(r), uint8(g), uint8(b), uint8(a)})
}

func (e *Experiment) InitEnvironment(f func(environment *core.Environment)) {
	f(e.env)
}

func (e *Experiment) InitNTurtles(f func() core.Actor, num int) {
	numSuccess := 0
	for i := e.initializedTurtles; i < e.initializedTurtles+num; i++ {
		t := core.NewTurtle(e.env, f())
		if t != nil {
			e.turtles = append(e.turtles, t)
			numSuccess++
		}
	}
	e.initializedTurtles += numSuccess
}

func (e *Experiment) InitTurtles(f func() core.Actor) {
	e.InitNTurtles(f, (*numTurtles))
}

func (e *Experiment) AddPlot(title string, plotfunc func(environment *core.Environment, turtles []*core.Turtle) *chart.Chart) {
	plot := graphics.NewPlot(title, int32(*width), int32(*height), int32(*pxsize))
	plot.GeneratePlot = plotfunc
	go plot.RenderAsync()
	e.plots = append(e.plots, plot)
}

func (e *Experiment) Stop() {
	e.running = false
}

func (e *Experiment) Restart() {
	e.running = false
}

func (e *Experiment) Run() {
	wait := sync.WaitGroup{}

	e.running = true
	wait.Add(1)
	e.env.ResetQuadtree()
	go func() {
		step := 0
		for e.running {
			e.env.Step = step
			e.graphics.Render(e.env, e.turtles)
			e.OnStep(e.env, e.turtles, step)

			for _, plot := range e.plots {
				plot.Render(e.env, e.turtles)
			}

			newTurtles := make([]*core.Turtle, 0)

			for _, t := range e.turtles {
				t.Run(e.env)
				if t.GetAttribute("status") != "dead" {
					newTurtles = append(newTurtles, t)
				} else {
					e.env.Leave(t.Pos()) // Dead turtles occupy no space
				}
			}

			e.turtles = newTurtles
			e.env.ResetQuadtree()
			for _, t := range e.turtles {
				e.env.InsertIntoQuadTree(t)
			}

			step++
		}
		wait.Done()
	}()

	wait.Add(1)
	for e.running {
		for event := sdl.PollEvent(); event != nil; event = sdl.PollEvent() {
			switch event.(type) {
			case *sdl.QuitEvent:
				e.running = false
				break
			}
		}
	}
	wait.Done()
	wait.Wait()
}