import math
import numpy as np
import os
if 'ADFHOME' in os.environ:
    adfhome = os.environ['ADFHOME']
else:
    print('error! ADFHOME NOT FOUND!')


class GammaResults(Results):

    def get_difference(self, job, jobplus):
        """Calculate the difference between HOMO energy and Ionization Potential. *jobplus* should be the counterpart of *job* with one less electron."""
        prop = job.results.get_properties()
        if 'HOMO' in prop:
            homo = prop['HOMO']
        else:
            homos = job.results.awk_output('/^ HOMO :   /{print $5}')
            homos = list(map(float, homos))
            restricted = (job.results.readkf('General', 'nspin') == 1)
            homo = homos[-1] if restricted else max(homos[-2:])
        IP = jobplus.results.get_energy() - job.results.get_energy()
        return IP + homo

    def get_j(self):
        """Calculate the value of the J function (J = sqrt(j0^2 + j1^2)) where j0 and j1 are differences between HOMO and IP for, respectively, neutral species and anion.
        """
        j0 = self.get_difference(self.job.children[1], self.job.children[2])
        j1 = self.get_difference(self.job.children[0], self.job.children[1])
        return math.sqrt((j0*j0)+(j1*j1))


class GammaJob(MultiJob):
    _result_type = GammaResults

    def __init__(self, molecule, gamma, charge=0, spins=(1,0,1), **kwargs):
        MultiJob.__init__(self, **kwargs)
        self.molecule = molecule
        self.charge = charge
        self.spins = spins
        self.gamma = gamma


    def prerun(self):
        for ch,sp in zip(range(self.charge-1, self.charge+2), self.spins):
            newjob = ADFJob(name=self.name+'_'+str(ch), molecule=self.molecule, settings=self.settings)
            newjob.settings.input.charge = '{} {}'.format(ch,sp)
            if sp != 0:
                newjob.settings.input.unrestricted = True
            newjob.settings.input.xc.rangesep = "gamma={:f}".format(self.gamma)
            self.children.append(newjob)

#===========================================================================

def gamma_scan(gammas, settings, molecule, name='gammascan', charge=0, spins=(1,0,1)):
    """Calculate values of J function for given range of gammas.

    Arguments:
    gammas   - list of gamma values to calculate the J function for
    settings - Settings object compatible with ADFJob
    molecule - Molecule object with the system of interest
    name     - base name of all the jobs
    charge   - base charge of the system of interest. The J function is going to be calculated based on two systems: with charge, and charge-1
    spins    - values of spin polarization (see keyword CHARGE of ADF) for jobs with, respectively, charge-1, charge and charge +1

    In other words, if charge=X and spins=(a,b,c) the three resulting jobs are going to have the following values of CHARGE keyword:

    CHARGE X-1 a
    CHARGE  X  b
    CHARGE X+1 c

    Returns: a list of pairs (gamma, J) of the same lenght as the parameter *gammas*
    """
    

    jobs = [GammaJob(molecule=molecule, settings=settings, gamma=g, charge=charge, spins=spins, name=name+str(g)) for g in gammas]
    results = [j.run() for j in jobs]
    js = [r.get_j() for r in results]
    return list(zip(gammas, js))


#===========================================================================

# first do a coarse scan

s = Settings()
s.input.basis.type = 'TZP'
s.input.basis.core = 'None'
s.input.numericalquality = 'Good'
s.input.Relativistic = 'Scalar ZORA'
s.input.xc.gga = 'pbe'
s.input.xc.xcfun = True

molfile = os.path.join(adfhome, 'atomicdata', 'Molecules', 'TestMols', 'p-Nitroaniline.xyz')
mol = Molecule(molfile)
gammas = np.arange(0.2, 0.51, 0.1)

#===========================================================================

results = gamma_scan(gammas, s, mol, name='test')
log('gamma \t J')
for g,j in results:
    log('{:.4f} \t {:.8f}'.format(g,j))
optimal_gamma = min(results,key=lambda x:x[1])[0]       
log('Optimal gamma value: {:.4f}'.format(optimal_gamma))

# then a finer grid

gammas = np.arange(optimal_gamma - 0.06, optimal_gamma + 0.07 , 0.03)

results = gamma_scan(gammas, s, mol, name='test')
log('gamma \t J')
for g,j in results:
    log('{:.4f} \t {:.8f}'.format(g,j))
optimal_gamma = min(results,key=lambda x:x[1])[0]       
log('Refined optimal gamma value: {:.4f}'.format(optimal_gamma))

# then an even finer grid

gammas = np.arange(optimal_gamma - 0.02, optimal_gamma + 0.03 , 0.01)

results = gamma_scan(gammas, s, mol, name='test')
log('gamma \t J')
for g,j in results:
    log('{:.4f} \t {:.8f}'.format(g,j))
optimal_gamma = min(results,key=lambda x:x[1])[0]       
log('Refined optimal gamma value: {:.4f}'.format(optimal_gamma))