#!/usr/bin/env python3
"""Create minimal synthetic RNA-seq data for strandness parameter validation."""
from __future__ import annotations

import gzip
import random
import shutil
from pathlib import Path

try:
    import yaml  # type: ignore
except Exception as exc:  # pragma: no cover
    raise SystemExit(f"PyYAML is required: {exc}")


PROJECT_ROOT = Path(__file__).resolve().parents[1]
BASE = PROJECT_ROOT / "项目记录" / "rnaseqUpdate" / "testdata" / "strandness_minimal"


def revcomp(seq: str) -> str:
    return seq.translate(str.maketrans("ACGTN", "TGCAN"))[::-1]


def write_fastq(path: Path, records: list[tuple[str, str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(path, "wt", encoding="utf-8") as handle:
        for name, seq in records:
            handle.write(f"@{name}\n{seq}\n+\n{'I' * len(seq)}\n")


def sample_yaml(strandness: str, mixed: bool = False) -> dict:
    samples = {
        "S1": {
            "run": "SYN001",
            "r1": "S1_1.clean.fastq.gz",
            "r2": "S1_2.clean.fastq.gz",
            "condition": "CK",
            "replicate": 1,
            "library_strandness": strandness,
        }
    }
    if mixed:
        samples["S2"] = {
            "run": "SYN002",
            "r1": "S2_1.clean.fastq.gz",
            "r2": "S2_2.clean.fastq.gz",
            "condition": "TRT",
            "replicate": 1,
            "library_strandness": "reverse",
        }
    return {
        "species": {"code": "Test", "speciesName": "Synthetic test", "keggOrganism": ""},
        "rnaseq": {"output_prefix": "Test"},
        "samples": samples,
        "groups": {sample: meta["condition"] for sample, meta in samples.items()},
        "experiment": {"design": "two_group", "control_group": "CK"},
    }


def main() -> int:
    if BASE.exists():
        shutil.rmtree(BASE)
    input_dir = BASE / "input"
    final_dir = BASE / "workspace" / "job_synthetic_strandness" / "final"
    input_dir.mkdir(parents=True, exist_ok=True)
    final_dir.mkdir(parents=True, exist_ok=True)

    rng = random.Random(7)
    genome = "".join(rng.choice("ACGT") for _ in range(1000))
    (input_dir / "Test.genome.fasta").write_text(f">chr1\n{genome}\n", encoding="utf-8")
    (input_dir / "Test.gtf").write_text(
        'chr1\tsynthetic\texon\t80\t320\t.\t+\t.\tgene_id "gene1"; transcript_id "tx1";\n',
        encoding="utf-8",
    )

    r1_records = []
    r2_records = []
    for idx, start in enumerate((100, 120, 140, 160), start=1):
        frag_start = start - 1
        r1 = genome[frag_start : frag_start + 50]
        r2 = revcomp(genome[frag_start + 100 : frag_start + 150])
        r1_records.append((f"S1_{idx}/1", r1))
        r2_records.append((f"S1_{idx}/2", r2))
    write_fastq(final_dir / "S1_1.clean.fastq.gz", r1_records)
    write_fastq(final_dir / "S1_2.clean.fastq.gz", r2_records)
    write_fastq(final_dir / "S2_1.clean.fastq.gz", r1_records)
    write_fastq(final_dir / "S2_2.clean.fastq.gz", r2_records)

    for strandness in ("non_stranded", "forward", "reverse"):
        payload = sample_yaml(strandness)
        (input_dir / f"samples_{strandness}.yaml").write_text(
            yaml.safe_dump(payload, allow_unicode=True, sort_keys=False),
            encoding="utf-8",
        )
    (input_dir / "samples_mixed.yaml").write_text(
        yaml.safe_dump(sample_yaml("forward", mixed=True), allow_unicode=True, sort_keys=False),
        encoding="utf-8",
    )
    print(BASE)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())