from __future__ import annotations

import csv
import copy
import json

import numpy as np
import pytest

from lgasieve.clinical_depth import (
    LEGACY_SCHEMA,
    AnchoredDepthModel,
    DepthTarget,
)
from lgasieve.clinical_depth_cli import main as depth_cli_main


def _targets() -> tuple[DepthTarget, ...]:
    controls = tuple(
        DepthTarget(f"C{index:02d}", "CONTROL", str(index), index, False)
        for index in range(30)
    )
    reportable = tuple(
        DepthTarget(f"B{index}", "BRCA1", str(index), index, True)
        for index in range(1, 6)
    )
    return controls + reportable


def _normal_counts(
    n: int, *, seed: int, noisy_reportable: bool = False
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    efficiency = np.linspace(0.75, 1.25, 35)
    depth = np.exp(rng.normal(np.log(380.0), 0.18, n))
    batch = rng.normal(0.0, 0.05, (n, 1))
    loading = np.linspace(-0.5, 0.5, 35)[None, :]
    rate = depth[:, None] * efficiency[None, :] * np.exp(batch @ loading)
    if noisy_reportable:
        rate[:, 32] *= np.exp(rng.normal(0.0, 0.45, n))
    return rng.poisson(rate)


def _model(*, noisy_reportable: bool = False) -> AnchoredDepthModel:
    return AnchoredDepthModel.fit(
        _normal_counts(
            90, seed=11, noisy_reportable=noisy_reportable
        ),
        _normal_counts(400, seed=12),
        _targets(),
        alpha=0.05,
        n_components=2,
        max_target_sigma=0.18,
    )


def test_single_exon_deletion_is_sent_to_reflex() -> None:
    model = _model()
    sample = _normal_counts(1, seed=20)[0]
    sample[31] = int(round(sample[31] * 0.5))
    result = model.screen(sample, sample="carrier")
    assert result.status == "REFLEX_POSITIVE"
    assert any(
        call.kind == "DEL" and "B2" in call.target_ids
        for call in result.calls
    )
    assert all(
        call.disposition == "REFLEX_POSITIVE" for call in result.calls
    )


def test_windows_do_not_bridge_an_uncallable_target() -> None:
    model = _model(noisy_reportable=True)
    missing_index = next(
        index
        for index, target in enumerate(model.targets)
        if target.target_id == "B3"
    )
    assert "B3" in {
        target.target_id
        for target, callable_target in zip(
            model.targets, model.callable_mask
        )
        if target.reportable and not callable_target
    }
    assert all(
        not (
            min(window.indices) < missing_index < max(window.indices)
        )
        for window in model.windows
    )
    sample = _normal_counts(1, seed=21)[0]
    sample[30:] = np.rint(sample[30:] * 0.5).astype(int)
    result = model.screen(sample, sample="whole_gene_deletion")
    assert result.status == "REFLEX_POSITIVE"
    assert all(
        not (
            {"B1", "B2"}.intersection(call.target_ids)
            and {"B4", "B5"}.intersection(call.target_ids)
        )
        for call in result.calls
    )


def test_windows_do_not_bridge_a_same_gene_nonreportable_target() -> None:
    targets = list(_targets())
    targets[32] = DepthTarget("B3", "BRCA1", "3", 3, False)
    model = AnchoredDepthModel.fit(
        _normal_counts(90, seed=25),
        _normal_counts(30, seed=26),
        tuple(targets),
        alpha=0.05,
        n_components=2,
    )
    gap_index = 32
    assert all(
        not (min(window.indices) < gap_index < max(window.indices))
        for window in model.windows
    )


def test_low_depth_is_no_call() -> None:
    model = _model()
    result = model.screen(np.ones(35), sample="low")
    assert result.status == "NO_CALL_QC"
    assert any("median control depth" in reason for reason in result.reasons)


def test_excluded_reportable_target_has_partial_coverage_vocabulary() -> None:
    model = _model(noisy_reportable=True)
    result = model.screen(_normal_counts(1, seed=23)[0], sample="incomplete")
    assert result.status == "NO_CANDIDATE_PARTIAL_COVERAGE"
    assert "B3" in result.no_call_targets
    assert "does not establish absence" in result.claims_boundary
    assert "negative" not in result.claims_boundary.lower()


def test_complete_clean_screen_uses_no_reflex_vocabulary() -> None:
    result = _model().screen(
        _normal_counts(1, seed=24)[0], sample="clean"
    )
    assert result.status == "NO_REFLEX_FLAG"
    assert "NEGATIVE" not in result.status


def test_model_json_round_trip(tmp_path) -> None:
    model = _model()
    destination = tmp_path / "model.json"
    model.save(destination)
    restored = AnchoredDepthModel.load(destination)
    sample = _normal_counts(1, seed=22)[0]
    sample[33] = int(round(sample[33] * 1.5))
    assert restored.screen(sample).to_json() == model.screen(sample).to_json()


def test_calibration_cohort_cannot_change_fitted_nuisance_structure() -> None:
    fit_counts = _normal_counts(90, seed=41)
    first = AnchoredDepthModel.fit(
        fit_counts,
        _normal_counts(80, seed=42),
        _targets(),
        alpha=0.05,
        n_components=2,
    )
    second = AnchoredDepthModel.fit(
        fit_counts,
        _normal_counts(80, seed=43, noisy_reportable=True),
        _targets(),
        alpha=0.05,
        n_components=2,
    )

    assert np.array_equal(first.callable_mask, second.callable_mask)
    assert np.array_equal(first.residual_centre, second.residual_centre)
    assert np.array_equal(first.residual_sigma, second.residual_sigma)
    assert np.array_equal(first.components, second.components)
    assert first.noise_threshold == second.noise_threshold
    assert tuple(
        (window.gene, window.indices, window.weights)
        for window in first.windows
    ) == tuple(
        (window.gene, window.indices, window.weights)
        for window in second.windows
    )


def test_joint_empirical_resolution_is_enforced_exactly() -> None:
    fit_counts = _normal_counts(90, seed=44)
    with pytest.raises(ValueError, match="need at least 99"):
        AnchoredDepthModel.fit(
            fit_counts,
            _normal_counts(98, seed=45),
            _targets(),
            alpha=0.01,
            n_components=2,
        )

    model = AnchoredDepthModel.fit(
        fit_counts,
        _normal_counts(99, seed=46),
        _targets(),
        alpha=0.01,
        n_components=2,
    )
    assert model.n_calibration == 99
    assert np.isclose(
        model.family_threshold,
        model._conformal_threshold(
            model.calibration_family_max, model.alpha
        ),
    )
    assert "family_thresholds" not in model.to_json()


def test_core_rejects_identical_rows_across_fit_and_calibration() -> None:
    fit_counts = _normal_counts(90, seed=47)
    calibration = _normal_counts(30, seed=48)
    calibration[0] = fit_counts[0]
    with pytest.raises(ValueError, match="identical count rows"):
        AnchoredDepthModel.fit(
            fit_counts,
            calibration,
            _targets(),
            alpha=0.05,
            n_components=2,
        )


def test_core_rejects_overlapping_sample_ids() -> None:
    fit_counts = _normal_counts(90, seed=49)
    calibration = _normal_counts(30, seed=50)
    fit_ids = tuple(f"fit_{index}" for index in range(len(fit_counts)))
    calibration_ids = (
        fit_ids[0],
        *(f"cal_{index}" for index in range(1, len(calibration))),
    )
    with pytest.raises(ValueError, match="sample IDs must be disjoint"):
        AnchoredDepthModel.fit(
            fit_counts,
            calibration,
            _targets(),
            alpha=0.05,
            n_components=2,
            fit_sample_ids=fit_ids,
            calibration_sample_ids=calibration_ids,
        )


def test_model_v2_rejects_legacy_and_corrupt_artifacts() -> None:
    payload = _model().to_json()

    legacy = copy.deepcopy(payload)
    legacy["schema"] = LEGACY_SCHEMA
    with pytest.raises(ValueError, match="refit"):
        AnchoredDepthModel.from_json(legacy)

    wrong_shape = copy.deepcopy(payload)
    wrong_shape["centre"] = [wrong_shape["centre"]]
    with pytest.raises(ValueError, match="shape"):
        AnchoredDepthModel.from_json(wrong_shape)

    numeric_string = copy.deepcopy(payload)
    numeric_string["centre"][0] = str(numeric_string["centre"][0])
    with pytest.raises(ValueError, match="JSON numeric"):
        AnchoredDepthModel.from_json(numeric_string)

    wrong_window = copy.deepcopy(payload)
    wrong_window["windows"][0]["indices"][0] = len(payload["targets"])
    with pytest.raises(ValueError, match="indices"):
        AnchoredDepthModel.from_json(wrong_window)

    numeric_string_weight = copy.deepcopy(payload)
    numeric_string_weight["windows"][0]["weights"][0] = str(
        numeric_string_weight["windows"][0]["weights"][0]
    )
    with pytest.raises(ValueError, match="weights"):
        AnchoredDepthModel.from_json(numeric_string_weight)

    wrong_threshold = copy.deepcopy(payload)
    wrong_threshold["family_threshold"] += 1.0
    with pytest.raises(ValueError, match="calibration maxima"):
        AnchoredDepthModel.from_json(wrong_threshold)

    impossible_sigma = copy.deepcopy(payload)
    impossible_sigma["residual_sigma"][0] = 0.01
    with pytest.raises(ValueError, match="variance floor"):
        AnchoredDepthModel.from_json(impossible_sigma)

    duplicate_provenance = copy.deepcopy(payload)
    duplicate_provenance["cohort_provenance"][
        "calibration_count_rows_sha256"
    ] = duplicate_provenance["cohort_provenance"][
        "fit_count_rows_sha256"
    ]
    with pytest.raises(ValueError, match="count-row hashes"):
        AnchoredDepthModel.from_json(duplicate_provenance)

    unresolved = copy.deepcopy(payload)
    unresolved["n_calibration"] = 1
    unresolved["calibration_family_max"] = unresolved[
        "calibration_family_max"
    ][:1]
    with pytest.raises(ValueError, match="cannot attain"):
        AnchoredDepthModel.from_json(unresolved)


def test_brca_only_matrix_is_rejected() -> None:
    targets = tuple(
        DepthTarget(f"B{index}", "BRCA2", str(index), index, True)
        for index in range(1, 6)
    )
    counts = np.full((100, len(targets)), 300)
    with pytest.raises(ValueError, match="control targets"):
        AnchoredDepthModel.fit(counts, counts, targets)


def _write_targets(path, targets: tuple[DepthTarget, ...]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle, delimiter="\t")
        writer.writerow(
            ("target_id", "gene", "exon", "order", "reportable")
        )
        for target in targets:
            writer.writerow(
                (
                    target.target_id,
                    target.gene,
                    target.exon,
                    target.order,
                    int(target.reportable),
                )
            )


def _write_counts(path, targets, samples, matrix) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle, delimiter="\t")
        writer.writerow(("sample", *(target.target_id for target in targets)))
        for sample, row in zip(samples, matrix):
            writer.writerow((sample, *row))


def test_fit_and_screen_cli(tmp_path) -> None:
    targets = _targets()
    target_path = tmp_path / "targets.tsv"
    fit_path = tmp_path / "fit.tsv"
    calibration_path = tmp_path / "calibration.tsv"
    screen_path = tmp_path / "screen.tsv"
    model_path = tmp_path / "depth-model.json"
    output_path = tmp_path / "screen-results.json"
    _write_targets(target_path, targets)
    fit_counts = _normal_counts(90, seed=31)
    calibration_counts = _normal_counts(400, seed=32)
    _write_counts(
        fit_path,
        targets,
        (f"fit_{index}" for index in range(len(fit_counts))),
        fit_counts,
    )
    _write_counts(
        calibration_path,
        targets,
        (f"cal_{index}" for index in range(len(calibration_counts))),
        calibration_counts,
    )
    assert (
        depth_cli_main(
            [
                "fit",
                "--targets",
                str(target_path),
                "--fit-counts",
                str(fit_path),
                "--calibration-counts",
                str(calibration_path),
                "--out",
                str(model_path),
                "--alpha",
                "0.05",
                "--components",
                "2",
            ]
        )
        == 0
    )
    carrier = _normal_counts(1, seed=33)
    carrier[0, 34] = int(round(carrier[0, 34] * 0.5))
    _write_counts(screen_path, targets, ("carrier",), carrier)
    assert (
        depth_cli_main(
            [
                "screen",
                "--model",
                str(model_path),
                "--counts",
                str(screen_path),
                "--out",
                str(output_path),
            ]
        )
        == 0
    )
    payload = json.loads(output_path.read_text(encoding="utf-8"))
    assert payload["results"][0]["status"] == "REFLEX_POSITIVE"
    assert payload["results"][0]["calls"]
    model_payload = json.loads(model_path.read_text(encoding="utf-8"))
    assert model_payload["schema"].endswith(".v2")
    assert model_payload["cohort_provenance"]["fit_sample_ids_sha256"]
    assert model_payload["cohort_provenance"][
        "calibration_sample_ids_sha256"
    ]
    assert model_payload["fit_parameters"] == {
        "requested_n_components": 2,
        "min_median_control_depth": 100.0,
        "min_target_median_depth": 100.0,
        "max_target_sigma": 0.13,
        "max_window_targets": 60,
    }


def test_fit_cli_rejects_overlapping_sample_ids(tmp_path, capsys) -> None:
    targets = _targets()
    target_path = tmp_path / "targets.tsv"
    fit_path = tmp_path / "fit.tsv"
    calibration_path = tmp_path / "calibration.tsv"
    model_path = tmp_path / "depth-model.json"
    _write_targets(target_path, targets)
    fit_counts = _normal_counts(90, seed=51)
    calibration_counts = _normal_counts(30, seed=52)
    fit_ids = tuple(f"fit_{index}" for index in range(len(fit_counts)))
    calibration_ids = (
        fit_ids[0],
        *(f"cal_{index}" for index in range(1, len(calibration_counts))),
    )
    _write_counts(fit_path, targets, fit_ids, fit_counts)
    _write_counts(
        calibration_path,
        targets,
        calibration_ids,
        calibration_counts,
    )

    assert (
        depth_cli_main(
            [
                "fit",
                "--targets",
                str(target_path),
                "--fit-counts",
                str(fit_path),
                "--calibration-counts",
                str(calibration_path),
                "--out",
                str(model_path),
                "--alpha",
                "0.05",
            ]
        )
        == 2
    )
    assert "sample IDs overlap" in capsys.readouterr().err
    assert not model_path.exists()