#!/usr/bin/env python
#JSON {"lot": "UKS/6-31G(d)",
#JSON  "scf": "CDIISSCFSolver",
#JSON  "er": "cholesky",
#JSON  "difficulty": 7,
#JSON  "description": "Basic UKS DFT example with hybrid MGGA exhange-correlation functional (M05)"}

import numpy as np
from horton import *  # pylint: disable=wildcard-import,unused-wildcard-import


# Load the coordinates from file.
# Use the XYZ file from HORTON's test data directory.
fn_xyz = context.get_fn('test/methyl.xyz')
mol = IOData.from_file(fn_xyz)

# Create a Gaussian basis set
obasis = get_gobasis(mol.coordinates, mol.numbers, '6-31g(d)')

# Compute Gaussian integrals
olp = obasis.compute_overlap()
kin = obasis.compute_kinetic()
na = obasis.compute_nuclear_attraction(mol.coordinates, mol.pseudo_numbers)
er_vecs = obasis.compute_electron_repulsion_cholesky()

# Define a numerical integration grid needed the XC functionals
grid = BeckeMolGrid(mol.coordinates, mol.numbers, mol.pseudo_numbers)

# Create alpha orbitals
orb_alpha = Orbitals(obasis.nbasis)
orb_beta = Orbitals(obasis.nbasis)

# Initial guess
guess_core_hamiltonian(olp, kin + na, orb_alpha, orb_beta)

# Construct the restricted HF effective Hamiltonian
external = {'nn': compute_nucnuc(mol.coordinates, mol.pseudo_numbers)}
libxc_term_x = ULibXCHybridMGGA('x_m05')
libxc_term_c = ULibXCMGGA('c_m05')
terms = [
    UTwoIndexTerm(kin, 'kin'),
    UDirectTerm(er_vecs, 'hartree'),
    UGridGroup(obasis, grid, [libxc_term_x, libxc_term_c]),
    UExchangeTerm(er_vecs, 'x_hf', libxc_term_x.get_exx_fraction()),
    UTwoIndexTerm(na, 'ne'),
]
ham = UEffHam(terms, external)

# Decide how to occupy the orbitals (5 alpha electrons, 4 beta electrons)
occ_model = AufbauOccModel(5, 4)

# Converge WFN with CDIIS SCF
# - Construct the initial density matrix (needed for CDIIS).
occ_model.assign(orb_alpha, orb_beta)
dm_alpha = orb_alpha.to_dm()
dm_beta = orb_beta.to_dm()
# - SCF solver
scf_solver = CDIISSCFSolver(1e-6)
scf_solver(ham, olp, occ_model, dm_alpha, dm_beta)

# Derive orbitals (coeffs, energies and occupations) from the Fock and density
# matrices. The energy is also computed to store it in the output file below.
fock_alpha = np.zeros(olp.shape)
fock_beta = np.zeros(olp.shape)
ham.reset(dm_alpha, dm_beta)
ham.compute_energy()
ham.compute_fock(fock_alpha, fock_beta)
orb_alpha.from_fock_and_dm(fock_alpha, dm_alpha, olp)
orb_beta.from_fock_and_dm(fock_beta, dm_beta, olp)

# Assign results to the molecule object and write it to a file, e.g. for
# later analysis. Note that the CDIIS algorithm can only really construct an
# optimized density matrix and no orbitals.
mol.title = 'UKS computation on methyl'
mol.energy = ham.cache['energy']
mol.obasis = obasis
mol.orb_alpha = orb_alpha
mol.orb_beta = orb_beta
mol.dm_alpha = dm_alpha
mol.dm_beta = dm_beta

# useful for post-processing (results stored in double precision):
mol.to_file('methyl.h5')

# CODE BELOW IS FOR horton-regression-test.py ONLY. IT IS NOT PART OF THE EXAMPLE.
rt_results = {
    'energy': ham.cache['energy'],
    'orb_alpha': orb_alpha.energies,
    'orb_beta': orb_beta.energies,
    'nn': ham.cache["energy_nn"],
    'kin': ham.cache["energy_kin"],
    'ne': ham.cache["energy_ne"],
    'grid': ham.cache["energy_grid_group"],
    'ex': ham.cache["energy_x_hf"],
    'hartree': ham.cache["energy_hartree"],
}
# BEGIN AUTOGENERATED CODE. DO NOT CHANGE MANUALLY.
rt_previous = {
    'energy': -39.785947853471306,
    'ex': -1.7225865319035798,
    'orb_alpha': np.array([
        -10.23828545833133, -0.67858423851559102, -0.41160132632670871,
        -0.41021975233062052, -0.2304072733693352, 0.10564107678753258,
        0.17361036059448601, 0.17475384959943174, 0.5703614324584384, 0.60006398922443194,
        0.6019504171018244, 0.70677515185472173, 0.86421064444925833, 0.87393058602376417,
        0.94089134365277749, 1.622965745696175, 1.6314015490418585, 1.9621478712905653,
        2.0859187871191329, 2.1014348229240505
    ]),
    'orb_beta': np.array([
        -10.229951836398778, -0.63473610512484957, -0.40501128236445993,
        -0.39623017626304829, -0.034403316878747814, 0.12790741841493219,
        0.18852605259256111, 0.19455240647895397, 0.59406665317865992,
        0.60159605396061111, 0.61472346327942695, 0.71136977385060918,
        0.85570424077503049, 0.92731867786853817, 0.99074789531487739, 1.7283129823914731,
        1.7804498840106648, 2.030083671908538, 2.126106888785416, 2.2130467516595327
    ]),
    'grid': -4.671970776001134,
    'hartree': 28.105878950392608,
    'kin': 39.35093706265106,
    'ne': -109.9279915012739,
    'nn': 9.0797849426636361,
}
