module neuroevolution

import rand
import math

fn random_clamped() f64 {
    return rand.f64() * 2 - 1
}

pub fn activation(a f64) f64 {
    ap := (-a) / 1
    return 1 / (1 + math.exp(ap))
}

fn round(a int, b f64) int {
    return int(math.round(f64(a) * b))
}

struct Neuron {
mut:
    value   f64
    weights []f64
}

fn (mut n Neuron) populate(nb int) {
    for _ in 0 .. nb {
        n.weights << random_clamped()
    }
}

struct Layer {
    id int
mut:
    neurons []Neuron
}

fn (mut l Layer) populate(nb_neurons int, nb_inputs int) {
    for _ in 0 .. nb_neurons {
        mut n := Neuron{}
        n.populate(nb_inputs)
        l.neurons << n
    }
}

pub struct Network {
mut:
    layers []Layer
}

fn (mut n Network) populate(network []int) {
    assert network.len >= 2
    input := network[0]
    hiddens := network[1..network.len - 1]
    output := network[network.len - 1]
    mut index := 0
    mut previous_neurons := 0
    mut input_layer := Layer{
        id: index
    }
    input_layer.populate(input, previous_neurons)
    n.layers << input_layer
    previous_neurons = input
    index++
    for hidden in hiddens {
        mut hidden_layer := Layer{
            id: index
        }
        hidden_layer.populate(hidden, previous_neurons)
        previous_neurons = hidden
        n.layers << hidden_layer
        index++
    }
    mut output_layer := Layer{
        id: index
    }
    output_layer.populate(output, previous_neurons)
    n.layers << output_layer
}

fn (n Network) get_save() Save {
    mut save := Save{}
    for layer in n.layers {
        save.neurons << layer.neurons.len
        for neuron in layer.neurons {
            for weight in neuron.weights {
                save.weights << weight
            }
        }
    }
    return save
}

fn (mut n Network) set_save(save Save) {
    mut previous_neurons := 0
    mut index := 0
    mut index_weights := 0
    n.layers = []
    for save_neuron in save.neurons {
        mut layer := Layer{
            id: index
        }
        layer.populate(save_neuron, previous_neurons)
        for mut neuron in layer.neurons {
            for i in 0 .. neuron.weights.len {
                neuron.weights[i] = save.weights[index_weights]
                index_weights++
            }
        }
        previous_neurons = save_neuron
        index++
        n.layers << layer
    }
}

pub fn (mut n Network) compute(inputs []f64) []f64 {
    assert n.layers.len > 0
    assert inputs.len == n.layers[0].neurons.len
    for i, input in inputs {
        n.layers[0].neurons[i].value = input
    }
    mut prev_layer := n.layers[0]
    for i in 1 .. n.layers.len {
        for j, neuron in n.layers[i].neurons {
            mut sum := f64(0)
            for k, prev_layer_neuron in prev_layer.neurons {
                sum += prev_layer_neuron.value * neuron.weights[k]
            }
            n.layers[i].neurons[j].value = activation(sum)
        }
        prev_layer = n.layers[i]
    }
    mut outputs := []f64{}
    mut last_layer := n.layers[n.layers.len - 1]
    for neuron in last_layer.neurons {
        outputs << neuron.value
    }
    return outputs
}

struct Save {
mut:
    neurons []int
    weights []f64
}

fn (s Save) clone() Save {
    mut save := Save{}
    save.neurons << s.neurons
    save.weights << s.weights
    return save
}

struct Genome {
    score   int
    network Save
}

struct Generation {
mut:
    genomes []Genome
}

fn (mut g Generation) add_genome(genome Genome) {
    mut i := 0
    for gg in g.genomes {
        if genome.score > gg.score {
            break
        }
        i++
    }
    g.genomes.insert(i, genome)
}

fn (g1 Genome) breed(g2 Genome, nb_child int) []Save {
    mut datas := []Save{}
    for _ in 0 .. nb_child {
        mut data := g1.network.clone()
        for i, weight in g2.network.weights {
            if rand.f64() <= 0.5 {
                data.weights[i] = weight
            }
        }
        for i, _ in data.weights {
            if rand.f64() <= 0.1 {
                data.weights[i] += (rand.f64() * 2 - 1) * 0.5
            }
        }
        datas << data
    }
    return datas
}

fn (g Generation) next(population int) []Save {
    mut nexts := []Save{}
    if population == 0 {
        return nexts
    }
    keep := round(population, 0.2)
    for i in 0 .. keep {
        if nexts.len < population {
            nexts << g.genomes[i].network.clone()
        }
    }
    random := round(population, 0.2)
    for _ in 0 .. random {
        if nexts.len < population {
            mut n := g.genomes[0].network.clone()
            for k, _ in n.weights {
                n.weights[k] = random_clamped()
            }
            nexts << n
        }
    }
    mut max := 0
    out: for {
        for i in 0 .. max {
            mut childs := g.genomes[i].breed(g.genomes[max], 1)
            for c in childs {
                nexts << c
                if nexts.len >= population {
                    break out
                }
            }
        }
        max++
        if max >= g.genomes.len - 1 {
            max = 0
        }
    }
    return nexts
}

pub struct Generations {
pub:
    population int
    network    []int
mut:
    generations []Generation
}

fn (mut gs Generations) first() []Save {
    mut out := []Save{}
    for _ in 0 .. gs.population {
        mut nn := Network{}
        nn.populate(gs.network)
        out << nn.get_save()
    }
    gs.generations << Generation{}
    return out
}

fn (mut gs Generations) next() []Save {
    assert gs.generations.len > 0
    gen := gs.generations[gs.generations.len - 1].next(gs.population)
    gs.generations << Generation{}
    return gen
}

fn (mut gs Generations) add_genome(genome Genome) {
    assert gs.generations.len > 0
    gs.generations[gs.generations.len - 1].add_genome(genome)
}

fn (mut gs Generations) restart() {
    gs.generations = []
}

pub fn (mut gs Generations) generate() []Network {
    saves := if gs.generations.len == 0 { gs.first() } else { gs.next() }
    mut nns := []Network{}
    for save in saves {
        mut nn := Network{}
        nn.set_save(save)
        nns << nn
    }
    if gs.generations.len >= 2 {
        gs.generations.delete(0)
    }
    return nns
}

pub fn (mut gs Generations) network_score(network Network, score int) {
    gs.add_genome(Genome{
        score:   score
        network: network.get_save()
    })
}