#!/usr/bin/env python3
"""Create a small, fully synthetic SpectraREML input bundle."""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

PROJECT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT / "python"))
import spectra_reml  # noqa: E402


def main() -> None:
    output = Path(__file__).resolve().parent / "synthetic_work"
    output.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(20260809)
    n = 12

    z = rng.normal(size=(n, 5))
    grm = z @ z.T / z.shape[1]
    diagonal = np.sqrt(np.diag(grm))
    grm /= np.outer(diagonal, diagonal)
    packed = np.asarray(
        [grm[row, col] for row in range(n) for col in range(row + 1)],
        dtype="<f4",
    )
    packed.tofile(output / "example.grm.bin")
    (output / "example.grm.id").write_text(
        "".join("sample{0}\tsample{0}\n".format(i + 1) for i in range(n)),
        encoding="utf-8",
    )

    base_x = np.column_stack((np.ones(n), np.linspace(-1.0, 1.0, n)))
    extra = rng.normal(size=(2, n)).astype("<f4")
    phenotypes = np.vstack(
        (
            0.3 + 0.5 * base_x[:, 1] + rng.normal(scale=0.2, size=n),
            -0.2 + 0.8 * extra[0] + rng.normal(scale=0.2, size=n),
            0.1 + 0.4 * extra[0] - 0.3 * extra[1] + rng.normal(scale=0.2, size=n),
        )
    )
    np.asarray(base_x, dtype="<f8").tofile(output / "base_x.f64.bin")
    np.asarray(phenotypes, dtype="<f8").tofile(output / "phenotypes.f64.bin")
    extra.tofile(output / "extra_covariates.f32.bin")
    (output / "tasks.tsv").write_text(
        "task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
        "0\ttrait_0\t0\t0\n"
        "1\ttrait_1\t1\t1\n"
        "2\ttrait_2\t2\t2\n",
        encoding="utf-8",
    )
    np.asarray([0, 0, 1, 3], dtype="<i8").tofile(output / "extra_offsets.i64.bin")
    np.asarray([0, 0, 1], dtype="<i4").tofile(output / "extra_indices.i32.bin")

    manifest = output / "manifest.json"
    spectra_reml.main(
        [
            "make-manifest",
            "--manifest", str(manifest),
            "--grm-bin", str(output / "example.grm.bin"),
            "--grm-id", str(output / "example.grm.id"),
            "--base-x", str(output / "base_x.f64.bin"),
            "--phenotypes", str(output / "phenotypes.f64.bin"),
            "--extra-covariates", str(output / "extra_covariates.f32.bin"),
            "--tasks", str(output / "tasks.tsv"),
            "--extra-offsets", str(output / "extra_offsets.i64.bin"),
            "--extra-indices", str(output / "extra_indices.i32.bin"),
            "--output-dir", str(output / "blocks"),
            "--n-samples", str(n),
            "--n-base-covariates", str(base_x.shape[1]),
            "--n-phenotype-rows", str(phenotypes.shape[0]),
            "--n-extra-covariate-rows", str(extra.shape[0]),
        ]
    )
    print("\nNext:")
    print("python python/spectra_reml.py run --manifest {} --engine /path/to/spectra_reml --threads 4".format(manifest))


if __name__ == "__main__":
    main()