class GammaResults(Results):

    @staticmethod
    def get_difference(job, jobplus):
        """Calculate the difference between HOMO and IP.
        *jobplus* should be the counterpart of *job* with one less electron."""
        homo = job.results.readrkf('Properties','HOMO', file='engine')
        IP = jobplus.results.get_energy() - job.results.get_energy()
        return IP + homo

    def get_J(self):
        N = GammaResults.get_difference(self.job.children[1], self.job.children[2])
        A = GammaResults.get_difference(self.job.children[0], self.job.children[1])
        return (N**2 + A**2)**0.5


class GammaJob(MultiJob):
    _result_type = GammaResults

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

    def prerun(self):
        charges = [self.charge-1, self.charge, self.charge+1]
        for charge, spin in zip(charges, self.spins):
            name = '{}_charge_{}'.format(self.name, charge)
            newjob = AMSJob(name=name, molecule=self.molecule, settings=self.settings)
            newjob.molecule.properties.charge = charge
            newjob.settings.input.adf.xc.rangesep = "gamma={:f}".format(self.gamma)
            if spin != 0:
                newjob.settings.input.adf.unrestricted = True
                newjob.settings.input.adf.SpinPolarization = spin

            self.children.append(newjob)


def gamma_scan(gammas, settings, molecule, name='scan', 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 for an ADF calculation
    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 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 for charge and spin:

    Charge=X-1  SpinPolarization=a
    Charge=X    SpinPolarization=b
    Charge=X+1  SpinPolarization=c

    Returns a list of pairs (gamma, J) of the same length as the parameter *gammas*
    """
    jobs = [GammaJob(molecule=molecule, settings=settings, gamma=g,
            charge=charge, spins=spins, name=name+'_gamma_'+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))


# =============================================================
# Now we simply use the gamma_scan function to find the optimal 
# gamma value for a toy system (H2)
# =============================================================


import numpy as np
import multiprocessing

# Run as many jobs in parallel as there are cores:
config.default_jobrunner = JobRunner(parallel=True, maxjobs=multiprocessing.cpu_count())

# Settings of the ADF calculations
# ================================

s = Settings()
s.input.ams.task = 'SinglePoint'
s.input.adf.basis.type = 'DZP'
s.input.adf.basis.core = 'None'
s.input.adf.xc.gga = 'PBE'
s.input.adf.xc.xcfun = True
s.runscript.nproc = 1

# The molecule (here we just use H2)
# ==================================

mol = Molecule()
mol.add_atom(Atom(symbol='H', coords=(0,0,-0.3540)))
mol.add_atom(Atom(symbol='H', coords=(0,0, 0.3540)))

# The list of gamma values
# ========================

# Here we scan just a few values for gamma. 
# In practice, you want to scan a wider range and smaller step.
gammas = np.around(np.arange(1.2, 1.9, 0.2), decimals=3)

results = gamma_scan(gammas, s, mol)

print('== Results ==')
print('gamma \t J')
for g,j in results:
    print('{:.4f} \t {:.8f}'.format(g,j))
print('Optimal gamma value: {:.4f}'.format(min(results,key=lambda x:x[1])[0]))