package main

import (
	"fmt"
	"git.openprivacy.ca/sarah/microworlds/core"
	"git.openprivacy.ca/sarah/microworlds/experiments"
	"github.com/faiface/pixel/pixelgl"
	"github.com/wcharczuk/go-chart"
	"image/color"
	"math/rand"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
)

type Player struct {
	dna  [40]int
	Team int
}

func (player *Player) Setup(env *core.Environment, t *core.Turtle) {
	for i := 0; i < len(player.dna); i++ {
		player.dna[i] = rand.Intn(NumGenes)
	}
	t.SetAttribute("team", strconv.Itoa(player.Team))
	if player.Team == 1 {
		t.SetColor(color.RGBA{0xFF, 0x00, 0x00, 0xFF})
	} else {
		t.SetColor(color.RGBA{0x00, 0x00, 0xFF, 0xFF})
	}

}

const (
	Wiggle = iota
	FollowMyTeam
	FollowOtherTeam
	FollowBall
	Drop
	RetreatMyTeam
	RetreatOtherTeam
	RetreatBall
	TurnAround
	HeadTowardsOurGoal
	HeadTowardsOtherGoal
)

var NumGenes = 11

func (player *Player) Run(env *core.Environment, t *core.Turtle) {

	MyTeam := strconv.Itoa(player.Team)
	OtherTeam := "1"
	if player.Team == 1 {
		OtherTeam = "2"
	}

	ahead := t.Check(env)
	if ahead != nil {
		ball, ok := ahead.GetActor().(*Ball)
		if ok {
			ahead.SetHeading(t.Heading())
			if rand.Intn(5) == 0 {
				ahead.Wiggle()
			}
			ahead.Step(env)
			ball.Check(env, ahead)
		}
	}
	t.Step(env)

	for i := 0; i < len(player.dna); i++ {
		switch player.dna[i] {
		case TurnAround:
			t.TurnAround()
		case Wiggle:
			t.Wiggle()
		case FollowMyTeam:
			if i+2 < len(player.dna) {
				t.FollowGradient(env, player.dna[i+1], float32(player.dna[i+2]), MyTeam)
				i += 2
			}
		case FollowOtherTeam:
			if i+2 < len(player.dna) {
				t.FollowGradient(env, player.dna[i+1], float32(player.dna[i+2]), OtherTeam)
				i += 2
			}
		case FollowBall:
			t.FollowGradient(env, 1, 0, "ball")
		case Drop:
			if i+1 < len(player.dna) {
				t.Drop(env, float32(player.dna[i+1]), strconv.Itoa(player.Team))
			}
		case RetreatMyTeam:
			if i+2 < len(player.dna) {
				t.AvoidAverageGradient(env, player.dna[i+1], float32(player.dna[i+2]), MyTeam)
				i += 2
			}
		case RetreatOtherTeam:
			if i+2 < len(player.dna) {
				t.AvoidAverageGradient(env, player.dna[i+1], float32(player.dna[i+2]), OtherTeam)
				i += 2
			}
		case RetreatBall:
			if i+2 < len(player.dna) {
				t.AvoidAverageGradient(env, player.dna[i+1], float32(player.dna[i+2]), "ball")
				i += 2
			}
		case HeadTowardsOurGoal:
			x, _ := t.Pos()
			if player.Team == 1 {
				if x > 1 {
					t.SetHeading(7)
				}
			} else {
				if x < 299 {
					t.SetHeading(3)
				}
			}
		case HeadTowardsOtherGoal:
			x, _ := t.Pos()
			if player.Team == 2 {
				if x > 1 {
					t.SetHeading(7)
				}
			} else {
				if x < 299 {
					t.SetHeading(3)
				}
			}
		}
	}
}

func (player *Player) Mutate() {
	mutatePoint := rand.Intn(len(player.dna))
	fmt.Printf("\t\t Mutating %v at %v\n", player.Team, mutatePoint)
	player.dna[mutatePoint] = rand.Intn(NumGenes)
}

func (player *Player) Clone(parentA *Player, parentB *Player) {
	fmt.Printf("\t\tCrossing Over team %v with:%v %v\n", player.Team, parentA.dna, parentB.dna)
	crossoverPoint := rand.Intn(len(player.dna))
	copy(player.dna[0:crossoverPoint], parentA.dna[0:crossoverPoint])
	copy(player.dna[crossoverPoint:], parentB.dna[crossoverPoint:])
}

type Ball struct {
	Team1Score int
	Team2Score int
}

func (b *Ball) Setup(env *core.Environment, t *core.Turtle) {
	t.SetXY(50, 50)
	t.SetColor(color.RGBA{0xFF, 0xFF, 0xFF, 0xFF})
}

func (b *Ball) Run(env *core.Environment, t *core.Turtle) {
	b.Check(env, t)
	t.Drop(env, 10, "ball")
}

func (b *Ball) Check(env *core.Environment, t *core.Turtle) {
	x, _ := t.Pos()
	if x > 98 {
		fmt.Printf("\t\tTeam 1 Scored!!! %v\n", env.Step)
		b.Reset(env, t)
		b.Team1Score++
	} else if x < 1 {
		fmt.Printf("\t\tTeam 2 Scored!!! %v\n", env.Step)
		b.Reset(env, t)
		b.Team2Score++
	}
}

func (b *Ball) Reset(env *core.Environment, t *core.Turtle) {
	x, y := t.Pos()
	t.SetHeading(rand.Intn(8))
	env.Leave(x, y)
	t.SetXY(50, 50)
	env.Occupy(t, 50, 50)
}

func mainrun() {
	ball := new(Ball)
	experiment := new(experiments.Experiment)
	experiment.InitializeExperiment()
	experiment.InitNTurtles(func() core.Actor {
		player := new(Player)
		player.Team = 1
		return player
	}, 25)
	experiment.InitNTurtles(func() core.Actor {
		player := new(Player)
		player.Team = 2
		return player
	}, 25)
	experiment.InitNTurtles(func() core.Actor {
		return ball
	}, 1)
	experiment.InitPheromone("1", color.RGBA{0xFF, 0x00, 0x00, 0x00})
	experiment.InitPheromone("2", color.RGBA{0x00, 0x00, 0xFF, 0x00})
	experiment.InitPheromone("ball", color.RGBA{0xff, 0xff, 0xff, 0xff})

	x := []float64{-2}
	team1 := []float64{0}
	team2 := []float64{1}

	var graph chart.Chart
	experiment.AddPlot("Goals Scored", func(environment *core.Environment, turtles []*core.Turtle) *chart.Chart {
		x = append(x, float64(environment.Step))
		team1 = append(team1, float64(ball.Team1Score))
		team2 = append(team2, float64(ball.Team2Score))
		graph = chart.Chart{
			Background: chart.Style{
				Padding: chart.Box{
					Top: 50,
				},
			},
			XAxis: chart.XAxis{Name: "Time Step", NameStyle: chart.Style{Show: true}, Style: chart.Style{Show: true, TextRotationDegrees: 90}, ValueFormatter: func(v interface{}) string {
				return fmt.Sprintf("%d", int(v.(float64)))
			}},
			YAxis: chart.YAxis{Name: "Goals Scored", NameStyle: chart.Style{Show: true}, Style: chart.Style{Show: true}, ValueFormatter: func(v interface{}) string {
				return fmt.Sprintf("%d", int(v.(float64)))
			}},
			Series: []chart.Series{
				chart.ContinuousSeries{
					XValues: x,
					YValues: team1,
				},
				chart.ContinuousSeries{
					XValues: x,
					YValues: team2,
				},
			},
		}
		return &graph
	})

	gamelengths := []int{}
	gamewins := []int{}
	gamescores := []int{}
	lastGameStart := 0

	experiment.OnStep = func(environment *core.Environment, turtles []*core.Turtle, i int) {

		environment.EvaporateAndDiffuse(0.95, "ball")
		environment.EvaporateAndDiffuse(0.95, "1")
		environment.EvaporateAndDiffuse(0.95, "2")

		if environment.Step-lastGameStart > 4000 {
			fmt.Printf("Time!\n")

			if ball.Team1Score == ball.Team2Score {
				fmt.Printf("Draw!\n")
				gamelengths = append(gamelengths, environment.Step-lastGameStart)
				gamewins = append(gamewins, 0)
				gamescores = append(gamescores, ball.Team1Score, ball.Team2Score)
				lastGameStart = environment.Step

				ball.Team2Score = 0
				ball.Team1Score = 0

				red, blue := GetTeams(turtles)
				for _, r := range red {
					player := r.GetActor().(*Player)
					player.Mutate()
					player.Mutate()
					player.Mutate()
					player.Mutate()
					player.Mutate()
				}
				for _, b := range blue {
					player := b.GetActor().(*Player)
					player.Mutate()
					player.Mutate()
					player.Mutate()
					player.Mutate()
					player.Mutate()

				}
				for i, lengths := range gamelengths {
					fmt.Printf("Generation: %v, %v time steps Team %v Won (%v - %v)\n", i, lengths, gamewins[i], gamescores[i*2], gamescores[(i*2)+1])
				}
			} else {
				ball.Team1Score++
				ball.Team2Score++
			}
		}

		if ball.Team1Score >= 3 && ball.Team2Score < 3 {
			fmt.Printf("Team 1 Win")
			red, blue := GetTeams(turtles)
			for _, b := range blue {
				player := b.GetActor().(*Player)
				parentA := rand.Intn(len(red))
				parentB := rand.Intn(len(red))
				player.Clone(red[parentA].GetActor().(*Player), red[parentB].GetActor().(*Player))
				player.Mutate()
			}
			adjust := 0
			if environment.Step-lastGameStart > 4000 {
				adjust = (environment.Step - lastGameStart) - 4000
			}

			gamelengths = append(gamelengths, environment.Step-lastGameStart)
			gamewins = append(gamewins, 1)
			gamescores = append(gamescores, ball.Team1Score-adjust, ball.Team2Score-adjust)
			lastGameStart = environment.Step

			for i, lengths := range gamelengths {
				fmt.Printf("Generation: %v, %v time steps Team %v Won (%v - %v)\n", i, lengths, gamewins[i], gamescores[i*2], gamescores[(i*2)+1])
			}

			ball.Team2Score = 0
			ball.Team1Score = 0

		} else if ball.Team2Score >= 3 && ball.Team1Score < 3 {
			fmt.Printf("Team 2 Win")
			red, blue := GetTeams(turtles)
			for _, r := range red {
				player := r.GetActor().(*Player)
				parentA := rand.Intn(len(blue))
				parentB := rand.Intn(len(blue))
				player.Clone(blue[parentA].GetActor().(*Player), blue[parentB].GetActor().(*Player))
				player.Mutate()
			}

			adjust := 0
			if environment.Step-lastGameStart > 4000 {
				adjust = (environment.Step - lastGameStart) - 4000
			}

			gamelengths = append(gamelengths, environment.Step-lastGameStart)
			gamewins = append(gamewins, 2)
			gamescores = append(gamescores, ball.Team1Score-adjust, ball.Team2Score-adjust)
			lastGameStart = environment.Step

			ball.Team2Score = 0
			ball.Team1Score = 0

			for i, lengths := range gamelengths {
				fmt.Printf("Generation: %v, %v time steps Team %v Won (%v - %v)\n", i, lengths, gamewins[i], gamescores[i*2], gamescores[(i*2)+1])
			}
		}

	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for sig := range c {
			fmt.Printf("Got Signal %v", sig)
			pprof.StopCPUProfile()
			os.Exit(0)
		}
	}()

	experiment.Run()
}

func main() {
	pixelgl.Run(mainrun)
}

func GetTeams(turtles []*core.Turtle) (redTeam, blueTeam []*core.Turtle) {
	for _, turtle := range turtles {
		if turtle.GetAttribute("team") == "1" {
			redTeam = append(redTeam, turtle)
		} else if turtle.GetAttribute("team") == "2" {
			blueTeam = append(blueTeam, turtle)
		}
	}
	return
}