"""
Radial Distribution Function (RDF) analysis utilities.

This module provides single-frame and multi-frame RDF calculations for
ReaxFF trajectories using either the FREUD or OVITO backends, as well as
helper utilities for extracting RDF-derived structural descriptors.

Typical use cases include:

- computing total or partial RDFs for selected frames
- averaging RDFs across multiple frames
- extracting RDF-based properties such as peak positions or integrated areas
"""


from __future__ import annotations
from typing import Iterable, Optional, Sequence, Tuple, List, Dict, Union
import numpy as np
import pandas as pd

# FREUD backend (pip install freud-analysis)
import freud

from reaxkit.io.handlers.xmolout_handler import XmoloutHandler


# ==========================================================
# =============== Single-frame RDF backends ================
# ==========================================================

def _rdf_freud_frame(
    handler: XmoloutHandler,
    frame_index: int,
    *,
    types_a: Optional[Iterable[str]] = None,
    types_b: Optional[Iterable[str]] = None,
    r_max: Optional[float] = None,
    bins: int = 200,
    drop_first_bin: bool = True,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute the RDF for a single frame using the FREUD backend.

    Works on
    --------
    XmoloutHandler — ``xmolout``

    Parameters
    ----------
    handler : XmoloutHandler
        Parsed trajectory handler.
    frame_index : int
        Frame index to analyze.
    types_a, types_b : iterable of str, optional
        Atom types defining A–B RDF. If None, all atoms are used.
    r_max : float, optional
        Maximum radius. Defaults to half the smallest box length.
    bins : int, default=200
        Number of RDF bins.
    drop_first_bin : bool, default=True
        Drop the first bin to avoid self-counting artifacts.

    Returns
    -------
    r : np.ndarray
        Radial distances.
    g : np.ndarray
        RDF values g(r).

    Examples
    --------
    >>> r, g = _rdf_freud_frame(xh, frame_index=0, types_a=["Al"], types_b=["N"])
    """
    df = handler.dataframe()
    fr = handler.frame(int(frame_index))

    coords = fr["coords"]
    atom_types = np.asarray(fr["atom_types"], dtype=str)

    # Box from sim DF (supports non-orthogonal if available)
    row = df.iloc[int(frame_index)]
    a, b, c = float(row["a"]), float(row["b"]), float(row["c"])
    alpha = float(row["alpha"]) if "alpha" in df.columns else 90.0
    beta  = float(row["beta"])  if "beta"  in df.columns else 90.0
    gamma = float(row["gamma"]) if "gamma" in df.columns else 90.0

    box = freud.box.Box.from_box_lengths_and_angles(
        a, b, c, np.radians(alpha), np.radians(beta), np.radians(gamma)
    )

    # A/B masks
    A_mask = np.ones(len(atom_types), dtype=bool) if types_a is None else np.isin(atom_types, list(types_a))
    B_mask = np.ones(len(atom_types), dtype=bool) if types_b is None else np.isin(atom_types, list(types_b))
    A = coords[A_mask]
    B = coords[B_mask]

    if r_max is None:
        r_max = 0.5 * float(min(a, b, c))

    rdf = freud.density.RDF(bins=int(bins), r_max=float(r_max))
    rdf.compute((box, B), query_points=A)

    r = rdf.bin_centers.copy()
    g = rdf.rdf.copy()
    if drop_first_bin and len(r) > 1:
        r, g = r[1:], g[1:]
    return r, g


def _write_xyz_temp(coords: np.ndarray, types: Sequence[str]) -> str:
    import tempfile, os
    def _norm(sym: str) -> str:
        s = str(sym).strip()
        return s[:1].upper() + s[1:].lower()
    n = coords.shape[0]
    fd, path = tempfile.mkstemp(suffix=".xyz", prefix="reaxkit_rdf_")
    os.close(fd)
    with open(path, "w") as f:
        f.write(f"{n}\n")
        f.write("generated by reaxkit RDF (single-frame)\n")
        for t, (x, y, z) in zip(types, coords):
            f.write(f"{_norm(t)} {x:.9f} {y:.9f} {z:.9f}\n")
    return path


def _rdf_ovito_total_frame(
    handler: XmoloutHandler,
    frame_index: int,
    *,
    r_max: float = 4.0,
    bins: int = 200,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute the total RDF for a single frame using the OVITO backend.

    Works on
    --------
    XmoloutHandler — ``xmolout``

    Parameters
    ----------
    handler : XmoloutHandler
        Parsed trajectory handler.
    frame_index : int
        Frame index to analyze.
    r_max : float, default=4.0
        RDF cutoff radius.
    bins : int, default=200
        Number of RDF bins.

    Returns
    -------
    r : np.ndarray
        Radial distances.
    g : np.ndarray
        RDF values g(r).

    Examples
    --------
    >>> r, g = _rdf_ovito_total_frame(xh, frame_index=10)
    """
    fr = handler.frame(int(frame_index))
    coords = fr["coords"]
    types = fr["atom_types"]

    import os
    xyz = _write_xyz_temp(coords, types)

    try:
        from ovito.io import import_file
        from ovito.modifiers import CoordinationAnalysisModifier
    except Exception as e:
        raise ImportError(
            "OVITO is required for RDF analysis. "
            "Install reaxkit[viz] and run in an environment with EGL/Qt."
        ) from e

    try:
        pipe = import_file(xyz)
        pipe.modifiers.append(CoordinationAnalysisModifier(cutoff=float(r_max), number_of_bins=int(bins), partial=False))
        data = pipe.compute()
        table = data.tables["coordination-rdf"]
        arr = table.xy()
        r = np.asarray(arr[:, 0])
        g = np.asarray(arr[:, 1])
        return r, g
    finally:
        try: os.remove(xyz)
        except OSError: pass


def _rdf_ovito_partial_frame(
    handler: XmoloutHandler,
    frame_index: int,
    *,
    r_max: float = 4.0,
    bins: int = 200,
    type_a: str,
    type_b: str,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute a partial A–B RDF for a single frame using the OVITO backend.

    Works on
    --------
    XmoloutHandler — ``xmolout``

    Parameters
    ----------
    handler : XmoloutHandler
        Parsed trajectory handler.
    frame_index : int
        Frame index to analyze.
    r_max : float, default=4.0
        RDF cutoff radius.
    bins : int, default=200
        Number of RDF bins.
    type_a, type_b : str
        Atom types defining the partial RDF (e.g. ``Al``–``N``).

    Returns
    -------
    r : np.ndarray
        Radial distances.
    g : np.ndarray
        RDF values g(r).

    Examples
    --------
    >>> r, g = _rdf_ovito_partial_frame(xh, 0, type_a="Al", type_b="N")
    """
    fr = handler.frame(int(frame_index))
    coords = fr["coords"]
    types = fr["atom_types"]

    import os
    xyz = _write_xyz_temp(coords, types)

    try:
        from ovito.io import import_file
        from ovito.modifiers import CoordinationAnalysisModifier
    except Exception as e:
        raise ImportError(
            "OVITO is required for RDF analysis. "
            "Install reaxkit[viz] and run in an environment with EGL/Qt."
        ) from e

    try:
        pipe = import_file(xyz)
        pipe.modifiers.append(CoordinationAnalysisModifier(cutoff=float(r_max), number_of_bins=int(bins), partial=True))
        data = pipe.compute()
        base = "coordination-rdf"
        if base not in data.tables:
            raise KeyError(f"'{base}' table not found. Available: {list(data.tables.keys())}")
        table = data.tables[base]

        # locate the pair component
        try:
            names = list(table.y.component_names)
        except Exception:
            names = list(getattr(table.y, "components", []))
        want = f"{type_a}-{type_b}"
        if want not in names and f"{type_b}-{type_a}" in names:
            want = f"{type_b}-{type_a}"
        idx = names.index(want)

        arr = table.xy()
        r = np.asarray(arr[:, 0])
        g = np.asarray(arr[:, 1 + idx])
        return r, g
    finally:
        try: os.remove(xyz)
        except OSError: pass


# ==========================================================
# ================= RDF utilities (no smoothing) ===========
# ==========================================================

def _dominant_peak(r: np.ndarray, g: np.ndarray) -> tuple[float, float]:
    """Return the position and height of the dominant (global) RDF peak.

    Parameters
    ----------
    r : np.ndarray
        Radial distances.
    g : np.ndarray
        RDF values.

    Returns
    -------
    r_peak : float
        Radius at the global maximum of g(r).
    g_peak : float
        Height of the global maximum.

    Examples
    --------
    >>> r_peak, g_peak = _dominant_peak(r, g)
    """
    if len(g) == 0:
        return float("nan"), float("nan")
    k = int(np.argmax(g))
    return float(r[k]), float(g[k])


def _first_local_max(r: np.ndarray, g: np.ndarray) -> tuple[float, float]:
    """Return the first local maximum of an RDF curve.

    Parameters
    ----------
    r : np.ndarray
        Radial distances.
    g : np.ndarray
        RDF values.

    Returns
    -------
    r_peak : float
        Radius of the first local maximum.
    g_peak : float
        Height of the first local maximum.

    Notes
    -----
    Falls back to the dominant peak if no local maximum is found.

    Examples
    --------
    >>> r_first, g_first = _first_local_max(r, g)
    """
    n = len(g)
    for i in range(1, n - 1):
        if g[i] > g[i - 1] and g[i] > g[i + 1]:
            return float(r[i]), float(g[i])
    return _dominant_peak(r, g)


# ==========================================================
# ============== Multi-frame RDF + properties ==============
# ==========================================================

def rdf_using_freud(
    handler: XmoloutHandler,
    *,
    frames: Optional[Iterable[int]] = None,
    types_a: Optional[Iterable[str]] = None,
    types_b: Optional[Iterable[str]] = None,
    r_max: Optional[float] = None,
    bins: int = 200,
    average: bool = True,
    return_stack: bool = False,
) -> Tuple[np.ndarray, Union[np.ndarray, List[np.ndarray]]]:
    """Compute RDFs across multiple frames using the FREUD backend.

    Works on
    --------
    XmoloutHandler — ``xmolout``

    Parameters
    ----------
    frames : iterable of int, optional
        Frame indices to include.
    types_a, types_b : iterable of str, optional
        Atom types defining A–B RDF.
    average : bool, default=True
        If True, return the average RDF across frames.
    return_stack : bool, default=False
        If True and ``average=False``, return all per-frame RDFs.

    Returns
    -------
    r : np.ndarray
        Radial distances.
    g : np.ndarray or list[np.ndarray]
        Averaged RDF, last RDF, or stack of RDFs.

    Examples
    --------
    >>> r, g = rdf_using_freud(xh, frames=range(100), average=True)
    """
    df = handler.dataframe()
    if frames is None:
        frames = range(len(df))
    frames = list(frames)

    r_ref: np.ndarray | None = None
    stack: List[np.ndarray] = []

    for i in frames:
        r, g = _rdf_freud_frame(
            handler, frame_index=int(i),
            types_a=types_a, types_b=types_b, r_max=r_max, bins=bins, drop_first_bin=True
        )
        if r_ref is None:
            r_ref = r
        else:
            if len(r) != len(r_ref) or np.max(np.abs(r - r_ref)) > 1e-10:
                raise ValueError("R grids differ between frames; fix bins/r_max.")
        stack.append(g)

    if r_ref is None:
        raise ValueError("No frames selected.")

    if average:
        return r_ref, np.mean(np.vstack(stack), axis=0)
    if return_stack:
        return r_ref, stack
    return r_ref, stack[-1]


def rdf_using_ovito(
    handler: XmoloutHandler,
    *,
    frames: Optional[Iterable[int]] = None,
    r_max: float = 4.0,
    bins: int = 200,
    types_a: Optional[Iterable[str]] = None,  # if both provided => partial
    types_b: Optional[Iterable[str]] = None,
    average: bool = True,
    return_stack: bool = False,
) -> Tuple[np.ndarray, Union[np.ndarray, List[np.ndarray]]]:
    """Compute RDFs across multiple frames using the OVITO backend.

    Works on
    --------
    XmoloutHandler — ``xmolout``

    Parameters
    ----------
    frames : iterable of int, optional
        Frame indices to include.
    types_a, types_b : iterable of str, optional
        If provided, compute partial RDF.
    average : bool, default=True
        Return the average RDF across frames.

    Returns
    -------
    r : np.ndarray
        Radial distances.
    g : np.ndarray or list[np.ndarray]
        RDF data.

    Examples
    --------
    >>> r, g = rdf_using_ovito(xh, types_a=["Al"], types_b=["N"])
    """
    df = handler.dataframe()
    if frames is None:
        frames = range(len(df))
    frames = list(frames)

    do_partial = (types_a is not None) and (types_b is not None)
    if do_partial:
        ta = next(iter(types_a))
        tb = next(iter(types_b))

    r_ref: np.ndarray | None = None
    stack: List[np.ndarray] = []

    for i in frames:
        if do_partial:
            r, g = _rdf_ovito_partial_frame(handler, int(i), r_max=r_max, bins=bins, type_a=str(ta), type_b=str(tb))
        else:
            r, g = _rdf_ovito_total_frame(handler, int(i), r_max=r_max, bins=bins)

        if r_ref is None:
            r_ref = r
        else:
            if len(r) != len(r_ref) or np.max(np.abs(r - r_ref)) > 1e-10:
                raise ValueError("R grids differ between frames; fix bins/r_max.")
        stack.append(g)

    if r_ref is None:
        raise ValueError("No frames selected.")

    if average:
        return r_ref, np.mean(np.vstack(stack), axis=0)
    if return_stack:
        return r_ref, stack
    return r_ref, stack[-1]


def rdf_property_over_frames(
    handler: XmoloutHandler,
    *,
    backend: str = "ovito",
    frames: Optional[Iterable[int]] = None,
    property: str = "first_peak",  # 'first_peak' | 'dominant_peak' | 'area' | 'excess_area'
    r_max: Optional[float] = None,  # for OVITO total/partial also accepts float
    bins: int = 200,
    types_a: Optional[Iterable[str]] = None,
    types_b: Optional[Iterable[str]] = None,
) -> pd.DataFrame:
    """Compute a single RDF-derived property for each simulation frame.

    Works on
    --------
    XmoloutHandler — ``xmolout``

    Parameters
    ----------
    backend : {"ovito", "freud"}, default="ovito"
        RDF backend to use.
    property : {"first_peak", "dominant_peak", "area", "excess_area"}
        RDF-based property to compute.
    frames : iterable of int, optional
        Frame indices to include.
    types_a, types_b : iterable of str, optional
        Atom types defining partial RDF.

    Returns
    -------
    pandas.DataFrame
        One row per frame with frame index, iteration, and property values.

    Examples
    --------
    >>> df = rdf_property_over_frames(xh, property="first_peak")
    """
    df_sim = handler.dataframe()
    if frames is None:
        frames = range(len(df_sim))
    frames = list(frames)

    prop = property.strip().lower()
    allowed = {"first_peak", "dominant_peak", "area", "excess_area"}
    if prop not in allowed:
        raise ValueError(f"property must be one of {allowed}")

    rows: List[Dict[str, float]] = []

    for i in frames:
        if backend.lower() == "freud":
            r, g = _rdf_freud_frame(handler, int(i), types_a=types_a, types_b=types_b, r_max=r_max, bins=bins)
        elif backend.lower() == "ovito":
            if (types_a is not None) and (types_b is not None):
                ta = next(iter(types_a)); tb = next(iter(types_b))
                r, g = _rdf_ovito_partial_frame(handler, int(i), r_max=float(r_max or 4.0), bins=bins, type_a=str(ta), type_b=str(tb))
            else:
                r, g = _rdf_ovito_total_frame(handler, int(i), r_max=float(r_max or 4.0), bins=bins)
        else:
            raise ValueError("backend must be 'freud' or 'ovito'")

        if prop == "first_peak":
            rp, gp = _first_local_max(r, g)
            out = {"r_first_peak": rp, "g_first_peak": gp}
        elif prop == "dominant_peak":
            rp, gp = _dominant_peak(r, g)
            out = {"r_peak": rp, "g_peak": gp}
        elif prop == "area":
            out = {"area": float(np.trapezoid(g, r)) if len(r) else np.nan}
        else:  # "excess_area"
            out = {"excess_area": float(np.trapezoid(g - 1.0, r)) if len(r) else np.nan}

        rows.append({
            "frame_index": int(i),
            "iter": int(df_sim.iloc[i]["iter"]) if "iter" in df_sim.columns else int(i),
            **out,
        })

    return pd.DataFrame(rows).sort_values("frame_index").reset_index(drop=True)
