"""Shared utilities for UMA-ASE workflows."""

from __future__ import annotations

import argparse
import logging
import os
import sys
from typing import Iterable, Mapping, Optional

import numpy as np
from ase import Atoms

# for UMA Fairchem
try:
    from fairchem.core import FAIRChemCalculator, pretrained_mlip
except ModuleNotFoundError:  # pragma: no cover - environment dependent
    pretrained_mlip = None
    FAIRChemCalculator = None


def Sum_of_atomic_energies(
    counts: Mapping[str, int],
    uma_checkpoint: str,
    uma_model: str,
    device: str,
    logger: Optional[logging.Logger] = None,
) -> float:
    """
    Compute the total energy by summing the energy of each atom type multiplied
    by its number of occurrences.

    Parameters
    ----------
    counts : dict[str, int]
        Dictionary mapping atom symbol -> number of occurrences.
    uma_checkpoint : str
        UMA checkpoint identifier.
    uma_model : str
        UMA task/model identifier.
    device : str
        Execution device (e.g. 'cpu', 'cuda').
    logger : logging.Logger, optional
        Logger used to report per-element energies.

    Returns
    -------
    float
        Total summed energy.
    """
    if pretrained_mlip is None or FAIRChemCalculator is None:
        raise RuntimeError("UMA FairChem libraries are unavailable; cannot compute atomic energies.")

    sum_atomic_energies = 0.0
    for atom_type, count in counts.items():
        # Generate a single-atom periodic cell inside the cutoff so the
        # calculator has neighbours to build the graph representation.
        # UMA default cutoff = 6.0, so a=5.0
        a = 5.0  # choose a < cutoff
        atom = Atoms(atom_type, positions=[[0, 0, 0]], cell=[a, a, a], pbc=True)
        atom.info["charge"] = 0
        atom.info["spin"] = 1

        try:
            predictor = pretrained_mlip.get_predict_unit(uma_checkpoint, device=device)
        except KeyError as exc:
            raise KeyError(f"UMA checkpoint '{uma_checkpoint}' not found.") from exc
        atom.calc = FAIRChemCalculator(predictor, task_name=uma_model)
        energy = atom.get_total_energy()
        message = f"Atomic energy for atom type {atom_type}: {energy} eV"
        if logger:
            logger.info(message)
        else:
            print(message)
        sum_atomic_energies += energy * count
    return sum_atomic_energies


def sum_of_atomic_energies(
    counts: Mapping[str, int],
    uma_checkpoint: str,
    uma_model: str,
    device: str,
    logger: Optional[logging.Logger] = None,
) -> float:
    """PEP 8 alias for :func:`Sum_of_atomic_energies`."""
    return Sum_of_atomic_energies(counts, uma_checkpoint, uma_model, device, logger=logger)


class Tee:
    """
    A file-like object that writes simultaneously to a file and stdout.

    Usage:
        from log_utils import Tee
        tee = Tee("output.log")
        print("Hello", file=tee)
        tee.close()
    """
    def __init__(self, filename: str, mode: str = "w", stream=sys.stdout):
        """
        Args:
            filename (str): Path to the log file.
            mode (str): File open mode ('w' to overwrite, 'a' to append).
            stream: Output stream to duplicate (default: sys.stdout).
        """
        self.file = open(filename, mode)
        self.stream = stream

    def write(self, message: str) -> None:
        self.stream.write(message)
        self.file.write(message)

    def flush(self) -> None:
        self.stream.flush()
        self.file.flush()

    def close(self) -> None:
        """Close the underlying file."""
        self.file.close()


def rmsd_simple(atoms1: Atoms, atoms2: Atoms) -> float:
    """Compute RMSD between two ASE Atoms objects without alignment."""
    pos1 = atoms1.get_positions()
    pos2 = atoms2.get_positions()
    diff = pos1 - pos2
    return np.sqrt((diff ** 2).sum() / len(pos1))


def get_filename_only(filepath: str) -> str:
    """
    Given a file path like 'folder/filename.ext', return only 'filename'.

    Args:
        filepath (str): The full file name or path.

    Returns:
        str: The filename without its extension.
    """
    # Extract the base name (remove path)
    base = os.path.basename(filepath)
    # Split into (filename, extension)
    name, _ = os.path.splitext(base)
    return name


def delete_file_if_exists(filepath: str) -> bool:
    """
    Check if a file exists, and delete it if it does.

    Args:
        filepath (str): Path to the file to delete.

    Returns:
        bool: True if file existed and was deleted, False otherwise.
    """
    try:
        if os.path.isfile(filepath):
            os.remove(filepath)
            print(f"Deleted file: {filepath}")
            return True
        else:
            print(f"File not found: {filepath}")
            return False
    except PermissionError:
        print(f"Permission denied: {filepath}")
        return False
    except Exception as exc:
        print(f"Error deleting {filepath}: {exc}")
        return False


def parse_cli_args(argv: Optional[Iterable[str]] = None) -> argparse.Namespace:
    """
    Parse command-line arguments shared across UMA utilities.

    Required arguments
    ------------------
    - input : Path to the input geometry readable by ASE.
    - chg   : Total molecular charge.
    - spin  : Spin multiplicity.

    Optional arguments
    ------------------
    - iter             : Maximum number of optimisation cycles (default 250).
    - grad             : Convergence threshold for maximum force (default 0.01).
    - mlff-chk         : UMA checkpoint identifier (default 'uma-sm').
    - mlff-task        : UMA task/model identifier (default 'omol').
    - visualize        : Open the final trajectory in an interactive viewer.
    - optimizer        : ASE optimizer to use (default 'BFGS').
    - run-type         : Workflow to execute ('sp', 'geoopt', 'freqs'; default 'sp').
    - temperature      : Thermochemistry temperature in Kelvin (default 298.15).
    - pressure         : Thermochemistry pressure in Pascals (default 101325.0).
    """
    parser = argparse.ArgumentParser(
        prog="UMA-ASE",
        description="Please, provide required -input and -chg values. Other options:"
    )

    # --- Required arguments ---
    parser.add_argument(
        "-input",
        type=str,
        required=True,
        help="Filename of input geometry, usually in xyz format. ASE atoms object. Required"
    )

    parser.add_argument(
        "-chg",
        type=int,
        required=True,
        help="Molecular Charge. Required",
    )

    # --- Optional arguments ---
    parser.add_argument(
        "-spin",
        type=int,
        default=1,
        help="Spin multiplicity. Default 1."
    )

    parser.add_argument(
        "-run-type",
        nargs="+",
        choices=("sp", "geoopt", "freqs"),
        default=["sp"],
        help="Run type(s) to execute: 'sp', 'geoopt', 'freqs', or any sequence thereof (default: sp).",
    )
    parser.add_argument(
        "-iter",
        type=int,
        default=250,
        help="Maximum number of geometry optimization cycles. Default=250"
    )

    parser.add_argument(
        "-grad",
        type=float,
        default=0.01,
        help="Max grad for convergence. Default=0.01 eV/A"
    )

    parser.add_argument(
        "-optimizer",
        default="LBFGS",
        metavar="NAME",
        help="ASE optimizer to use (e.g. BFGS, LBFGS, FIRE, BFGSLineSearch,MDMin). Default='LBFGS'."
    )

    parser.add_argument(
        "-mlff-chk",
        default="uma-sm",
        metavar="CHECKPOINT",
        help="UMA checkpoint identifier to load. Default='uma-sm'."
    )

    parser.add_argument(
        "-mlff-task",
        default="omol",
        metavar="TASK",
        help="UMA task/model identifier. Default='omol'."
    )

    parser.add_argument(
        "-visualize",
        action="store_true",
        help="Open the trajectory of the geoopt process in an interactive viewer."
    )

    parser.add_argument(
        "-temperature",
        type=float,
        metavar="T",
        default=298.15,
        help="Temperature in Kelvin for vibrational analysis (default 298.15 K)."
    )

    parser.add_argument(
        "-pressure",
        type=float,
        metavar="P",
        default=101325.0,
        help="Pressure in Pascals for vibrational analysis (default 101325.0 Pa)."
    )

    return parser.parse_args(argv)
