"""
Electrostatics workflow for ReaxKit.

This workflow provides tools for computing and analyzing electrostatic properties
from ReaxFF simulations, including dipole moments, polarizations, and
polarization–electric-field hysteresis behavior.

It supports:
- Dipole moment and polarization calculations for single frames, at both
  total-system and local (core-atom cluster) levels.
- Polarization–field hysteresis analysis using time-dependent electric fields
  (via fort.78 and control files), including extraction of coercive fields and
  remnant polarizations.
- Visualization of local electrostatics through 3D scatter plots and 2D
  projected heatmaps based on atomic coordinates.

The workflow integrates xmolout, fort.7, fort.78, and control files, and is
designed to bridge atomistic ReaxFF simulations with experiment-facing
electrostatic observables.
"""


from __future__ import annotations
import argparse
from pathlib import Path
from typing import Optional, Sequence, Tuple, Union

import numpy as np
import pandas as pd

from reaxkit.io.handlers.xmolout_handler import XmoloutHandler
from reaxkit.io.handlers.fort7_handler import Fort7Handler
from reaxkit.io.handlers.fort78_handler import Fort78Handler
from reaxkit.io.handlers.control_handler import ControlHandler
from reaxkit.analysis.per_file.xmolout_analyzer import get_atom_trajectories
from reaxkit.utils.media.plotter import scatter3d_points, heatmap2d_from_3d
from reaxkit.utils.frame_utils import parse_frames, resolve_indices
from reaxkit.utils.alias import resolve_alias_from_columns
from reaxkit.utils.path import resolve_output_path
from reaxkit.utils.media.plotter import single_plot

from reaxkit.utils.alias import (
    normalize_choice,
    _resolve_alias,
)
from reaxkit.analysis.composed.electrostatics_analyzer import (
    dipoles_polarizations_over_multiple_frames,
    polarization_field_analysis,
)


# -------------------------------------------------------------------------
# Tasks
# -------------------------------------------------------------------------


def _dipole_task(args: argparse.Namespace) -> int:
    """
    reaxkit elect dipole
        --xmolout xmolout --fort7 fort.7
        --frame X --scope {total,local}
        [--core Al,Mg] --export out.csv [--polarization]

    total scope:
        exports total dipole or polarization for a single frame.

    local scope:
        exports local dipole or polarization for each core atom cluster for a single frame.
        NOTE: local polarization uses bbox volume by default (fast). Use hull only if you add an option.
    """
    # Handlers
    xh = XmoloutHandler(args.xmolout)
    f7 = Fort7Handler(args.fort7)

    # Mode: dipole vs polarization
    mode = "polarization" if args.polarization else "dipole"

    # Scope/core types
    scope = args.scope
    core_types: Optional[list[str]] = None
    if scope == "local":
        if not args.core:
            raise ValueError("When --scope local is used, --core must be provided (e.g. --core Al,Mg).")
        core_types = [c.strip() for c in args.core.split(",") if c.strip()]

    # Volume method policy (per your new analyzer design)
    # - total polarization: hull (reasonable)
    # - local polarization: bbox (fast default)
    volume_method = None
    if mode == "polarization":
        volume_method = "bbox" if scope == "local" else "hull"

    # Compute for the requested single frame (by passing frames=[...])
    df = dipoles_polarizations_over_multiple_frames(
        xh,
        f7,
        frames=[args.frame],
        scope=scope,
        core_types=core_types,
        mode=mode,
        volume_method=volume_method,
    )

    workflow_name = args.kind

    # Export
    out = resolve_output_path(args.export, workflow_name)
    df.to_csv(out, index=False)
    print(f"\n[Done] Exported data to {out}")

    return 0


def _hyst_task(args: argparse.Namespace) -> int:
    """
    reaxkit elect hyst
        --xmolout xmolout --fort7 fort.7 --fort78 fort.78 --control control
        [--plot] [--save hyst.png]
        [--yaxis pol_z] [--xaxis time|field_z]
        [--aggregate mean|max|min|last]
        [--export hysteresis.csv]
        [--summary hysteresis_summary.txt]
        [--roots]

    - Uses polarization_field_analysis to build polarization + field dataset.
    - Aggregates according to --aggregate.
    - Exports aggregated joint DataFrame (if --export is given).
    - Writes coercive fields and remnant polarizations to a text file (summary).
    - If --roots is present, also prints these values to the terminal.
    - Makes a hysteresis / time-series plot according to --xaxis/--yaxis
      and plot/save flags, using alias_utils to map names.
    """
    # Handlers
    xh = XmoloutHandler(args.xmolout)
    f7 = Fort7Handler(args.fort7)
    f78 = Fort78Handler(args.fort78)
    ctrl = ControlHandler(args.control)

    # Run basic polarization vs field analysis; keep x/y for roots as field_z vs P_z
    full_df, agg_df, coercive_fields, remnant_pols = polarization_field_analysis(
        xh,
        f7,
        f78,
        ctrl,
        field_var="field_z",
        aggregate=args.aggregate,
        x_variable="field_z",
        y_variable="P_z (uC/cm^2)",
    )

    # ------------------------------------------------------------------
    # Optionally add time information to full_df (for plotting vs time)
    # ------------------------------------------------------------------
    try:
        sim_df = xh.dataframe()
        iter_col = _resolve_alias(sim_df, "iter")
        time_col_name = _resolve_alias(sim_df, "time")
        sim_small = sim_df[[iter_col, time_col_name]].copy()
        sim_small = sim_small.rename(columns={iter_col: "iter", time_col_name: "time"})
        full_df = full_df.merge(sim_small, on="iter", how="left")
    except Exception:
        # If no time is resolvable, we simply skip it; plotting vs time may fail later.
        pass

    # ------------------------------------------------------------------
    # Export aggregated joint DataFrame
    # ------------------------------------------------------------------
    workflow_name = args.kind
    if args.export:
        # Export aggregated data (agg_df)
        out_csv = resolve_output_path(args.export, workflow_name)
        agg_df.to_csv(out_csv, index=False)
        print(f"\n[Done] Exported aggregated joint hysteresis data to {out_csv}")

        # Export full dataframe in same directory
        full_save_dir = out_csv.parent  # Directory where --export was saved
        full_save_dir.mkdir(parents=True, exist_ok=True)  # Ensure directory exists
        full_save_path = full_save_dir / "hysteresis_full_data.csv"
        full_df.to_csv(full_save_path, index=False)
        print(f"[Done] Exported full hysteresis dataset to {full_save_path}\n")

    # ------------------------------------------------------------------
    # Write coercive fields & remnant polarizations to a text file
    # ------------------------------------------------------------------
    summary_path = Path(args.summary or "hysteresis_summary.txt")

    def _fmt_list(values) -> str:
        if not values:
            return "None found"
        return ", ".join(f"{v:.6g}" for v in values)

    with summary_path.open("w", encoding="utf-8") as f:
        f.write("Hysteresis Analysis Summary\n")
        f.write("===========================\n\n")
        f.write("Coercive fields (where polarization crosses zero vs field_z)\n")
        f.write("Units: MV/cm\n")
        f.write(f"Values: {_fmt_list(coercive_fields)}\n\n")
        f.write("Remnant polarizations (where field_z crosses zero vs P_z)\n")
        f.write("Units: µC/cm^2\n")
        f.write(f"Values: {_fmt_list(remnant_pols)}\n")

    print(f"[Done] Wrote hysteresis summary to {summary_path}")

    # Also print to terminal if --roots is requested
    if args.roots:
        print("\n[Hysteresis roots]")
        print("  Coercive fields (MV/cm):", _fmt_list(coercive_fields))
        print("  Remnant polarizations (µC/cm^2):", _fmt_list(remnant_pols))

    # ------------------------------------------------------------------
    # Plot: use alias_utils to map --xaxis / --yaxis
    # - If xaxis == time → use per-frame data (full_df)
    # - Else             → use aggregated data (agg_df) for hysteresis
    # ------------------------------------------------------------------
    if args.plot or args.save:
        # Canonical keys
        canonical_x = normalize_choice(args.xaxis or "field_z", domain="xaxis")
        canonical_y = normalize_choice(args.yaxis or "P_z (uC/cm^2)", domain="yaxis")

        # Choose DataFrame for plotting
        if canonical_x == "time":
            df_plot = full_df
        else:
            df_plot = agg_df

        # Resolve actual column names present in df_plot
        x_col = _resolve_alias(df_plot, canonical_x)
        y_col = _resolve_alias(df_plot, canonical_y)

        x = df_plot[x_col]
        y = df_plot[y_col]

        title = f"{y_col} vs {x_col}"
        xlabel = x_col
        ylabel = y_col

        single_plot(x, y, title=title, xlabel=xlabel, ylabel=ylabel, save=args.save, figsize=(6,4))

    return 0

# ------------------------- internals -------------------------

def _indices_from_spec(xh: XmoloutHandler, frames_spec) -> list[int]:
    """Turn a frames spec (slice or list) into concrete frame indices for xh."""
    return resolve_indices(xh, frames=frames_spec, iterations=None, step=None)


def _local_pol_with_coords(
    xh: XmoloutHandler,
    f7: Fort7Handler,
    *,
    core_types: Sequence[str],
    frames_spec,
    mode: str = "polarization",  # or "dipole"
) -> pd.DataFrame:
    """
    Build a DataFrame with local electrostatics + core-atom coordinates:

    columns include:
      frame_index, iter, core_atom_type, core_atom_id,
      x, y, z,
      mu_x (debye), mu_y (debye), mu_z (debye),
      P_x (uC/cm^2), P_y (uC/cm^2), P_z (uC/cm^2), volume (angstrom^3)
    """
    # 1) local dipole/polarization over all frames
    df_local = dipoles_polarizations_over_multiple_frames(
        xh,
        f7,
        scope="local",
        core_types=core_types,
        mode=mode,
    )
    if df_local.empty:
        raise ValueError("No local polarization/dipole data found. Check core types and files.")

    # Restrict to requested frames
    idx_list = _indices_from_spec(xh, frames_spec)
    df_local = df_local[df_local["frame_index"].isin(idx_list)].copy()
    if df_local.empty:
        raise ValueError("No local data in the requested frames.")

    # 2) coordinates of atoms from xmolout (long format)
    traj = get_atom_trajectories(
        xh,
        frames=idx_list,
        every=1,
        atoms=None,
        atom_types=None,
        dims=("x", "y", "z"),
        format="long",
    )
    if traj.empty:
        raise ValueError("No trajectory data returned for requested frames.")

    traj = traj[["frame_index", "atom_id", "x", "y", "z"]].copy()
    traj = traj.rename(columns={"atom_id": "core_atom_id"})

    # 3) join local pol/dipole with core-atom coordinates
    df = pd.merge(
        df_local,
        traj,
        on=["frame_index", "core_atom_id"],
        how="inner",
    )
    if df.empty:
        raise ValueError("Could not match local polarization entries to core atom coordinates.")

    return df


# ========================= 3D SCATTER =========================

def _local_pol_plot3d_task(args: argparse.Namespace) -> int:
    """
    3D scatter of local polarization/dipole (one point per core atom) across frames.

    Example:
      reaxkit elect plot3d --core Al,Mg --component pol_z --frames 0:200:20 --save figs/local_pol3d
    """
    xh = XmoloutHandler(args.xmolout)
    f7 = Fort7Handler(args.fort7)

    frames_spec = parse_frames(args.frames)
    core_types = [c.strip() for c in args.core.split(",") if c.strip()]

    df = _local_pol_with_coords(
        xh,
        f7,
        core_types=core_types,
        frames_spec=frames_spec,
        mode="polarization" if args.polarization else "dipole",
    )

    # Resolve which column to color by (mu_z, P_z, etc.) using alias map
    col = resolve_alias_from_columns(df.columns, args.component)
    if col is None:
        raise ValueError(
            f"Component '{args.component}' not found. "
            f"Available columns include: {list(df.columns)[:12]} ..."
        )

    # Output dir
    save_dir = Path(args.save) if args.save else None
    if save_dir:
        save_dir.mkdir(parents=True, exist_ok=True)

    # Per-frame 3D plot
    for fi in sorted(df["frame_index"].unique()):
        sub = df[df["frame_index"] == fi]
        coords = sub[["x", "y", "z"]].to_numpy(float)
        vals = sub[col].to_numpy(float)

        # skip frames with all-NaN or empty
        if coords.size == 0 or not np.isfinite(vals).any():
            continue

        vmin = args.vmin if args.vmin is not None else float(np.nanmin(vals))
        vmax = args.vmax if args.vmax is not None else float(np.nanmax(vals))

        title = f"{col}_local_3D_frame_{fi}"
        scatter3d_points(
            coords,
            vals,
            title=title,
            s=args.size,
            alpha=args.alpha,
            cmap=args.cmap,
            vmin=vmin,
            vmax=vmax,
            elev=args.elev,
            azim=args.azim,
            save=(save_dir / f"{title}.png" if save_dir else None),
            show_message=False,
        )

    if args.save:
        print(f"[Done] All 3D local polarization plots saved in {args.save}")
    return 0


# ========================= 2D HEATMAP =========================

def _parse_bins(bins: str) -> Union[int, Tuple[int, int]]:
    if "," in bins:
        nx, ny = [int(x) for x in bins.split(",")]
        return (nx, ny)
    return int(bins)


def _local_pol_heatmap2d_task(args: argparse.Namespace) -> int:
    """
    2D heatmap of local polarization/dipole on core-atom positions.

    Example:
      reaxkit elect heatmap2d --core Al,Mg --component pol_z --plane xz --bins 40 --frames 0:200:20 --save figs/local_pol_heat
    """
    xh = XmoloutHandler(args.xmolout)
    f7 = Fort7Handler(args.fort7)

    frames_spec = parse_frames(args.frames)
    core_types = [c.strip() for c in args.core.split(",") if c.strip()]

    df = _local_pol_with_coords(
        xh,
        f7,
        core_types=core_types,
        frames_spec=frames_spec,
        mode="polarization" if args.polarization else "dipole",
    )

    col = resolve_alias_from_columns(df.columns, args.component)
    if col is None:
        raise ValueError(
            f"Component '{args.component}' not found. "
            f"Available columns include: {list(df.columns)[:12]} ..."
        )

    bins = _parse_bins(args.bins)

    save_dir = Path(args.save) if args.save else None
    if save_dir:
        save_dir.mkdir(parents=True, exist_ok=True)

    for fi in sorted(df["frame_index"].unique()):
        sub = df[df["frame_index"] == fi]
        coords = sub[["x", "y", "z"]].to_numpy(float)
        vals = sub[col].to_numpy(float)

        if coords.size == 0 or not np.isfinite(vals).any():
            continue

        m = np.isfinite(vals)
        coords = coords[m]
        vals = vals[m]

        vmin = args.vmin if args.vmin is not None else float(np.nanmin(vals))
        vmax = args.vmax if args.vmax is not None else float(np.nanmax(vals))

        title = f"{col}_local_{args.plane}_frame_{fi}"
        out_path = (save_dir / f"{title}.png") if save_dir else None

        heatmap2d_from_3d(
            coords,
            vals,
            plane=args.plane,
            bins=bins,
            agg=args.agg,
            vmin=vmin,
            vmax=vmax,
            cmap=args.cmap,
            title=title,
            save=out_path,
            show_message=False,
        )

    if args.save:
        print(f"[Done] All 2D local polarization heatmaps saved in {args.save}")
    return 0


# -------------------------------------------------------------------------
# Registration
# -------------------------------------------------------------------------

def register_tasks(subparsers: argparse._SubParsersAction) -> None:
    """
    Register coordination-related CLI subcommands.

    This function defines the task-level interface for the
    `reaxkit file` workflow and attaches its tasks (i.e., subcommands).

    Each task may share common input arguments and defines task-specific options as needed.
    """
    # ---------------------- dipole ----------------------
    p_dip = subparsers.add_parser(
        "dipole",
        help="Compute dipole moment or polarization for a single frame\n",
        description=(
            "Examples:\n"
            "  reaxkit elect dipole --frame 10 --scope total --export dipole_pol_total_frame10.csv --polarization \n"
            "  reaxkit elect dipole --frame 10 --scope local --core Al --export dipole_pol_local_frame10.csv --polarization \n"
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    p_dip.add_argument("--xmolout", default="xmolout", help="Path to xmolout file")
    p_dip.add_argument("--fort7", default="fort.7", help="Path to fort.7 file")
    p_dip.add_argument("--frame", type=int, required=True, help="0-based frame index in xmolout")
    p_dip.add_argument("--scope", choices=["total", "local"],
                       default="total", help="Electrostatics scope: total or local")
    p_dip.add_argument("--core", default=None,
                       help="Comma-separated core atom types for local scope (e.g. Al,Mg)")
    p_dip.add_argument("--export", required=True, help="CSV file to export the dipole/polarization data")
    p_dip.add_argument("--polarization", action="store_true",
                       help="If present, compute polarization instead of dipole")
    p_dip.set_defaults(_run=_dipole_task)

    # ---------------------- hysteresis ----------------------
    p_hyst = subparsers.add_parser(
        "hyst",
        help="Polarization-field hysteresis analysis\n",
        description=(
            "Examples:\n"
            "  reaxkit elect hyst --plot --yaxis pol_z --xaxis field_z --aggregate mean --roots"
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    p_hyst.add_argument("--xmolout", default="xmolout", help="Path to xmolout file")
    p_hyst.add_argument("--fort7", default="fort.7", help="Path to fort.7 file")
    p_hyst.add_argument("--fort78", default="fort.78", help="Path to fort.78 file")
    p_hyst.add_argument("--control", default="control", help="Path to control file")
    p_hyst.add_argument("--plot", action="store_true", help="Show the hysteresis or time-series plot")
    p_hyst.add_argument("--save", default='hysteresis_aggregated.png',
                        help="If set, save plot to file (e.g. hyst.png)")
    p_hyst.add_argument("--export", default='hysteresis_aggregated.csv',
                        help="CSV file to export aggregated hysteresis data")
    p_hyst.add_argument("--summary", default="hysteresis_summary.txt",
                        help="Text file to write coercive fields and remnant polarizations")
    p_hyst.add_argument("--yaxis", default="pol_z", help="Quantity for y-axis (e.g. pol_z, mu_z, time, P_z)")
    p_hyst.add_argument("--xaxis", default="field_z", help="Quantity for x-axis (e.g. field_z, time, iter)")
    p_hyst.add_argument("--aggregate", choices=["mean", "max", "min", "last"],
                        default="mean", help="Aggregation method")

    p_hyst.add_argument("--roots", action="store_true", help="Also print coercive and remnant values to terminal")
    p_hyst.set_defaults(_run=_hyst_task)

    # ---------------------- 3D local polarization scatter ----------------------
    p3d = subparsers.add_parser(
        "plot3d",
        help="3D scatter of local polarization/dipole on core-atom positions\n",
        description=(
            "Examples:\n"
            "  reaxkit elect plot3d --core Al --component mu_z --frames 0:3:1 "
            "--save reaxkit_outputs/elect/local_mu_3d/"
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    p3d.add_argument("--xmolout", default="xmolout", help="Path to xmolout file")
    p3d.add_argument("--fort7", default="fort.7", help="Path to fort.7 file")
    p3d.add_argument("--core", required=True, help="Comma-separated core atom types (e.g. Al,Mg)")
    p3d.add_argument("--component", default="pol_z",
                     help="Which component to color by (e.g. pol_z, P_z (uC/cm^2), mu_z)")
    p3d.add_argument("--frames", default=None, help='Frames: "0,10,20" or "0:100:5"')
    p3d.add_argument("--polarization", action="store_true",
                     help="Use polarization components (P_x/y/z) instead of dipole")
    p3d.add_argument("--save", default=None, help="Directory to save PNGs (one per frame)")
    p3d.add_argument("--vmin", type=float, default=None, help="Color scale min (auto if not set)")
    p3d.add_argument("--vmax", type=float, default=None, help="Color scale max (auto if not set)")
    p3d.add_argument("--size", type=float, default=20.0, help="Marker size")
    p3d.add_argument("--alpha", type=float, default=0.9, help="Marker transparency")
    p3d.add_argument("--cmap", default="coolwarm", help="Matplotlib colormap")
    p3d.add_argument("--elev", type=float, default=22.0, help="3D view elevation")
    p3d.add_argument("--azim", type=float, default=38.0, help="3D view azimuth")
    p3d.set_defaults(_run=_local_pol_plot3d_task)

    # ---------------------- 2D local polarization heatmap ----------------------
    p2d = subparsers.add_parser(
        "heatmap2d",
        help="2D heatmap of local polarization/dipole on core-atom positions",
        description=(
            "Examples:\n"
            "  reaxkit elect heatmap2d --core Al --component mu_z --plane xz --bins 10 --frames 0:2:1 "
            "--save reaxkit_outputs/elect/local_mu_2d"
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    p2d.add_argument("--xmolout", default="xmolout", help="Path to xmolout file")
    p2d.add_argument("--fort7", default="fort.7", help="Path to fort.7 file")
    p2d.add_argument("--core", required=True, help="Comma-separated core atom types (e.g. Al,Mg)")
    p2d.add_argument("--component", default="pol_z",
                     help="Which component to aggregate (e.g. pol_z, P_z (uC/cm^2), mu_z)")
    p2d.add_argument("--frames", default=None, help='Frames: "0,10,20" or "0:100:5"')
    p2d.add_argument("--polarization", action="store_true",
                     help="Use polarization components (P_x/y/z) instead of dipole")

    # Projection/grid
    p2d.add_argument("--plane", default="xy", choices=["xy", "xz", "yz"], help="Projection plane")
    p2d.add_argument("--bins", default="40", help='Grid bins: "N" or "Nx,Ny" (e.g., "10,25")')
    p2d.add_argument("--agg", default="mean", help="Aggregation: mean|max|min|sum|count")

    # Viz
    p2d.add_argument("--vmin", type=float, default=None, help="Color scale min (auto if not set)")
    p2d.add_argument("--vmax", type=float, default=None, help="Color scale max (auto if not set)")
    p2d.add_argument("--cmap", default="viridis", help="Matplotlib colormap")
    p2d.add_argument("--save", default=None, help="Directory to save PNGs (one per frame)")
    p2d.set_defaults(_run=_local_pol_heatmap2d_task)


