# coding: utf-8

# # This notebook runs genetic algorithm varying rates, thermo and rates+thermo

# ## Dependencies

# * Install DEAP using any package manager
#     * `pip install deap==1.0.2` 

# In[1]:

import IPython
from IPython.display import display
import sys
import os
import re
sys.path.insert(1,os.getenv('RMGpy', default=os.path.expanduser('~/Code/RMG-Py')))
import rmgpy
print "Loading RMG-Py from {!s}".format(rmgpy.__file__)
from rmgpy.molecule import Molecule
import rmgpy.kinetics
import numpy
import cPickle as pickle
import json
from collections import Counter, defaultdict
import operator
import multiprocessing


# In[2]:

master = 'CombFlame2012/2028-Sarathy' # this should be what was used for '--names=' when generating the .pkl files


# ## Change the directory paths

# In[5]:

# Find and read the chemkin file
with open(os.path.join('/home/sirumalla.s/Code/butanol-comparison/',master,'import.sh')) as infile:
    shellscript = infile.read()
reactions_filename = re.search('--reactions\s+(\S+)',shellscript).group(1)
reactions_filepath = os.path.join('/home/sirumalla.s/Code/butanol-comparison/',master,reactions_filename)
thermo_filename = re.search('--thermo\s+(\S+)',shellscript).group(1)
thermo_filepath = os.path.join('/home/sirumalla.s/Code/butanol-comparison/',master,thermo_filename)
print(reactions_filepath)
print(thermo_filepath)
with open(reactions_filepath) as infile:
    chemkin = infile.readlines()
#for i,line in enumerate(chemkin):
    #print i, line.strip()


# In[6]:

# If you have an up to date cantera, then
#from cantera import ck2cti
# otherwise, do the next two cells
# no, use the hacked version of ck2cti.py in this repo


# In[7]:

#%%bash
#curl https://raw.githubusercontent.com/Cantera/cantera/master/interfaces/cython/cantera/ck2cti.py > ck2cti.py


# In[8]:

import ck2cti


# In[9]:

parser = ck2cti.Parser()
surfaces = parser.convertMech(inputFile=reactions_filepath,
                              thermoFile=thermo_filepath,
                              transportFile=None,
                                surfaceFile=None,
                              phaseName=None,
                              outName='master.cti',
                                  permissive=True)
parser.reactions[0]


# In[10]:

original_reactions = list()
for i, r in enumerate(parser.reactions):
    original_reactions.append(r)
    assert original_reactions[i] == parser.reactions[i]
len(original_reactions)


# In[11]:

original_thermo=list()
for i, s in enumerate(parser.speciesList):
    original_thermo.append(s.thermo)
    assert original_thermo[i] == parser.speciesList[i].thermo
len(original_thermo)


# ## Use the pickle files generated by `get-alternative-rates-thermo.ipynb`

# In[15]:

with open("../../get_alternatives/alternatives_rates.pkl") as infile:
    alternatives_rates = pickle.load(infile)


# In[17]:

with open("../../get_alternatives/alternatives_thermo.pkl") as infile:
    alternatives_thermo = pickle.load(infile)


# In[18]:

import os, sys
from collections import defaultdict
import cantera as ct
import numpy as np
import pandas as pd
import itertools
import random


# In[19]:

from deap import base, creator, algorithms, tools


# In[21]:

NBR_REACTIONS = len(parser.reactions)  # 2346
NBR_SPECIES = len(parser.speciesList)  #  431


# In[22]:

ALL_OPTIONS_THERMO = [tuple(range(len(alternatives_thermo[i])+1)) for i in range(NBR_SPECIES) ]
ALL_OPTIONS_THERMO[:10]


# In[23]:

ALL_OPTIONS_RATES = [tuple(range(len(alternatives_rates[i])+1)) for i in range(NBR_REACTIONS) ]
ALL_OPTIONS_RATES[:10]


# In[24]:

ALL_OPTIONS =  ALL_OPTIONS_THERMO + ALL_OPTIONS_RATES
len(ALL_OPTIONS)


# ## Change the `+1.0` to `-1.0` to minimize

# In[36]:

creator.create('Fitness', base.Fitness , weights = ((+1.0),))  # minimize or maximize a single objective
creator.create('Individual', list , fitness = creator.Fitness) # represent things as a list


# In[25]:

def mutMechanismAll(individual, indpb):
    """Mutate a mechanism. EVERY rate is mutated with probability indpb"""
    for i, choices in enumerate(ALL_OPTIONS):
        if random.random() < indpb:
            individual[i] = random.choice(choices)
    return individual,


# In[26]:

def mutMechanismOne(individual):
    """Mutate a mechanism. Only one rate is mutated."""
    variable_rates = [i for i, choices in enumerate(ALL_OPTIONS) if choices > (0,)]
    i = random.choice(variable_rates)
    individual[i] = random.choice(ALL_OPTIONS[i])
    return individual,


# In[27]:

outdir = 'cantera-files-GA'
os.path.exists(outdir) or os.mkdir(outdir)
for f in os.listdir(outdir):
    os.unlink(os.path.join(outdir, f))


# ## Convert the following cell into `code` to run GA varying just rates
def makeCanteraFile(list_of_choices):
    """
    Makes a cantera file for the provided list of rate choices.
    Returns the path to the file.
    """
    for i, choice in enumerate(list_of_choices):
        if choice == 0:
            parser.reactions[i] = original_reactions[i]
        else:
            parser.reactions[i] = alternatives[i][choice-1]

    worker_name = multiprocessing.current_process().name
    output_filename = worker_name+'.cti'
    output_file_path = os.path.join(outdir, output_filename)
    header = ["#"*80,
             "## {:74s} ##".format("Model built from this genome:")]
    s = ''.join([str(i) for i in list_of_choices])
    import textwrap
    header.extend(["## {:74s} ##".format(line) for line in textwrap.wrap(s, width=74)])
    header.append("#"*80)
    parser.writeCTI(header=header, outName=output_file_path)
    return output_file_path
# ## Convert the following cell into `code` to run GA varying just thermo
def makeCanteraFile(list_of_choices):
    """
    Makes a cantera file for the provided list of rate choices.
    Returns the path to the file.
    """
    for i, choice in enumerate(list_of_choices):
        if choice == 0:
            parser.speciesList[i].thermo = original_thermo[i]
        else:
            replacement, models = alternatives[i][choice-1]
            parser.speciesList[i].thermo = replacement
            parser.speciesList[i].note = "Replaced with option {} as seen in {}".format(j+1, ', '.join(models))
    worker_name = multiprocessing.current_process().name
    output_filename = worker_name+'.cti'
    output_file_path = os.path.join(outdir, output_filename)
    header = ["#"*80,
             "## {:74s} ##".format("Model built from this genome:")]
    s = ''.join([str(i) for i in list_of_choices])
    import textwrap
    header.extend(["## {:74s} ##".format(line) for line in textwrap.wrap(s, width=74)])
    header.append("#"*80)
    parser.writeCTI(header=header, outName=output_file_path)
    return output_file_path
# ## The following cells runs GA while varying rates and thermo

# In[28]:

def makeCanteraFile(list_of_choices):
    """
    Makes a cantera file for the provided list of parameter choices.
    Returns the path to the file.
    """
    n_thermo = len(ALL_OPTIONS_THERMO)
    list_of_choices_thermo = list_of_choices[:n_thermo]
    list_of_choices_rates = list_of_choices[n_thermo:]
    # Thermo substitutions
    for i, choice in enumerate(list_of_choices_thermo):
        if choice == 0:
            parser.speciesList[i].thermo = original_thermo[i]
            parser.speciesList[i].note = None
        else:
            replacement, models = alternatives_thermo[i][choice-1]
            parser.speciesList[i].thermo = replacement
            parser.speciesList[i].note = "Replaced with option {} as seen in {}".format(choice, ', '.join(models))
    # Rate substitutions
    for i, choice in enumerate(list_of_choices_rates):
        if choice == 0:
            parser.reactions[i] = original_reactions[i]
        else:
            parser.reactions[i] = alternatives_rates[i][choice-1]

    worker_name = multiprocessing.current_process().name
    process_id = multiprocessing.current_process().pid
    output_filename = "{}.{}.cti".format(worker_name,process_id)
    output_file_path = os.path.join(outdir, output_filename)
    header = ["#"*80,
             "## {:74s} ##".format("Model built from this genome:")]
    s = ''.join([str(int(i)) for i in list_of_choices])
    import textwrap
    header.extend(["## {:74s} ##".format(line) for line in textwrap.wrap(s, width=74)])
    header.append("#"*80)
    parser.writeCTI(header=header, outName=output_file_path)
    return output_file_path


# In[31]:

def get_ignition_delay(cantera_file_path, temperature , pressure, stoichiometry=1.0, plot=False, isomer='n'):
    """
    Get the ignition delay at temperature (K) and pressure (bar) and stochiometry (phi),
    for the butanol isomer (n,s,t,i)
    """
    try:
        ct.suppress_thermo_warnings(True)
    except AttributeError:
        print("Sorry about the warnings...")
    gas = ct.Solution(cantera_file_path)
    assert isomer in ['n','s','t','i'], "Expecting isomer n,s,t, or i not {}".format(isomer)
    oxygen_mole = 1.0
    argon_mole = 96./4.*oxygen_mole
    butanol_mole = stoichiometry * oxygen_mole/6.
    X_string = isomer + 'c4h9oh:{0}, o2:{1}, ar:{2}'.format(butanol_mole, oxygen_mole, argon_mole)
    gas.TPX = temperature, pressure*1e5, X_string
    reactor=ct.IdealGasReactor(gas)
    reactor_network=ct.ReactorNet([reactor])
    time=0.0
    end_time=1000e-3
    times=[]
    concentrations=[]
    
    pressures=[]
    temperatures=[]
    print_data = True
    while time < end_time:
        time=reactor_network.time
        times.append(time)
        temperatures.append(reactor.T)
        pressures.append(reactor.thermo.P)
        concentrations.append(reactor.thermo.concentrations)
        reactor_network.step(end_time)
    print("reached end time {0:.4f} ms in {1} steps ". format(times[-1]*1e3, len(times)))
    concentrations = np.array(concentrations)
    times = np.array(times)
    pressures = np.array(pressures)
    temperatures = np.array(temperatures)
    
    dTdt = (temperatures[1:] - temperatures[:-1]) / (times[1:] - times[:-1])
    
    if plot:
        plt.subplot(2,1,1)
        plt.plot(times,temperatures)
        plt.subplot(2,1,2)
        plt.plot(times[1:], dTdt)
        plt.show()
        
    step_with_fastest_T_rise = dTdt.argmax()
    if step_with_fastest_T_rise > 1 and step_with_fastest_T_rise < len(times)-2:
        ignition_time_ms = 1e3 * times[step_with_fastest_T_rise]
        print("At {0} K {1} bar, ignition delay time is {2} ms for {3}-butanol".format(temperature,pressure,ignition_time_ms,isomer))
        return ignition_time_ms
    else:
        print("At {0} K {1} bar, no ignition is detected for {2}-butanol" .format(temperature, pressure, isomer))
        return np.infty


# In[42]:

class DummyFile(object):
    # could write to this instead of a log file, to suppress output entirely
    def write(self, x): pass

class LogToFile(object):
    def __init__(self, filename):
        self.filename = filename
    def __enter__(self):
        self.stdout = sys.stdout
        self.output = open(self.filename,'w')
        sys.stdout = self.output
    def __exit__(self, error_type, value, traceback):
        sys.stdout = self.stdout
        self.output.close()


# ## Rename the `.log` and `.pkl` files to avoid conflicts but the default is rates+thermo variation

# In[43]:

def evaluate(individual):
    """
    Evaluate a mechanism
    
    Takes in an individaul, which is a list of choices, one for each parameter (thermo then rates).
    With this, it should generate a cantera file, and run a simulation, and return 
    the ignition delay.
    """
    cantera_file_path = makeCanteraFile(individual)
    temperature = 1500.
    pressure = 43 * 1.01325
    stoichiometry=1.0
    with LogToFile(cantera_file_path+'.log'):
        try:
            idt = get_ignition_delay(cantera_file_path, temperature, pressure, stoichiometry, plot=False, isomer='n')
        except Exception as e:
            print "ERROR", e.message
            print "using nominal value to continue optimization"
            return NOMINAL_IDT
    with open('evaluations-rates-thermo-new-conditions-max.log','a') as output:
        output.write("{!r}: {!r}\n".format(''.join([str(int(i)) for i in individual]), idt))
    return idt,


# In[44]:

# The ignition delay time of the original model
NOMINAL_IDT = evaluate(np.zeros(NBR_SPECIES+NBR_REACTIONS))
print "Nominal IDT used in case of errors:", NOMINAL_IDT


# In[34]:

def generator():
    "To generate a random kinetic model"
    for i, options in enumerate(ALL_OPTIONS):
        yield random.choice(options)


# In[37]:

toolbox = base.Toolbox()
toolbox.register("indices", generator)
toolbox.register("individual", tools.initIterate, creator.Individual, toolbox.indices)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)


# In[38]:

toolbox.register("mate", tools.cxOnePoint)
# toolbox.register("mutate", mutMechanismAll, indpb=0.01)
toolbox.register("mutate", mutMechanismOne)
toolbox.register("select", tools.selTournament, tournsize=3)
toolbox.register("evaluate", evaluate)


# In[45]:

print multiprocessing.cpu_count()
if multiprocessing.cpu_count() < 10 :
    print "Running on a laptop. Not using multiprocessing, to simplify debugging"
    toolbox.register("map", map)
else:
    pool = multiprocessing.Pool(maxtasksperchild=100) # estimate ran ~800 per worker before crashing on largemem node
    toolbox.register("map", pool.map)


# In[46]:

# Monkey-patch the logbook so it saves it to disk at each generation
import deap, deap.tools
if 'old_record' in locals():
    "looks like this cell has already been run. Do nothing!"
else:
    old_record = deap.tools.Logbook.record
    def new_record(self, **infos):
        old_record(self, **infos)
        with open("logbook-rates-thermo-1500k-43atm-max.pkl","w") as outfile:
            pickle.dump(self, outfile)
        with open("logbook-rates-thermo-1500k-43atm-max.log","w") as outfile:
            outfile.write(str(self))
    deap.tools.Logbook.record = new_record


# In[ ]:

def main():
    random.seed(64)
    NGEN = 1000
    MU = 100
    LAMBDA = 200
    CXPB = 0.7
    MUTPB = 0.2
    
    if multiprocessing.cpu_count() < 10 :
        print "Running on a laptop."
        NGEN = 2
        MU = 3
        LAMBDA = 6

    
    pop = toolbox.population(n=MU)
    hof = tools.ParetoFront()
    stats = tools.Statistics(lambda ind: ind.fitness.values)
    stats.register("avg", np.mean, axis=0)
    stats.register("std", np.std, axis=0)
    stats.register("min", np.min, axis=0)
    stats.register("max", np.max, axis=0)
    
    pop, logbook = algorithms.eaMuPlusLambda(pop, toolbox, MU, LAMBDA, CXPB, MUTPB, NGEN,
                                             stats=stats, halloffame=hof, verbose=True
                             )
    record = stats.compile(pop)
    
    # pickle.dump(logbook, open("logbook.pkl","w"))
    logbook.header = "gen", "avg", "std", "min", "max"
    print(logbook)
    
    #print (tools.selBest(pop, k=1))
    
    return pop, stats, hof


# In[ ]:

if __name__ == "__main__":
    final_pop, stats, hall_of_fame = main()
    print final_pop
    with open("final-pop-thermo-rates-1500k-43atm-max.pkl","w") as outfile:
        pickle.dump(final_pop, outfile)