#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
from typing import Any

from assaylens.data.access import iter_limited_rows
from assaylens.data.hf_upload import auto_upload_folder
from assaylens.data.schemas import EVE_HF_ID
from assaylens.splits.assay_frame import make_assay_frame_holdout_split
from assaylens.splits.release import make_release_time_split
from assaylens.splits.scaffold import make_scaffold_split
from assaylens.splits.target import make_target_split
from assaylens.splits.target_class import make_target_class_holdout_split
from assaylens.reports.markdown import write_markdown


def random_split(n_rows: int, seed: int = 20260505) -> dict[str, list[int]]:
    indices = list(range(n_rows))
    random.Random(seed).shuffle(indices)
    train_end = int(0.70 * n_rows)
    val_end = int(0.85 * n_rows)
    return {
        "train": sorted(indices[:train_end]),
        "validation": sorted(indices[train_end:val_end]),
        "test": sorted(indices[val_end:]),
    }


def split_rows(split: dict[str, list[int]]) -> list[dict[str, Any]]:
    return [
        {"row_index": idx, "split": split_name}
        for split_name, indices in split.items()
        for idx in indices
    ]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--out-dir", default="splits")
    args = parser.parse_args()
    rows = list(iter_limited_rows(EVE_HF_ID, "train", args.limit))
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    splitters = {
        "random_split": random_split(len(rows)),
        "scaffold_split": make_scaffold_split(rows),
        "target_split": make_target_split(rows),
        "target_class_split": make_target_class_holdout_split(rows),
        "release_time_split": make_release_time_split(rows),
        "assay_frame_holdout": make_assay_frame_holdout_split(rows),
    }
    import pandas as pd

    for name, split in splitters.items():
        (out / f"{name}.json").write_text(json.dumps(split, indent=2), encoding="utf-8")
        pd.DataFrame(split_rows(split)).to_parquet(out / f"{name}.parquet", index=False)
    report_lines = [
        f"- {name}: train={len(split['train'])}, validation={len(split['validation'])}, test={len(split['test'])}"
        for name, split in splitters.items()
    ]
    reports_dir = Path("/tmp/assaylens_outputs/reports")
    write_markdown(
        reports_dir / "split_report.md",
        "Split Report",
        [("Rows", f"{len(rows)} EvE rows assigned per split."), ("Definitions", "\n".join(report_lines))],
    )
    uploaded = auto_upload_folder(out, "splits")
    uploaded_reports = auto_upload_folder(reports_dir, "reports")
    if uploaded:
        print(f"uploaded split artifacts to {uploaded}/splits")
    if uploaded_reports:
        print(f"uploaded split report to {uploaded_reports}/reports")
    print(f"wrote {len(splitters)} split files to {out}")


if __name__ == "__main__":
    main()