params = [    
    { "start": 34.77,  "attr": "Ra",        "low": 5.0, "high": 100.0, "lists": ["all"] },
    { "start": 2.706,   "attr": "cm",        "low": 0.1, "high": 4.0, "lists": ["all"] },
    { "start": 49.95,  "attr": "ena",       "low": 40.0, "high": 50.0, "lists": ["all"] },
    { "start": -70.03,  "attr": "ek",        "low": -100.0, "high": -70.0, "lists": ["all"] },
    { "start": -64.42,  "attr": "e_pas",     "low": -70.0, "high": -50.0, "lists": ["all"] },
    { "start": 0.0005955, "attr": "g_pas",     "low": 0, "high": 0.00003, "lists": ["all"] },    
    { "start": 0.5955,  "attr": "sh_Na",     "low": 0, "high": 10, "lists": ["all"] },    
    { "start": 10,  "attr": "tau_CaPool",     "low": 1, "high": 500, "lists": ["all"] },    

    { "start":  0.87485,  "attr": "gbar_Na",     "low": 0, "high": 0.05, "lists": ["all"] },
    { "start": 0.0297,  "attr": "gbar_Kd",     "low": 0, "high": 0.04, "lists": ["all"] },
    { "start": 0.000264,  "attr": "gbar_Kslow",  "low": 0, "high": 0.004, "lists": ["all"] },
    { "start": 0.07215,  "attr": "gbar_KA",     "low": 0, "high": 0.005, "lists": ["all"] },
    { "start": 0.001,  "attr": "gbar_KCa",    "low": 0, "high": 0.004, "lists": ["all"] },
    { "start": 0.00081441,  "attr": "gbar_LCa",  "low": 0, "high": 0.001, "lists": ["all"] },
    
    { "start": -30.805,  "attr": "eh",       "low": -40.0, "high": -25.0, "lists": ["apical"] },
    { "start": 0.00335,  "attr": "gbar_Ih",    "low": 0, "high": 0.00003, "lists": ["apical"] },
    { "start": 0.000107,  "attr": "gbar_CaT",    "low": 0, "high": 18e-3, "lists": ["apical"] },
]

top = [[34.40577448714667,
  1.7746337170486868,
  48.600061083126825,
  -73.37824874831657,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  6.513661312706145,
  0.01835806177392501,
  0.012685911771409467,
  0.001009786028927451,
  0.003485358049971202,
  0.0009111182551995163,
  0.0006405838020076363,
  -36.191181538089175,
  2.6308088272967942e-06,
  0.009520566435955282],
 [54.371057327671025,
  1.343586187918468,
  44.40845548134062,
  -81.17397852577881,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  12.492102024675049,
  0.01835806177392501,
  0.012243090987314832,
  0.001009786028927451,
  0.0025920677224618473,
  0.0006458839218994715,
  0.0006437525063628547,
  -36.474771746065244,
  1.480641964136122e-06,
  0.009384267411892713],
 [34.40577448714667,
  1.7746337170486868,
  48.25121832219984,
  -85.83988610793635,
  -52.149512474246556,
  2.9121194317970837e-05,
  7.600064936952291,
  11.469065054901318,
  0.01835806177392501,
  0.012685911771409467,
  0.001009786028927451,
  0.003485358049971202,
  0.0009111182551995163,
  0.0006405838020076363,
  -36.191181538089175,
  4.9750908402217285e-06,
  0.009520566435955282],
 [34.40577448714667,
  1.7746337170486868,
  48.25121832219984,
  -73.37824874831657,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  11.99833703626091,
  0.01835806177392501,
  0.012685911771409467,
  0.001009786028927451,
  0.003485358049971202,
  0.0009111182551995163,
  0.0006405838020076363,
  -36.191181538089175,
  4.9750908402217285e-06,
  0.009520566435955282],
 [34.40577448714667,
  1.7746337170486868,
  48.25121832219984,
  -73.37824874831657,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  12.492102024675049,
  0.01835806177392501,
  0.012685911771409467,
  0.001009786028927451,
  0.003485358049971202,
  0.0009111182551995163,
  0.0006405838020076363,
  -36.191181538089175,
  4.9750908402217285e-06,
  0.009520566435955282],
 [30.15666549418986,
  1.7746337170486868,
  44.40845548134062,
  -81.17397852577881,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  12.492102024675049,
  0.01835806177392501,
  0.012243090987314832,
  0.001009786028927451,
  0.0025920677224618473,
  0.0006458839218994715,
  0.0006437525063628547,
  -36.474771746065244,
  4.9750908402217285e-06,
  0.009384267411892713],
 [30.15666549418986,
  1.7746337170486868,
  47.11501790497679,
  -81.17397852577881,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  12.492102024675049,
  0.01835806177392501,
  0.012243090987314832,
  0.001009786028927451,
  0.0025920677224618473,
  0.0006458839218994715,
  0.0006437525063628547,
  -36.474771746065244,
  4.9750908402217285e-06,
  0.009384267411892713],
 [34.40577448714667,
  1.7746337170486868,
  48.25121832219984,
  -85.83988610793635,
  -52.149512474246556,
  2.8689454891603497e-05,
  7.600064936952291,
  12.492102024675049,
  0.01835806177392501,
  0.012685911771409467,
  0.001009786028927451,
  0.003485358049971202,
  0.0009111182551995163,
  0.0006405838020076363,
  -36.191181538089175,
  4.9750908402217285e-06,
  0.009520566435955282]]


from olfactorybulb.database import *
import os,sys
from neuronunit.tests.olfactory_bulb.publications import *
from neuronunit.tests.olfactory_bulb.tests import *
from neuronunit.models.neuron_cell import NeuronCellModel
from sciunit.suites import TestSuite
from pandas import DataFrame
import quantities as pq
from neuronunit.tests.olfactory_bulb.utilities import cache
from linetimer import CodeTimer
import string, math
from scipy import optimize
import numpy as np
import matplotlib.pyplot as plt
#import smt
#from smt.surrogate_models.genn import GENN, load_smt_data
from scipy import optimize
import linetimer
import multiprocessing
from multiprocessing import Pool, TimeoutError
from sciunit.scores import ZScore

cell_type = 'mc'

# Load tests
measurements = Measurement\
    .select(Measurement)\
    .join(Property)\
    .switch(Measurement)\
    .join(Source)\
    .where((Measurement.property.type == "Electrophysiology") & (Measurement.property.id.startswith(cell_type+'_')))\
    .order_by(Measurement.property.id)

properties = {}

for m in measurements:
    test_generic = str(m.property.test_class_generic)
    pub = str(m.source.publication_class).strip()
    class_name = test_generic+pub
    
    if test_generic not in properties:
        properties[test_generic] = []
    
    globals()[class_name] = type(class_name, 
                                 (eval(pub), eval(test_generic)), 
                                 {})
    
    print('creating specific test class: ' + class_name)
    
    test_instance = eval(class_name)(observation={"mean": m.mean * eval(m.property.units), "std": m.std * eval(m.property.units), "n": m.n})
    
    properties[test_generic].append(test_instance)
    
    
# Load model classes
model_classes = list(CellModel\
                    .select(CellModel)\
                    .where(CellModel.cell_type == cell_type.upper())             
                     )
for i, m in enumerate(model_classes):
    nmsp = string.join(m.isolated_model_class.split('.')[:-1],'.')
    cls = m.isolated_model_class.split('.')[-1]

    import_cmd = 'from '+nmsp+' import '+cls+' as Model'+str(i)
    print(import_cmd)
    exec(import_cmd)
    
# Create work item list
work_items = []

for model in model_classes:
    work_items.append({"model_class": model.isolated_model_class })
        
def get_suite_score(item):
    results = item
    results["properties"] = {}
    results["model_score"] = 0
    
    import prev_ob_models
    exec('cell = '+ str(item["model_class"]) +'()')
    
    from neuron import h
    ind = item["param_values"]
    
    for pi, pv in enumerate(ind):
        if params[pi]["attr"] == "tau_CaPool":
            setattr(h, params[pi]["attr"], pv)
        else:
            for param_list in params[pi]["lists"]:
                for sec in getattr(cell.cell, param_list):
                    setattr(sec, params[pi]["attr"], pv)

    model = NeuronCellModel(cell.soma(0.5),name=cell.__class__.__module__+'.'+cell.__class__.__name__+'|'+str(ind))


    for prop in properties.keys():
        
        if prop not in results["properties"]:
            results["properties"][prop] = { "tests": {}, "total_n": 0, "z_score_combined": None}
        
        prop_tests = properties[prop]
        
        for prop_test in prop_tests:
            
            prop_test_result = {}
            results["properties"][prop]["tests"][prop_test.__class__.__name__] = prop_test_result          
            
            try:
                #print('Starting', item, prop_test)
                prediction = prop_test.generate_prediction(model)
                #print('Finished', item, prop_test, prediction)

            except:
                import traceback
                prediction = traceback.format_exc()
                print(prediction)
                
            prop_test_result["observation"] = prop_test.observation
            prop_test_result["prediction"] = prediction
            
            if type(prediction) != str:
                z_score = (prediction - prop_test.observation["mean"])/prop_test.observation["std"]
                z_score = z_score.simplified
            else:
                z_score = 6.0 # errors are treated as 6 std deviation
                
            z_weighed = z_score * prop_test.observation["n"]

            prop_test_result["z_score"] = z_score
            prop_test_result["z_score_weighed"] = z_weighed
            
            results["properties"][prop]["total_n"] += prop_test.observation["n"]
            
        results["properties"][prop]["z_score_combined"] = sum([i["z_score_weighed"] for i in results["properties"][prop]["tests"].values()])
        results["properties"][prop]["z_score_combined"] /= results["properties"][prop]["total_n"]
        
        results["model_score"] += results["properties"][prop]["z_score_combined"].magnitude**2
        
    import math
    results["model_score"] = math.sqrt(results["model_score"])

    return results

def evaluate(ind):  
    def do_work():
        from prev_ob_models.Birgiolas2020.isolated_cells import MC
        score = get_suite_score({
            "model_class": "prev_ob_models.Birgiolas2020.isolated_cells.MC",
            "param_values": ind
        })
        
        return score["model_score"],
    
    from multiprocess import Pool, TimeoutError
    pool = Pool(processes = 1)
    
    try:
        result = pool.apply_async(do_work).get(timeout=4*60)
    except TimeoutError:
        print('Simulation timed out')
        pool.terminate()
        result = 9*10.0,
    except:
        print('Error in simulation')
        import traceback
        print(traceback.format_exc())
        pool.terminate()
        result = 9*10.999,
        
    return result

class NoDaemonProcess(multiprocessing.Process):
    # make 'daemon' attribute always return False
    def _get_daemon(self):
        return False
    def _set_daemon(self, value):
        pass
    daemon = property(_get_daemon, _set_daemon)

# We sub-class multiprocessing.pool.Pool instead of multiprocessing.Pool
# because the latter is only a wrapper function, not a proper class.
class MyPool(multiprocessing.pool.Pool):
    Process = NoDaemonProcess
    
    
def GA(suggested_pop=None, n=30, NGEN=30):
    #genn.options['print_global'] = False

    from deap import base, creator
    import math

    creator.create("FitnessMin", base.Fitness, weights=(-1,))
    creator.create("Individual", list, fitness=creator.FitnessMin)

    import random
    from deap import tools

    def random_indiv():
        result = [random.random()] * len(params)
        for i, pv in enumerate(result):
            result[i] = (params[i]["high"]-params[i]["low"])*pv+params[i]["low"]
        
        return creator.Individual(result)

    toolbox = base.Toolbox()
    toolbox.register("individual", random_indiv)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)

    def evaluate_surrogate(individual):
        return genn.predict_values(np.array([individual])).flatten()[0],

    lows = [p["low"] for p in params]
    highs = [p["high"] for p in params]
    
    toolbox.register("mate", tools.cxSimulatedBinaryBounded, eta=0.1, low=lows, up=highs)
    toolbox.register("mutate", tools.mutPolynomialBounded, eta=0.1, low=lows, up=highs, indpb=0.1)
    toolbox.register("evaluate", evaluate)



    toolbox.register("select", tools.selNSGA2, k=int(n*0.2))
    if suggested_pop is None:
        pop = toolbox.population(n=n)
    else:
        pop = [creator.Individual(i) for i in suggested_pop]
        
    CXPB, MUTPB = 1, 1
    F_DIVERSITY = 0.5

    # Evaluate the entire population
    processes = 15   
    pool = MyPool(processes = processes, maxtasksperchild=1)    
    fitnesses = pool.map(toolbox.evaluate, pop)
    pool.terminate()
    pool.join()
    
    for ind, fit in zip(pop, fitnesses):
        ind.fitness.values = fit

    for g in range(NGEN):
        # Select the parents
        elite = toolbox.select(pop)      

        random_offspring = toolbox.population(n=int(n*F_DIVERSITY/2.0))
        diversity_offspring = random_offspring + tools.selRandom(pop, int(n*F_DIVERSITY/2.0))        
        elite_offspring = tools.selRandom(elite, n-len(elite)-len(diversity_offspring))

        offspring = random_offspring + diversity_offspring + elite_offspring

        # Clone the selected individuals
        offspring = map(toolbox.clone, offspring)

        # Apply crossover and mutation on the offspring
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < CXPB:
                toolbox.mate(child1, child2)
                del child1.fitness.values
                del child2.fitness.values

        for mutant in offspring:
            if random.random() < MUTPB:
                toolbox.mutate(mutant)
                del mutant.fitness.values

        # Evaluate the individuals with an invalid fitness
        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        
        processes = max(1, multiprocessing.cpu_count()-1)
        pool = MyPool(processes = processes, maxtasksperchild=1)    
        fitnesses = pool.map(toolbox.evaluate, invalid_ind)
        pool.terminate()
        pool.join()        

        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit

        # The population is entirely replaced by the parents + offspring
        pop[:] = elite + offspring
        
        print("Generation",g,"out of",NGEN,"COMPLETE")
        print(np.min([i.fitness.values[0] for i in pop]))
        
    

    return pop, pop[0], toolbox.evaluate(pop[0])


pop, pop0, sur_fit = GA(top, n=max(5,multiprocessing.cpu_count()-1), NGEN=1)

import deap
top = deap.tools.selNSGA2(pop,8)
top


