# -------------------------------------------------------------
# Bulk configuration
# -------------------------------------------------------------

# Set up lattice
lattice = Hexagonal(9.8448*Angstrom, 13.418*Angstrom)

# Define elements
elements = [Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon, Carbon,
            Lithium]

# Define coordinates
fractional_coordinates = [[ 0.        ,  0.        ,  0.125     ],
                          [ 0.        ,  0.        ,  0.625     ],
                          [ 0.        ,  0.25      ,  0.125     ],
                          [ 0.        ,  0.25      ,  0.625     ],
                          [ 0.        ,  0.5       ,  0.125     ],
                          [ 0.        ,  0.5       ,  0.625     ],
                          [ 0.        ,  0.75      ,  0.125     ],
                          [ 0.        ,  0.75      ,  0.625     ],
                          [ 0.25      ,  0.        ,  0.125     ],
                          [ 0.25      ,  0.        ,  0.625     ],
                          [ 0.25      ,  0.25      ,  0.125     ],
                          [ 0.25      ,  0.25      ,  0.625     ],
                          [ 0.25      ,  0.5       ,  0.125     ],
                          [ 0.25      ,  0.5       ,  0.625     ],
                          [ 0.25      ,  0.75      ,  0.125     ],
                          [ 0.25      ,  0.75      ,  0.625     ],
                          [ 0.5       ,  0.        ,  0.125     ],
                          [ 0.5       ,  0.        ,  0.625     ],
                          [ 0.5       ,  0.25      ,  0.125     ],
                          [ 0.5       ,  0.25      ,  0.625     ],
                          [ 0.5       ,  0.5       ,  0.125     ],
                          [ 0.5       ,  0.5       ,  0.625     ],
                          [ 0.5       ,  0.75      ,  0.125     ],
                          [ 0.5       ,  0.75      ,  0.625     ],
                          [ 0.75      ,  0.        ,  0.125     ],
                          [ 0.75      ,  0.        ,  0.625     ],
                          [ 0.75      ,  0.25      ,  0.125     ],
                          [ 0.75      ,  0.25      ,  0.625     ],
                          [ 0.75      ,  0.5       ,  0.125     ],
                          [ 0.75      ,  0.5       ,  0.625     ],
                          [ 0.75      ,  0.75      ,  0.125     ],
                          [ 0.75      ,  0.75      ,  0.625     ],
                          [ 0.        ,  0.        ,  0.375     ],
                          [ 0.        ,  0.        ,  0.875     ],
                          [ 0.        ,  0.25      ,  0.375     ],
                          [ 0.        ,  0.25      ,  0.875     ],
                          [ 0.        ,  0.5       ,  0.375     ],
                          [ 0.        ,  0.5       ,  0.875     ],
                          [ 0.        ,  0.75      ,  0.375     ],
                          [ 0.        ,  0.75      ,  0.875     ],
                          [ 0.25      ,  0.        ,  0.375     ],
                          [ 0.25      ,  0.        ,  0.875     ],
                          [ 0.25      ,  0.25      ,  0.375     ],
                          [ 0.25      ,  0.25      ,  0.875     ],
                          [ 0.25      ,  0.5       ,  0.375     ],
                          [ 0.25      ,  0.5       ,  0.875     ],
                          [ 0.25      ,  0.75      ,  0.375     ],
                          [ 0.25      ,  0.75      ,  0.875     ],
                          [ 0.5       ,  0.        ,  0.375     ],
                          [ 0.5       ,  0.        ,  0.875     ],
                          [ 0.5       ,  0.25      ,  0.375     ],
                          [ 0.5       ,  0.25      ,  0.875     ],
                          [ 0.5       ,  0.5       ,  0.375     ],
                          [ 0.5       ,  0.5       ,  0.875     ],
                          [ 0.5       ,  0.75      ,  0.375     ],
                          [ 0.5       ,  0.75      ,  0.875     ],
                          [ 0.75      ,  0.        ,  0.375     ],
                          [ 0.75      ,  0.        ,  0.875     ],
                          [ 0.75      ,  0.25      ,  0.375     ],
                          [ 0.75      ,  0.25      ,  0.875     ],
                          [ 0.75      ,  0.5       ,  0.375     ],
                          [ 0.75      ,  0.5       ,  0.875     ],
                          [ 0.75      ,  0.75      ,  0.375     ],
                          [ 0.75      ,  0.75      ,  0.875     ],
                          [ 0.08333333,  0.16666667,  0.125     ],
                          [ 0.08333333,  0.16666667,  0.625     ],
                          [ 0.08333333,  0.41666667,  0.125     ],
                          [ 0.08333333,  0.41666667,  0.625     ],
                          [ 0.08333333,  0.66666667,  0.125     ],
                          [ 0.08333333,  0.66666667,  0.625     ],
                          [ 0.08333333,  0.91666667,  0.125     ],
                          [ 0.08333333,  0.91666667,  0.625     ],
                          [ 0.33333333,  0.16666667,  0.125     ],
                          [ 0.33333333,  0.16666667,  0.625     ],
                          [ 0.33333333,  0.41666667,  0.125     ],
                          [ 0.33333333,  0.41666667,  0.625     ],
                          [ 0.33333333,  0.66666667,  0.125     ],
                          [ 0.33333333,  0.66666667,  0.625     ],
                          [ 0.33333333,  0.91666667,  0.125     ],
                          [ 0.33333333,  0.91666667,  0.625     ],
                          [ 0.58333333,  0.16666667,  0.125     ],
                          [ 0.58333333,  0.16666667,  0.625     ],
                          [ 0.58333333,  0.41666667,  0.125     ],
                          [ 0.58333333,  0.41666667,  0.625     ],
                          [ 0.58333333,  0.66666667,  0.125     ],
                          [ 0.58333333,  0.66666667,  0.625     ],
                          [ 0.58333333,  0.91666667,  0.125     ],
                          [ 0.58333333,  0.91666667,  0.625     ],
                          [ 0.83333333,  0.16666667,  0.125     ],
                          [ 0.83333333,  0.16666667,  0.625     ],
                          [ 0.83333333,  0.41666667,  0.125     ],
                          [ 0.83333333,  0.41666667,  0.625     ],
                          [ 0.83333333,  0.66666667,  0.125     ],
                          [ 0.83333333,  0.66666667,  0.625     ],
                          [ 0.83333333,  0.91666667,  0.125     ],
                          [ 0.83333333,  0.91666667,  0.625     ],
                          [ 0.16666667,  0.08333333,  0.375     ],
                          [ 0.16666667,  0.08333333,  0.875     ],
                          [ 0.16666667,  0.33333333,  0.375     ],
                          [ 0.16666667,  0.33333333,  0.875     ],
                          [ 0.16666667,  0.58333333,  0.375     ],
                          [ 0.16666667,  0.58333333,  0.875     ],
                          [ 0.16666667,  0.83333333,  0.375     ],
                          [ 0.16666667,  0.83333333,  0.875     ],
                          [ 0.41666667,  0.08333333,  0.375     ],
                          [ 0.41666667,  0.08333333,  0.875     ],
                          [ 0.41666667,  0.33333333,  0.375     ],
                          [ 0.41666667,  0.33333333,  0.875     ],
                          [ 0.41666667,  0.58333333,  0.375     ],
                          [ 0.41666667,  0.58333333,  0.875     ],
                          [ 0.41666667,  0.83333333,  0.375     ],
                          [ 0.41666667,  0.83333333,  0.875     ],
                          [ 0.66666667,  0.08333333,  0.375     ],
                          [ 0.66666667,  0.08333333,  0.875     ],
                          [ 0.66666667,  0.33333333,  0.375     ],
                          [ 0.66666667,  0.33333333,  0.875     ],
                          [ 0.66666667,  0.58333333,  0.375     ],
                          [ 0.66666667,  0.58333333,  0.875     ],
                          [ 0.66666667,  0.83333333,  0.375     ],
                          [ 0.66666667,  0.83333333,  0.875     ],
                          [ 0.91666667,  0.08333333,  0.375     ],
                          [ 0.91666667,  0.08333333,  0.875     ],
                          [ 0.91666667,  0.33333333,  0.375     ],
                          [ 0.91666667,  0.33333333,  0.875     ],
                          [ 0.91666667,  0.58333333,  0.375     ],
                          [ 0.91666667,  0.58333333,  0.875     ],
                          [ 0.91666667,  0.83333333,  0.375     ],
                          [ 0.91666667,  0.83333333,  0.875     ],
                          [ 0.43147197,  0.46855409,  0.4987717 ]]

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

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

# Initialize a new potential set.
potentialSet = TremoloXPotentialSet(name = 'Tersoff_LJ_C_Li')

# Add particles for C and Li
potentialSet.addParticleType(ParticleType(symbol='C',
                                          mass=12.0107*atomic_mass_unit,
                                          charge=None,
                                          # Sigma and epsilon from Ref [1].
                                          sigma=3.3611*Ang,
                                          epsilon=0.004207*eV))
potentialSet.addParticleType(ParticleType(symbol='Li',
                                          mass=6.941*atomic_mass_unit,
                                          charge=None,
                                          # Sigma and epsilon from Ref [1].
                                          sigma=0.826*Ang,
                                          epsilon=0.271115*eV))

# Add the Tersoff potential: No modifications here.
potentialSet.addPotential(TersoffSingleTypePotential(particleType = 'C',
                                                     A = 2019.8449*eV,
                                                     B = 175.426651*eV,
                                                     R = 1.85*Angstrom,
                                                     S = 2.15*Angstrom,
                                                     l = 4.18426232*1/Ang,
                                                     mu = 1.93090093*1/Ang,
                                                     alpha = 0.0*1/Ang,
                                                     beta = 1.0,
                                                     omega = 0.11233,
                                                     chi = 1.0,
                                                     chiR = 1.0,
                                                     m = 0,
                                                     n = 1.0,
                                                     c = 181.91,
                                                     d = 6.28433,
                                                     h = -0.5556))

# Add the Lennard-Jones potential between C-Li and Li-Li. No LJ between C-C !
# The epsilon and sigma parameters have been defined with the particle types.
# Contributions between different atoms are employed via the
# Lorenz-Berthelot mixing rules.
potentialSet.addPotential(LennardJonesPotential(particleType1='C',
                                                particleType2='Li',
                                                r_cut=9.0*Ang))
potentialSet.addPotential(LennardJonesPotential(particleType1='Li',
                                                particleType2='Li',
                                                r_cut=9.0*Ang))

# Create a new calculator based on this potential set
calculator = TremoloXCalculator(parameters=potentialSet)

bulk_configuration.setCalculator(calculator)
bulk_configuration.update()

# -------------------------------------------------------------
# Molecular Dynamics
# -------------------------------------------------------------

initial_velocity = MaxwellBoltzmannDistribution(
    temperature=300.0*Kelvin,
    remove_center_of_mass_momentum=True
)

method = Langevin(
    time_step=1*femtoSecond,
    reservoir_temperature=300*Kelvin,
    friction=0.01*femtoSecond**-1,
    initial_velocity=initial_velocity
)

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=[],
    trajectory_filename='C_li_traj.nc',
    steps=5000,
    log_interval=100,
    method=method
)

bulk_configuration = md_trajectory.lastImage()
