from noodles import (gather, schedule)
from qmworks import (adf, dftb, molkit, run, templates, Settings)

import numpy as np


def main():
    backbone = '{}C1=CC=C(C=C1)C#CC2=CC=C({})C=C2'

    # smiles of the electronDonors
    # electronDonors = ['N', 'P', 'S', 'CO']
    electronDonors = ['N']
    # smiles of the electronAcceptors
    # electronAcceptors = ['B', 'C(O)=O', 'N(=O)=O', 'C(F)(F)(F)']
    electronAcceptors = ['O']

    # Genereate all the molecules
    molecules, names = create_molecules(backbone, electronDonors, electronAcceptors)

    # Optimize with DFTB
    dftb_jobs = [dftb(templates.geometry, m, job_name='dftb_{}'.format(n)) for m, n in zip(molecules, names)]

    # Schedule functions
    schedule_if_succesful = schedule(run_if_succesful)
    schedule_if_failure = schedule(run_if_failure)
    schedule_choose = schedule(choose_job)

    # Optimize with DFT in gas phase
    s = create_settings_dft_gas()
    gas_jobs = [schedule_if_failure(adf, s, job, mol, name, prefix='adf_')
                for job, mol, name in zip(dftb_jobs, molecules, names)]

    preoptimized_jobs = [schedule_choose(j1, j2) for j1, j2 in zip(dftb_jobs, gas_jobs)]

    #Optimize with COSMO and compute frequency
    s = create_settings_dft_cosmo_freq()
    adf_cosmo_freq_jobs = [schedule_if_succesful(adf, s, job, name, prefix='solvation')
                           for job, name in zip(preoptimized_jobs, names)]

    # Filter jobs which frequencies are greater than zero
    schedule_check_freq = schedule(check_frequencies)
    candidates = [schedule_check_freq(j) for j in adf_cosmo_freq_jobs]

    # Schedule a TD-DFT job with non-negative frequencies
    s = create_settings_TD_DFT()
    td_dft_jobs = [adf(s, job, job_name='td_dft_'.format(job.job_name)) for job in candidates]

    schedule_getattr = schedule(extract_property)
    energies = [schedule_getattr(j, 'energy') for j in td_dft_jobs]

    rs = run(gather(*energies), folder='biphenyl')

    print(rs)

def create_molecules(backbone, electronDonors, electronAcceptors):
    """
    Create all the possible combination by changing one para substituent
    by an electronAcceptors and the other para position by an electronDonors
    """
    # Smiles representing the molecules to simulate
    smiles = [backbone.format(donor, acceptor)
              for donor in electronDonors for acceptor in electronAcceptors]
    # Plams molecules
    geometries = [molkit.from_smiles(s, forcefield='uff') for s in smiles]

    # Name of the molecules
    invalid_chars = ['=', '(', ')', '[', ']', '#']
    names = [''.join(list(filter(lambda c: c not in invalid_chars, name))) for name in smiles]

    return geometries, names



def adf_general_settings():
    """
    ADF settings use for the calculations
    """
    s = Settings()

    # Basis
    s.specific.adf.basis.createoutput = None

    # SCF
    s.specific.adf.SCF.iterations = 200
    s.specific.adf.SCF.convergence = '1.0e-6 1.0e-5'

    # zlmfit
    s.specific.adf.zlmfit.quality = 'verygood'

    # BeckeGrid
    s.specific.adf.BeckeGrid.quality = 'verygood'

    # realtivistic
    s.specific.adf.Relativistic = 'SCALAR ZORA'

    # Solvation
    with open('radii', 'r') as f:
        radii = f.read()

    s.specific.adf.Solvation.div  = 'ndiv=3 min=0.5 ofac=0.8'
    s.specific.adf.Solvation.charged = 'conv=1e-10 iter=1000 corr'
    s.specific.adf.Solvation.CSMRSP = ''
    s.specific.adf.Solvation.Radii = '\n{}'.format(radii)

    return s

def create_settings_TD_DFT(basis='TZ2P'):
    """
    create input for a TD-DFT calculation
    """
    s = adf_general_settings()

     # Basis
    s.specific.adf.basis.type = '{}'.format(basis)
    s.specific.adf.basis.core = 'none'


    #COSMO
    s.specific.adf.Solvation.solvent = 'name=Dichloromethane emp=0.0 neql=2.028'

    # functional
    s.specific.adf.xc.hybrid = 'CAMYB3LYP'
    s.specific.adf.xc.xcfun = ''
    s.specific.adf.xc.rangesep = ''

    # excitations
    s.specific.adf.excitations.Davidson = ''
    s.specific.adf.excitations.lowest  = 7
    s.specific.adf.eprint.sfo = "eig ovl"

    return s

def create_settings_dft_gas(basis='DZP'):
    """
    Create input for an optimizatiion with ADF in gas phase
    """
    s = Settings()

    # Basis
    s.specific.adf.basis.createoutput = None

    # SCF
    s.specific.adf.SCF.iterations = 200
    s.specific.adf.SCF.convergence = '1.0e-6 1.0e-5'

    # zlmfit
    s.specific.adf.zlmfit.quality = 'good'

    # BeckeGrid
    s.specific.adf.BeckeGrid.quality = 'good'

    # realtivistic
    s.specific.adf.Relativistic = 'SCALAR ZORA'


    # Basis
    s.specific.adf.basis.type = '{}'.format(basis)
    s.specific.adf.basis.core = 'small'

    # functional
    s.specific.adf.xc.gga = 'BP86'

    # geometry
    s.specific.adf.geometry.optim = 'delocalized'
    s.specific.adf.geometry.iterations = 99
    s.specific.adf.geometry.convergence = 'e=0.00001 grad=0.0001'

    # SCF
    s.specific.adf.SCF.mix = 0.2

    return s



def create_settings_dft_cosmo_freq(basis='TZ2P'):
    """
    Create input for an optimization and posterior frequencies calculation with ADF
    """
    s = adf_general_settings()
    s.specific.adf.analyticalfreq
    s.specific.adf.scanfreq = '-500 100'
    #COSMO
    s.specific.adf.Solvation.solvent = 'name=Dichloromethane emp=0.0'

    # Basis
    s.specific.adf.basis.type = '{}'.format(basis)
    s.specific.adf.basis.core = 'small'

    # functional
    s.specific.adf.xc.gga = 'BP86'

    # geometry
    s.specific.adf.geometry.optim = 'delocalized'
    s.specific.adf.geometry.iterations = 99
    s.specific.adf.geometry.convergence = 'e=0.000001 grad=0.00001'

    # SCF
    s.specific.adf.SCF.mix = 0.2

    return s

def filter_None(jobs):
    """
    Remove the non-valid jobs
    """
    return list(filter(lambda j: j is not None, jobs))


def check_frequencies(job):
    """
    Filter the jobs which frequencies are non-negative
    """
    predicate = np.any(np.array(job.frequencies) < 0)
    if predicate:
        print("Job: ", job.job_name, " has at least a negative frequency")
        return None
    else:
        return job


def extract_property(job, prop):
    """
    wrapper around getattr
    """
    if job is None:
        return None
    else:
        return getattr(job, prop)


def choose_job(job1, job2):
    """
    Try to get an attribute from job1 otherwise from job2 else None
    """
    if job1 is not None and job1.status not in ['failed', 'crashed']:
        return  job1
    elif job2 is not None and job2.status not in ['failed', 'crashed']:
        return  job2
    else:
        return None


def run_if_succesful(fun, settings, job, name, prefix='job'):
    """
    Schedule a job if its depencies finished succesfully
    """
    if job is not None and job.status not in ['failed', 'crashed']:
        job_name = '{}_{}'.format(prefix, name)
        return fun(settings, job.molecule, job_name=job_name)
    else:
        return None

def run_if_failure(fun, settings, job, mol, name, prefix='job'):
    """
    Schedule a job if its depencies finished succesfully
    """
    if job is None or job.status in ['failed', 'crashed']:
        job_name = '{}_{}'.format(prefix, name)
        return fun(settings, mol, job_name=job_name)
    else:
        return None

if __name__ == "__main__":
    main()