#!/usr/bin/env python

from embodied_ising import ising
from embodied_ising import food
from embodied_ising import CriticalLearning, EvolutionLearning
import numpy as np
import argparse
import matplotlib.pyplot as plt
from sys import argv

import pickle



def create_settings():

    args = parse()
    # --- CONSTANTS ----------------------------------------------------------------+
    settings = {}
    
    # ENVIRONMENT SETTINGS
    settings['pop_size'] = 150  # number of organisms #150
    settings['numKill'] = int(settings['pop_size'] / 1.66)
    settings['food_num'] = 100  # number of food particles
    settings['food_radius'] = 0.03
    settings['org_radius'] = 0.05
    settings['ANN'] = False  # Use ANN or Ising?
    settings['BoidOn'] = False  # Only use Boid model? #True
    
    # SIMULATION SETTINGS
    settings['TimeSteps'] = args.time_steps # number of timesteps per iteration #2000

    # number of system-wide spin updates per unit of time  (multiplies computation time)
    settings['thermalTime'] = 3
    
    settings['evolution_toggle'] = False  # only toggles for CriticalLearning
    settings['evolution_rate'] = 1  # only with critical learning number of iterations to skip to kill/mate (gives more time to eat before evolution)
    
    settings['dt'] = 0.2  # kinetic time step      (dt)
    settings['r_max'] = 720
    settings['dr_max'] = 90  # max rotational speed      (degrees per second)
    settings['v_max'] = 999  # 0.5 max velocity              (units per second)
    settings['dv_max'] = 0.05  # max acceleration (+/-)    (units per second^2)
    
    settings['x_min'] = 0.0  # arena eastern border
    settings['x_max'] = 8.0  # arena western border
    settings['y_min'] = 0.0  # arena southern border
    settings['y_max'] = 8.0  # arena northern border
    
    settings['save_data'] = args.save_data
    settings['plot'] = args.plot # make plots?
    settings['plot_n_last_generations'] = 5  # switches on plotting for n last generations
    # iterations. Also begins saving figures after this many iterations if 'plot' setting is 'False'
    settings['plotLive'] = False  # live updates of figures
    settings['frameRate'] = 1
    settings['animation_fps'] = 20
    
    settings['size'] = 10
    settings['nSensors'] = 3
    settings['nMotors'] = 4
    settings['learningrate'] = 0.01  # 0.01
    # how many hidden neurons are not connected to each other
    settings['numDisconnectedNeurons'] = 0 #  int((settings['size'] - settings['nSensors'] - settings['nMotors']) / 1.2)
    # how should organisms repopulate, duplicate or mate?
    settings['mateDupRatio'] = 0.5
    settings['mutationRateDup'] = 0.1  # DUPLICATION mutation rate
    
    settings['mutateB'] = True  # toggle to allow Beta (temperature) mutations (toggle off if critical learning is on)
    settings['sigB'] = 0.02  # std for Beta mutation
    
    settings['LoadIsings'] = True
    settings['loadfile'] = 'sim-20191114-000009_server'
    settings['iter'] = '500'
    
    #Seasons
    settings['seasons'] = False #BOO; Activates seasons
    settings['years_per_iteration'] = 1 #INT amount of seasonal changes per iteration
    settings['min_food_winter'] = 0.5 #FLOAT [0,1]; relative decimation of food in winter
    
    settings['parallel_computing'] = False #BOO
    settings['cores'] = 4 #INT if 0 number is determined automatically
    
    settings['energy_model'] = True #BOO
    settings['v_min'] = 0.05  # FLOAT [0,1)
    settings['cost_speed'] = 0.05  # FLOAT [0,1] energy cost of speed as a factor of speed #before 0.025
    settings['initial_energy'] = 0  # Energy that each organism starts with in each simulation

    
    Iterations = args.iterations




    return settings, Iterations

def parse():
    parser = argparse.ArgumentParser(description=
                                     '''Agent-based simulation of artificial organisms 
                                     controlled by a statistical neural net (ising model)''')
    parser.add_argument('-s', '--save', dest='save_data', action='store_true', help = "Save data of simulation")
    parser.add_argument('-p', '--plot', dest='plot', action='store_true',
                        help="Plot simulation data, drastically decreases runtime")
    parser.add_argument('-i', '--iterations', type=int, dest = 'iterations',
                        help='Number of Iterations in simulation')
    parser.add_argument('-t', '--timesteps', type=int, dest = 'time_steps',
                        help='Number of Iterations in simulation')
    parser.set_defaults(save_data=False, plot=False, iterations=2000, time_steps=2000)
    args = parser.parse_args()
    return args

# --- MAIN ---------------------------------------------------------------------+

def run(settings):



    size = settings['size']
    nSensors = settings['nSensors']
    nMotors = settings['nMotors']
    # LOAD ISING CORRELATIONS
    # filename = 'correlations-ising2D-size400.npy'
    filename2 = 'correlations-ising-generalized-size83.npy'
    settings['Cdist'] = np.load(filename2)

    # --- POPULATE THE ENVIRONMENT WITH FOOD ---------------+
    foods = []
    for i in range(0, settings['food_num']):
        foods.append(food(settings))

    # Food is only created uniformly distributed at the very beginning.
    # For a new iteration the placement of the food is kept.


    # --- POPULATE THE ENVIRONMENT WITH ORGANISMS ----------+
    if settings['LoadIsings']:
        loadfile = 'save/' + settings['loadfile'] + '/isings/gen[' + settings['iter'] + ']-isings.pickle'
        startstr = 'Loading simulation:' + loadfile + ' (' + str(settings['TimeSteps']) + \
                   ' timesteps) x (' + str(Iterations) + ' iterations)'
        print(startstr)
        isings = pickle.load(open(loadfile, 'rb'))
    else:
        startstr = 'Starting simulation: (' + str(settings['TimeSteps']) + \
                   ' timesteps) x (' + str(Iterations) + ' iterations)'
        print(startstr)
        isings = []
        for i in range(0, settings['pop_size']):
            isings.append(ising(settings, size, nSensors, nMotors, name='gen[0]-org[' + str(i) + ']'))
        

    # --- CYCLE THROUGH EACH GENERATION --------------------+
    # Choose between CriticalLearning (which has both inverse-ising and GA with toggle)
    # or EvolutionLearning which is only GA. The functions are fairly similar, should find a
    # better way to call them than this.
    # ------------------------------------------------------+
    
    #No critical learning:
    # CriticalLearning(isings, foods, settings, Iterations)

    EvolutionLearning(isings, foods, settings, Iterations)

    pass


# --- RUN ----------------------------------------------------------------------+

if __name__ ==  '__main__':
    settings, Iterations = create_settings()
    run(settings)

# --- END ----------------------------------------------------------------------+
