# -*- coding: utf-8 -*-
# -------------------------------------------------------------
# Bulk Configuration
# -------------------------------------------------------------

# Set up lattice
vector_a = [11.0976166667, 0.0, 0.0]*Angstrom
vector_b = [0.0, 11.0976166667, 0.0]*Angstrom
vector_c = [0.0, 0.0, 19.809]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum, Platinum,
            Platinum, Platinum, Platinum, Platinum, Platinum]

# Define coordinates
fractional_coordinates = [[ 0.            ,  0.            ,  0.099035791812],
                          [ 0.25          , -0.            ,  0.099035791812],
                          [ 0.5           , -0.            ,  0.099035791812],
                          [ 0.75          , -0.            ,  0.099035791812],
                          [ 0.            ,  0.25          ,  0.099035791812],
                          [ 0.25          ,  0.25          ,  0.099035791812],
                          [ 0.5           ,  0.25          ,  0.099035791812],
                          [ 0.75          ,  0.25          ,  0.099035791812],
                          [ 0.            ,  0.5           ,  0.099035791812],
                          [ 0.25          ,  0.5           ,  0.099035791812],
                          [ 0.5           ,  0.5           ,  0.099035791812],
                          [ 0.75          ,  0.5           ,  0.099035791812],
                          [ 0.            ,  0.75          ,  0.099035791812],
                          [ 0.25          ,  0.75          ,  0.099035791812],
                          [ 0.5           ,  0.75          ,  0.099035791812],
                          [ 0.75          ,  0.75          ,  0.099035791812],
                          [ 0.125         ,  0.125         ,  0.198071583624],
                          [ 0.375         ,  0.125         ,  0.198071583624],
                          [ 0.625         ,  0.125         ,  0.198071583624],
                          [ 0.875         ,  0.125         ,  0.198071583624],
                          [ 0.125         ,  0.375         ,  0.198071583624],
                          [ 0.375         ,  0.375         ,  0.198071583624],
                          [ 0.625         ,  0.375         ,  0.198071583624],
                          [ 0.875         ,  0.375         ,  0.198071583624],
                          [ 0.125         ,  0.625         ,  0.198071583624],
                          [ 0.375         ,  0.625         ,  0.198071583624],
                          [ 0.625         ,  0.625         ,  0.198071583624],
                          [ 0.875         ,  0.625         ,  0.198071583624],
                          [ 0.125         ,  0.875         ,  0.198071583624],
                          [ 0.375         ,  0.875         ,  0.198071583624],
                          [ 0.625         ,  0.875         ,  0.198071583624],
                          [ 0.875         ,  0.875         ,  0.198071583624],
                          [ 0.            ,  0.            ,  0.297107375435],
                          [ 0.25          , -0.            ,  0.297107375435],
                          [ 0.5           , -0.            ,  0.297107375435],
                          [ 0.75          , -0.            ,  0.297107375435],
                          [ 0.            ,  0.25          ,  0.297107375435],
                          [ 0.25          ,  0.25          ,  0.297107375435],
                          [ 0.5           ,  0.25          ,  0.297107375435],
                          [ 0.75          ,  0.25          ,  0.297107375435],
                          [ 0.            ,  0.5           ,  0.297107375435],
                          [ 0.25          ,  0.5           ,  0.297107375435],
                          [ 0.5           ,  0.5           ,  0.297107375435],
                          [ 0.75          ,  0.5           ,  0.297107375435],
                          [ 0.            ,  0.75          ,  0.297107375435],
                          [ 0.25          ,  0.75          ,  0.297107375435],
                          [ 0.5           ,  0.75          ,  0.297107375435],
                          [ 0.75          ,  0.75          ,  0.297107375435],
                          [ 0.125         ,  0.125         ,  0.396143167247],
                          [ 0.375         ,  0.125         ,  0.396143167247],
                          [ 0.625         ,  0.125         ,  0.396143167247],
                          [ 0.875         ,  0.125         ,  0.396143167247],
                          [ 0.125         ,  0.375         ,  0.396143167247],
                          [ 0.375         ,  0.375         ,  0.396143167247],
                          [ 0.625         ,  0.375         ,  0.396143167247],
                          [ 0.875         ,  0.375         ,  0.396143167247],
                          [ 0.125         ,  0.625         ,  0.396143167247],
                          [ 0.375         ,  0.625         ,  0.396143167247],
                          [ 0.625         ,  0.625         ,  0.396143167247],
                          [ 0.875         ,  0.625         ,  0.396143167247],
                          [ 0.125         ,  0.875         ,  0.396143167247],
                          [ 0.375         ,  0.875         ,  0.396143167247],
                          [ 0.625         ,  0.875         ,  0.396143167247],
                          [ 0.875         ,  0.875         ,  0.396143167247],
                          [ 0.25          ,  0.25          ,  0.495178959059]]

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

# Add tags
bulk_configuration.addTags('Selection 0', [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12,
                                           13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,
                                           26, 27, 28, 29, 30, 31])

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------

potentialSet = EAM_Pt_2004()
calculator = TremoloXCalculator(parameters=potentialSet)
calculator.setVerletListsDelta(0.25*Angstrom)

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('akmc.nc', bulk_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
fix_atom_indices_0 = [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12,
                      13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,
                      26, 27, 28, 29, 30, 31]
constraints = [FixAtomConstraints(fix_atom_indices_0)]

bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.01*eV/Ang,
    max_steps=200,
    max_step_length=0.2*Ang,
    constraints=constraints,
    trajectory_filename=None,
    disable_stress=True,
    optimizer_method=LBFGS(),
)
nlsave('akmc.nc', bulk_configuration)
nlprint(bulk_configuration)

# -------------------------------------------------------------
# Adaptive Kinetic Monte Carlo
# -------------------------------------------------------------

saddle_search_parameters = SaddleSearchParameters(
    max_neb_images=5,
    neb_climbing_image=True,
    )

if os.path.isfile('akmc_markov_chain.nc'):
    markov_chain = nlread('akmc_markov_chain.nc')[-1]
else:
    markov_chain = MarkovChain(
        configuration=bulk_configuration,
        configuration_energy=TotalEnergy(bulk_configuration).evaluate(),
        )

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

constraints = [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12,
               13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,
               26, 27, 28, 29, 30, 31]

akmc = AdaptiveKineticMonteCarlo(
    markov_chain=markov_chain,
    kmc_temperature=300.0*Kelvin,
    md_temperature=2000.0*Kelvin,
    calculator=bulk_configuration.calculator(),
    kmc=kmc,
    saddle_search_parameters=saddle_search_parameters,
    constraints=constraints,
    confidence=0.95,
    filename_prefix='akmc',
    )

akmc.run(max_searches=100, max_kmc_steps=10000)