# %% ForceFieldCalculator

# This script uses a forcefield calculator for simply showing how to set up the workflow
# but for real studies of defect dynamics and other complex phenomena a GGA DFT calculator
# is advisable.
potentialSet = ReaxFF_CHONSSiCaCsKSrNaMgAl_2014(strict_bondpairs=False)
calculator = TremoloXCalculator(parameters=potentialSet)


# %% Load TrainingSet
# Load a premade training set. This simple training set has been created with
# respectively pure silicon carbide cells of varying sizes and some simple
# defect configurations created with the Defects object with an interstitial
# magnesium atom included.
training_set_data_results_training_set_0 = nlread(
    filename=r'C:/full/path/to/folder/TrainingSet_data_results.hdf5',
    object_id='TrainingSet_0',
)[0]


# %% Table

if 'table' not in locals():
    table = Table('AKMC-AL_SiC_Mg-interstitial.hdf5', object_id='table')
    table.addInstanceColumn(key='training_set', types=TrainingSet)
    table.setMetatext('Table')

table.append(training_set_data_results_training_set_0)


# %% NonLinearCoefficientsParameters

non_linear_coefficients_parameters = NonLinearCoefficientsParameters(random_seed=1)


# %% MomentTensorPotentialFittingParameters

moment_tensor_potential_fitting_parameters = MomentTensorPotentialFittingParameters(
    basis_size=500,
    outer_cutoff_radii=4.5 * Angstrom,
    forces_cap=500.0 * eV / Angstrom,
    non_linear_coefficients_parameters=non_linear_coefficients_parameters,
    use_element_specific_coefficients=True,
)


# %% ActiveLearningSimulation

extrapolation_selection_parameters = ExtrapolationSelectionParameters(
    extrapolation_grade_algorithm=QueryByCommitteeForces,
    write_atomic_error_estimates=True,
)

active_learning_simulation = ActiveLearningSimulation(
    fitting_parameters=moment_tensor_potential_fitting_parameters,
    initial_training_data=table,
    mtp_study_filename='AKMC-AL_SiC_Mg-interstitial.hdf5',
    mtp_study_object_id='mtp_active_learning',
    reference_calculator=calculator,
    retrain_threshold=2.0,
    candidate_trajectory_filename='AKMC-AL_SiC_Mg-interstitial.hdf5',
    candidate_trajectory_object_id='al_candidates',
    max_iterations=100,
    extrapolation_selection_parameters=extrapolation_selection_parameters,
)


# %% ConfigurationTable(List)

# %% Configuration

# Set up lattice
lattice = Hexagonal(9.24153*Angstrom, 10.0848*Angstrom)

# Define elements
elements = [Magnesium, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon, Silicon, Carbon, Silicon, Carbon, Silicon, Carbon,
            Silicon, Carbon]

# Define coordinates
fractional_coordinates = [[-0.000001796922, -0.000000562305,  0.596070368757],
                          [-0.000000491118, -0.000001423676,  0.000842041026],
                          [-0.000000205884, -0.000001383541,  0.188233014251],
                          [ 0.110811058927,  0.221621814395,  0.249969311004],
                          [ 0.109829059508,  0.2196583359  ,  0.439589629801],
                          [ 0.245293333655,  0.122645489506,  0.488142170214],
                          [ 0.225402718954,  0.112700398204,  0.687221178349],
                          [ 0.116653980832,  0.233307899771,  0.757331706657],
                          [ 0.11146277856 ,  0.222925272715,  0.940751317959],
                          [-0.000182451102,  0.334335956139,  0.001401835037],
                          [ 0.000126658693,  0.333610745075,  0.188919306969],
                          [ 0.111713775738,  0.55585572834 ,  0.250483969249],
                          [ 0.111211131028,  0.555604200648,  0.437760793502],
                          [ 0.222330915802,  0.444660951181,  0.500401792859],
                          [ 0.223660337491,  0.447319977989,  0.688113927628],
                          [ 0.110764068664,  0.555380822675,  0.75124517968 ],
                          [ 0.111094766015,  0.555546249638,  0.938001717188],
                          [-0.000182413478,  0.665479510391,  0.001401800378],
                          [ 0.000126722726,  0.666513720997,  0.188919311512],
                          [ 0.110811253577,  0.889186932759,  0.2499694597  ],
                          [ 0.109829395363,  0.890168437617,  0.439589668067],
                          [ 0.222330695434,  0.777667199888,  0.500401732683],
                          [ 0.223660088609,  0.776337844776,  0.688113886822],
                          [ 0.116654009306,  0.883343889406,  0.75733141284 ],
                          [ 0.111462832867,  0.888535342786,  0.940751158024],
                          [ 0.334518658582,  0.000180596666,  0.001401815154],
                          [ 0.333484463299, -0.000128552864,  0.188919261249],
                          [ 0.446812059424,  0.223404856537,  0.245228012251],
                          [ 0.456543201939,  0.228270381015,  0.431675884813],
                          [ 0.554798395266,  0.109596872961,  0.500145533664],
                          [ 0.555556707125,  0.111113437972,  0.687570795315],
                          [ 0.446612498081,  0.223305123252,  0.751441590001],
                          [ 0.443672532677,  0.221835220216,  0.937622702914],
                          [ 0.334518665183,  0.334335953788,  0.001401886277],
                          [ 0.333484416638,  0.333610760408,  0.188919345063],
                          [ 0.444142309997,  0.555855766271,  0.250483916962],
                          [ 0.444393617909,  0.555604225616,  0.437760642442],
                          [ 0.55479867684 ,  0.445199205235,  0.500145623278],
                          [ 0.555557011734,  0.444440938365,  0.687570890175],
                          [ 0.444617255708,  0.55538068838 ,  0.751245236818],
                          [ 0.444451784801,  0.555546238995,  0.938001772576],
                          [ 0.333332812795,  0.666665348345,  0.000498637507],
                          [ 0.333332807161,  0.666665295376,  0.188195233904],
                          [ 0.444142341554,  0.888284277919,  0.250483948497],
                          [ 0.444393593349,  0.888786637092,  0.437760784015],
                          [ 0.55533696638 ,  0.777667004673,  0.500401762112],
                          [ 0.552677949783,  0.776337601652,  0.68811393903 ],
                          [ 0.444617087112,  0.8892338381  ,  0.751245172959],
                          [ 0.444451782942,  0.888903289278,  0.938001710532],
                          [ 0.665662186868,  0.000180653452,  0.001401825274],
                          [ 0.666387391661, -0.000128521614,  0.188919250589],
                          [ 0.776593048989,  0.223404954227,  0.245227802701],
                          [ 0.771726970503,  0.22827093797 ,  0.431675605268],
                          [ 0.877351901913,  0.122645867165,  0.488141315009],
                          [ 0.887297100201,  0.112700891892,  0.687221327326],
                          [ 0.776692560289,  0.223305456976,  0.75144168897 ],
                          [ 0.778162860019,  0.221835322908,  0.937622793709],
                          [ 0.666666109383,  0.333332111268, -0.002869186472],
                          [ 0.666666128427,  0.333332135772,  0.182762928989],
                          [ 0.776593232543,  0.553186114161,  0.2452280084  ],
                          [ 0.771727165243,  0.543453991573,  0.431675654892],
                          [ 0.89040092211 ,  0.44519944496 ,  0.500145529853],
                          [ 0.888884545611,  0.444441257865,  0.687570775483],
                          [ 0.776692914778,  0.553385577665,  0.751441571953],
                          [ 0.778162915485,  0.556325532925,  0.937622740482],
                          [ 0.665662181617,  0.6654794412  ,  0.001401871408],
                          [ 0.666387398713,  0.666513773324,  0.188919337018],
                          [ 0.778376349524,  0.88918711857 ,  0.249969284304],
                          [ 0.780339644533,  0.890168850272,  0.439589560342],
                          [ 0.877353113378,  0.754706111706,  0.488142521702],
                          [ 0.887297635426,  0.774595359534,  0.68722101625 ],
                          [ 0.766690016129,  0.883343927058,  0.757331621642],
                          [ 0.777072777134,  0.888535332651,  0.940751367245]]

# First AKMC simulation: Initial configuration setup for an AKMC simulation
# with active learning that should be started from scratch

# Set up configuration
configuration_0 = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# Add tags - constraints needed for AKMC
configuration_0.addTags('fixed', [36, 37])

configuration_name_0 = "optimized_configuration_sic_mginterstitial"

configuration_table = Table()
configuration_table.addInstanceColumn(key='configuration', types=BulkConfiguration)
configuration_table.append(configuration_0)

configuration_names_table = Table()
configuration_names_table.addStringColumn(key='configuration_name')
configuration_names_table.append(configuration_name_0)


# %% ActiveLearningOptimizeGeometry

optimized_configuration_table = runActiveLearningOptimizeGeometry(
    active_learning_simulation=active_learning_simulation,
    configurations=configuration_table,
    max_forces=0.01 * eV / Angstrom,
    constraints=[BravaisLatticeConstraint()],
    trajectory_filename='AKMC-AL_SiC_Mg-interstitial.hdf5',
    trajectory_object_id='alopt',
    log_filename_prefix='AKMC-AL_pre-optimize_geometry',
)

nlsave(
    'AKMC-AL_SiC_Mg-interstitial.hdf5', optimized_configuration_table, object_id='alopt_1'
)

training_set_table = active_learning_simulation.trainingSetTable()


# Custom block for extraction the configuration from the table


def custom(table):
    # This custom script supports atkpython syntax
    # and can perform almost any procedure.

    configuration = table[0, 0]
    configuration.update()
    return configuration


configuration = custom(optimized_configuration_table)

nlsave('AKMC-AL_SiC_Mg-interstitial.hdf5', configuration)


# %% TotalEnergy

total_energy = TotalEnergy(configuration=configuration)
nlsave('AKMC-AL_SiC_Mg-interstitial.hdf5', total_energy)


## Restarting the AKMC simulation: Initial configuration setup for an AKMC simulation
## with active learning that should be restarted
# %% Load BulkConfiguration
#
# configuration = nlread(
#     filename=r'C:/full/path/to/folder/AKMC-AL_SiC_Mg-interstitial.hdf5',
#     object_id='BulkConfiguration_0',
# )[0]
#
#
# %% Load TotalEnergy
#
# total_energy = nlread(
#     filename=r'C:/full/path/to/folder/AKMC-AL_SiC_Mg-interstitial.hdf5',
#     object_id='TotalEnergy_0',
# )[0]


# %% AdaptiveKineticMonteCarlo

saddle_search_parameters = SaddleSearchParameters(neb_climbing_image=True)

if os.path.isfile('akmc-al_markov_chain.nc'):
    markov_chain = nlread('akmc-al_markov_chain.nc')[-1]
else:
    markov_chain = MarkovChain(
        configuration=configuration, configuration_energy=total_energy.evaluate()
    )

if os.path.isfile('akmc-al_kmc.nc'):
    kmc = nlread('akmc-al_kmc.nc')[-1]
else:
    kmc = None

saddle_search_method = MDSaddleSearch(
    temperature=2900.0 * Kelvin,
    saddle_search_parameters=saddle_search_parameters,
    active_learning=active_learning_simulation,
    constraints=configuration.indicesFromTags(['fixed']),
)

akmc = AdaptiveKineticMonteCarlo(
    markov_chain=markov_chain,
    kmc_temperature=300.0 * Kelvin,
    calculator=configuration.calculator(),
    kmc=kmc,
    filename_prefix='akmc-al',
    saddle_search_method=saddle_search_method,
)

akmc.run(max_searches=50, max_kmc_steps=50)
