from pathlib import Path
from typing import Optional
import os
import logging

logger = logging.getLogger(__name__)

DEFAULTS = {
    "structure": None,
    "pme": True,
    "cutoff": 9.0,
    "switching": True,
    "switchdistance": 7.5,
    "implicitsolvent": False,
    "igb": 2,
    "extforces": None,
    "fbrefcoor": None,
    "plumedfile": None,
    "coordinates": None,
    "boxsize": None,
    "velocities": 298.15,
    "timestep": 4.0,
    "slowperiod": 1,
    "thermostat": False,
    "thermostattemperature": 298.15,
    "thermostatdamping": 1,  # 1/ps is the recommended value from OMM
    "barostat": False,
    "barostatpressure": 1.0,
    "barostatanisotropic": False,
    "barostatconstratio": False,
    "barostatconstxy": False,
    "trajectoryfile": "output.xtc",
    "trajvelocityfile": None,
    "trajforcefile": None,
    "trajectoryperiod": 25000,
    "restart": False,
    "parameters": None,
    "run": 0,
    "minimize": 0,
    "stepzero": False,
    "nnp": None,
    "hydrogenmass": 4.032,
    "hmr": None,
    "hbondconstr": None,
    "rigidwater": None,
    "molecules": None,
    "protonate": False,
}
DEPRECATIONS = {  # The order is important. If multiple keys point to the same value the last non-null one wins
    "bincoordinates": "coordinates",
    "temperature": "velocities",  # Temperature is overridden by velocities and velocities by binvelocities
    "velocities": "velocities",  # Not a real deprectation. Just to specify the order priority
    "binvelocities": "velocities",
    "parmfile": "structure",
    "celldimension": "boxsize",
    "extendedsystem": "boxsize",
    "thermostattemp": "thermostattemperature",
    "switchdist": "switchdistance",
    "trajectoryfreq": "trajectoryperiod",
    "implicit": "implicitsolvent",
    "atomrestraint": "extforces",
    "grouprestraint": "extforces",
}
DEPRECATIONS_MSG = {
    "temperature": "The `temperature` option is deprecated in favor of the `velocities` option which can accept a temperature value to generate initial velocities from a Maxwell-Boltzmann distribution",
    "extendedsystem": "The `extendedsystem` option is deprecated in favor of the `boxsize` option which can accept a list of 3 numbers or a filename of a NAMD XSC file",
}


def _handle_deprecations(_input_args, error=False):
    to_delete = []
    _input_args = {key.lower(): val for key, val in _input_args.items()}
    input_args = _input_args.copy()
    for key in DEPRECATIONS:
        if key in _input_args and _input_args[key] is not None:
            val = _input_args[key]
            newkey = DEPRECATIONS[key].lower()

            if key in DEPRECATIONS_MSG:
                warn = f"# Deprecation: {DEPRECATIONS_MSG[key]}"
            else:
                warn = f"# Deprecation: `{key}` is deprecated in favor of `{newkey}`"

            if key == "parmfile" and _input_args.get("structure", None) is not None:
                if error:
                    raise ValueError(warn + ". Ignoring the parmfile option")
                else:
                    logger.warning(warn + ". Ignoring the parmfile option")
                to_delete.append(key)
                continue  # Ignore the parmfile option

            if _input_args.get(newkey, None) is not None and _input_args[newkey] != val:
                warn += (
                    f". Replacing {newkey} value '{_input_args[newkey]}' with '{val}'"
                )

            if key == newkey:  # No deprecation, just making sure ordering is maintained
                input_args[newkey] = val
                continue

            if error:
                raise ValueError(warn)
            else:
                logger.warning(warn)

            to_delete.append(key)

            if key in ("atomrestraint", "grouprestraint"):
                if newkey not in input_args or input_args[newkey] is None:
                    input_args[newkey] = []
                if isinstance(val, str):
                    input_args[newkey].append(val)
                elif isinstance(val, list):
                    input_args[newkey].extend(val)
                else:
                    raise ValueError(f"Invalid restraint type: {type(val)}")
            else:
                input_args[newkey] = val

    for key in to_delete:
        del input_args[key]
    return input_args


def _parse_runtime(runtime, timestep):
    if isinstance(runtime, int):
        return runtime
    if runtime is None:
        return 0
    if isinstance(runtime, str):
        if runtime.endswith("fs"):
            runtime = float(runtime.replace("fs", ""))
        elif runtime.endswith("ps"):
            runtime = float(runtime.replace("ps", "")) * 1e3
        elif runtime.endswith("ns"):
            runtime = float(runtime.replace("ns", "")) * 1e6
        elif runtime.endswith("us"):
            runtime = float(runtime.replace("us", "")) * 1e9
        else:
            raise ValueError(
                f"Invalid runtime format: {runtime}. Must be in fs, ps, ns or us"
            )
    return int(float(runtime) / timestep)


def detect_input_file(directory: Path, input_file: Optional[Path]):
    if input_file is None:
        fnames = ["input.yaml", "input.yml", "input.json", "input"]
        for fname in fnames:
            if os.path.isfile(os.path.join(directory, fname)):
                return os.path.join(directory, fname)
        raise FileNotFoundError(
            f"No input file found in directory: {directory} (valid names: `{'`, `'.join(fnames)}`). Please specify an input file."
        )
    elif os.path.isfile(os.path.join(directory, input_file)):
        return os.path.join(directory, input_file)
    elif os.path.isfile(input_file):
        return input_file
    else:
        raise FileNotFoundError(f"Input file {input_file} not found")


def parse_input_file(input_file: Optional[Path], directory: Path):
    from acemd.restraints import PositionalRestraint
    import yaml
    import json
    import re

    new_format = os.path.splitext(input_file)[-1] in (".yaml", ".yml", ".json")
    if not new_format:
        logger.warning(
            "# Parsing input file with old input file parser. We recommend using the new YAML or JSON inputs instead."
        )
        _input_args = {}
        with open(input_file, "r") as f:
            for line in f:
                line = line.strip()
                if len(line) == 0 or line.startswith("#"):
                    continue
                pieces = line.split()
                key = pieces[0].lower()

                # Remove the key from the line to get the rest
                pattern = re.compile(f"^{key}", re.IGNORECASE)
                val = pattern.sub("", line).strip()

                try:
                    val = int(val)
                except ValueError:
                    try:
                        val = float(val)
                    except ValueError:
                        pass

                if key.lower() in ("atomrestraint", "grouprestraint"):
                    val = key.lower() + " " + val
                    if key not in _input_args or _input_args[key] is None:
                        _input_args[key] = []
                    _input_args[key].append(val)
                else:
                    _input_args[key] = val
    elif input_file.endswith(".yaml") or input_file.endswith(".yml"):
        with open(input_file, "r") as f:
            _input_args = yaml.load(f, Loader=yaml.FullLoader)
            _input_args = {k.lower(): v for k, v in _input_args.items()}
    elif input_file.endswith(".json"):
        with open(input_file, "r") as f:
            _input_args = json.load(f)
            _input_args = {k.lower(): v for k, v in _input_args.items()}

    # Handle the deprecations
    input_args = _handle_deprecations(_input_args, error=new_format)

    if (
        input_args.get("boxsize", None) is not None
        and isinstance(input_args["boxsize"], str)
        and not input_args["boxsize"].endswith(".xsc")
    ):
        input_args["boxsize"] = [float(x) for x in input_args["boxsize"].split()]

    # Convert string restraints to dict
    if input_args.get("extforces", None) is not None:
        new_extfrc = []
        for rr in input_args["extforces"]:
            if new_format and isinstance(rr, str):
                raise RuntimeError(
                    f"# Deprecation: restraint '{rr}' is using the old format. Please convert it to the new yaml format."
                )
            if new_format and rr["type"].lower() != "positionalrestraint":
                raise ValueError(
                    f"# Error: Currently only positional restraints are supported with type: positionalRestraint. You used type {rr['type']}"
                )
            new_extfrc.append(PositionalRestraint.parse_restraint(rr))
        input_args["extforces"] = new_extfrc

    # Defaults
    for key, val in input_args.items():
        if isinstance(val, str) and val.lower() in ["on", "off"]:
            val = {"on": True, "off": False}[val.lower()]
        if isinstance(val, str) and val.lower() in ["true", "false"]:
            val = {"true": True, "false": False}[val.lower()]
        input_args[key] = val

    # Create input.yaml.new to ease user transition from old to new format
    if not new_format:
        logger.info(
            "# Writing input file in the new YAML format to input.yaml.new. Feel free to use it in future runs by renaming it to input.yaml"
        )
        # Convert the extforces to dict
        input_args_c = input_args.copy()
        if "extforces" in input_args_c:
            input_args_c["extforces"] = [
                rr.to_dict() for rr in input_args_c["extforces"]
            ]
        with open(os.path.join(directory, "input.yaml.new"), "w") as f:
            yaml.dump(input_args_c, f)

    args = DEFAULTS.copy()
    args.update(input_args)
    args["run"] = _parse_runtime(args["run"], args["timestep"])
    return args


def _print_input_args(args):
    logger.info("#")
    logger.info("# Input arguments")
    for key, value in sorted(args.items()):
        if key == "igb" and not args["implicitsolvent"]:
            continue
        if key == "switchdistance" and not args["switching"]:
            continue
        if (
            key in ("thermostattemperature", "thermostatdamping")
            and not args["thermostat"]
        ):
            continue
        if (
            key
            in (
                "barostatpressure",
                "barostatanisotropic",
                "barostatconstratio",
                "barostatconstxy",
            )
            and not args["barostat"]
        ):
            continue
        logger.info(f"#   {key}: {value}")
