#imports
import numpy as np
from ase.build import molecule
from ase.calculators.emt import EMT
from ase.optimize import QuasiNewton, LBFGS
from ase.vibrations import Vibrations
from ase.thermochemistry import HarmonicThermo
from ase.calculators.cp2k import CP2K
from ase.units import Rydberg
from ase.io import read
from ase import io

inp=  """
!  &EXT_RESTART
 !   RESTART_FILE_NAME gluc_C3N4_ads_21_1-1.restart
!  &END

&GLOBAL
 RUN_TYPE GEO_OPT
 EXTENDED_FFT_LENGTHS .TRUE.         !Get more libraries
 FFTW_PLAN_TYPE ESTIMATE             !Quick estimate, no runtime measurements.
 PREFERRED_FFT_LIBRARY FFTW3         !Preferred library
 &TIMINGS SILENT                     !Specify % of CPUTIME above which the contribution will be inserted in the final timing report
  THRESHOLD 1e-08
 &END TIMINGS
&END GLOBAL

&MOTION

 &GEO_OPT
  MAX_ITER 400
  OPTIMIZER LBFGS                       ! Better than LBFGS for small systems!!
  MAX_DR 5e-04                 !Convergence criterion for the maximum geometry change between the current and the last optimizer iteration.
 &END GEO_OPT

 &PRINT
  &TRAJECTORY
   FORMAT XYZ
  &END TRAJECTORY
 &END PRINT

 &CONSTRAINT
  &FIXED_ATOMS
   COMPONENTS_TO_FIX XYZ
   LIST 25 26 27..150
  &END FIXED_ATOMS
 &END CONSTRAINT

&END MOTION

&FORCE_EVAL
 &PRINT
  &STRESS_TENSOR
  &END STRESS_TENSOR
 &END PRINT
 &DFT
  UKS .TRUE.                      !Requests a spin-polarized calculation using alpha and beta orbitals, i.e. no spin restriction is applied
  CHARGE 0                   !Total charge
  MULTIPLICITY 1             !  Two times the total spin plus one. 2 is for a doublet for an odd number of electrons
  WFN_RESTART_FILE_NAME gluc_C3N4_ads_21_1-RESTART.kp
  &MGRID
   NGRIDS 5
   RELATIVE_CUTOFF 50
  &END MGRID
  &QS
   EPS_DEFAULT 1.0E-12
   METHOD GPW
   EXTRAPOLATION USE_GUESS
  &END QS
  &SCF
   EPS_SCF 1.0E-7
   SCF_GUESS RESTART            !Use the RESTART file as an initial guess (and ATOMIC if not present).
   ADDED_MOS  200
   CHOLESKY  OFF
   &OUTER_SCF  F
    EPS_SCF 1.0E-7
    MAX_SCF 50
   &END OUTER_SCF
   &SMEAR  T
    METHOD  FERMI_DIRAC
    ELECTRONIC_TEMPERATURE     3.0000000000000000E+02
   &END SMEAR
   &MIXING  T
    METHOD  BROYDEN_MIXING
    ALPHA     4.0000000000000002E-01
    BETA     1.5000000000000000E+00
    NMIXING  5
    NBUFFER  8
   &END MIXING
  &END SCF
  &XC
   FUNCTIONAL_ROUTINE NEW
   DENSITY_CUTOFF 1.0e-12
   GRADIENT_CUTOFF 1.0e-12
   TAU_CUTOFF 1.0e-12
   &XC_FUNCTIONAL
    &PBE
     PARAMETRIZATION ORIG
    &END PBE
   &END XC_FUNCTIONAL
   &VDW_POTENTIAL                            ! ... dispersion interactions
         POTENTIAL_TYPE PAIR_POTENTIAL           
         &PAIR_POTENTIAL
            TYPE DFTD3                          ! computed with the DFTD3 method 
            REFERENCE_FUNCTIONAL PBE
            PARAMETER_FILE_NAME dftd3.dat     
            R_CUTOFF 8                         ! cutoff raddius for the dispersion interactions
         &END PAIR_POTENTIAL
   &END VDW_POTENTIAL
   &XC_GRID
    USE_FINER_GRID  .TRUE.
   &END XC_GRID
  &END XC
  &POISSON
   POISSON_SOLVER  PERIODIC
   PERIODIC  XYZ
  &END POISSON
  &KPOINTS
   SCHEME  MONKHORST-PACK  1 1 1
   FULL_GRID  .TRUE.
  &END KPOINTS
 &END DFT
 &SUBSYS
 &END SUBSYS
&END FORCE_EVAL

 """


atoms=io.read('input.xyz')
atoms.set_cell([21.31370,21.31370,30.00000,90,90,120])
atoms.set_pbc(True)
atoms.set_calculator(CP2K(print_level='MEDIUM',cutoff=450 * Rydberg,max_scf=100,poisson_solver='None',potential_file='GTH_POTENTIALS',inp=inp))
potentialenergy = atoms.get_potential_energy()
print('energy is', potentialenergy)

indices = []    #specify which are adsorbate atoms by index number
for i in range(0,24):
    indices.append(i)
vib = Vibrations(atoms,indices=indices,name='vib')
vib.run()
vib.summary(log='vib.txt')
for mode in range(len(indices)*3):
                vib.write_mode(mode)
vib_energies = vib.get_energies()
#this section inputs the necessary input to calculate thermodynamic properties

thermo = HarmonicThermo(vib_energies=vib_energies[1:],  potentialenergy=potentialenergy)
F = thermo.get_helmholtz_energy(temperature=298.15)