package grow

import "io"
import "math/rand"
import "sort"

import "gopkg.in/yaml.v2"

import "github.com/firestuff/subcoding/vm"

type Definition struct {
	GlobalMemorySize   uint64 `yaml:"global_memory_size"`
	FunctionMemorySize uint64 `yaml:"function_memory_size"`
	InstructionLimit   uint64 `yaml:"instruction_limit"`

	InstructionsPerFunctionMean   uint64 `yaml:"instructions_per_function_mean"`
	InstructionsPerFunctionStdDev uint64 `yaml:"instructions_per_function_std_dev"`

	ReturnToBestPeriodMean uint64 `yaml:"return_to_best_period_mean"`

	Samples []*Sample `yaml:"samples"`

	// Sample indices ranked by each output dimension
	SampleRanks [][]int
}

type Score struct {
	Current uint64
	Total   uint64
}

func NewDefinition(r io.Reader) (*Definition, error) {
	dec := yaml.NewDecoder(r)
	dec.SetStrict(true)

	def := &Definition{}

	err := dec.Decode(def)
	if err != nil {
		return nil, err
	}

	// TODO: Test & handle non-consistent In and Out dimensions

	def.buildSampleRanks()

	return def, nil
}

func (def *Definition) Grow(statusChan chan<- Status) (*vm.Program, error) {
	if statusChan != nil {
		defer close(statusChan)
	}

	status := Status{}

	if statusChan != nil {
		statusChan <- status
	}

	prog := &vm.Program{
		GlobalMemorySize:   def.GlobalMemorySize,
		FunctionMemorySize: def.FunctionMemorySize,
		InstructionLimit:   def.InstructionLimit,
		Functions: []*vm.Function{
			&vm.Function{},
		},
	}

	for {
		status.Attempts++

		if def.ReturnToBestPeriodMean != 0 && rand.Uint64()%def.ReturnToBestPeriodMean == 0 {
			prog = status.BestProgram.Copy()
		}

		Mutate(def, prog)

		scores, err := def.score(prog)
		if err != nil {
			// Can never get best score
			continue
		}

		if !def.scoreIsBetter(status.BestScores, scores) {
			continue
		}

		err = def.minifyProgram(prog)
		if err != nil {
			return nil, err
		}

		status.BestScores = scores
		status.BestProgram = prog.Copy()

		if statusChan != nil {
			statusChan <- status
		}

		if def.scoreIsPerfect(status.BestScores) {
			return prog, nil
		}
	}
}

func (def *Definition) buildSampleRanks() {
	for col := 0; col < len(def.Samples[0].Out); col++ {
		rank := []int{}

		for i := 0; i < len(def.Samples); i++ {
			rank = append(rank, i)
		}

		sort.SliceStable(rank, func(i, j int) bool {
			return def.Samples[i].Out[col] < def.Samples[j].Out[col]
		})

		def.SampleRanks = append(def.SampleRanks, rank)
	}
}

func (def *Definition) score(prog *vm.Program) ([]*Score, error) {
	outputs := [][]uint64{}

	for _, sample := range def.Samples {
		state, err := vm.NewState(prog)
		if err != nil {
			return nil, err
		}

		sample.SetInputs(state)

		err = state.Execute()
		if err != nil {
			return nil, err
		}

		output := []uint64{}
		for i := 0; i < len(def.Samples[0].Out); i++ {
			// TODO: Handle signedness?
			output = append(output, state.GlobalMemory().MustReadUnsigned(uint64(i)))
		}
		outputs = append(outputs, output)
	}

	// TODO: Points for proximity to target values?
	// TODO: Points for correlation coeficient with target values across samples?

	return []*Score{
		def.scoreMatching(outputs),
		def.scoreRank(outputs),
	}, nil
}

func (def *Definition) scoreMatching(outputs [][]uint64) *Score {
	ret := &Score{}

	for s, sample := range def.Samples {
		for o, out := range sample.Out {
			ret.Total++

			if outputs[s][o] == out {
				ret.Current++
			}
		}
	}

	return ret
}

func (def *Definition) scoreRank(outputs [][]uint64) *Score {
	ranks := [][]int{}

	for col := 0; col < len(outputs[0]); col++ {
		rank := []int{}

		for i := 0; i < len(def.Samples); i++ {
			rank = append(rank, i)
		}

		sort.SliceStable(rank, func(i, j int) bool {
			return outputs[i][col] < outputs[j][col]
		})

		ranks = append(ranks, rank)
	}

	ret := &Score{}

	for col, vals := range ranks {
		for i, val := range vals {
			ret.Total++

			if val == def.SampleRanks[col][i] {
				ret.Current++
			}
		}
	}

	return ret
}

func (def *Definition) scoreIsBetter(old, new []*Score) bool {
	if old == nil {
		return true
	}

	for i, score := range new {
		best := old[i]

		switch {
		case score.Current == best.Current:
			continue

		case score.Current > best.Current:
			return true

		case score.Current < best.Current:
			return false
		}
	}

	// Unchanged
	return false
}

func (def *Definition) scoreIsPerfect(scores []*Score) bool {
	for _, score := range scores {
		if score.Current < score.Total {
			return false
		}
	}

	return true
}

func (def *Definition) minifyProgram(prog *vm.Program) error {
	for f := 0; f < len(prog.Functions); f++ {
		err := def.minifyFunction(prog, f)
		if err != nil {
			return err
		}
	}

	return nil
}

func (def *Definition) minifyFunction(prog *vm.Program, f int) error {
	baseScores, err := def.score(prog)
	if err != nil {
		return err
	}

	for loop := true; loop; {
		loop = false

		for i := 0; i < len(prog.Functions[f].Instructions); i++ {
			origInstructions := prog.Functions[f].Instructions

			tmp := make([]*vm.Instruction, len(prog.Functions[f].Instructions))
			copy(tmp, prog.Functions[f].Instructions)
			prog.Functions[f].Instructions = append(tmp[:i], tmp[i+1:]...)

			newScores, err := def.score(prog)
			// XXX: Use all scores
			if err == nil && newScores[0].Current >= baseScores[0].Current {
				loop = true
				break
			} else {
				prog.Functions[f].Instructions = origInstructions
			}
		}
	}

	return nil
}