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) -> 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")
        manifest_path = root / "manifest.json"
        args = cli.build_parser().parse_args(
            [
                "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",
            ]
        )
        cli.make_manifest(args)
        return manifest_path, cli.validate_manifest(manifest_path)

    @staticmethod
    def write_blocks(root: Path) -> None:
        root.mkdir(parents=True, exist_ok=True)
        (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\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"
            "\tsatterthwaite\tboundary_conditional\t2\t1\t8\t2\t0.2\t\t\n",
            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"
        )
        (root / "block_000000.complete").write_text(
            "format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n"
            "fixed_test_elements\t5\n".format(cli.BLOCK_FORMAT),
            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",
            ):
                self.assertEqual(command.count(option), 1)

    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()