#!/usr/bin/env python3

# Please feel free to modify any part of the code or add any functions, or modify the argument of the given functions. But please keep the name of the given functions

# Please feel free to import any libraries you need.

# You are required to finish the genetic_algorithm function, and you may need to complete crossover, mutate and select.

import random
import matplotlib.pyplot as plt

def crossover(old_gen, probability_crossover):
    #TODO START
    new_population = []
    for individual in old_gen:
        new_population.append(individual)
    random.shuffle(new_population)
    
    sites = []
    
    for i in range(int(len(new_population) / 2)):
        choices = [x for x in range(1, 31)]
        random.shuffle(choices)
        site_1 = min(choices[0], choices[-1])
        site_2 = max(choices[0], choices[-1])
        if random.random() > probability_crossover:
            site_1 = 0
            site_2 = 0
        sites.append((site_1, site_2))
    
    for i in range(len(sites)):
        new_population[2 * i] = new_population[2 * i][: sites[i][0]] + new_population[2 * i + 1][sites[i][0]: sites[i][1]] + new_population[2 * i][sites[i][1] :]
        new_population[2 * i + 1] = new_population[2 * i + 1][: sites[i][0]] + new_population[2 * i][sites[i][0]: sites[i][1]] + new_population[2 * i + 1][sites[i][1] : ]
        
        assert(len(new_population[2 * i]) == 30)
        assert(len(new_population[2 * i + 1]) == 30)
        
    #TODO END
    return new_population

def mutate(old_gen, probability_mutation):
    #TODO START
    new_gen = ""
    for i in range(len(old_gen)):
        if random.random() < probability_mutation:
            if i % 3 == 0:
                new_gene = old_gen[i]
                while new_gene == old_gen[i]:
                    new_gene = str(random.randint(1, 4))
            else:
                new_gene = old_gen[i]
                while new_gene == old_gen[i]:
                    new_gene = str(random.randint(0, 9)) 
            new_gen += new_gene
        else:
            new_gen += old_gen[i]
    
    #TODO END
    return new_gen

def select(old_gen, fitness):
    #TODO START
    new_population = []
    
    standardize_fitness = [i / sum(fitness) for i in fitness]
    selected = [random.random() for i in range(len(fitness))]
    sum_fitness = 0
    
    for i in range(len(fitness)):
        sum_fitness += standardize_fitness[i]
        for j in range(len(selected)):
            if selected[j] < sum_fitness:
                selected[j] = i + 1
    
    for selection in selected:
        new_population.append(old_gen[selection - 1])
    
    #TODO END
    return new_population

def genetic_algorithm(population, food_map_file_name, max_generation, probability_crossover, probability_mutation):
    #TODO START
    food_map, map_size = get_map(food_map_file_name)
    stats = []
    
    for i in range(max_generation):
        fitness = []
        for individual in population:
            trial, individual_fitness = ant_simulator(food_map, map_size, individual)
            fitness.append(individual_fitness)
        stats.append([max(fitness), min(fitness), sum(fitness) / len(fitness)])
        print(len(stats), stats[-1])
        
        new_population = select(population, fitness)
        
        new_population = crossover(new_population, probability_crossover)
        
        for i in range(len(new_population)):
            new_population[i] = mutate(new_population[i], probability_mutation)
        
        population = new_population
    
    max_fitness = -1
    max_individual = ""
    max_trial = ""
    for individual in population:
        trial, individual_fitness = ant_simulator(food_map, map_size, individual)
        if individual_fitness > max_fitness:
            max_fitness = individual_fitness
            max_individual = individual
            max_trial = trial
    
    return max_fitness, max_individual, max_trial, stats, population
    #TODO END    

def _initialize_individual():
    individual = [""] *  30
    
    for i in range(30):
        if i % 3 == 0:
            individual[i] = str(random.randint(1, 4))
        else:
            individual[i] = str(random.randint(0, 9))

    return "".join(individual)

def initialize_population(num_population):
    #TODO START
    population = []
    
    for i in range(num_population):
        population.append(_initialize_individual())
        
    #TODO END
    return population
    
def ant_simulator(food_map, map_size, ant_genes):
    """
    parameters:
        food_map: a list of list of strings, representing the map of the environment with food
            "1": there is a food at the position
            "0": there is no food at the position
            (0, 0) position: the top left corner of the map
            (x, y) position: x is the row, and y is the column
        map_size: a list of int, the dimension of the map. It is in the format of [row, column]
        ant_genes: a string with length 30. It encodes the ant's genes, for more information, please refer to the handout.
    
    return:
        trial: a list of list of strings, representing the trials
            1: there is food at that position, and the spot was not visited by the ant
            0: there is no food at that position, and the spot was not visited by the ant
            empty: the spot has been visited by the ant
    
    It takes in the food_map and its dimension of the map and the ant's gene information, and return the trial in the map
    """
    
    step_time = 200
    
    trial = []
    for i in food_map:
        line = []
        for j in i:
            line.append(j)
        trial.append(line)

    position_x, position_y = 0, 0
    orientation = [(1, 0), (0, -1), (-1, 0), (0, 1)] # face down, left, up, right
    fitness = 0
    state = 0
    orientation_state = 3
    gene_list = [ant_genes[i : i + 3] for i in range(0, len(ant_genes), 3)]
    
    for i in range(step_time):
        if trial[position_x][position_y] == "1":
            fitness += 1
        trial[position_x][position_y] = " "
        
        sensor_x = (position_x + orientation[orientation_state][0]) % map_size[0]
        sensor_y = (position_y + orientation[orientation_state][1]) % map_size[1]
        sensor_result = trial[sensor_x][sensor_y]
        
        if sensor_result == "1":
            state = int(gene_list[state][2])
        else:
            state = int(gene_list[state][1])
        
        action = gene_list[state][0]
        
        if action == "1": # move forward
            position_x = (position_x + orientation[orientation_state][0]) % map_size[0]
            position_y = (position_y + orientation[orientation_state][1]) % map_size[1]
        elif action == "2": # turn right
            orientation_state = (orientation_state + 1) % 4
        elif action == "3": # turn left
            orientation_state = (orientation_state - 1) % 4
        elif action == "4": # do nothing
            pass
        else:
            raise Exception("invalid action number!")
    
    return trial, fitness
        

def get_map(file_name):
    """
    parameters:
        file_name: a string, the name of the file which stored the map. The first line of the map is the dimension (row, column), the rest is the map
            1: there is a food at the position
            0: there is no food at the position
    
    return:
        food_map: a list of list of strings, representing the map of the environment with food
            "1": there is a food at the position
            "0": there is no food at the position
            (0, 0) position: the top left corner of the map
            (x, y) position: x is the row, and y is the column
        map_size: a list of int, the dimension of the map. It is in the format of [row, column]
    
    It takes in the file_name of the map, and return the food_map and the dimension map_size
    """
    food_map = []
    map_file = open(file_name, "r")
    first_line = True
    map_size = []
    
    for line in map_file:
        line = line.strip()
        if first_line:
            first_line = False
            map_size = line.split()
            continue
        if line:
            food_map.append(line.split())
    
    map_file.close()
    return food_map, [int(i) for i in map_size]

def display_trials(trials, target_file):
    """
    parameters:
        trials: a list of list of strings, representing the trials
            1: there is food at that position, and the spot was not visited by the ant
            0: there is no food at that position, and the spot was not visited by the ant
            empty: the spot has been visited by the ant
        taret_file: a string, the name the target_file to be saved
    
    It takes in the trials, and target_file, and saved the trials in the target_file. You can open the target_file to take a look at the ant's trial.
    """
    trial_file = open(target_file, "w")
    for line in trials:
        trial_file.write(" ".join(line))
        trial_file.write("\n")
    trial_file.close()

if __name__ == "__main__":
    #TODO START
    #You will need to modify the parameters below.
    #The parameters are for references, please feel free add more or delete the ones you do not intend to use in your genetic algorithm
    random.seed(3)
    
    population_size = 100
    population = initialize_population(population_size)
    food_map_file_name = "muir.txt"
    max_generation = 200
    probability_crossover = 0.1
    probability_mutation = 0.05
    max_fitness, max_individual, max_trial, stats, population = genetic_algorithm(population, food_map_file_name, max_generation, probability_crossover, probability_mutation)
    
    display_trials(max_trial, "max_trial.txt")
    
    plt.figure(1)
    plt.plot([i for i in range(len(stats))], [i[0] for i in stats], marker = "o")
    plt.xlabel("generation")
    plt.xlim((0, 200))
    plt.ylim((0, max(i[0] for i in stats) + 10))
    plt.ylabel("most fit individual")
    plt.savefig("max fitness of each generation.png")
    
    plt.figure(2)
    muir_fitness = []
    santafe_fitness = []
    muir_food_map, muir_map_size = get_map("muir.txt")
    santafe_food_map, santafe_map_size = get_map("santafe.txt")
    for individual in population:
        trial, individual_muir_fitness = ant_simulator(muir_food_map, muir_map_size, individual)
        trial, individual_santafe_fitness = ant_simulator(santafe_food_map, santafe_map_size, individual)
        muir_fitness.append(individual_muir_fitness)
        santafe_fitness.append(individual_santafe_fitness)
    
    plt.plot([i for i in range(len(muir_fitness))], muir_fitness, marker = "o", color = "blue", label = "muir")
    plt.plot([i for i in range(len(santafe_fitness))], santafe_fitness, marker = "o", color = "green", label = "santa fe")
    plt.xlabel("individuals in the last generation")
    plt.xlim((0, population_size))
    plt.ylim((0, max(muir_fitness + santafe_fitness) + 10))
    plt.ylabel("fitness")
    plt.legend()
    plt.savefig("fitness of last generation on muir and santa fe map.png")    

    #TODO END
    
    # Example of how to use get_map, ant_simulator and display trials function
    """
    food_map, map_size = get_map("muir.txt")
    ant_genes = "335149249494173115455311387263"
    trial, fitness = ant_simulator(food_map, map_size, ant_genes)
    display_trials(trial, "trial.txt")
    """