import matplotlib

matplotlib.use("Agg")  # no display needed for plot smoke tests

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from giant.analysis import (
    RAW_TARGET_NAMES,
    SampleCollection,
    constraint_report,
    constraint_report_pl,
    correlation_matrices,
    direction_alignment,
    load_predicted_local,
    marginal_table,
    marginal_table_pl,
    plot_constraint_violations,
    plot_correlation_matrices,
    plot_direction_alignment,
    plot_marginals,
    plot_pairwise,
)
from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
)
from giant.data.transforms import log_transform


def _unit_vectors(rng, n):
    v = rng.standard_normal((n, 3)).astype(np.float32)
    return v / np.linalg.norm(v, axis=1, keepdims=True)


def _make_collection(n=200, seed=0, gen_offset=0.0) -> SampleCollection:
    rng = np.random.default_rng(seed)
    real = np.column_stack(
        [
            rng.uniform(0.1, 5.0, n),  # step_length
            rng.uniform(0.1, 5.0, n),  # delta_e
            rng.uniform(0.1, 5.0, n),  # edep
            _unit_vectors(rng, n),  # post_dir
            _unit_vectors(rng, n),  # travel_dir
        ]
    ).astype(np.float32)
    gen = real + gen_offset

    pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
    cond_cont_raw = np.column_stack(
        [
            rng.standard_normal((n, 3)),
            pre_E,
            rng.standard_normal((n, 3)),
            rng.integers(0, 5, n),
            rng.integers(0, 3, n),
        ]
    ).astype(np.float32)

    return SampleCollection(
        cond_cont_raw=cond_cont_raw,
        pdg=rng.choice([11, -11, 22], size=n),
        material=rng.choice(["W", "Pb"], size=n),
        real_raw=real,
        gen_raw=gen,
        real_norm=real,
        gen_norm=gen,
    )


def test_marginal_table_aggregate_has_all_dims():
    table = marginal_table(_make_collection())
    assert set(table["dim"]) == set(RAW_TARGET_NAMES)
    assert (table["group"] == "all").all()


@pytest.mark.parametrize("group_by", ["pdg", "material", "energy"])
def test_marginal_table_grouped_covers_all_rows(group_by):
    collection = _make_collection()
    table = marginal_table(collection, group_by=group_by)
    assert table["n"].groupby(table["group"]).first().sum() == len(collection.pdg)


def test_marginal_table_identical_distributions_have_zero_kl():
    collection = _make_collection(gen_offset=0.0)
    table = marginal_table(collection)
    np.testing.assert_allclose(table["kl_real_gen"], 0.0, atol=1e-6)


def test_marginal_table_shifted_distribution_has_positive_kl():
    collection = _make_collection(gen_offset=3.0)
    table = marginal_table(collection)
    assert (table["kl_real_gen"] > 0).all()


def test_correlation_matrices_are_symmetric_unit_diagonal():
    real_corr, gen_corr = correlation_matrices(_make_collection())
    for corr in (real_corr, gen_corr):
        np.testing.assert_allclose(np.diag(corr), 1.0, atol=1e-5)
        np.testing.assert_allclose(corr, corr.T, atol=1e-5)


def test_direction_alignment_real_data_is_unit_norm_dot_product():
    real_cos, gen_cos = direction_alignment(_make_collection())
    assert np.all(real_cos >= -1.0 - 1e-5) and np.all(real_cos <= 1.0 + 1e-5)
    assert np.all(gen_cos >= -1.0 - 1e-5) and np.all(gen_cos <= 1.0 + 1e-5)


def test_constraint_report_clean_data_has_no_violations():
    report = constraint_report(_make_collection(gen_offset=0.0))
    assert (report["violation_rate"] == 0.0).all()


def test_constraint_report_flags_negative_log_dims_and_bad_norms():
    collection = _make_collection(gen_offset=0.0)
    collection.gen_raw[:, 0] = -1.0  # negative step_length
    collection.gen_raw[:, 3:6] *= 2.0  # post_dir no longer unit norm
    report = constraint_report(collection)
    violations = dict(zip(report["check"], report["violation_rate"]))
    assert violations["step_length >= 0"] == 1.0
    assert violations["post_dir unit norm"] == 1.0


def test_plot_marginals_runs_without_error():
    fig = plot_marginals(_make_collection())
    assert fig is not None


def test_plot_marginals_grouped_runs_without_error():
    fig = plot_marginals(_make_collection(), group_by="material")
    assert fig is not None


def test_plot_correlation_matrices_runs_without_error():
    fig = plot_correlation_matrices(_make_collection())
    assert fig is not None


def test_plot_pairwise_runs_without_error():
    fig = plot_pairwise(_make_collection())
    assert fig is not None


def test_plot_direction_alignment_runs_without_error():
    fig = plot_direction_alignment(_make_collection())
    assert fig is not None


def test_plot_constraint_violations_runs_without_error():
    fig = plot_constraint_violations(_make_collection())
    assert fig is not None


def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
    """Mimic `giant predict --coord local`'s output schema for the loader tests."""
    rng = rng or np.random.default_rng(0)
    true_log_local = rng.standard_normal((n, 9)).astype(np.float32)
    true_log_local[:, :3] = log_transform(
        rng.uniform(0.1, 5.0, (n, 3)).astype(np.float32)
    )
    pred_log_local = true_log_local + rng.normal(0, 0.01, (n, 9)).astype(np.float32)

    table = pa.table(
        {
            "event_id": rng.integers(0, 10, n),
            "pdg": rng.choice([11, -11, 22], n),
            "pre_x": rng.standard_normal(n).astype(np.float32),
            "pre_y": rng.standard_normal(n).astype(np.float32),
            "pre_z": rng.standard_normal(n).astype(np.float32),
            "pre_E": rng.uniform(1.0, 100.0, n).astype(np.float32),
            "pre_dx": rng.standard_normal(n).astype(np.float32),
            "pre_dy": rng.standard_normal(n).astype(np.float32),
            "pre_dz": rng.standard_normal(n).astype(np.float32),
            "material": rng.choice(["W", "Pb"], n),
            "layer_id": rng.integers(0, 10, n).astype(np.int32),
            "n_sec": rng.integers(0, 3, n).astype(np.int32),
            **{
                f"pred_{name}": pred_log_local[:, j]
                for j, name in enumerate(LOCAL_TARGET_NAMES)
            },
            **{
                f"true_{name}": true_log_local[:, j]
                for j, name in enumerate(LOCAL_TARGET_NAMES)
            },
        }
    )
    if metadata is not None:
        table = table.replace_schema_metadata(metadata)
    pq.write_table(table, path)
    return true_log_local, pred_log_local


def test_load_predicted_local_round_trips_values(tmp_path):
    path = tmp_path / "predicted_local.parquet"
    true_log_local, pred_log_local = _write_predicted_local_parquet(
        path,
        metadata={
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        },
    )

    collection = load_predicted_local(path)

    expected_real = true_log_local.copy()
    expected_real[:, :3] = np.exp(expected_real[:, :3]) - 1e-8
    expected_gen = pred_log_local.copy()
    expected_gen[:, :3] = np.exp(expected_gen[:, :3]) - 1e-8
    np.testing.assert_allclose(collection.real_raw, expected_real, atol=1e-4)
    np.testing.assert_allclose(collection.gen_raw, expected_gen, atol=1e-4)
    assert collection.real_norm is None
    assert collection.gen_norm is None


def test_load_predicted_local_usable_by_downstream_plots(tmp_path):
    path = tmp_path / "predicted_local.parquet"
    _write_predicted_local_parquet(
        path,
        metadata={
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        },
    )
    collection = load_predicted_local(path)
    assert marginal_table(collection) is not None
    assert plot_marginals(collection) is not None


def test_load_predicted_local_rejects_missing_metadata(tmp_path):
    path = tmp_path / "no_metadata.parquet"
    _write_predicted_local_parquet(path, metadata=None)
    with pytest.raises(ValueError, match="no '.*' parquet metadata"):
        load_predicted_local(path)


def test_load_predicted_local_rejects_global_coord(tmp_path):
    path = tmp_path / "global.parquet"
    _write_predicted_local_parquet(
        path,
        metadata={
            PREDICT_COORD_METADATA_KEY: "global",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        },
    )
    with pytest.raises(ValueError, match="coord=local"):
        load_predicted_local(path)


def test_load_predicted_local_rejects_mismatched_schema_version(tmp_path):
    path = tmp_path / "old_version.parquet"
    _write_predicted_local_parquet(
        path,
        metadata={
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: "999",
        },
    )
    with pytest.raises(ValueError, match="schema version"):
        load_predicted_local(path)


def _predicted_local_path(tmp_path, n=200, seed=0):
    path = tmp_path / "predicted_local.parquet"
    _write_predicted_local_parquet(
        path,
        n=n,
        rng=np.random.default_rng(seed),
        metadata={
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        },
    )
    return path


@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
    path = _predicted_local_path(tmp_path)
    collection = load_predicted_local(path)

    expected = marginal_table(collection, group_by=group_by).sort_values(
        ["group", "dim"]
    )
    actual = (
        marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
    )

    assert list(expected["group"]) == list(actual["group"])
    assert list(expected["n"]) == list(actual["n"])
    for col in ["real_mean", "gen_mean", "real_std", "gen_std", "kl_real_gen"]:
        np.testing.assert_allclose(
            expected[col].to_numpy(),
            actual[col].to_numpy(),
            atol=1e-4,
            rtol=1e-4,
        )


def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
    path = tmp_path / "no_metadata.parquet"
    _write_predicted_local_parquet(path, metadata=None)
    with pytest.raises(ValueError, match="no '.*' parquet metadata"):
        marginal_table_pl(path)


def test_constraint_report_pl_matches_numpy_version(tmp_path):
    path = _predicted_local_path(tmp_path)
    collection = load_predicted_local(path)

    expected = constraint_report(collection)
    actual = constraint_report_pl(path).to_pandas()

    assert list(expected["check"]) == list(actual["check"])
    np.testing.assert_allclose(
        expected["violation_rate"].to_numpy(),
        actual["violation_rate"].to_numpy(),
        atol=1e-6,
    )
    np.testing.assert_allclose(
        expected["mean_abs_error"].to_numpy(),
        actual["mean_abs_error"].to_numpy(),
        atol=1e-4,
    )