# wrapper for optking's optimize function for input by psi4API
# creates a moleuclar system from psi4s and generates optkings options from psi4's lsit of options
import copy
import inspect
import json
import logging
import pprint

import qcelemental
from qcelemental.util.serialization import json_dumps
from pydantic import ValidationError
from pydantic.v1.error_wrappers import ValidationError as v1ValidationError

import optking

from . import caseInsensitiveDict, molsys
from .compute_wrappers import ComputeWrapper, Psi4Computer, QCEngineComputer, UserComputer
from .exceptions import OptError
from .optimize import optimize
from .printTools import welcome
from . import log_name
from . import op

logger = logging.getLogger(f"{log_name}{__name__}")


def optimize_psi4(calc_name, program="psi4", dertype=None, **xtra_opt_params):
    """Wrapper for optimize.optimize() Looks for an active psi4 molecule and optimizes.
    This is the written warning that Optking will try to use psi4 if no program is provided

    Parameters
    ----------
    calc_name: str
        level of theory for optimization. eg MP2
    program: str
        program used for gradients, hessians...
    dertype: ?
        hack to try to get finite differences working in psi4
    xtra_opt_params: dictionary
        extra keywords currently forbidden by psi4's read_options, but supported by optking

    Returns
    -------
    opt_output: dict
        dictionary serialized MolSSI OptimizationResult.
        If Psi4 supports QCSchema v2 (~v1.11 Spring 2026), *opt_output* will be v2: https://molssi.github.io/QCElemental/next/model_opt.html .
        Otherwise, *opt_output* will be v1: https://molssi.github.io/QCElemental/dev/api/qcelemental.models.OptimizationInput.html#qcelemental.models.OptimizationInput .
    """

    opt_input = {}
    opt_output = {}

    try:
        op.Params, oMolsys, computer, opt_input = initialize_from_psi4(
            calc_name, program, computer_type="psi4", dertype=dertype, **xtra_opt_params
        )
        opt_output = optimize(oMolsys, computer)
    except (ValidationError, v1ValidationError) as error:
        logger.critical("A ValidationError has occured: %s", error, exc_info=True)
        opt_output = {
            "success": False,
            "error": {"error_type": "ValidationError", "error_message": str(error)}
        }
    except TypeError as error:
        logger.critical("A TypeError has occured: %s", error, exc_info=True)
        logger.critical("This TypeError is likely related to option validation.")
        opt_output = {
            "success": False,
            "error": {"error_type": "ValidationError", "error_message": str(error)}
        }
    except Exception as error:
        logger.critical(
            "A critical exception has occured:\n%s - %s",
            type(error),
            str(error),
            exc_info=True
        )
        opt_output = {
            "success": False,
            "error": {"error_type": type(error), "error_message": str(error)},
        }
    finally:
        if computer.dtype == 1:
            opt_input.update({"provenance": optking._optking_provenance_stamp})
            opt_input["provenance"]["routine"] = "optimize_psi4"
            opt_input.update(opt_output)

            return opt_input

        elif computer.dtype == 2:
            opt_output["input_data"] = opt_input
            opt_output["provenance"]["routine"] = "optimize_psi4"

            return opt_output


def initialize_from_psi4(calc_name, program, computer_type, dertype=None, **xtra_opt_params):
    """Gathers information from an active psi4 instance. to cleanly run optking from a
    psithon or psi4api input file

    Parameters
    ----------
    calc_name: str
    computer_type: str
    dertype: Union[int, None]
    program: str
    **xtra_opt_params
        extra keywords which are not recognized by psi4

    Returns
    -------
    params: op.OptParams
    o_molsys: molsys.Molsys
    computer: ComputeWrapper
    opt_input: qcelemental.models.OptimizationInput
    """
    import psi4

    # Get Molecule and Options from Psi4 as QCSchema
    logger.debug("Setting up optimization from Psi4's current state: p4util.state_to_atomicinput")

    psi4_can_v2 = "dtype" in inspect.signature(psi4.driver.p4util.state_to_atomicinput).parameters  # PR 3341 pre-v1.11

    if psi4_can_v2:
        full_atomic_input = psi4.driver.p4util.state_to_atomicinput(
            dtype=2,
            driver="gradient",
            method=calc_name,
        ).model_dump()
        atomic_input = full_atomic_input["specification"]
    else:
        atomic_input = psi4.driver.p4util.state_to_atomicinput(
            driver="gradient",
            method=calc_name,
        ).dict()

    opt_keys = {}

    # QCEngine doesn't expect information to already have been validated
    if psi4_can_v2:
        o_molsys = molsys.Molsys.from_schema(full_atomic_input.get("molecule"))
        atomic_input["program"] = program
    else:
        atomic_input.pop("id")
        atomic_input.pop("provenance")
        atomic_input.pop("protocols")
        o_molsys = molsys.Molsys.from_schema(atomic_input.get("molecule"))
        opt_keys["program"] = program

    optking_canon = op.OptParams().to_dict(by_alias=True).keys()

    # Psi4 options can get mixed in with optking's options in prepare_options_for_module anyway.
    # Atomic_input will contain all set options so only accept options that are present in
    # options dict. Assume everything else is for Psi4.
    for opt, optval in atomic_input.get("keywords").items():
        if opt.upper() in optking_canon:
            opt_keys[opt] = optval

    if xtra_opt_params:
        for xtra_key, xtra_value in xtra_opt_params.items():
            opt_keys[xtra_key] = xtra_value

    # Make a qcSchema OptimizationInput
    if psi4_can_v2:
        opt_input = {
            "initial_molecule": full_atomic_input.pop("molecule"),
            "specification": {
                "specification": atomic_input,
                "keywords": opt_keys,
            },
        }
    else:
        opt_input = {
            "keywords": opt_keys,
            "initial_molecule": atomic_input.pop("molecule"),
            "input_specification": atomic_input,
        }

    # in case user has selected a specific dertype
    if dertype:
        if psi4_can_v2:
            opt_input.get("specification").get("specification").get("keywords").update("dertype", dertype)
        else:
            opt_input.get("input_specification").get("keywords").update("dertype", dertype)

    # any Psi4 that is v2-capable requires QCElemental >=0.50 where QCSchema v2 is RTG
    if psi4_can_v2:
        from qcelemental.models.v2 import OptimizationInput
    else:
        from qcelemental.models import OptimizationInput

    try:
        logger.debug("Creating OptimizationInput")
        opt_input = OptimizationInput(**opt_input)
    except (ValidationError, v1ValidationError) as error:
        logger.critical("A Validation Error has occured while initializing optking: %s", error)
        logger.critical("Could not create an OptimizationInput")
        logger.critical(
            """Note: `optimize_psi4` is not recommended for users. Consider calling `psi4.optimize`
            or see `tests/test_opthelper` for an example of using the OptHelper interface."""
    )

    # Remove numpy elements that can appear from psi4 to allow at will json serialization
    # Json cannot handle numpy types only python builtins
    opt_input = json.loads(json_dumps(opt_input))
    params = initialize_options(opt_keys)
    computer = make_computer(opt_input, computer_type)
    return params, o_molsys, computer, opt_input


def optimize_qcengine(opt_input, computer_type="qc"):
    """Try to optimize, find TS, or find IRC of the system as specifed by a QCSchema
    OptimizationInput.

    Parameters
    ----------
    opt_input: Union[OptimizationInput, dict]
        Pydantic Schema of the OptimizationInput model.
        see https://github.com/MolSSI/QCElemental/blob/master/qcelemental/models/procedures.py
    computer_type: str

    Returns
    -------
    dict

    """

    try:
        allowed = (qcelemental.models.v1.OptimizationInput, qcelemental.models.v2.OptimizationInput)
    except AttributeError:
        allowed = (qcelemental.models.OptimizationInput)

    if isinstance(opt_input, allowed):
        # Remove numpy elements and turn into dictionary
        opt_input = json.loads(json_dumps(opt_input))
    opt_output = copy.deepcopy(opt_input)

    dtype = 2 if "specification" in opt_output.keys() else 1

    # Make basic optking molecular system
    oMolsys = molsys.Molsys.from_schema(opt_input["initial_molecule"])
    try:
        if dtype == 1:
            initialize_options(opt_input["keywords"])
        elif dtype == 2:
            initialize_options(opt_input["specification"]["keywords"])
        computer = make_computer(opt_input, computer_type)
        opt_output = optimize(oMolsys, computer)
    except Exception as error:
        logger.critical("A critical error has occured: %s - %s", type(error), error, exc_info=True)
        opt_output = {
            "success": False,
            "error": {"error_type": type(error), "error_message": str(error)},
        }
        logger.critical(f"Error placed in qcschema: {opt_output}")
    finally:
        if dtype == 1:
            opt_input.update(opt_output)
            opt_input.update({"provenance": optking._optking_provenance_stamp})
            opt_input["provenance"]["routine"] = "optimize_qcengine"
            return opt_input
        elif dtype == 2:
            opt_output["input_data"] = opt_input
            opt_output["provenance"] = optking._optking_provenance_stamp
            opt_output["provenance"]["routine"] = "optimize_qcengine"
            return opt_output

    # QCEngine.procedures.optking.py takes 'output_data', unpacks and creates Optimization Schema
    # from qcel.models.v1.procedures.py or qcel.models.v2.optimization.py


def make_computer(opt_input: dict, computer_type):
    logger.debug("Creating a Compute Wrapper")
    program = op.Params.program

    # This gets updated so it shouldn't be a reference
    molecule = copy.deepcopy(opt_input["initial_molecule"])

    schver = 2 if "specification" in opt_input.keys() else 1

    # Sorting by spec_schema_name isn't foolproof b/c opt_input might not be a
    #   constructed model at this point if it's not arriving through QCEngine.
    if schver == 1:
        spec_schema_name = opt_input["input_specification"].get("schema_name", "qcschema_input")
        if spec_schema_name == "qcschema_manybodyspecification":
            # route is defunct
            model = "(proc_spec_in_options)"
            options = opt_input["input_specification"]
            protocols = {}
        else:
            qc_input = opt_input["input_specification"]
            options = qc_input["keywords"]
            model = qc_input["model"]
            protocols = {}  # no field for protocols in v1 QCInputSpecification
    elif schver == 2:
        spec_schema_name = opt_input["specification"]["specification"].get("schema_name", "qcschema_atomic_specification")
        if spec_schema_name == "qcschema_many_body_specification":
            model = "(proc_spec_in_options)"
            options = opt_input["specification"]["specification"]
            if "model" in (qc := options["specification"]):  # single spec
                protocols = qc.get("protocols", {})
            else:  # mapping spec
                protocols = next(iter(qc.values())).get("protocols", {})
            program = opt_input["specification"]["specification"].get("program") or "qcmanybody"
        else:
            qc_input = opt_input["specification"]["specification"]
            options = qc_input["keywords"]
            model = qc_input["model"]
            protocols = qc_input.get("protocols", {})
            program = qc_input.get("program", program)

    if computer_type == "psi4":
        # Please note that program is not actually used here
        return Psi4Computer(molecule, model, options, program, protocols, dtype=schver)
    elif computer_type == "qc":
        return QCEngineComputer(molecule, model, options, program, protocols, dtype=schver)
    elif computer_type == "user":
        logger.info("Creating a UserComputer")
        return UserComputer(molecule, model, options, program, protocols, dtype=schver)
    else:
        raise OptError("computer_type is unknown")


def initialize_options(opt_keys, silent=False):
    if not silent:
        logger.info(welcome())

    # Turn keys but not vals to upper. Optparams will handle lowercase values
    opt_keys = {key.upper(): val for key, val in opt_keys.items()}

    # Create full list of parameters from user options plus defaults.
    try:
        params = op.OptParams(**opt_keys)
    except (KeyError, ValueError, AttributeError, ValidationError) as e:
        error_key = ''
        if isinstance(e, ValidationError):
            for key in opt_keys.keys():
                if key in str(e):
                    error_key = key
        logger.critical("\nEncountered an %s error parsing options %s", type(e), e)
        logger.critical(
            "\nFor valdidation errors please check the documenation to find an example of the desired\n"
            "keyword being uitilized: https://optking.readthedocs.io/en/latest/options.html#options\n"
        )
        if error_key:
            raise OptError(
                f"unable to parse user provided options. \n{e}"
                f"\n{error_key}:\t{opt_keys[error_key]}"
            ) from e
        else:
            raise OptError(f"unable to parse provided options. \n{e}") from e

    # TODO we should make this just be a normal object
    #  we should return it to the optimize method
    if not silent:
        logger.info(params)
    op.Params = params

    return params
