"""End-to-end Seminario validation using ethane .fchk Hessians.

Validates the full Seminario pipeline against stored reference values:
  .fchk → Hessian → Seminario projection → ForceField parameters

Uses both ground state (staggered) and transition state (eclipsed)
ethane at B3LYP/6-31G(d).  Reference fixture was generated by the
same pipeline and is stored as ethane_reference.json.

Covers issue #73.
"""

import json
from pathlib import Path

import numpy as np
import pytest

from q2mm.models.forcefield import ForceField
from q2mm.models.molecule import Q2MMMolecule
from q2mm.models.seminario import (
    estimate_force_constants,
    seminario_angle_fc,
    seminario_bond_fc,
)
from q2mm.optimizers.objective import ReferenceData

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
REPO = Path(__file__).resolve().parents[2]
EXAMPLES = REPO / "examples" / "ethane"
FIXTURES = REPO / "test" / "fixtures" / "seminario_parity"

GS_FCHK = EXAMPLES / "GS.fchk"
TS_FCHK = EXAMPLES / "TS.fchk"
REFERENCE = FIXTURES / "ethane_reference.json"


@pytest.fixture(scope="module")
def reference():
    with REFERENCE.open() as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# .fchk parsing
# ---------------------------------------------------------------------------
class TestFchkParsing:
    """Verify the .fchk parser extracts correct geometry and Hessian."""

    def test_gs_atom_count(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        assert len(mol.symbols) == 8

    def test_gs_symbols(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        assert mol.symbols.count("C") == 2
        assert mol.symbols.count("H") == 6

    def test_gs_hessian_shape(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        assert mol.hessian.shape == (24, 24)

    def test_gs_hessian_symmetric(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        np.testing.assert_allclose(mol.hessian, mol.hessian.T, atol=1e-12)

    def test_ts_atom_count(self):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        assert len(mol.symbols) == 8

    def test_ts_hessian_shape(self):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        assert mol.hessian.shape == (24, 24)


# ---------------------------------------------------------------------------
# Bond detection
# ---------------------------------------------------------------------------
class TestBondDetection:
    """Verify bond detection finds the expected connectivity."""

    def test_gs_bond_count(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        assert len(mol.bonds) == reference["GS"]["n_atoms"] - 1  # 7 for ethane

    def test_gs_has_cc_bond(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        cc_bonds = [b for b in mol.bonds if set(b.elements) == {"C"}]
        assert len(cc_bonds) == 1

    def test_gs_has_six_ch_bonds(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ch_bonds = [b for b in mol.bonds if set(b.elements) == {"C", "H"}]
        assert len(ch_bonds) == 6

    def test_gs_angle_count(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        assert len(mol.angles) == reference["GS"]["n_atoms"] + 4  # 12 for ethane

    def test_ts_bond_count(self, reference):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        assert len(mol.bonds) == reference["TS"]["n_atoms"] - 1


# ---------------------------------------------------------------------------
# Individual bond projections
# ---------------------------------------------------------------------------
class TestIndividualBonds:
    """Verify each Seminario bond projection matches the fixture."""

    def test_gs_individual_bond_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        for ref_bond in reference["GS"]["individual_bonds"]:
            k = seminario_bond_fc(
                ref_bond["atom_i"],
                ref_bond["atom_j"],
                mol.geometry,
                mol.hessian,
            )
            np.testing.assert_allclose(
                k,
                ref_bond["force_constant"],
                atol=1e-8,
                err_msg=f"Bond {ref_bond['atom_i']}-{ref_bond['atom_j']}",
            )

    def test_ts_individual_bond_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        for ref_bond in reference["TS"]["individual_bonds"]:
            k = seminario_bond_fc(
                ref_bond["atom_i"],
                ref_bond["atom_j"],
                mol.geometry,
                mol.hessian,
            )
            np.testing.assert_allclose(
                k,
                ref_bond["force_constant"],
                atol=1e-8,
                err_msg=f"Bond {ref_bond['atom_i']}-{ref_bond['atom_j']}",
            )

    def test_gs_ch_bonds_near_equivalent(self):
        """All 6 C-H bonds in staggered ethane should have similar force constants."""
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ch_fcs = []
        for b in mol.bonds:
            if set(b.elements) == {"C", "H"}:
                k = seminario_bond_fc(b.atom_i, b.atom_j, mol.geometry, mol.hessian)
                ch_fcs.append(k)
        assert len(ch_fcs) == 6
        # All C-H bonds should be within 0.1 mdyn/A of each other
        assert max(ch_fcs) - min(ch_fcs) < 0.1


# ---------------------------------------------------------------------------
# Individual angle projections
# ---------------------------------------------------------------------------
class TestIndividualAngles:
    """Verify each Seminario angle projection matches the fixture."""

    def test_gs_individual_angle_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        for ref_angle in reference["GS"]["individual_angles"]:
            k = seminario_angle_fc(
                ref_angle["atom_i"],
                ref_angle["atom_center"],
                ref_angle["atom_j"],
                mol.geometry,
                mol.hessian,
            )
            np.testing.assert_allclose(
                k,
                ref_angle["force_constant"],
                atol=1e-8,
                err_msg=f"Angle {ref_angle['atom_i']}-{ref_angle['atom_center']}-{ref_angle['atom_j']}",
            )

    def test_ts_individual_angle_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        for ref_angle in reference["TS"]["individual_angles"]:
            k = seminario_angle_fc(
                ref_angle["atom_i"],
                ref_angle["atom_center"],
                ref_angle["atom_j"],
                mol.geometry,
                mol.hessian,
            )
            np.testing.assert_allclose(
                k,
                ref_angle["force_constant"],
                atol=1e-8,
                err_msg=f"Angle {ref_angle['atom_i']}-{ref_angle['atom_center']}-{ref_angle['atom_j']}",
            )


# ---------------------------------------------------------------------------
# Full pipeline (averaged parameters)
# ---------------------------------------------------------------------------
class TestFullPipeline:
    """End-to-end: .fchk → estimate_force_constants → ForceField."""

    def test_gs_averaged_bond_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ff = estimate_force_constants(mol)
        for ref_bond, param in zip(reference["GS"]["averaged_bonds"], ff.bonds):
            np.testing.assert_allclose(
                param.force_constant,
                ref_bond["force_constant"],
                atol=1e-8,
                err_msg=f"Bond {ref_bond['elements']}",
            )

    def test_gs_averaged_bond_equilibria(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ff = estimate_force_constants(mol)
        for ref_bond, param in zip(reference["GS"]["averaged_bonds"], ff.bonds):
            np.testing.assert_allclose(
                param.equilibrium,
                ref_bond["equilibrium"],
                atol=1e-8,
                err_msg=f"Bond {ref_bond['elements']}",
            )

    def test_gs_averaged_angle_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ff = estimate_force_constants(mol)
        for ref_angle, param in zip(reference["GS"]["averaged_angles"], ff.angles):
            np.testing.assert_allclose(
                param.force_constant,
                ref_angle["force_constant"],
                atol=1e-8,
                err_msg=f"Angle {ref_angle['elements']}",
            )

    def test_gs_averaged_angle_equilibria(self, reference):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ff = estimate_force_constants(mol)
        for ref_angle, param in zip(reference["GS"]["averaged_angles"], ff.angles):
            np.testing.assert_allclose(
                param.equilibrium,
                ref_angle["equilibrium"],
                atol=1e-8,
                err_msg=f"Angle {ref_angle['elements']}",
            )

    def test_ts_averaged_bond_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        ff = estimate_force_constants(mol)
        for ref_bond, param in zip(reference["TS"]["averaged_bonds"], ff.bonds):
            np.testing.assert_allclose(
                param.force_constant,
                ref_bond["force_constant"],
                atol=1e-8,
                err_msg=f"Bond {ref_bond['elements']}",
            )

    def test_ts_averaged_angle_force_constants(self, reference):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        ff = estimate_force_constants(mol)
        for ref_angle, param in zip(reference["TS"]["averaged_angles"], ff.angles):
            np.testing.assert_allclose(
                param.force_constant,
                ref_angle["force_constant"],
                atol=1e-8,
                err_msg=f"Angle {ref_angle['elements']}",
            )

    def test_forcefield_returns_correct_type(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ff = estimate_force_constants(mol)
        assert isinstance(ff, ForceField)

    def test_torsions_zeroed_by_default(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        ff = estimate_force_constants(mol)
        for t in ff.torsions:
            assert t.force_constant == 0.0


# ---------------------------------------------------------------------------
# GS vs TS comparison (chemical sanity checks)
# ---------------------------------------------------------------------------
class TestGSvsTSComparison:
    """Verify GS and TS ethane produce chemically sensible differences."""

    @pytest.fixture(scope="class")
    def gs_ff(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        return estimate_force_constants(mol)

    @pytest.fixture(scope="class")
    def ts_ff(self):
        _, mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)
        return estimate_force_constants(mol)

    def test_ts_cc_bond_longer(self, gs_ff, ts_ff):
        """Eclipsed TS has longer C-C bond due to steric repulsion."""
        gs_cc = next(b for b in gs_ff.bonds if b.elements == ("C", "C"))
        ts_cc = next(b for b in ts_ff.bonds if b.elements == ("C", "C"))
        assert ts_cc.equilibrium > gs_cc.equilibrium

    def test_ts_cc_bond_weaker(self, gs_ff, ts_ff):
        """Eclipsed TS has weaker C-C bond force constant."""
        gs_cc = next(b for b in gs_ff.bonds if b.elements == ("C", "C"))
        ts_cc = next(b for b in ts_ff.bonds if b.elements == ("C", "C"))
        assert ts_cc.force_constant < gs_cc.force_constant

    def test_ch_bond_fc_similar(self, gs_ff, ts_ff):
        """C-H bonds should be similar between GS and TS (rotation doesn't affect them much)."""
        gs_ch = next(b for b in gs_ff.bonds if set(b.elements) == {"C", "H"})
        ts_ch = next(b for b in ts_ff.bonds if set(b.elements) == {"C", "H"})
        assert abs(gs_ch.force_constant - ts_ch.force_constant) < 0.2

    def test_bond_fc_positive(self, gs_ff, ts_ff):
        """All bond force constants should be positive for ethane."""
        for b in gs_ff.bonds:
            assert b.force_constant > 0, f"GS bond {b.elements} has negative FC"
        for b in ts_ff.bonds:
            assert b.force_constant > 0, f"TS bond {b.elements} has negative FC"

    def test_angle_fc_positive(self, gs_ff, ts_ff):
        """All angle force constants should be positive for ethane."""
        for a in gs_ff.angles:
            assert a.force_constant > 0, f"GS angle {a.elements} has negative FC"
        for a in ts_ff.angles:
            assert a.force_constant > 0, f"TS angle {a.elements} has negative FC"

    def test_cc_bond_length_reasonable(self, gs_ff):
        """C-C single bond should be ~1.53 Å."""
        cc = next(b for b in gs_ff.bonds if b.elements == ("C", "C"))
        assert 1.4 < cc.equilibrium < 1.6

    def test_ch_bond_length_reasonable(self, gs_ff):
        """C-H bond should be ~1.09 Å."""
        ch = next(b for b in gs_ff.bonds if set(b.elements) == {"C", "H"})
        assert 1.0 < ch.equilibrium < 1.2

    def test_hch_angle_reasonable(self, gs_ff):
        """H-C-H angle should be ~107-108° (slightly less than tetrahedral)."""
        hch = next(a for a in gs_ff.angles if a.elements == ("H", "C", "H"))
        assert 105 < hch.equilibrium < 112

    def test_cch_angle_reasonable(self, gs_ff):
        """C-C-H angle should be ~111° (slightly greater than tetrahedral)."""
        cch = next(a for a in gs_ff.angles if a.elements == ("C", "C", "H"))
        assert 108 < cch.equilibrium < 115


# ---------------------------------------------------------------------------
# Multi-molecule averaging
# ---------------------------------------------------------------------------
class TestMultiMoleculeAveraging:
    """Verify averaging across GS and TS molecules."""

    def test_averaged_cc_between_gs_and_ts(self):
        """C-C force constant averaged across GS+TS should be between individual values."""
        _, gs_mol = ReferenceData.from_fchk(GS_FCHK)
        _, ts_mol = ReferenceData.from_fchk(TS_FCHK, bond_tolerance=1.4)

        # Get individual values
        gs_ff = estimate_force_constants(gs_mol)
        ts_ff = estimate_force_constants(ts_mol)
        gs_cc = next(b for b in gs_ff.bonds if b.elements == ("C", "C"))
        ts_cc = next(b for b in ts_ff.bonds if b.elements == ("C", "C"))

        # Create shared force field for averaging
        shared_ff = ForceField.create_for_molecule(gs_mol, name="shared")
        avg_ff = estimate_force_constants([gs_mol, ts_mol], forcefield=shared_ff)
        avg_cc = next(b for b in avg_ff.bonds if b.elements == ("C", "C"))

        lo = min(gs_cc.force_constant, ts_cc.force_constant)
        hi = max(gs_cc.force_constant, ts_cc.force_constant)
        assert lo <= avg_cc.force_constant <= hi


# ---------------------------------------------------------------------------
# Independent literature validation (not self-referential)
# ---------------------------------------------------------------------------
class TestLiteratureValidation:
    """Validate Seminario results against independent literature values.

    Reference ranges come from experimental/high-level QM sources,
    NOT from the Seminario pipeline itself.

    Sources:
    - C-C bond: ~4.3-5.0 mdyn/Å (MM3 force field, Allinger et al.)
    - C-H bond: ~4.7-5.0 mdyn/Å (IR spectroscopy / MM3)
    - H-C-H angle: ~0.3-0.6 mdyn·Å/rad² (MM3)
    - C-C-H angle: ~0.4-0.7 mdyn·Å/rad² (MM3)
    """

    @pytest.fixture(scope="class")
    def gs_ff(self):
        _, mol = ReferenceData.from_fchk(GS_FCHK)
        return estimate_force_constants(mol)

    def test_cc_fc_in_literature_range(self, gs_ff):
        """C-C stretch FC should be ~3-6 mdyn/Å (Seminario from DFT can be softer than MM3)."""
        cc = next((b for b in gs_ff.bonds if b.elements == ("C", "C")), None)
        assert cc is not None, "No C-C bond found in force field"
        assert 2.5 < cc.force_constant < 6.5, f"C-C FC {cc.force_constant:.2f} outside lit range"

    def test_ch_fc_in_literature_range(self, gs_ff):
        """C-H stretch FC should be ~4.5-5.5 mdyn/Å."""
        ch = next((b for b in gs_ff.bonds if set(b.elements) == {"C", "H"}), None)
        assert ch is not None, "No C-H bond found in force field"
        assert 4.0 < ch.force_constant < 6.0, f"C-H FC {ch.force_constant:.2f} outside lit range"

    def test_hch_angle_fc_in_literature_range(self, gs_ff):
        """H-C-H bend FC should be ~0.3-0.7 mdyn·Å/rad²."""
        hch = next((a for a in gs_ff.angles if a.elements == ("H", "C", "H")), None)
        assert hch is not None, "No H-C-H angle found in force field"
        assert 0.2 < hch.force_constant < 0.8, f"H-C-H FC {hch.force_constant:.2f} outside lit range"

    def test_cch_angle_fc_in_literature_range(self, gs_ff):
        """C-C-H bend FC should be ~0.4-0.8 mdyn·Å/rad²."""
        cch = next((a for a in gs_ff.angles if a.elements == ("C", "C", "H")), None)
        assert cch is not None, "No C-C-H angle found in force field"
        assert 0.3 < cch.force_constant < 0.9, f"C-C-H FC {cch.force_constant:.2f} outside lit range"
