"""
/******************************************************************************
  This source file is part of the Avogadro project.

  This source code is released under the New BSD License, (the "License").
******************************************************************************/
"""

import argparse
import json
import sys

# Some globals:
targetName = 'MOLPRO'
extension = 'inp'
debug = False


def getOptions():
    userOptions = {}

    userOptions['Title'] = {}
    userOptions['Title']['type'] = 'string'
    userOptions['Title']['default'] = ''

    userOptions['Calculation Type'] = {}
    userOptions['Calculation Type']['type'] = 'stringList'
    userOptions['Calculation Type']['default'] = 1
    userOptions['Calculation Type']['values'] = \
        ['Single Point', 'Equilibrium Geometry', 'Frequencies']

    userOptions['Theory'] = {}
    userOptions['Theory']['type'] = 'stringList'
    userOptions['Theory']['default'] = 2
    userOptions['Theory']['values'] = [
        'RHF', 'MP2', 'B3LYP', 'CCSD', 'CCSD(T)']

    userOptions['Basis'] = {}
    userOptions['Basis']['type'] = 'stringList'
    userOptions['Basis']['default'] = 2
    userOptions['Basis']['values'] = \
        ['STO-3G', '3-21G', '6-31G', '6-31G(d)', '6-31G(d,p)', '6-31+G(d)',
         '6-311G(d)', 'cc-pVDZ', 'cc-pVTZ', 'AUG-cc-pVDZ', 'AUG-cc-pVTZ']

    userOptions['Filename Base'] = {}
    userOptions['Filename Base']['type'] = 'string'
    userOptions['Filename Base']['default'] = 'job'

    userOptions['Charge'] = {}
    userOptions['Charge']['type'] = 'integer'
    userOptions['Charge']['default'] = 0
    userOptions['Charge']['minimum'] = -9
    userOptions['Charge']['maximum'] = 9

    userOptions['Multiplicity'] = {}
    userOptions['Multiplicity']['type'] = 'integer'
    userOptions['Multiplicity']['default'] = 1
    userOptions['Multiplicity']['minimum'] = 1
    userOptions['Multiplicity']['maximum'] = 6

    # TODO Coordinate format (need zmatrix)

    userOptions['Use Pre-2009.1 Format'] = {}
    userOptions['Use Pre-2009.1 Format']['type'] = 'boolean'
    userOptions['Use Pre-2009.1 Format']['default'] = False

    opts = {'userOptions': userOptions}
    opts['inputMoleculeFormat'] = 'cjson'

    return opts


def generateInputFile(cjson, opts):
    # Extract options:
    title = opts['Title']
    calculate = opts['Calculation Type']
    theory = opts['Theory']
    basis = opts['Basis']
    charge = opts['Charge']
    multiplicity = opts['Multiplicity']
    oldVersion = opts['Use Pre-2009.1 Format']

    # Convert to code-specific strings
    basisStr = ''
    if basis in ['STO-3G', '3-21G', '6-31G', '6-31G(d)', '6-31G(d,p)',
                 '6-31+G(d)', '6-311G(d)']:
        basisStr = basis
    elif basis == 'cc-pVDZ':
        basisStr = 'vdz'
    elif basis == 'cc-pVTZ':
        basisStr = 'vtz'
    elif basis == 'AUG-cc-pVDZ':
        basisStr = 'avdz'
    elif basis == 'AUG-cc-pVTZ':
        basisStr = 'avtz'
    else:
        raise Exception('Unhandled basis type: %s' % basis)

    # MOLPRO needs us to calculate some rough wavefunction parameters:
    numElectrons = -charge
    try:
        for z in cjson['atoms']['elements']['number']:
            numElectrons += z
    except KeyError:
        numElectrons = 0
    wavefnStr = 'wf,%d,1,%d' % (numElectrons, multiplicity - 1)

    theoryStr = ''
    if theory != 'B3LYP':
        theoryStr += '{rhf\n%s}\n' % wavefnStr
    # Intentionally not using elif here:
    if theory != 'RHF':
        theoryKey = ''
        if theory in ['MP2', 'CCSD', 'CCSD(T)']:
            theoryKey = theory.lower()
        elif theory == 'B3LYP':
            theoryKey = 'uks,b3lyp'
        else:
            raise Exception('Unhandled theory type: %s' % theory)
        theoryStr += f'{{{theoryKey}\n{wavefnStr}}}\n'

    calcStr = ''
    if calculate == 'Single Point':
        pass
    elif calculate == 'Equilibrium Geometry':
        calcStr = '{optg}\n\n'
    elif calculate == 'Frequencies':
        calcStr = '{optg}\n{frequencies}\n\n'
    else:
        raise Exception('Unhandled calculation type: %s' % calculate)

    # Create input file
    output = ''

    output += '*** %s\n\n' % title
    output += 'gprint,basis\n'
    output += 'gprint,orbital\n\n'

    output += 'basis, %s\n\n' % basisStr

    if oldVersion:
        output += 'geomtyp=xyz\n'
    output += 'geometry={\n'
    if oldVersion:
        numAtoms = 0
        try:
            numAtoms = len(cjson['atoms']['elements']['number'])
        except KeyError:
            numAtoms = 0
        output += '%d\n\n' % numAtoms

    output += '$$coords:Sxyz$$\n'
    output += '}\n\n'

    output += '%s\n' % theoryStr

    output += '%s' % calcStr

    output += "---\n"

    return output


def generateInput():
    # Read options from stdin
    stdinStr = sys.stdin.read()

    # Parse the JSON strings
    opts = json.loads(stdinStr)

    # Generate the input file
    inp = generateInputFile(opts['cjson'], opts['options'])

    # Basename for input files:
    baseName = opts['options']['Filename Base']

    # Prepare the result
    result = {}
    # Input file text -- will appear in the same order in the GUI as they are
    # listed in the array:
    files = []
    files.append({'filename': f'{baseName}.{extension}', 'contents': inp})
    if debug:
        files.append({'filename': 'debug_info', 'contents': stdinStr})
    result['files'] = files
    # Specify the main input file. This will be used by MoleQueue to determine
    # the value of the $$inputFileName$$ and $$inputFileBaseName$$ keywords.
    result['mainFile'] = f'{baseName}.{extension}'
    return result

if __name__ == "__main__":
    parser = argparse.ArgumentParser('Generate a %s input file.' % targetName)
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--print-options', action='store_true')
    parser.add_argument('--generate-input', action='store_true')
    parser.add_argument('--display-name', action='store_true')
    parser.add_argument('--lang', nargs='?', default='en')
    args = vars(parser.parse_args())

    debug = args['debug']

    if args['display_name']:
        print(targetName)
    if args['print_options']:
        print(json.dumps(getOptions()))
    elif args['generate_input']:
        print(json.dumps(generateInput()))
