#
# @BEGIN LICENSE
#
# Psi4: an open-source quantum chemistry software package
#
# Copyright (c) 2007-2022 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
#

__all__ = ["task_planner"]

import os
import copy
import logging
from typing import Dict, Tuple, Union

from qcelemental.models import DriverEnum
from psi4.driver import p4util, pp
from psi4.driver.task_base import AtomicComputer
from psi4.driver.driver_findif import FiniteDifferenceComputer
from psi4.driver.driver_nbody import ManyBodyComputer
from psi4.driver.driver_cbs import CompositeComputer, composite_procedures, cbs_text_parser
from psi4.driver.driver_util import negotiate_derivative_type, negotiate_convergence_criterion

logger = logging.getLogger(__name__)

TaskComputers = Union[AtomicComputer, CompositeComputer, FiniteDifferenceComputer, ManyBodyComputer]


def _expand_cbs_methods(method: str, basis: str, driver: DriverEnum, **kwargs) -> Tuple[str, str, Dict]:
    """Sort out the user input method string into recognized fields. Handles cases like:

    * (i) `"mp2"` --> passes through;
    * (ii) `"mp2/cc-pvdz"` --> broken into method and basis fields;
    * (iii) `"mp2/cc-pv[d,t]z"` --> processed into method="cbs" and CBSMetadata spec;
    * (iv) `method="cbs", cbsmeta=CBSMetadata` --> passes through.

    """
    if method == 'cbs' and kwargs.get('cbsmeta', False):
        return method, basis, kwargs['cbsmeta']

    # Expand CBS methods
    if "/" in method:
        kwargs["ptype"] = driver
        cbsmeta = cbs_text_parser(method, **kwargs)

        # Single call detected
        if "cbs_metadata" not in cbsmeta:
            method = cbsmeta["method"]
            basis = cbsmeta["basis"]
        else:
            method = "cbs"
    else:
        cbsmeta = {}

    return method, basis, cbsmeta


def task_planner(driver: DriverEnum, method: str, molecule: "psi4.core.Molecule", **kwargs) -> TaskComputers:
    """Plans a task graph of a complex computation.

    Canonical Task layering:
     - ManyBody
     - FiniteDifference
     - Composite
     - Atomic

    Parameters
    ----------
    driver
        The resulting type of computation: e/g/h. Note for finite difference that this should be the target driver, not the means driver.
    method
        A string representation of the method such as "HF" or "B3LYP". Special cases are:
        "cbs"
    molecule
        A Psi4 base molecule to use
    kwargs
        User keyword arguments, often used to configure task computers.

    Returns
    -------
    computer
        A simple (AtomicComputer) or layered (CompositeComputer, FiniteDifferenceComputer, ManyBodyComputer) task object.
        Layered objects contain many and multiple types of computers in a graph.

    """

    # Only pull the changed options
    keywords = p4util.prepare_options_for_set_options()

    keywords["function_kwargs"] = {}
    if "external_potentials" in kwargs:
        keywords["function_kwargs"].update({"external_potentials": kwargs.pop("external_potentials")})

    # Need to add full path to pcm file
    if "PCM__PCMSOLVER_PARSED_FNAME" in keywords.keys():
        fname = keywords["PCM__PCMSOLVER_PARSED_FNAME"]
        keywords["PCM__PCMSOLVER_PARSED_FNAME"] = os.path.join(os.getcwd(), fname)

    # Pull basis out of kwargs, override globals if user specified
    basis = kwargs.pop("basis", keywords.pop("BASIS", "(auto)"))
    method = method.lower()

    # Expand CBS methods
    method, basis, cbsmeta = _expand_cbs_methods(method, basis, driver, **kwargs)
    if method in composite_procedures:
        kwargs.update({'cbs_metadata': composite_procedures[method](**kwargs)})
        method = 'cbs'

    pertinent_findif_kwargs = ['findif_irrep', 'findif_stencil_size', 'findif_step_size', 'findif_verbose']
    current_findif_kwargs = {kw: kwargs.pop(kw) for kw in pertinent_findif_kwargs if kw in kwargs}
    # explicit: 'findif_mode'

    # Build a packet
    packet = {"molecule": molecule, "driver": driver, "method": method, "basis": basis, "keywords": keywords}

    # First check for BSSE type
    if kwargs.get("bsse_type", None) is not None:
        levels = kwargs.pop('levels', None)

        plan = ManyBodyComputer(**packet, **kwargs)
        original_molecule = packet.pop("molecule")

        # Add tasks for every nbody level requested
        if levels is None:
            levels = {plan.max_nbody: method}
        else:
            # rearrange bodies in order with supersystem last lest body count fail in organization loop below
            levels = dict(sorted(levels.items(), key=lambda item: 1000 if item[0] == "supersystem" else item[0]))

            # We define cp as being a correction to only interaction energies
            # If only doing cp, we need to ignore any user-specified 1st (monomer) level
            if 'cp' in kwargs.get("bsse_type", None) and 'nocp' not in kwargs.get("bsse_type", None): 
                if 1 in levels.keys():
                    removed_level = levels.pop(1)
                    logger.info("NOTE: User specified exclusively 'cp' correction, but provided level 1 details") 
                    logger.info(f"NOTE: Removing level {removed_level}")
                    logger.info("NOTE: For total energies, add 'nocp' to bsse_list")


        # Organize nbody calculations into modelchem levels
        # * expand keys of `levels` into full lists of nbodies covered. save to plan, resetting max_nbody accordingly
        # * below, process values of `levels`, which are modelchem strings, into kwargs specs
        nbodies_per_mc_level = []
        prev_body = 0
        for nb in levels:
            nbodies = []
            if nb == "supersystem":
                nbodies.append(nb)
            elif nb != (prev_body + 1):
                for m in range(prev_body + 1, nb + 1):
                    nbodies.append(m)
            else:
                nbodies.append(nb)
            nbodies_per_mc_level.append(nbodies)
            prev_body += 1

        plan.max_nbody = max(nb for nb in levels if nb != "supersystem")
        plan.nbodies_per_mc_level = nbodies_per_mc_level

        for mc_level_idx, mtd in enumerate(levels.values()):
            method, basis, cbsmeta = _expand_cbs_methods(mtd, basis, driver, cbsmeta=cbsmeta, **kwargs)
            packet.update({'method': method, 'basis': basis})

            # Tell the task builder which level to add a task list for
            # * see https://github.com/psi4/psi4/pull/1351#issuecomment-549948276 for discussion of where build_tasks logic should live
            if method == "cbs":
                # This CompositeComputer is discarded after being used for dermode.
                simplekwargs = copy.deepcopy(kwargs)
                simplekwargs.pop('dertype', None)
                simplecbsmeta = copy.deepcopy(cbsmeta)
                simplecbsmeta['verbose'] = 0
                dummyplan = CompositeComputer(**packet, **simplecbsmeta, molecule=original_molecule, **simplekwargs)

                methods = [sr.method for sr in dummyplan.task_list]
                # TODO: pass more info, so fn can use for managed_methods -- ref, qc_module, fc/ae, conv/df
                dermode = negotiate_derivative_type(driver, methods, kwargs.pop('dertype', None), verbose=1)

                if dermode[0] == dermode[1]:  # analytic
                    logger.info("PLANNING MB(CBS):  {mc_level_idx=} {packet=} {cbsmeta=} kw={kwargs}")
                    plan.build_tasks(CompositeComputer, **packet, mc_level_idx=mc_level_idx, **cbsmeta, **kwargs)

                else:
                    logger.info(
                        f"PLANNING MB(FD(CBS):  {mc_level_idx=} {packet=} {cbsmeta=} findif_kw={current_findif_kwargs} kw={kwargs}"
                    )
                    plan.build_tasks(FiniteDifferenceComputer,
                                     **packet,
                                     mc_level_idx=mc_level_idx,
                                     findif_mode=dermode,
                                     computer=CompositeComputer,
                                     **cbsmeta,
                                     **current_findif_kwargs,
                                     **kwargs)

            else:
                dermode = negotiate_derivative_type(driver, method, kwargs.pop('dertype', None), verbose=1)
                if dermode[0] == dermode[1]:  # analytic
                    logger.info(f"PLANNING MB:  {mc_level_idx=} {packet=}")
                    plan.build_tasks(AtomicComputer, **packet, mc_level_idx=mc_level_idx, **kwargs)
                else:
                    logger.info(
                        f"PLANNING MB(FD):  {mc_level_idx=} {packet=} findif_kw={current_findif_kwargs} kw={kwargs}"
                    )
                    plan.build_tasks(FiniteDifferenceComputer,
                                     **packet,
                                     mc_level_idx=mc_level_idx,
                                     findif_mode=dermode,
                                     **current_findif_kwargs,
                                     **kwargs)

        return plan

    # Check for CBS
    elif method == "cbs":
        kwargs.update(cbsmeta)
        # This CompositeComputer is discarded after being used for dermode. Could have used directly for analytic except for excess printing with FD
        simplekwargs = copy.deepcopy(kwargs)
        simplekwargs['verbose'] = 0
        dummyplan = CompositeComputer(**packet, **simplekwargs)

        methods = [sr.method for sr in dummyplan.task_list]
        # TODO: pass more info, so fn can use for managed_methods -- ref, qc_module, fc/ae, conv/df
        dermode = negotiate_derivative_type(driver, methods, kwargs.pop('dertype', None), verbose=1)

        if dermode[0] == dermode[1]:  # analytic
            logger.info('PLANNING CBS:  packet={packet} kw={kwargs}')
            plan = CompositeComputer(**packet, **kwargs)
            return plan
        else:
            # For FD(CBS(Atomic)), the CompositeComputer above is discarded after being used for dermode.
            logger.info(
                f'PLANNING FD(CBS):  dermode={dermode} packet={packet} findif_kw={current_findif_kwargs} kw={kwargs}')
            plan = FiniteDifferenceComputer(**packet,
                                            findif_mode=dermode,
                                            computer=CompositeComputer,
                                            **current_findif_kwargs,
                                            **kwargs)
            return plan

    # Done with Wrappers -- know we want E, G, or H -- but may still be FD or AtomicComputer
    else:
        dermode = negotiate_derivative_type(driver, method, kwargs.pop('dertype', None), verbose=1)
        convcrit = negotiate_convergence_criterion(dermode, method, return_optstash=False)

        if dermode[0] == dermode[1]:  # analytic
            logger.info(f'PLANNING Atomic:  keywords={keywords}')
            return AtomicComputer(**packet, **kwargs)
        else:
            keywords.update(convcrit)
            logger.info(
                f'PLANNING FD:  dermode={dermode} keywords={keywords} findif_kw={current_findif_kwargs} kw={kwargs}')
            return FiniteDifferenceComputer(**packet,
                                            findif_mode=dermode,
                                            **current_findif_kwargs,
                                            **kwargs)
