from ase.calculators.cp2k import CP2K
from ase.optimize import BFGS
from ase.io import read, write
from ase.constraints import FixAtoms

system = 1 # 1: Ag2Au2 loaded on Au
           # 2: Pt8O14 loaded on TiO2

if system == 1:

    xyz_file = 'Ag2Au2.xyz'
    sys_a = 11.539982668964457
    sys_b = 9.993918150555366
    sys_c = 20.0
    n_atom_fix = 4

elif system == 2:

    xyz_file = 'Pt8O14.xyz'
    sys_a = 17.741
    sys_b = 19.6334
    sys_c = 30.8281
    n_atom_fix = 22

else:

    print('You select a wrong system number, exit')
    exit()

atoms = read(xyz_file)
atoms.set_cell([sys_a, sys_b, sys_c])
atoms.set_pbc((True, True, True))
atoms.set_constraint(FixAtoms(mask=(len(atoms)-n_atom_fix) * [True] + n_atom_fix * [False]))

cutoff_Ry = 400
atoms.calc = CP2K(
        basis_set = 'DZVP-MOLOPT-SR-GTH',
        charge = 0.,
        cutoff = 13.605662285137*cutoff_Ry,
        force_eval_method = 'quickstep',
        pseudo_potential = 'GTH-PBE',
        stress_tensor = True,
        uks = False,
        xc = 'PBE',
        max_scf = 50,
        print_level = 'MEDIUM',
        inp = '''
&FORCE_EVAL
    &DFT
	&QS
            METHOD xTB
            &xTB
                DO_EWALD .TRUE.
                CHECK_ATOMIC_CHARGES .FALSE.
                &PARAMETER
                    DISPERSION_PARAMETER_FILE dftd3.dat
                &END PARAMETER
            &END xTB
	&END QS
	&POISSON
            &EWALD
                EWALD_TYPE SPME
            &END EWALD
	&END POISSON
        &MGRID
            REL_CUTOFF 30
            NGRIDS 5
        &END MGRID
        &SCF
            SCF_GUESS ATOMIC
            EPS_SCF 1.0E-5
            &OUTER_SCF
                EPS_SCF 1.0E-5
                MAX_SCF 10
            &END OUTER_SCF
            &OT
                ALGORITHM IRAC
                ENERGY_GAP 0.5
                PRECONDITIONER FULL_SINGLE_INVERSE
                MINIMIZER DIIS
            &END OT
        &END SCF
    &END DFT
&END FORCE_EVAL
        '''
        )

dyn = BFGS(atoms)
dyn.run(fmax = 0.05)
print(atoms.get_total_energy())
write('relaxed.xyz', atoms)