%% Assignment 04
clear all;
clc;

%% Initial parameters
max_generations=500;
generation=1;
target_fitness = 1000;
num_i_nodes = 6;
num_b_nodes = 1;
num_o_nodes = 1;
num_h_nodes = 0;
pop_size = 150;
fixed_topology=1;

stagnation.threshold=1e-2;
stagnation.number_generation=15;
refocus.threshold=1e-2;
refocus.number_generation=20;

initial.kill_percentage=0.2;
initial.number_for_kill=5;
initial.number_copy=5;

selection.pressure=2;

% crossover.percentage=0.8;
crossover.percentage=0;
crossover.probability_interspecies=0.0;
crossover.probability_multipoint=0.6;
crossover.rate = 0.1;
% Mutation related parameters
mutation.rate = 0.2;
mutation.perturbation_rate = 0.9;
mutation.perturbation_range=0.1;
mutation.cap = 5;
mutation.range = 2;
mutation.mutate_weight=0.9;
mutation.add_node_rate=0.03;
mutation.add_connection_rate=0.05;
mutation.recurrency_rate=0.0;
mutation.reenabled_rate=0.25;

speciation.c1 = 0;
speciation.c2 = 0;
speciation.c3 = 4;
speciation.threshold=5;

species_record(1).ID=0;
species_record(1).number_individuals=0;
species_record(1).generation_record=[];

% if fixed_topology == 1
%
%     num_nodes = num_i_nodes+num_b_nodes+num_o_nodes+num_h_nodes;
%     num_conections = num_i_nodes+num_b_nodes;
%     connections_from = 1:num_i_nodes+num_b_nodes;
%     connections_to = num_nodes*ones(1,num_conections);
%     global_innovation_number=num_conections+1;
% end

num_nodes = num_i_nodes+num_b_nodes+num_o_nodes+num_h_nodes;
num_conections = num_i_nodes+num_b_nodes;
connections_from = 1:num_i_nodes+num_b_nodes;
connections_to = num_nodes*ones(1,num_conections);
global_innovation_number=num_conections+1;
innovation_record = [];
max_gen_fit = [];
%% Initialize population
for ind=1:pop_size

    %node ID's
    %node type: 1=input 2=output 3=hidden 4=bias
    %node input state
    %node output state
    population(ind).node_genes=[1:(num_i_nodes+num_b_nodes+num_o_nodes+num_h_nodes);
        ones(1,num_i_nodes),4*ones(1,num_b_nodes),2*ones(1,num_o_nodes),3*ones(1,num_h_nodes);
        zeros(1,num_i_nodes),ones(1,num_b_nodes),zeros(1,num_o_nodes),zeros(1,num_h_nodes);
        zeros(1,num_nodes)];

    %innovation number
    %connection from,
    %connection to
    %weights (uniformly distributed in [-2 +2], all connections enabled)
    %enable bit
    population(ind).connection_genes=[1:num_conections;
        connections_from;
        connections_to;
        rand(1,num_conections)*2*mutation.range-mutation.range;
        ones(1,num_conections)];
    population(ind).wMat = weight_matrix(population(ind).node_genes,population(ind).connection_genes);
    population(ind).fitness=0;
    population(ind).species=0;
end
%% Innovation reccord for initial population
%Innovation ID
%Connections from
%Connection to
%New node(Highest)
%Generation this innovation occured
innovation_record=[population(ind).connection_genes(1:3,:);zeros(size(population(ind).connection_genes(1:2,:)))];
innovation_record(4,size(innovation_record,2))=max(population(1).node_genes(1,:));

%% Initial speciation
population(1).species=1;
%fill in first individual in species record and propogating species matrix

mat=population(1).connection_genes(4,:);
species_record(1).ID=1;
species_record(1).number_individuals=1;

for ind=2:size(population,2);

    added = false
    sp=1;

    while ~added
        %species compatibility distance

        distance=speciation.c3*sum(abs(population(ind).connection_genes(4,:)-mat(sp,:)))/num_conections;
        if distance<speciation.threshold %If within threshold, assign to the existing species
            population(ind).species=sp;
            species_record(sp).number_individuals=species_record(sp).number_individuals+1;
            added=true;
        end
        sp=sp+1; % check with all species

        %Outside of species references, must be new species
        if sp>size(mat,1) & ~added
            population(ind).species=sp;
            species_record(sp).ID=sp;
            species_record(sp).number_individuals=1;
            mat=[mat;population(ind).connection_genes(4,:)];
            added=true;
        end
    end
end

%% Start generation loop
maximal_fit=0;
sol_found=false;
while generation<=max_generations & ~sol_found
    cur_max_fit=zeros(1,size(species_record,2));

    % Fitness of population
    population = get_pop_fitness(population, @RNNet, target_fitness);

    for sp=1:size(species_record,2)
        if species_record(sp).number_individuals>0
            % max fitness
            [max_fit,ind_max]=max(([population.species]==sp).*[population.fitness]);
            %mean fitness
            mean_fit=sum(([population.species]==sp).*[population.fitness])/species_record(sp).number_individuals;
            % Compute stagnation vector (last stagnation.number_generation-1 max fitnesses plus current fitness

            if size(species_record(sp).generation_record,2)>stagnation.number_generation-2
                stag_vec=[species_record(sp).generation_record(3,size(species_record(sp).generation_record,2)-stagnation.number_generation+2:size(species_record(sp).generation_record,2)),max_fit];
                if sum(abs(stag_vec-mean(stag_vec))<stagnation.threshold)==stagnation.number_generation %Check for stagnation
                    mean_fit=0.01; %set mean fitness to small value to eliminate species (cannot be set to 0, if only one species is present, we would have divide by zero in fitness sharing. anyways, with only one species present, we have to keep it)
                end
            end
            % update generation record
            species_record(sp).generation_record=[species_record(sp).generation_record,[generation;mean_fit;max_fit;ind_max]];
            cur_max_fit(1,sp)=max_fit;
        end
    end

    c=[];
    for sp=1:size(species_record,2)
        c=[c,species_record(sp).generation_record(1:4,size(species_record(sp).generation_record,2))];
    end
    maximal_fit=max(c(3,:).*(c(1,:)==generation));
    max_gen_fit=[max_gen_fit,[maximal_fit;generation]];
    if maximal_fit>=target_fitness
        sol_found = true;
    else
        % Function reproduce goes here
        [population,species_record,innovation_record]=...
            reproduce_pop(population, species_record, innovation_record, initial, selection, ...
            crossover, mutation, speciation, generation, pop_size,@RNNet,target_fitness);
    end
    generation
    generation=generation+1;
    subplot(2,2,3);
    plot(max_gen_fit(2,:),max_gen_fit(1,:));
    ylabel('max fitness');
end

[max_fit, max_fit_individual] = max([population(:).fitness]);

% simulate = twoPole_test( population(max_fit_individual).wMat, @RNNet, target_fitness,'vis');


%% Function implementing crossover, mutation, and speciation
function [new_population,updated_species_record,updated_innovation_record]=...
    reproduce_pop(population, species_record, innovation_record, initial, selection, ...
    crossover, mutation, speciation, generation, pop_size,RNNet,target_fitness)



% new_population(1) = population(1);
%new_population = rank_based_selection(population,species_record);

% Sum of avg from every species
avg_fit=0;
for s_ind=1:size(species_record,2)
    avg_fit=avg_fit+species_record(s_ind).generation_record(2,size(species_record(s_ind).generation_record,2))...
        *(species_record(s_ind).number_individuals>0);
end

% Matrix containing iformation about species that currently exist and
% will be propagating.
propagating_species=[];
overflow=0;
index_individual=0;
for s_ind = 1:size(species_record,2)
    if species_record(s_ind).number_individuals>0
        avg_specie_fitness = species_record(s_ind).generation_record(2,size(species_record(s_ind).generation_record,2));
        offsprings = (avg_specie_fitness/avg_fit)*pop_size;
        overflow=overflow+offsprings-floor(offsprings);
        if overflow>=1
            offsprings=ceil(offsprings);
            overflow=overflow-1;
        else
            offsprings=floor(offsprings);
        end

        if offsprings>0
            propagating_species=[propagating_species,[species_record(s_ind).ID;offsprings;0]];
            if species_record(s_ind).number_individuals>=initial.number_copy %check for condition for objective 2
                index_individual=index_individual+1;
                new_population(index_individual)=population(species_record(s_ind).generation_record(4,size(species_record(s_ind).generation_record,2))); % Objective 2
                propagating_species(3,size(propagating_species,2))=1; %Update matrix_existing_and_propagating_species
            end
        end

    end
end
% propagating_species
% Reference individuals from every species
ref=0;
for sp=1:size(species_record,2)
    if sum([population.species]==sp)>0
        ref=ref+1;
        % take any of the individuals from all existing species in population
        [discard,ind]=max(([population.species]==sp).*rand(1,size(population,2)));
        pop_ref(ref)=population(ind);
    end
end

pop_ind = 0;
for sp=1:size(propagating_species,2)
    count_individuals_species=0;
    species_ID=propagating_species(1,sp);

    ind_species=find([population.species]==species_ID);

    fitnesses_species=[population(ind_species).fitness];
    [discard,sorted_fitnesses]=sort(fitnesses_species);
    ranking=zeros(1,size(fitnesses_species,2));
    ranking(sorted_fitnesses)=1:size(fitnesses_species,2);
    %normalized fitness
    if size(fitnesses_species,2)>1
        FitnV=(2-selection.pressure+2*(selection.pressure-1)/(size(fitnesses_species,2)-1)*(ranking-1))';
    else
        FitnV=2;
    end

    number_overall=propagating_species(2,sp)-propagating_species(3,sp);
    number_crossover=round(crossover.percentage*number_overall);
    number_mutate=number_overall-number_crossover;
    Nind=size(fitnesses_species,2);
    Nsel=2*number_crossover+number_mutate;

    if Nsel==0
        Nsel=1;
    end

    cumfit = cumsum(FitnV);
    trials = cumfit(Nind) / Nsel * (rand + (0:Nsel-1)');
    Mf = cumfit(:, ones(1, Nsel));
    Mt = trials(:, ones(1, Nind))';
    [NewChrIx, ans] = find(Mt < Mf & [ zeros(1, Nsel); Mf(1:Nind-1, :) ] <= Mt);
    [ans, shuf] = sort(rand(Nsel, 1));
    NewChrIx = NewChrIx(shuf);
    NewChrIx = ind_species(NewChrIx);



    species_individuals = population(find([population(:).species]==propagating_species(1,sp)));
    while propagating_species(3,find([propagating_species(1,:)]==species_ID)) < propagating_species(2,find([propagating_species(1,:)]==species_ID))
        %         species_individuals
        %         rand_offs = randi(size(species_individuals,2));
        %         offspring = species_individuals(rand_offs);
        pop_ind=pop_ind+1;
        count_individuals_species=count_individuals_species+1;
        offspring = population(NewChrIx(count_individuals_species));


        if rand < crossover.rate
            parent1 = population(NewChrIx(randi(size(NewChrIx,2))));
            parent2 = population(NewChrIx(randi(size(NewChrIx,2))));
            parent_nodes = [parent1.node_genes(:,:) parent2.node_genes(:,:)];
            parent_connection_genes = [parent1.connection_genes(:,:) parent2.connection_genes(:,:)];

            [unique_nodes, u_n_ind, ~] = ...
                unique(parent_nodes(1,:));
            [unique_innovation_num, u_i_ind, ~]  = ...
                unique(parent_connection_genes(1,:));
            offspring.node_genes = [];
            offspring.connection_genes = [];
            for node_id = 1:size(unique_nodes,2)
                offspring.node_genes = [offspring.node_genes ...
                    parent_nodes(:,u_n_ind(node_id))];
            end
            for innov_id = 1:size(unique_innovation_num,2)
                offspring.connection_genes = [offspring.connection_genes ...
                    parent_connection_genes(:,u_i_ind(innov_id))];
            end
            offspring.wMat = weight_matrix(offspring.node_genes,offspring.connection_genes);
            offspring.fitness = get_ind_fitness(offspring, RNNet, target_fitness);
        else
            offspring = population(NewChrIx(count_individuals_species));
        end








        % Weight value mutation
        if rand < mutation.mutate_weight
            offspring = weight_value_mutation(offspring,mutation);
        end

        % Add connection mutation
        if rand < mutation.add_node_rate
            [offspring,innovation_record] = ...
                add_connection(offspring,innovation_record,generation,mutation.range);
        end

        % Add node mutation
        if rand < mutation.add_connection_rate
            [offspring,innovation_record] = ...
                add_node(offspring,innovation_record,generation,mutation.range);

        end

        %         % Speciation
        species_assigned=0;
        ref=0;
        while species_assigned==0 & ref<size(pop_ref,2)
            ref=ref+1;
            ref_ind=pop_ref(ref);

            max_ref_inn = max(ref_ind.connection_genes(1,:));
            max_off_inn = max(offspring.connection_genes(1,:));


            if max_ref_inn > max_off_inn
                excess_cutoff = max_off_inn;
            else
                excess_cutoff = max_ref_inn;

            end

            ref_conn = ref_ind.connection_genes(1,:);
            off_conn = offspring.connection_genes(1,:);

            matching_genes = intersect(ref_conn,off_conn);
            all_conn = union(ref_conn,off_conn);

            disj_or_exs = setdiff(all_conn,matching_genes);

            excess_inn = disj_or_exs(disj_or_exs>excess_cutoff);

            disjoint_inn = setdiff(disj_or_exs,excess_inn);


            excess = size(excess_inn,2);
            disjoint = size(disjoint_inn,2);



            weight_diff =0;
            for weight_ind = 1:size(matching_genes,2)
                weight1 = offspring.connection_genes(4,(offspring.connection_genes(1,:)==matching_genes(weight_ind)));
                weight2 = ref_ind.connection_genes(4,(ref_ind.connection_genes(1,:)==matching_genes(weight_ind)));
                weight_diff = weight_diff + abs(weight1 - weight2);
            end

            average_weight_difference=weight_diff/size(matching_genes,2);

            max_genes=1;
            distance=speciation.c1*excess/max_genes+speciation.c2*disjoint/max_genes+speciation.c3*average_weight_difference;
            if distance<speciation.threshold
                % assign individual to same species as current reference individual
                offspring.species=ref_ind.species;
                species_assigned=1; %set flag indicating new_individual has been assigned to species
            end
        end
        % not compatible with any? well, then create new species
        if species_assigned==0
            new_species_ID=size(species_record,2)+1;
            % assign individual to new species
            offspring.species=new_species_ID;
            % update species_record
            species_record(new_species_ID).ID=new_species_ID;
            species_record(new_species_ID).number_individuals=1;
            species_record(new_species_ID).generation_record=[generation;offspring.fitness;offspring.fitness;pop_ind];
            % update population reference
            pop_ref(size(pop_ref,2)+1)=offspring;
        end


        new_population(pop_ind) = offspring;
        propagating_species(3,sp) = propagating_species(3,sp) + 1;
    end
end

for sp=1:size(species_record,2)
    species_record(sp).number_individuals=sum([new_population(:).species]==sp);
end

updated_species_record = species_record;
updated_innovation_record = innovation_record;
end

%% Rank based selection
function [new_population] = rank_based_selection(population,species_record)
pop_ind = 1;
new_population = population;
for species_ID=1:size(species_record,2)
    species_individuals = population(find([population(:).species]==species_ID));
    if size(species_individuals,2)<0
        num_ind_species = species_record(species_ID).number_individuals;
        competitors = randi([1,size(species_individuals,2)],2,num_ind_species);
        for species_ind=1:num_ind_species
            species_ind
            fitness_1 = species_individuals(competitors(1,species_ind)).fitness
            fitness_2 = species_individuals(competitors(2,species_ind)).fitness;

            if fitness_1 >= fitness_2
                new_population(pop_ind) = species_individuals(competitors(1,species_ind));
            else
                new_population(pop_ind) = species_individuals(competitors(2,species_ind));
            end
            pop_ind = pop_ind + 1;
        end
    end
end
% max_fitness=[];
% max_fitness = [max_fitness population(:).fitness];
% [max_v,max_id] = max(max_fitness);
% min_fitness=[];
% min_fitness = [min_fitness new_population(:).fitness];
% [min_v,min_id] = min(min_fitness);
% new_population(min_id) = population(max_id);
end

%% Weight value mutation
function individual = weight_value_mutation(individual,mutation)
num_connections = size(individual.connection_genes(1,:),2);
mutate_weights = rand(1,num_connections);
perturbate_weight = rand(1,num_connections);
for con_id=1:num_connections
    node_from = find(individual.node_genes(1,:)==individual.connection_genes(2,con_id));
    node_to = find(individual.node_genes(1,:)==individual.connection_genes(3,con_id));
    if mutate_weights(con_id) < mutation.rate
        if perturbate_weight(con_id) < mutation.perturbation_rate
            perturbation = rand(1)*2*mutation.perturbation_range-...
                mutation.perturbation_range;
            new_weight = individual.connection_genes(4,con_id) + perturbation;
            if new_weight < -mutation.cap
                new_weight = -mutation.cap;
            elseif new_weight > mutation.cap
                new_weight = mutation.cap;
            end
            individual.wMat(node_from,node_to) = new_weight;
            individual.connection_genes(4,con_id) = new_weight;
        else
            new_weight = rand(1)*2*mutation.range-mutation.range;
            individual.wMat(node_from,node_to) =  new_weight;
            individual.connection_genes(4,con_id) = new_weight;
        end
    end
end
end

%% Call simulator for every individual in population to get fitness
function population = get_pop_fitness(population, RNNet, target_fitness)
for ind=1:size(population(:),1)
    population(ind).fitness = twoPole_test( population(ind).wMat, RNNet, target_fitness);%,'vis')
end
end

%% Call simulator for one individual to get fitness
function steps = get_ind_fitness(individual, RNNet, target_fitness)
    steps = twoPole_test(individual.wMat, RNNet, target_fitness);%,'vis')

end

%% Built the weight matrix from connected genes
function wMat = weight_matrix(node_genes,connection_genes)
wMat = zeros(size(node_genes(1,:),2),size(node_genes(1,:),2));
for c = 1:size(connection_genes(1,:),2)
    enabled = connection_genes(5,c);
    from_node = find(node_genes(1,:)==connection_genes(2,c));
    to_node = find(node_genes(1,:)==connection_genes(3,c));
    weight=connection_genes(4,c);
    if enabled
        wMat(from_node, to_node) = weight;
    else
        wMat(from_node, to_node) = 0;
    end
end
end

%% Add a connection to the network
function [individual,innovation_record] = ...
    add_connection(individual,innovation_record,generation, mutation_range)

% individual.connection_genes
%get all possible connection from enabled ones
inn_old=max((innovation_record(5,:)<generation).*innovation_record(1,:));

pos_conn=individual.connection_genes(2:3,(individual.connection_genes(5,:)==1)  & (individual.connection_genes(1,:)<=inn_old));

%choose a random 'from node' and 'to node' connection

rand_node_connection=pos_conn(:,round(rand*size(pos_conn,2)+0.5));
%
% rand_from_node = rand_node_connection(1);
% rand_to_node = rand_node_connection(2);
rand_from_node = individual.node_genes(1,round(rand*size(individual.node_genes,2)+0.5));
rand_to_node = individual.node_genes(1,round(rand*size(individual.node_genes,2)+0.5));

%check innovation duplicacy
current_gen_changes = innovation_record(:,innovation_record(5,:)==generation);
dup = 0;
if (~isempty(current_gen_changes))

    for cc=1:size(current_gen_changes,2)
        if current_gen_changes(2,cc) == rand_from_node & current_gen_changes(3,cc) == rand_to_node
            next_inn_num = current_gen_changes(1,cc);

            dup = 1;
        else
            next_inn_num = max(innovation_record(1,:))+1;
        end
    end
else
    next_inn_num = max(innovation_record(1,:))+1;
end

%choose a random 'from node' and 'to node'

weight = rand(1)*2*mutation_range-mutation_range;
new_connection = [next_inn_num; rand_from_node; rand_to_node; weight; 1];
individual.connection_genes = [individual.connection_genes new_connection];
if dup==0
    innovation_record = [innovation_record [next_inn_num rand_from_node rand_to_node 0 generation]'];
end
% innovation_record

individual.wMat = weight_matrix(individual.node_genes,individual.connection_genes);

end

%% Add node to network
function [individual,innovation_record] = ...
    add_node(individual,innovation_record,generation,mutation_range)

num_nodes = size(individual.node_genes(1,:),2);
num_connections = size(individual.connection_genes(1,:),2);
new_node_ID =max(innovation_record(4,:))+1;

%get all possible connection from enabled ones
inn_old=max((innovation_record(5,:)<generation).*innovation_record(1,:));

pos_conn=individual.connection_genes(2:3,(individual.connection_genes(5,:)==1)  & (individual.connection_genes(1,:)<=inn_old));

%choose a random 'from node' and 'to node' connection

rand_node_connection=pos_conn(:,round(rand*size(pos_conn,2)+0.5));
%
rand_from_node = rand_node_connection(1);
rand_to_node = rand_node_connection(2);

%sanity check
% rand_from_node = innovation_record(2,8);
% rand_to_node = innovation_record(3,9);

%check innovation duplicacy
current_gen_changes = innovation_record(:,innovation_record(5,:)==generation);
% current_gen_changes;
dup = 0;
current_gen_changes;
if (~isempty(current_gen_changes))
    next_inn_num = max(innovation_record(1,:))+1;
    for cc=1:size(current_gen_changes,2)-1
        if current_gen_changes(2,cc) == rand_from_node & current_gen_changes(3,cc+1) == rand_to_node
            next_inn_num = current_gen_changes(1,cc);
            new_node_ID =current_gen_changes(3,cc);
            dup = 1;

        end
    end
else
    next_inn_num = max(innovation_record(1,:))+1;
end
% new node gene's properties [assign node ID, type=hidden,
% inp_state = 0, o_state = 0
new_node_gene = [ new_node_ID ; 3; 0; 0];
individual.node_genes(:,num_nodes+1) = new_node_gene;



%check if connection it already exists
exist_from_con = find(individual.connection_genes(2,:)==rand_from_node);

if(isempty(exist_from_con))
    weight1=rand(1)*2*mutation_range-mutation_range;
    weight2=rand(1)*2*mutation_range-mutation_range;
else
    for i=1:length(exist_from_con)
        if(individual.connection_genes(3,exist_from_con(i)) == rand_to_node)
            individual.connection_genes(5,exist_from_con(i))=0;

            weight1=1;
            weight2=individual.connection_genes(4,exist_from_con(i));

            break;
        else
            weight1=rand(1)*2*mutation_range-mutation_range;
            weight2=rand(1)*2*mutation_range-mutation_range;
        end
    end
end
new_connection_gene_1 = [ next_inn_num; rand_from_node;new_node_ID;weight1;1];
new_connection_gene_2 = [ next_inn_num+1; new_node_ID;rand_to_node;weight2;1];

%return incremented global innovation number
individual.connection_genes(:,num_connections+1) = new_connection_gene_1;
individual.connection_genes(:,num_connections+2) = new_connection_gene_2;

individual.wMat = weight_matrix(individual.node_genes,individual.connection_genes);

%update innovation
if dup==0
    innovation_record = [innovation_record  [next_inn_num;rand_from_node;new_node_ID;new_node_ID;generation] [next_inn_num+1;new_node_ID;rand_to_node;0;generation]];
end
innovation_record;
% individual
end