from __future__ import annotations

from collections import Counter
from datetime import datetime
from pathlib import Path
import numpy as np
from pymatgen.analysis.ewald import EwaldSummation
from pymatgen.core import Structure


XRAY_WAVELENGTHS = {
    "CuKa1": 1.5405929,
    "CuKa2": 1.5444274,
    "CuKb": 1.392234,
    "AgKa1": 0.55942178,
    "AgKa2": 0.5638131,
    "MoKa1": 0.70931715,
    "MoKa2": 0.713607,
    "CoKa1": 1.788996,
    "CoKa2": 1.792835,
    "FeKa1": 1.936041,
    "FeKa2": 1.939973,
    "CrKa1": 2.289726,
    "CrKa2": 2.293651,
}


def _load_structure(input_path: str) -> Structure:
    path = Path(input_path)
    if not path.exists():
        raise FileNotFoundError(f"Structure file not found: {input_path}")
    return Structure.from_file(str(path))


def _parse_oxidation_states(oxidation_states: str | None) -> dict[str, float] | None:
    if not oxidation_states:
        return None
    ox_map: dict[str, float] = {}
    for token in oxidation_states.split(","):
        token = token.strip()
        if not token:
            continue
        if ":" not in token:
            raise ValueError(
                f"Invalid oxidation state token '{token}'. Use format 'Element:Value'."
            )
        elem, value = token.split(":", 1)
        elem = elem.strip()
        value = value.strip()
        if not elem:
            raise ValueError("Element symbol is empty in oxidation state string.")
        ox_map[elem] = float(value)
    if not ox_map:
        raise ValueError("No valid oxidation states parsed.")
    return ox_map


def xrd_from_poscar(
    poscar_path: str = "POSCAR",
    xray_type: str = "CuKa1",
    wavelength: float | None = None,
    output_prefix: str | None = None,
    two_theta_min: float = 0.0,
    two_theta_max: float = 120.0,
    plot_min: float = 10.0,
    plot_max: float = 80.0,
    sigma: float = 0.05,
    npoints: int = 5000,
    show: bool = False,
):
    import matplotlib.pyplot as plt
    from pymatgen.analysis.diffraction.core import AbstractDiffractionPatternCalculator
    from pymatgen.analysis.diffraction.xrd import XRDCalculator
    from pymatgen.core.composition import Composition
    from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

    poscar_file = Path(poscar_path)
    if not poscar_file.exists():
        raise FileNotFoundError(f"POSCAR file not found: {poscar_path}")

    if wavelength is None:
        if xray_type not in XRAY_WAVELENGTHS:
            valid = ", ".join(XRAY_WAVELENGTHS.keys())
            raise ValueError(f"Unsupported --xraytype '{xray_type}'. Supported: {valid}")
        wavelength = XRAY_WAVELENGTHS[xray_type]
    else:
        xray_type = f"custom-{wavelength:.6f}A"

    output_stem = output_prefix if output_prefix else poscar_file.name
    output_peak_dat = Path(f"{output_stem}_peaks.dat")
    output_profile_dat = Path(f"{output_stem}_profile.dat")
    output_sticks_dat = Path(f"{output_stem}_sticks.dat")
    output_pdf = Path(f"{output_stem}.pdf")

    structure = _load_structure(str(poscar_file))

    AbstractDiffractionPatternCalculator.SCALED_INTENSITY_TOL = 0
    calculator = XRDCalculator(wavelength=wavelength)
    pattern = calculator.get_pattern(structure, two_theta_range=(two_theta_min, two_theta_max))

    xrd_data = []
    for idx, theta in enumerate(pattern.x):
        h, k, l = 0, 0, 0
        multiplicity = 0
        if pattern.hkls[idx]:
            hkl_info = pattern.hkls[idx][0]
            h, k, l = hkl_info["hkl"]
            multiplicity = hkl_info.get("multiplicity", 0)
        d_hkl = pattern.d_hkls[idx]
        intensity = pattern.y[idx]
        xrd_data.append([h, k, l, multiplicity, d_hkl, theta, intensity])

    comp = Composition(structure.formula)
    reduced_formula = comp.reduced_formula
    space_group = SpacegroupAnalyzer(structure).get_space_group_symbol()

    header = (
        f"reduced_formula={reduced_formula}, space_group={space_group}, "
        f"xray={xray_type}, wavelength={wavelength} Angstrom\n"
        "h,k,l,multiplicity,d_hkl(Angstrom),two_theta(deg),intensity"
    )
    np.savetxt(
        output_peak_dat,
        np.array(xrd_data, dtype=float),
        fmt=["%d", "%d", "%d", "%d", "%.8f", "%.8f", "%.8f"],
        delimiter=",",
        header=header,
    )

    def _smeared_signal(x_vals):
        y_vals = np.zeros_like(x_vals, dtype=float)
        if sigma <= 0:
            return y_vals
        norm = sigma * np.sqrt(2.0 * np.pi)
        for peak_pos, peak_int in zip(pattern.x, pattern.y):
            y_vals += peak_int * np.exp(-0.5 * ((x_vals - peak_pos) / sigma) ** 2) / norm
        return y_vals

    x_plot = np.linspace(two_theta_min, two_theta_max, npoints)
    y_plot = _smeared_signal(x_plot)

    np.savetxt(
        output_profile_dat,
        np.column_stack([x_plot, y_plot]),
        fmt="%.8f",
        header="two_theta(deg),intensity",
    )

    np.savetxt(
        output_sticks_dat,
        np.column_stack([pattern.x, pattern.y]),
        fmt="%.8f",
        header="two_theta(deg),intensity",
    )

    fig = plt.figure()
    plt.plot(x_plot, y_plot, linestyle="solid", linewidth=0.8, color="red", label=poscar_file.name)
    plt.vlines(pattern.x, 0.0, pattern.y, linestyle="solid", color="black", linewidth=0.5, alpha=0.7)
    plt.xlim(plot_min, plot_max)
    plt.xlabel("2 theta (deg)")
    plt.ylabel("Intensity (arb. unit)")
    plt.legend()
    plt.tight_layout()
    fig.savefig(output_pdf)

    if show:
        plt.show()
    else:
        plt.close(fig)

    print(f"XRD peak data saved: {output_peak_dat}")
    print(f"XRD profile data saved: {output_profile_dat}")
    print(f"XRD sticks data saved: {output_sticks_dat}")
    print(f"XRD plot saved: {output_pdf}")


def calc_madelung_from_poscar(
    poscar_path: str = "POSCAR",
    output_log: str | None = None,
    oxidation_states: str | None = None,
):
    poscar_file = Path(poscar_path)
    if not poscar_file.exists():
        raise FileNotFoundError(f"POSCAR file not found: {poscar_path}")

    structure = _load_structure(str(poscar_file))

    oxidation_map = _parse_oxidation_states(oxidation_states)
    if oxidation_map:
        structure.add_oxidation_state_by_element(oxidation_map)
        oxidation_mode = f"manual ({oxidation_states})"
    else:
        guesses = structure.composition.oxi_state_guesses()
        guess_map = None
        if guesses:
            for candidate in guesses:
                if any(abs(float(v)) > 1e-12 for v in candidate.values()):
                    guess_map = {k: float(v) for k, v in candidate.items()}
                    break
            if guess_map is None:
                guess_map = {k: float(v) for k, v in guesses[0].items()}

        if guess_map:
            structure.add_oxidation_state_by_element(guess_map)
            oxidation_mode = f"guessed by composition ({guess_map})"
        else:
            structure.add_oxidation_state_by_guess()
            oxidation_mode = "guessed by pymatgen"

    if all(abs(float(site.specie.oxi_state)) <= 1e-12 for site in structure):
        raise ValueError(
            "Failed to assign non-zero oxidation states automatically. "
            "Please pass --oxidation-states (e.g., \"Mn:2,Cr:3,O:-2\")."
        )

    ewald = EwaldSummation(structure)
    total_energy = float(ewald.total_energy)
    natoms = len(structure)
    energy_per_atom = total_energy / natoms if natoms > 0 else np.nan
    species_counter = Counter(str(site.specie) for site in structure)
    species_order = []
    for site in structure:
        specie_str = str(site.specie)
        if specie_str not in species_order:
            species_order.append(specie_str)
    oxidation_summary = ", ".join(
        f"{species} x{species_counter[species]}" for species in species_order
    )

    if output_log is None:
        output_log = f"madelung-{poscar_file.name}.log"
    log_path = Path(output_log)

    lines = [
        f"Timestamp: {datetime.now().isoformat()}",
        f"Input POSCAR: {poscar_file.resolve()}",
        f"Oxidation assignment: {oxidation_mode}",
        f"Oxidation states summary: {oxidation_summary}",
        f"Number of atoms: {natoms}",
        f"Madelung energy (total): {total_energy:.12f} eV",
        f"Madelung energy (per atom): {energy_per_atom:.12f} eV/atom",
    ]
    log_path.write_text("\n".join(lines) + "\n", encoding="utf-8")

    for line in lines[1:]:
        print(line)
    print(f"Log saved: {log_path}")
