#
# @BEGIN LICENSE
#
# Psi4: an open-source quantum chemistry software package
#
# Copyright (c) 2007-2018 The Psi4 Developers.
#
# The copyrights for code used from other parties are included in
# the corresponding files.
#
# This file is part of Psi4.
#
# Psi4 is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, version 3.
#
# Psi4 is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with Psi4; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# @END LICENSE
#

import json

import numpy as np

from ..util import unnp
from ..physconst import psi_bohr2angstroms
from ..exceptions import *
from .to_string import formula_generator

def to_schema(molrec, dtype, units='Angstrom', return_type='json'):
    """Translate Psi4 json Molecule spec into json from other schemas.

    Parameters
    ----------
    molrec : dict
        Psi4 json Molecule spec.
    dtype : {'psi4', 1}
        Molecule schema format.
    units : {'Angstrom', 'Bohr'}
        Units in which to write string. There is not an option to write in
        intrinsic/input units. Some `dtype` may not allow all units.
    return_type : {'json', 'yaml'}
        Serialization format string to return.

    Returns
    -------
    qcschema : str
        String of the `return_type` repr of `molrec`.

    """
    qcschema = {}

    if molrec['units'] == 'Angstrom' and units == 'Angstrom':
        factor = 1.
    elif molrec['units'] == 'Angstrom' and units == 'Bohr':
        if 'input_units_to_au' in molrec:
            factor = molrec['input_units_to_au']
        else:
            factor = 1. / psi_bohr2angstroms
    elif molrec['units'] == 'Bohr' and units == 'Angstrom':
        factor = psi_bohr2angstroms
    elif molrec['units'] == 'Bohr' and units == 'Bohr':
        factor = 1.
    else:
        raise ValidationError("""units must be 'Angstrom'/'Bohr', not {}""".format(units))
    geom = np.array(molrec['geom']) * factor
    nat = geom.shape[0] // 3

    name = molrec.get('name', formula_generator(molrec['elem']))
#    tagline = """auto-generated by qcdb from molecule {}""".format(name)

    if dtype == 'psi4':
        qcschema = copy.deepcopy(molrec)
        qcschema['geom'] = geom
        qcschema['units'] = units
        qcschema['name'] = name

    elif dtype == 1:
        if units != 'Bohr':
            raise ValidationError("""QC_JSON_Schema {} allows only 'Bohr' coordinates, not {}.""".format(dtype, units))

        qcschema['symbols'] = np.array(molrec['elem'])
        qcschema['geometry'] = geom
        qcschema['masses'] = np.array(molrec['mass'])
        qcschema['name'] = name
        qcschema['molecular_charge'] = molrec['molecular_charge']
        qcschema['molecular_multiplicity'] = molrec['molecular_multiplicity']
        qcschema['real'] = np.array(molrec['real'])
        fidx = np.split(np.arange(nat), molrec['fragment_separators'])
        qcschema['fragments'] = [fr.tolist() for fr in fidx]
        qcschema['fragment_charges'] = np.array(molrec['fragment_charges'])
        qcschema['fragment_multiplicities'] = np.array(molrec['fragment_multiplicities'])
        qcschema['fix_com'] = molrec['fix_com']
        qcschema['fix_orientation'] = molrec['fix_orientation']
    else:
        raise TypeError("dtype not understood, valid options are {'psi4', 1}. Found {}.".format(dtype))

        # hmm, psi4/qcdb for provenance or does psi molrec need a passthrough field to preserve?
        #qcschema['provenance'] creator, version, routine

    qcschema = unnp(qcschema)

    if return_type == 'json':
        return json.dumps(qcschema)
    elif return_type == 'yaml':
        import yaml
        return yaml.dump(qcschema)
    else:
        raise ValidationError("""Return type ({}) not recognized.""".format(return_type))
