import multiprocessing
import numpy

import scm.plams
import scm.pyzacros as pz
import scm.pyzacros.models

import adaptiveDesignProcedure as adp
import warnings; warnings.simplefilter('ignore', UserWarning)

def get_rate( conditions ):

    print("")
    print("  Requesting:")
    for cond in conditions:
        print("    x_CO = ", cond[0])
    print("")

    #---------------------------------------
    # Zacros calculation
    #---------------------------------------
    zgb = pz.models.ZiffGulariBarshad()

    z_sett = pz.Settings()
    z_sett.random_seed = 953129
    z_sett.temperature = 500.0
    z_sett.pressure = 1.0
    z_sett.species_numbers = ('time', 0.1)
    z_sett.max_time = 10.0

    z_job = pz.ZacrosJob( settings=z_sett, lattice=zgb.lattice,
                          mechanism=zgb.mechanism,
                          cluster_expansion=zgb.cluster_expansion )

    #---------------------------------------
    # Steady-State calculation
    #---------------------------------------
    ss_sett = pz.Settings()
    ss_sett.turnover_frequency.nbatch = 20
    ss_sett.turnover_frequency.confidence = 0.96
    ss_sett.turnover_frequency.nreplicas = 4

    ss_params = pz.ZacrosSteadyStateJob.Parameters()
    ss_params.add( 'max_time', 'restart.max_time',
                       2*z_sett.max_time*( numpy.arange(10)+1 )**3 )

    ss_job = pz.ZacrosSteadyStateJob( settings=ss_sett, reference=z_job,
                                      parameters=ss_params )

    #---------------------------------------
    # Parameters scan calculation
    #---------------------------------------
    ps_params = pz.ZacrosParametersScanJob.Parameters()
    ps_params.add( 'x_CO', 'molar_fraction.CO', [ cond[0] for cond in conditions ] )
    ps_params.add( 'x_O2', 'molar_fraction.O2', lambda params: 1.0-params['x_CO'] )

    ps_job = pz.ZacrosParametersScanJob( reference=ss_job, parameters=ps_params )

    #---------------------------------------
    # Running the calculations
    #---------------------------------------
    results = ps_job.run()

    if not results.job.ok():
        print('Something went wrong!')

    #---------------------------------------
    # Collecting the results
    #---------------------------------------
    data = numpy.nan*numpy.empty((len(conditions),3))
    if( results.job.ok() ):
        results_dict = results.turnover_frequency()
        results_dict = results.average_coverage( last=20, update=results_dict )

        for i in range(len(results_dict)):
            data[i,0] = results_dict[i]['average_coverage']['O*']
            data[i,1] = results_dict[i]['average_coverage']['CO*']
            data[i,2] = results_dict[i]['turnover_frequency']['CO2']

    return data


scm.pyzacros.init()

maxjobs = multiprocessing.cpu_count()
scm.plams.config.default_jobrunner = scm.plams.JobRunner(parallel=True, maxjobs=maxjobs)
scm.plams.config.job.runscript.nproc = 1
print('Running up to {} jobs in parallel simultaneously'.format(maxjobs))

input_var = ( { 'name':'x_CO', 'min':0.2, 'max':0.8, 'num':5 }, )

output_var = ( {'name':'ac_O'},
               {'name':'ac_CO'},
               {'name':'TOF_CO2'}, )

adpML = adp.adaptiveDesignProcedure( input_var, output_var, get_rate,
                                     algorithmParams={'dth':0.01,'d2th':0.10},
                                     outputDir=scm.pyzacros.workdir()+'/adp.results',
                                     randomState=10 )

adpML.createTrainingDataAndML()


x_CO,ac_O,ac_CO,TOF_CO2 = adpML.trainingData.T

print( "-------------------------------------------------" )
print( "%4s"%"cond", "%8s"%"x_CO", "%10s"%"ac_O", "%10s"%"ac_CO", "%12s"%"TOF_CO2" )
print( "-------------------------------------------------" )
for i in range(len(x_CO)):
    print( "%4d"%i, "%8.3f"%x_CO[i], "%10.6f"%ac_O[i], "%10.6f"%ac_CO[i], "%12.6f"%TOF_CO2[i] )


import matplotlib.pyplot as plt

fig = plt.figure()

x_CO_model = numpy.linspace(0.2,0.8,201)
ac_O_model,ac_CO_model,TOF_CO2_model = adpML.predict( x_CO_model ).T

ax = plt.axes()
ax.set_xlabel('Molar Fraction CO', fontsize=14)

ax.set_ylabel("Coverage Fraction (%)", color="blue", fontsize=14)
ax.plot(x_CO_model, ac_O_model, color="blue", linestyle="-.", lw=2, zorder=1)
ax.plot(x_CO, ac_O, marker='$\u25EF$', color='blue', markersize=4, lw=0, zorder=1)
ax.plot(x_CO_model, ac_CO_model, color="blue", linestyle="-", lw=2, zorder=2)
ax.plot(x_CO, ac_CO, marker='$\u25EF$', color='blue', markersize=4, lw=0, zorder=1)
plt.text(0.3, 0.9, 'O', fontsize=18, color="blue")
plt.text(0.7, 0.9, 'CO', fontsize=18, color="blue")

ax2 = ax.twinx()
ax2.set_ylabel("TOF (mol/s/site)",color="red", fontsize=14)
ax2.plot(x_CO_model, TOF_CO2_model, color='red', linestyle='-', lw=2, zorder=0)
ax2.plot(x_CO, TOF_CO2, marker='$\u25EF$', color='red', markersize=4, lw=0, zorder=1)
plt.text(0.37, 1.5, 'CO$_2$', fontsize=18, color="red")

plt.show()