from __future__ import annotations

import csv
import gzip
import importlib.util
import json
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

import numpy as np


SCRIPT = Path(__file__).resolve().parents[2] / "python" / "spectra_reml.py"
SPEC = importlib.util.spec_from_file_location("spectra_reml", SCRIPT)
assert SPEC is not None and SPEC.loader is not None
cli = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(cli)


class SpectraRemlCliTests(unittest.TestCase):
    def fixture(self, root: Path, with_contrasts: bool = False) -> tuple[Path, dict]:
        n, p0 = 3, 2
        np.zeros(n * (n + 1) // 2, dtype="<f4").tofile(root / "grm.bin")
        (root / "grm.id").write_text("F1 I1\nF2 I2\nF3 I3\n", encoding="utf-8")
        np.asarray([[1, -1], [1, 0], [1, 1]], dtype="<f8").tofile(root / "base.bin")
        np.zeros((2, n), dtype="<f8").tofile(root / "phenotypes.bin")
        np.asarray([[0, 1, 2]], dtype="<f4").tofile(root / "extra.bin")
        (root / "tasks.tsv").write_text(
            "task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
            "0\ttrait_a\t0\t0\n"
            "1\ttrait_b\t1\t1\n",
            encoding="utf-8",
        )
        np.asarray([0, 0, 1], dtype="<i8").tofile(root / "offsets.bin")
        np.asarray([0], dtype="<i4").tofile(root / "indices.bin")
        if with_contrasts:
            np.asarray([[0, 1], [1, -1]], dtype="<f8").tofile(
                root / "contrasts.bin"
            )
            (root / "contrasts.tsv").write_text(
                "contrast_index\tcontrast_id\trhs\n"
                "0\tslope\t0\n"
                "1\tintercept_minus_slope\t0\n",
                encoding="utf-8",
            )
        manifest_path = root / "manifest.json"
        command = [
                "make-manifest",
                "--manifest", str(manifest_path),
                "--grm-bin", str(root / "grm.bin"),
                "--grm-id", str(root / "grm.id"),
                "--base-x", str(root / "base.bin"),
                "--phenotypes", str(root / "phenotypes.bin"),
                "--extra-covariates", str(root / "extra.bin"),
                "--tasks", str(root / "tasks.tsv"),
                "--extra-offsets", str(root / "offsets.bin"),
                "--extra-indices", str(root / "indices.bin"),
                "--output-dir", str(root / "blocks"),
                "--n-samples", "3",
                "--n-base-covariates", "2",
                "--n-phenotype-rows", "2",
                "--n-extra-covariate-rows", "1",
            ]
        if with_contrasts:
            command.extend(
                [
                    "--contrast-matrix", str(root / "contrasts.bin"),
                    "--contrast-metadata", str(root / "contrasts.tsv"),
                    "--n-contrasts", "2",
                ]
            )
        args = cli.build_parser().parse_args(command)
        cli.make_manifest(args)
        return manifest_path, cli.validate_manifest(manifest_path)

    @staticmethod
    def write_blocks(root: Path, contrast_count: int = 0) -> None:
        root.mkdir(parents=True, exist_ok=True)
        first_contrast_offset = 0 if contrast_count else -1
        second_contrast_offset = contrast_count if contrast_count else -1
        (root / "block_000000.summary.tsv").write_text(
            "\t".join(cli.SUMMARY_HEADER) + "\n"
            "0\ttrait_a\tconverged\t2\t0\t0\t0\t1\t1\t0.5\t-1\t4\t8\t1e-8"
            "\tsatterthwaite\tok\t0\t{}\t0\tnan\tnan\tnan\t\t\n"
            "1\ttrait_b\tconverged_boundary\t3\t1\t2\t3\t0\t1\t0\t-2\t3\t6\t1e-9"
            "\tordinary-least-squares\tboundary_ols\t2\t{}\t1\t8\t2\t0.2\t\t\n".format(
                first_contrast_offset, second_contrast_offset
            ),
            encoding="utf-8",
        )
        np.asarray([1, 2, 3, 4, 5], dtype="<f8").tofile(root / "block_000000.beta.f64.bin")
        np.arange(1, 10, dtype="<f8").tofile(root / "block_000000.cov.f64.bin")
        np.asarray([0.1, 0.2, 0.3, 0.4, 0.5], dtype="<f8").tofile(
            root / "block_000000.fixed_se.f64.bin"
        )
        np.asarray([1, 2, 3, 4, 5], dtype="<f8").tofile(
            root / "block_000000.fixed_stat.f64.bin"
        )
        np.asarray([10, 10, 8, 8, 8], dtype="<f8").tofile(
            root / "block_000000.fixed_ddf.f64.bin"
        )
        np.asarray([0.5, 0.2, 0.1, 0.05, 0.01], dtype="<f8").tofile(
            root / "block_000000.fixed_p.f64.bin"
        )
        contrast_elements = 2 * contrast_count
        contrast_arrays = {
            "contrast_estimate": np.arange(1, contrast_elements + 1, dtype="<f8"),
            "contrast_se": np.full(contrast_elements, 0.25, dtype="<f8"),
            "contrast_stat": np.arange(2, contrast_elements + 2, dtype="<f8"),
            "contrast_numdf": np.ones(contrast_elements, dtype="<f8"),
            "contrast_ddf": np.full(contrast_elements, 12.0, dtype="<f8"),
            "contrast_p": np.full(contrast_elements, 0.05, dtype="<f8"),
        }
        for suffix, values in contrast_arrays.items():
            values.tofile(root / "block_000000.{}.f64.bin".format(suffix))
        (root / "block_000000.complete").write_text(
            "format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n"
            "fixed_test_elements\t5\ncontrast_test_elements\t{}\n".format(
                cli.BLOCK_FORMAT, contrast_elements
            ),
            encoding="utf-8",
        )

    def test_manifest_and_generic_command(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            _, manifest = self.fixture(root)
            args = cli.build_parser().parse_args(
                ["run", "--manifest", str(root / "manifest.json"), "--engine", "spectra_reml"]
            )
            command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
            for option in (
                "--phenotypes",
                "--extra-covariates",
                "--extra-offsets",
                "--extra-indices",
                "--n-phenotype-rows",
                "--n-extra-covariate-rows",
                "--line-search-expansion",
                "--zoom-safeguard",
                "--n-contrasts",
            ):
                self.assertEqual(command.count(option), 1)

    def test_contrast_manifest_and_engine_command(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            _, manifest = self.fixture(root, with_contrasts=True)
            self.assertEqual(manifest["dimensions"]["contrast_count"], 2)
            args = cli.build_parser().parse_args(
                [
                    "run", "--manifest", str(root / "manifest.json"),
                    "--engine", "spectra_reml", "--fixed-effect-test",
                    "kenward-roger",
                ]
            )
            command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
            self.assertEqual(command.count("--contrast-matrix"), 1)
            self.assertEqual(command.count("--contrast-metadata"), 1)
            self.assertEqual(command[command.index("--n-contrasts") + 1], "2")
            self.assertEqual(command[command.index("--fixed-effect-test") + 1], "kenward-roger")

    def test_new_contract_rejects_missing_contrast_fields(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            manifest_path, _ = self.fixture(root)
            manifest = cli.read_json(manifest_path)
            del manifest["dimensions"]["contrast_count"]
            del manifest["paths"]["contrast_matrix_f64"]
            del manifest["paths"]["contrast_metadata_tsv"]
            cli.atomic_json(manifest_path, manifest)
            with self.assertRaises(RuntimeError):
                cli.validate_manifest(manifest_path)

    def test_both_fixed_effect_methods_are_forwarded(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            _, manifest = self.fixture(root, with_contrasts=True)
            for method in ("satterthwaite", "kenward-roger"):
                args = cli.build_parser().parse_args(
                    [
                        "run", "--manifest", str(root / "manifest.json"),
                        "--engine", "spectra_reml", "--fixed-effect-test", method,
                    ]
                )
                command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
                self.assertEqual(
                    command[command.index("--fixed-effect-test") + 1], method
                )

    def test_contrast_finalize_contract(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            manifest_path, manifest = self.fixture(root, with_contrasts=True)
            engine = root / "spectra_reml"
            engine.write_bytes(b"synthetic engine")
            args = cli.build_parser().parse_args(
                [
                    "run", "--manifest", str(manifest_path), "--engine",
                    str(engine), "--resume", "--fixed-effect-test",
                    "satterthwaite",
                ]
            )
            with mock.patch.object(
                cli.subprocess, "run", return_value=SimpleNamespace(returncode=0)
            ):
                cli.run_engine(manifest, args)
            output = Path(manifest["paths"]["output_directory"])
            self.write_blocks(output, contrast_count=2)
            result = root / "contrast_results.tsv.gz"
            cli.finalize(manifest, result)
            with gzip.open(result, "rt", encoding="utf-8", newline="") as handle:
                rows = list(csv.DictReader(handle, delimiter="\t"))
            self.assertEqual(json.loads(rows[0]["contrast_estimate_json"]), [1.0, 2.0])
            self.assertEqual(json.loads(rows[1]["contrast_estimate_json"]), [3.0, 4.0])
            self.assertEqual(json.loads(rows[0]["contrast_numerator_df_json"]), [1.0, 1.0])

    def test_signature_recovery_and_finalize(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            manifest_path, manifest = self.fixture(root)
            engine = root / "spectra_reml"
            engine.write_bytes(b"synthetic engine")
            args = cli.build_parser().parse_args(
                ["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
            )
            with mock.patch.object(cli.subprocess, "run", return_value=SimpleNamespace(returncode=0)):
                cli.run_engine(manifest, args)
            output = Path(manifest["paths"]["output_directory"])
            signature = cli.read_signature(output / "run.signature.json")
            self.assertEqual(signature["format"], cli.RUN_SIGNATURE_FORMAT)

            self.write_blocks(output)
            changed = cli.build_parser().parse_args(
                [
                    "run", "--manifest", str(manifest_path), "--engine", str(engine),
                    "--resume", "--max-iterations", "101",
                ]
            )
            with self.assertRaisesRegex(RuntimeError, "signature differs"):
                cli.run_engine(manifest, changed)

            result = root / "results.tsv.gz"
            cli.finalize(manifest, result)
            with gzip.open(result, "rt", encoding="utf-8", newline="") as handle:
                rows = list(csv.DictReader(handle, delimiter="\t"))
            self.assertEqual([row["task_id"] for row in rows], ["trait_a", "trait_b"])
            self.assertEqual(json.loads(rows[0]["beta_json"]), [1.0, 2.0])
            self.assertEqual(len(json.loads(rows[1]["covariance_packed_lower_json"])), 6)
            self.assertEqual(len(json.loads(rows[1]["fixed_effect_p_value_json"])), 3)

            altered = dict(manifest)
            altered["created_utc"] = "changed"
            with self.assertRaisesRegex(RuntimeError, "does not belong"):
                cli.finalize(altered, root / "unsafe.tsv")

    def test_force_invalidates_before_cleanup_and_dry_run_is_read_only(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            manifest_path, manifest = self.fixture(root)
            engine = root / "spectra_reml"
            engine.write_bytes(b"engine")
            output = Path(manifest["paths"]["output_directory"])
            self.write_blocks(output)
            resume_args = cli.build_parser().parse_args(
                ["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
            )
            cli.atomic_json(output / "run.signature.json", cli.make_signature(manifest, engine, resume_args))
            before = sorted(path.name for path in output.iterdir())
            dry = cli.build_parser().parse_args(
                ["run", "--manifest", str(manifest_path), "--engine", "missing", "--force", "--dry-run"]
            )
            cli.run_engine(manifest, dry)
            self.assertEqual(before, sorted(path.name for path in output.iterdir()))

            force = cli.build_parser().parse_args(
                ["run", "--manifest", str(manifest_path), "--engine", str(engine), "--force"]
            )
            with mock.patch.object(cli, "block_files", side_effect=RuntimeError("cleanup interrupted")):
                with self.assertRaisesRegex(RuntimeError, "cleanup interrupted"):
                    cli.run_engine(manifest, force)
            self.assertFalse((output / "run.signature.json").exists())


if __name__ == "__main__":
    unittest.main()