from openff.toolkit import Molecule
from openff.recharge.esp import ESPSettings
from openff.recharge.grids import LatticeGridSettings
from openff.recharge.grids import GridSettingsType, GridGenerator
from qcelemental.models.procedures import OptimizationInput, QCInputSpecification
from source.conformers.conformer_gen import Conformers
from openff.units import unit
from qcelemental.models.common_models import Model
from openff.recharge.utilities.molecule import smiles_to_molecule
import qcengine
import os
import numpy as np


CWD = os.getcwd()

#generate water test molecule as openff.toolkit.Molecule
test_mol =  smiles_to_molecule('[H]O[H]')
conformer_list = Conformers.generate(test_mol, generation_type='rdkit')
conformer_list[0]
qc_mol =  test_mol.to_qcschema(conformer=0)

#Setup geometry optimisation
hf_model = Model(method="hf", basis="6-31G*")
spec = QCInputSpecification(model=hf_model, keywords={}, driver="gradient")
opt_spec = OptimizationInput(
            initial_molecule=qc_mol,
            input_specification=spec,
            keywords={"coordsys": "dlc", 
                      "program": "psi4"}                                        
        )

opt = qcengine.compute_procedure(opt_spec, "geometric")
print(f'geometry optimiztion was {opt.error}')
#return optimized molecule
optmized_mol = opt.final_molecule

#Generate grid.dat file for grid_esp and grid_field
grid_settings = LatticeGridSettings(
        type="fcc", spacing=0.5, inner_vdw_scale=1.4, outer_vdw_scale=2.0
    )

grid = GridGenerator.generate(test_mol, optmized_mol.geometry*unit.angstrom, grid_settings)

grid = grid.to(unit.angstrom).m
np.savetxt("grid.dat", grid, delimiter=" ", fmt="%16.10f")
#compute one-electron properties.
opt_input_2 = { "molecule" : optmized_mol,
              "driver" : "energy",
              "model" : {"method":"scf","basis":"6-31G*"},
              "protocols":{"wavefunction":"all","stdout":True,"native_files":"all"},
              "keywords":{"scf_properties":["GRID_ESP", "GRID_FIELD","MULLIKEN_CHARGES", "LOWDIN_CHARGES", "DIPOLE", "QUADRUPOLE", "MBIS_CHARGES"]}                               
              }


opt_2 = qcengine.compute(opt_input_2, "psi4", task_config={"scratch_directory":CWD,"scratch_messy":True})

print(opt_2.dict())

esp = (np.loadtxt("grid_esp.dat").reshape(-1, 1) * unit.hartree / unit.e)

electric_field = (np.loadtxt("grid_field.dat")* unit.hartree/ (unit.e * unit.bohr)))