# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright (C) 2026 Suturina Group

"""Predict paramagnetic shifts and spectra from input data.

Loads molecular, susceptibility, and experimental inputs, computes shifts, and
writes tables and plots for selected temperatures.
"""

import copy
import logging
import os
import re
from pathlib import Path

import numpy as np

# Application layer
from simpnmr.app.loaders.dia_load import load_diamagnetic_shifts
from simpnmr.app.loaders.elstate_load import load_electronic_state
from simpnmr.app.loaders.exp_load import load_experiments
from simpnmr.app.loaders.hfc_load import load_hyperfines
from simpnmr.app.loaders.labels_load import load_chem_labels_from_csv
from simpnmr.app.loaders.mol_load import load_base_molecule
from simpnmr.app.loaders.paramag_centre_load import load_paramagnetic_centre
from simpnmr.app.loaders.sh_load import (
    load_g_tensor_ab_initio,
    load_g_tensor_dft,
)
from simpnmr.app.loaders.susc_load import load_susceptibilities
from simpnmr.app.params.options import PredictRunOptions
from simpnmr.app.policies.relax import resolve_relaxation_conditions
from simpnmr.app.policies.susc import resolve_susceptibility_source

# Core / domain
from simpnmr.core.const.gammas import NUCLEAR_GAMMAS
from simpnmr.core.const.physics import EGAMMA
from simpnmr.core.conv.ang_to_freq import angstrom_to_mhz
from simpnmr.core.domain.mol import Molecule
from simpnmr.core.relaxation.eval import evaluate_relaxation_rates

# Tools
from simpnmr.core.util import transform as tfm
from simpnmr.core.util.strings import remove_numbers

# IO layer
from simpnmr.io.csv.mol import save_molecule_to_csv
from simpnmr.io.csv.peaks import save_peak_data_to_csv
from simpnmr.io.csv.spec import read_spectrum
from simpnmr.io.csv.susc import save_susc
from simpnmr.io.qc.backends.orca.geom import read_orca5_output_xyz  # TODO: remove
from simpnmr.io.xyz import xyz_write

# Visualisation
from simpnmr.viz.plots.orb_dep import plot_orbital_shift_distance_dependence
from simpnmr.viz.plots.shifts import plot_shift_contrib, plot_shift_spread
from simpnmr.viz.plots.spect import plot_pred_spectrum, plot_raw_deconv_pred
from simpnmr.viz.style.theme import apply_profile

logger = logging.getLogger(__name__)


def run_predict(config, options: PredictRunOptions | None = None) -> int:
    """Run pNMR prediction from a YAML configuration file.

    Args:
        config: Prediction configuration loaded from YAML.

    Returns:
        Exit code: 0 on success.
    """

    # TODO(policy): introduce a unified InputSpec policy (file backend/format/section)
    #               to resolve HFC/SH/Susceptibility readers in one place and avoid
    #               duplicated format detection across pipelines and loaders.

    # Make output directory and file
    os.makedirs(config.project_name, exist_ok=True)

    if options is None:
        raise ValueError("PredictRunOptions is required")

    # Build the resolved plotting contract for this run.
    spec = apply_profile(options.runtime.plot_profile)

    # Load Molecule
    base_molecule = load_base_molecule(config)

    # Load canonical paramagnetic centre into the molecule domain container
    base_molecule = load_paramagnetic_centre(
        molecule=base_molecule,
        paramagnetic_centre=config.hyperfine_paramagnetic_centre,
    )

    # Load DFT g-tensor (if available)
    base_molecule.sh.g_tensor_dft = load_g_tensor_dft(
        config=config,
    )

    # Load ab-initio g-tensor (if available)
    base_molecule.sh.g_tensor_ab_initio = load_g_tensor_ab_initio(
        config=config,
    )

    # Load Hyperfines
    base_molecule = load_hyperfines(
        molecule=base_molecule,
        config=config,
    )

    # Load Electronic State
    base_molecule.electronic = load_electronic_state(
        spin_S=config.spin_S,
        orbit_L=config.orbit,
        total_J=config.total_momentum_J,
        hyperfine_file=config.hyperfine_file if config.spin_S is None else None,
        hyperfine_method=config.hyperfine_method if config.spin_S is None else None,
    )

    # Resolve susceptibility source for downstream operations when an explicit
    # susceptibility file is available.
    backend, section = None, None
    if config.susceptibility_file is not None:
        backend, section = resolve_susceptibility_source(
            config.susceptibility_file,
            config.susceptibility_format,
        )

    # TODO(app): Temporary ORCA-only chi-source geometry load for prediction.
    # Move this into the appropriate loader/builder layer once the chi-source
    # geometry flow is formalized outside the pipeline.
    if backend == "orca":
        try:
            base_molecule.chi_source_labels, base_molecule.chi_source_coords = (
                read_orca5_output_xyz(config.susceptibility_file)
            )
        except Exception as exc:
            logger.warning(
                "Failed to read chi-source geometry from ORCA susceptibility file: %s",
                exc,
            )

    # Load magnetic susceptibility objects.
    suscs = load_susceptibilities(
        config.susceptibility_file,
        config.susceptibility_format,
        electronic=base_molecule.electronic,
        g_tensor=base_molecule.sh.g_tensor_ab_initio,
    )

    suscs = [
        susc for susc in suscs if susc.temperature in config.susceptibility_temperatures
    ]

    if not suscs:
        raise ValueError("No susceptibility data found for specified temperature(s)")

    # Load chemical labels
    if len(config.chem_labels_file):
        al_to_cl, al_to_cml = load_chem_labels_from_csv(config.chem_labels_file)
        base_molecule.apply_chem_labels(al_to_cl, al_to_cml)

        # Save xyz file with chemical labels for chemcraft
        xyz_write.save_chemcraft_xyz(
            file_name=os.path.join(config.project_name, "chemcraft_structure.xyz"),
            labels=base_molecule.labels,
            coords=base_molecule.coords,
            chem_labels={nuc.label: nuc.chem_label for nuc in base_molecule.nuclei},
        )

    # Save xyz file with chemical labels for chemcraft
    xyz_write.save_xyz(
        file_name=os.path.join(config.project_name, "structure.xyz"),
        labels=base_molecule.labels,
        coords=base_molecule.coords,
        comment=f"Structure from {config.hyperfine_file}",
    )

    # Load diamagnetic shift file
    if len(config.diamagnetic_file):
        dia_by_key, key_kind, ref_avg_by_label_nn = load_diamagnetic_shifts(
            file_name=config.diamagnetic_file,
            file_type=config.diamagnetic_method,
            ref_file_name=config.diamagnetic_ref_file,
            ref_file_type=config.diamagnetic_ref_method,
        )
        base_molecule.apply_diamagnetic_shifts(
            dia_by_key=dia_by_key,
            key_kind=key_kind,
            ref_avg_by_label_nn=ref_avg_by_label_nn,
        )

    # Load experimental data from file into list of experiment objects
    if len(config.experiment_files):
        experiments = load_experiments(config.experiment_files)
        for susc, exp in zip(suscs, experiments):
            if susc.temperature != exp.temperature:
                logger.warning(
                    "Mismatch in Susceptibility (%.2f K) and "
                    "Experimental (%.2f K) temperatures. "
                    "Proceeding with susceptibility temperature (%.2f K) "
                    "as the active calculation temperature.",
                    susc.temperature,
                    exp.temperature,
                    susc.temperature,
                )
            if re.sub("[0-9]", "", exp.isotope) not in config.nuclei_include:
                logger.warning(
                    "Experimental isotope (%s) not requested in input file (%s)",
                    exp.isotope,
                    config.nuclei_include,
                )
    else:
        experiments = [None] * len(suscs)

    if len(config.experiment_spectrum_files):
        for experiment, spectrum in zip(experiments, config.experiment_spectrum_files):
            if experiment is None:
                continue
            spectrum_array = read_spectrum(spectrum)
            experiment.exp_reference = config.experiment_exp_reference
            experiment.spectrum = spectrum_array

    # Rotationally average hyperfines of user selected nuclei:
    if len(config.hyperfine_average):
        base_molecule.average_hyperfine(config.hyperfine_average)

    # TODO(app): Temporary ORCA-specific chi-frame wiring.
    # Move chi-source geometry loading, alignment susceptibility selection, and
    # frame preparation out of predict.py into the appropriate loader/builder
    # layer once chi-frame domain flow is formalized.

    # Use the first loaded susceptibility for chi-frame rotation
    rotation_susc = suscs[0]

    # Rotate hyperfine tensors from DFT frame into chi eigenframe (if provided)
    if backend == "orca":
        if "dft" in config.hyperfine_method:
            if base_molecule.chi_source_coords is None:
                logger.warning(
                    "Chi-source geometry is unavailable; skipping chi-frame rotation"
                )
            else:
                rot_mat, _ = tfm.get_rotation_and_transformation(
                    chi_tensor=rotation_susc.tensor,
                    temperature=rotation_susc.temperature,
                    chi_source_coords=base_molecule.chi_source_coords,
                    dft_coords=base_molecule.coords,
                )
                base_molecule.apply_frame_rotation(rot_mat)

                # Rotate DFT g-tensor into chi frame
                if base_molecule.sh.g_tensor_dft is not None:
                    base_molecule.sh.g_tensor_dft = (
                        rot_mat @ base_molecule.sh.g_tensor_dft @ rot_mat.T
                    )

                # Rotate chi-source coords into chi eigenframe and save the
                # transformed susceptibility-source structure.
                if base_molecule.chi_source_labels is not None:
                    tfm.rotate_coords_to_chi_frame(
                        config.project_name,
                        chi_tensor=rotation_susc.tensor,
                        chi_source_labels=base_molecule.chi_source_labels,
                        chi_source_coords=base_molecule.chi_source_coords,
                    )

    _terms = ["pc", "fc", "d"]

    if config.hyperfine_method == "pdip":
        _terms.pop(_terms.index("fc"))
    if not config.diamagnetic_file:
        _terms.pop(_terms.index("d"))

    # Create a molecule object which accompanies each experiment object
    molecules = [copy.deepcopy(base_molecule) for _ in range(len(experiments))]

    # Update susceptibility tensor of Molecule using model
    for molecule, susc, experiment in zip(molecules, suscs, experiments):
        molecule.susc = susc

        # Apply relaxation linewidth (relaxation-aware when inputs are available).
        _apply_relaxation_linewidths(config, molecule, experiment)

        # Calculate shifts using new susceptibility tensor and rotated hyperfines
        molecule.calculate_shifts()

        # Calculate average shifts
        molecule.average_shifts()

        # Plot theoretical shifts
        with spec.context():
            # Spread
            plot_shift_spread(
                molecule,
                experiment=experiment,
                spec=spec,
                save=True,
                show=options.runtime.show_plots,
                terms=_terms,
                save_name=os.path.join(
                    config.project_name,
                    f"pred_shift_spread_{molecule.susc.temperature:.2f}_K",
                ),
                verbose=True,
                window_title=f"Spread of predicted shifts at {susc.temperature:.2f} K",
                order="descending",
            )

            # Bar chart for means
            plot_shift_contrib(
                molecule,
                experiment=experiment,
                spec=spec,
                save=True,
                show=options.runtime.show_plots,
                save_name=os.path.join(
                    config.project_name,
                    f"pred_mean_components_{molecule.susc.temperature:.2f}_K",
                ),
                verbose=True,
                window_title=(
                    f"Predicted mean shifts and components at {susc.temperature:.2f} K"
                ),
                order="descending",
            )

            if molecule.metadata.get("hyperfine", {}).get("orbital_contribution") == (
                "available"
            ):
                plot_orbital_shift_distance_dependence(
                    molecule,
                    spec=spec,
                    save=True,
                    show=options.runtime.show_plots,
                    save_name=os.path.join(
                        config.project_name,
                        f"pred_orbital_distance_dependence_{molecule.susc.temperature:.2f}_K",
                    ),
                    verbose=True,
                    window_title=(
                        f"Orbital shift distance dependence at {susc.temperature:.2f} K"
                    ),
                    order="ascending",
                )

        shift_range = [
            np.min([nuc.shift.avg for nuc in molecule.nuclei]),
            np.max([nuc.shift.avg for nuc in molecule.nuclei]),
        ]

        with spec.context():
            if len(config.experiment_files):
                plot_raw_deconv_pred(
                    molecule=molecule,
                    isotope=molecule.nuclei[0].isotope,
                    shift_range=shift_range,
                    experiment=experiment,
                    spec=spec,
                    save=True,
                    show=options.runtime.show_plots,
                    save_name=os.path.join(
                        config.project_name,
                        f"pred_and_exp_spectrum_{molecule.susc.temperature:.2f}_K",
                    ),
                )

            plot_pred_spectrum(
                molecule,
                isotope=molecule.nuclei[0].isotope,
                shift_range=shift_range,
                spec=spec,
                save=True,
                show=options.runtime.show_plots,
                save_name=os.path.join(
                    config.project_name,
                    f"pred_spectrum_{molecule.susc.temperature:.2f}_K",
                ),
            )

    # TODO If more than one temperature, then make a stacked plot of spectra

    # Save susceptibility data to file
    save_susc(
        molecules,
        os.path.join(config.project_name, "susceptibility_tensor.csv"),
        comment=(
            "Data from spin-only fallback (no susceptibility file)"
            if config.susceptibility_file is None
            else "Data from {} ({})".format(
                Path(config.susceptibility_file).name,
                config.susceptibility_format
                if config.susceptibility_format is not None
                else (f"orca_{section}" if backend == "orca" else backend),
            )
        ),
        susc_units=getattr(config, "susc_units", "A3"),
    )

    # Write shift and peak data to file
    for molecule in molecules:
        save_molecule_to_csv(
            molecule=molecule,
            file_name=os.path.join(
                config.project_name,
                f"hyperfines_and_shifts_{molecule.susc.temperature:.2f}_K.csv",
            ),
            delimiter=options.runtime.csv_delimiter,
            comment=f"T = {molecule.susc.temperature:.2f} K",
            verbose=True,
        )

        save_peak_data_to_csv(
            molecule=molecule,
            file_name=os.path.join(
                config.project_name,
                f"peak_data_{molecule.susc.temperature:.2f}_K.csv",
            ),
            comment=f"T = {molecule.susc.temperature:.2f} K",
            verbose=True,
        )

    return 0


def _apply_relaxation_linewidths(
    config,
    base_molecule: Molecule,
    experiment,
):
    """
    Apply linewidths using a user-specified relaxation model.

    This function updates `base_molecule` in-place by storing the computed
    relaxation evaluation in the domain object and by setting `nuc.shift.lw`
    when relaxation inputs are provided in the config.

    Args:
        config (PredictConfig): Prediction configuration containing relaxation
            settings and physical parameters.
        base_molecule (Molecule): Molecule instance to update in-place.
        experiment: Optional experiment object used as fallback when relaxation
            temperature and magnetic field are not explicitly overridden in the
            config.

    Returns:
        None
    """

    if not getattr(config, "relaxation_model", None):
        logger.warning(
            "No relaxation model specified — linewidths will be fixed at 1 ppm"
        )
        base_molecule.relaxation = None
        return

    temperature, magnetic_field_tesla = resolve_relaxation_conditions(
        config,
        experiment,
    )

    if magnetic_field_tesla is None or temperature is None:
        base_molecule.relaxation = None
        return

    # Solomon linewidths if relaxation model is SBM
    nuclei_labels = (
        config.nuclei_include
        if isinstance(config.nuclei_include, list)
        else [config.nuclei_include]
    )

    # Use all nuclei in the molecule that match the requested element(s)
    nuclei_coords = {
        nuc.label: nuc.coord
        for nuc in base_molecule.nuclei
        if remove_numbers(nuc.label) in nuclei_labels
    }
    B0 = magnetic_field_tesla

    # Build Aiso, gamma and omega dictionaries for selected nuclei
    # Converts nuclear gyromagnetic ratios from MHz/T to rad/s/T
    # and multiplies Aiso by 1e6 to convert from MHz to Hz

    if config.hyperfine_method == "pdip":
        # In point-dipole (pdip) model, contact hyperfine A_iso = 0 for all nuclei.
        A_iso_dict = {label: 0.0 for label in nuclei_coords}
    else:
        # Convert domain A_iso (ppm Å^-3) back to MHz for relaxation formulas.
        # Note: conversion depends on the nuclear gyromagnetic ratio for each nucleus.
        A_iso_dict_MHz = {
            nuc.label: float(
                angstrom_to_mhz(
                    1.0 / 3.0 * np.trace(nuc.A.fc),
                    nuclear_gamma=NUCLEAR_GAMMAS[remove_numbers(nuc.label)],
                )
            )
            for nuc in base_molecule.nuclei
            if nuc.label in nuclei_coords
        }

        # Convert MHz -> Hz for relaxation routines.
        A_iso_dict = {label: val_mhz * 1e6 for label, val_mhz in A_iso_dict_MHz.items()}

    gamma_I_dict = {
        label: NUCLEAR_GAMMAS[remove_numbers(label)] * 2 * np.pi * 1e6
        for label in nuclei_coords
    }
    omega_I_dict = {label: gamma_I_dict[label] * B0 for label in nuclei_coords}
    omega_S = EGAMMA * B0 * 2 * np.pi * 1e6
    tau_c1 = 1 / ((1 / config.relaxation_tR) + (1 / config.relaxation_T1e))
    tau_c2 = 1 / ((1 / config.relaxation_tR) + (1 / config.relaxation_T2e))
    tau_e1 = config.relaxation_T1e
    tau_e2 = config.relaxation_T2e
    tau_R = config.relaxation_tR

    # Load electronic states
    spin = base_molecule.electronic.spin_S
    orbit = base_molecule.electronic.orbit_L
    total_momentum_J = base_molecule.electronic.total_J

    relaxation_eval = evaluate_relaxation_rates(
        relaxation_model=config.relaxation_model,
        nuclei_coords=nuclei_coords,
        electron_coords=base_molecule.paramagnetic_centre,
        gamma_I_dict=gamma_I_dict,
        omega_I_dict=omega_I_dict,
        omega_S=omega_S,
        spin=spin,
        orbit=orbit,
        total_momentum_J=total_momentum_J,
        A_iso_dict=A_iso_dict,
        temperature=temperature,
        tau_R=tau_R,
        tau_c1=tau_c1,
        tau_c2=tau_c2,
        tau_e1=tau_e1,
        tau_e2=tau_e2,
        compute_r1=True,
        compute_r2=True,
    )

    # Persist the computed relaxation evaluation on the molecule domain object.
    base_molecule.relaxation = relaxation_eval

    rates_r1 = base_molecule.relaxation.r1.total
    rates_r2 = base_molecule.relaxation.r2.total

    if rates_r1 is None or rates_r2 is None:
        raise ValueError("Shared relaxation evaluator returned incomplete R1/R2 rates")

    # Group R2 rates by chemical label for linewidth averaging in Hz.
    r2_by_chem_label = {}
    for nuc in base_molecule.nuclei:
        if nuc.label not in rates_r2:
            continue
        if nuc.chem_label not in r2_by_chem_label:
            r2_by_chem_label[nuc.chem_label] = []
        r2_by_chem_label[nuc.chem_label].append(rates_r2[nuc.label])

    # Calculate average linewidths for each chemical label (Hz)
    avg_lw_by_chem_label = {
        chem_label: np.mean([rate / np.pi for rate in rate_list])
        for chem_label, rate_list in r2_by_chem_label.items()
    }

    for nuc in base_molecule.nuclei:
        if nuc.chem_label in avg_lw_by_chem_label:
            nuc.shift.lw = (
                avg_lw_by_chem_label[nuc.chem_label]
                / (abs(omega_I_dict[nuc.label]) / (2 * np.pi))
                * 1e6
            )

    return
