# -*- coding: utf-8 -*-
setVerbosity(MinimalLog)

# -------------------------------------------------------------
# Bulk Configuration
# -------------------------------------------------------------

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

# Define elements
elements = [Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium,
            Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium,
            Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium]

# Define coordinates
fractional_coordinates = [[ 0.046849303132,  0.578048584239,  0.330834628862],
                          [ 0.747486215681,  0.435180554111,  0.527295513093],
                          [ 0.728887761335,  0.886922306732,  0.768860338684],
                          [ 0.414425990171,  0.535354936092,  0.752008448713],
                          [ 0.255345178522,  0.753701553286,  0.282131565008],
                          [ 0.90036373387 ,  0.145453898294,  0.91761587373 ],
                          [ 0.374066511603,  0.307760711155,  0.497068077455],
                          [ 1.001750691515,  0.826775399807,  0.767971573708],
                          [ 0.499595292484,  0.810872894696,  0.200633139672],
                          [ 0.73569352768 ,  0.733148719396,  0.412714298573],
                          [-0.040562971588,  0.573188171614,  0.660129132106],
                          [ 0.520051379459,  0.103549730374,  0.306573957993],
                          [ 0.754832828099,  0.760331884743,  0.032073457935],
                          [ 0.226495603727,  1.066519311917,  0.230294586588],
                          [ 0.073666311249,  0.222949067025,  0.422025179187],
                          [ 0.229250922081,  1.016061409098,  0.510825242303],
                          [ 0.351070436645,  0.22556170251 ,  0.045638330736],
                          [ 0.529866779197,  0.598981184289,  0.029685084869],
                          [ 0.218593337763,  0.390693512083,  0.279656897794],
                          [ 0.700494745201,  0.610231274156,  0.713194313872],
                          [ 0.032452381159,  0.257328382205,  0.138954792242],
                          [ 1.090090362588,  0.939570412957,  0.039249823067],
                          [ 0.493734371776,  0.918447728982,  0.534857432958],
                          [ 0.625707030726,  0.193798424052,  0.589531032883],
                          [ 0.891181246874,  0.95284347295 ,  0.46802753193 ],
                          [ 0.589992811466,  0.328089179471,  0.108967920543],
                          [ 0.9644318872  ,  0.280931870331,  0.652147203356],
                          [ 0.905813913574,  0.428478389797,  0.964274043418],
                          [ 0.400084250597,  0.821303330657,  0.874952965449],
                          [ 0.834679015026,  0.029420113918,  0.151709005136],
                          [ 0.832984558616,  0.472748135559,  0.223176653472],
                          [ 0.154519315133,  0.448726019407,  0.560259895709],
                          [ 0.255496737936,  0.708732865978,  0.597386689053],
                          [ 0.61936512292 ,  0.353847688099,  0.809562728881],
                          [ 0.821478501522,  0.211801165209,  0.335990745615],
                          [ 0.219490629981,  0.353621117151,  0.826266199085],
                          [ 0.488730235527,  0.512568404251,  0.350920527406],
                          [ 0.301190752318,  0.502371188879,  0.081521311454],
                          [ 0.233348375244,  0.080840790956,  0.782272506811],
                          [ 0.129249284668,  0.642179113642,  0.966219325888],
                          [ 0.262552811495,  0.566017036421,  0.407245798597],
                          [ 0.414951250419,  0.295557205543,  0.255790072536],
                          [ 0.330528046577,  0.720410601205,  0.070959634619],
                          [ 0.156232927199,  0.14176776161 ,  0.987993911627],
                          [ 0.6885651968  ,  0.056969405266,  0.446353090389],
                          [ 0.830958665112,  0.410041054826,  0.761145385667],
                          [ 0.321292346767,  0.904194231944,  0.676264889112],
                          [ 0.848524561724,  0.767786200934,  0.615072765218],
                          [ 0.063642600349,  0.485719269682,  0.125091384053],
                          [ 0.170185141422,  0.570617214913,  0.754224329446],
                          [ 0.803233430653,  0.256977991098,  0.104950251227],
                          [ 0.546998081622,  0.402095501559,  0.59679733679 ],
                          [ 0.369865336928,  0.942237013356,  0.341297492104],
                          [ 0.942213547077,  0.403206696145,  0.439938812898],
                          [ 0.664057568462,  0.647503486697,  0.229862535733],
                          [ 0.005501956376,  0.059369530507,  0.29825977379 ],
                          [ 0.183507487754,  0.227621365202,  0.613164544886],
                          [ 0.44074326472 ,  0.399616721589,  0.937065786032],
                          [ 0.889764293264,  0.923155092107,  0.940716933301],
                          [ 0.607199787647,  0.728431830883,  0.865394444287]]

# Define velocities
velocities = [[-1.696580250855e-02,  1.457001145269e-02, -7.184740254696e-03],
              [-3.131746995789e-03,  2.534983832653e-03, -9.882157986733e-03],
              [-5.769552749414e-03,  3.218205180341e-04, -2.219335824301e-02],
              [ 1.772054196800e-02, -1.383634874481e-02,  1.492380996053e-02],
              [-1.430712860988e-02,  7.077377533187e-03, -1.062838510734e-03],
              [ 1.683880846738e-02,  1.746790228741e-03, -8.635135529429e-04],
              [ 3.959646995992e-03, -6.706804791711e-04, -5.407956995707e-03],
              [-2.876649389428e-03,  1.075286773103e-03,  2.322664933880e-03],
              [-2.309507994169e-02,  8.411263540726e-03, -1.306859044801e-03],
              [ 6.630135840514e-03,  2.668440553874e-03,  1.862788244342e-02],
              [ 1.801351557124e-02, -7.515552779359e-03,  1.321834467776e-02],
              [ 2.571510517324e-02,  1.480949057662e-02,  3.886138043709e-03],
              [-1.451403976765e-02,  7.239859289853e-03,  5.101782050010e-05],
              [-1.816164286881e-02,  1.054766587589e-02, -1.368827808805e-02],
              [ 9.005527153282e-03, -1.212415532897e-02, -3.856324775962e-03],
              [-4.583058774760e-03, -7.886765988656e-03,  1.093365961048e-02],
              [-5.067098610516e-03, -7.962934985918e-03, -8.228774755095e-03],
              [-8.242725697357e-04,  2.068752643069e-03,  8.625565416900e-03],
              [ 1.374460986230e-02, -3.338244770370e-03, -1.413370987872e-02],
              [-6.883156257166e-03, -1.150058508020e-03,  5.485092723725e-03],
              [-6.487397786190e-03, -4.992052112256e-04, -2.559824514658e-04],
              [-3.267323757774e-02, -1.210340638781e-02, -1.740204202893e-02],
              [-2.301048959165e-02, -1.045577924621e-02, -1.102679467791e-02],
              [-2.460939149525e-02,  1.949342881719e-02,  1.819374410425e-03],
              [-8.252971336415e-03,  2.070679718519e-02,  9.518449565472e-03],
              [-1.552752839880e-02, -5.381355595456e-03, -1.284932778699e-02],
              [-2.454903512569e-02,  4.077285339012e-03, -8.961924431063e-03],
              [-6.719967593734e-03,  8.086739311696e-03,  1.704612879614e-02],
              [-1.257561716793e-02, -9.229093553670e-03, -1.120434238416e-04],
              [ 1.619720495135e-02, -3.277484356317e-02,  6.009512237017e-03],
              [ 1.645790232317e-03,  1.561816208047e-02, -1.268717629904e-02],
              [ 8.899058812628e-03, -1.223364599234e-02, -5.564443911452e-03],
              [-1.792658603537e-03,  5.420856989598e-03, -1.420688689700e-02],
              [ 1.544414164208e-02, -7.151286005892e-03,  3.998779171079e-03],
              [ 2.069452364627e-02,  3.272449032905e-03, -8.818055755476e-03],
              [ 2.556430963006e-03, -8.017668663144e-03,  1.225991855537e-02],
              [-1.895039250712e-02,  5.308140294287e-03, -8.160184370211e-03],
              [ 3.571475093748e-03,  1.409839644879e-02,  9.210704607059e-03],
              [ 1.884899793066e-02,  1.602263032151e-02, -2.382924948166e-03],
              [-3.571468829186e-03, -7.645291930146e-03,  4.511244990108e-03],
              [ 1.850769292865e-03, -6.019001358779e-03,  2.062524343668e-03],
              [ 2.509099644103e-03,  9.811918490498e-04, -4.329894855033e-03],
              [ 4.908760593883e-03,  2.873673465066e-03, -4.066930530408e-03],
              [-7.344679291964e-03, -8.490111461062e-03, -3.225747643865e-03],
              [-2.832678073286e-03,  2.298249110927e-04,  6.626791052081e-03],
              [ 1.394358142592e-03,  3.945568808851e-03,  1.971651011112e-03],
              [ 2.853681275203e-03, -8.748439946074e-04, -3.072415162012e-03],
              [ 3.187122769515e-03, -5.495272467071e-03,  4.056113041792e-03],
              [-1.529700971427e-03,  8.732656042273e-04, -2.496934183424e-03],
              [-1.146882669149e-04,  3.171546653427e-03, -1.850680819685e-03],
              [-3.632023976919e-03,  3.622308816317e-03, -4.251602415766e-03],
              [ 4.212637336031e-04,  4.740006218425e-03, -5.422969928190e-04],
              [-6.752743502465e-04, -7.535127585223e-04, -4.108305055283e-03],
              [-1.217419664633e-03,  9.842445642216e-03,  6.276308490177e-03],
              [-1.232235522555e-03,  5.827502916814e-03,  1.828381158681e-03],
              [-3.440044988692e-04, -5.197429198377e-03,  3.549952255471e-03],
              [ 6.099225631952e-03, -6.232499430705e-03, -1.265020037692e-04],
              [ 4.497075495982e-03,  3.485572654714e-03,  3.278846728088e-03],
              [-1.400031512446e-03, -1.818150678573e-03,  1.062487636660e-03],
              [ 1.154041792376e-03, -6.970979431843e-03,  1.641852492513e-03]]*Angstrom/fs

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

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

potentialSet = TremoloXPotentialSet(name = 'HfO2_MTP_10.2021')
potentialSet.addParticleType(ParticleType(
	symbol = 'O',
	mass = 15.9994*atomic_mass_unit,
	charge = None,
	sigma = 3.11814551349*Angstrom,
	sigma14 = 3.11814551349*Angstrom,
	epsilon = 0.06*kiloCaloriePerMol,
	epsilon14 = 0.06*kiloCaloriePerMol,
	atomicNumber = 8,
	tags = [],
))
potentialSet.addParticleType(ParticleType(
	symbol = 'Hf',
	mass = 178.49*atomic_mass_unit,
	charge = None,
	sigma = 2.79831287368*Angstrom,
	sigma14 = 2.79831287368*Angstrom,
	epsilon = 0.072*kiloCaloriePerMol,
	epsilon14 = 0.072*kiloCaloriePerMol,
	atomicNumber = 72,
	tags = [],
))


_potential = MTPPotential(
	file = './mtp_potential.mtp',
	suppress_intercept = False,
)
potentialSet.addPotential(_potential)
calculator = TremoloXCalculator(parameters=potentialSet)
calculator.setVerletListsDelta(1.0*Angstrom)

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('HfO2_MTP_3000K_md.hdf5', bulk_configuration)

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

initial_velocity = MaxwellBoltzmannDistribution(
    temperature=3000.0*Kelvin,
    remove_center_of_mass_momentum=True,
    random_seed=None,
    enforce_temperature=True,
)

method = Langevin(
    time_step=1*femtoSecond,
    reservoir_temperature=3000*Kelvin,
    friction=0.01*femtoSecond**-1,
    initial_velocity=initial_velocity,
    heating_rate=0*Kelvin/picoSecond,
)

constraints = [FixCenterOfMass()]

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=constraints,
    trajectory_filename='HfO2_MTP_3000K_md.hdf5',
    steps=100000,
    log_interval=500,
    method=method
)

bulk_configuration = md_trajectory.lastImage()
