"""Package benchmark tasks as the SimplexTasks-12 release artifact."""
from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import numpy as np
import yaml

import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from scripts.run_age_ldl import extract_image_features, get_age_predictions, load_utkface
from scripts.run_hyperspectral import load_hyperspectral, unmix_nmf
from scripts.run_topics import prepare_topic_data
from src.data import build_prediction_matrix, load_affective_text
from src.dgp.deconv import nnls_deconv
from src.dgp.discrete_groups import DiscreteGroupsDGP
from src.dgp.heavy_tail import HeavyTailDGP
from src.dgp.high_k import HighKDGP
from src.dgp.model_bias import ModelBiasDGP
from src.dgp.pseudobulk import generate_pseudobulk
from src.dgp.pure_scale import PureScaleDGP

PACKAGE_NAME = "SimplexTasks-12"
PACKAGE_SLUG = "simplextasks-12"
VERSION = "0.1.0"
SEED = 2026

DGP_MAP = {
    "pure_scale": PureScaleDGP,
    "discrete_groups": DiscreteGroupsDGP,
    "model_bias": ModelBiasDGP,
    "heavy_tail": HeavyTailDGP,
    "high_k": HighKDGP,
}

SYNTHETIC_SPECS = [
    ("d1_homogeneous", "configs/synthetic/d1_homogeneous.yaml", "Homogeneous negative control"),
    ("d2_pure_scale", "configs/synthetic/d2_pure_scale.yaml", "Smooth scale heterogeneity"),
    ("d3_discrete_groups_aligned", "configs/synthetic/d3_discrete_groups_aligned.yaml", "Aligned discrete groups"),
    ("d4_model_bias", "configs/synthetic/d4_model_bias.yaml", "Bias-type heterogeneity"),
    ("d5_heavy_tail", "configs/synthetic/d5_heavy_tail.yaml", "Heavy-tail robustness"),
    ("d6_high_k", "configs/synthetic/d6_high_k.yaml", "High-dimensional simplex"),
]


def _float32(x: np.ndarray) -> np.ndarray:
    return x.astype(np.float32) if isinstance(x, np.ndarray) and x.dtype.kind == "f" else x


def save_npz(path: Path, arrays: dict[str, np.ndarray]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, **{k: _float32(v) for k, v in arrays.items()})


def write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text.strip() + "\n", encoding="utf-8")


def build_cifar(root: Path) -> tuple[dict[str, np.ndarray], dict]:
    cache_path = root / "data/processed/cifar10_resnet18_softmax.npz"
    data = np.load(cache_path)
    arrays = {
        "Y": data["Y"],
        "U": data["U"],
        "class_names": data["class_names"],
        "label_index": np.argmax(data["Y"], axis=1).astype(np.int16),
    }
    meta = {
        "task_id": "cifar10_softmax",
        "task_name": "CIFAR-10 softmax",
        "subset": "real",
        "n_samples": int(arrays["Y"].shape[0]),
        "simplex_dim": int(arrays["Y"].shape[1]),
        "source_asset": "CIFAR-10",
        "predictor": "Frozen ResNet-18 softmax cache",
        "redistribution": "derived-only",
        "notes": "No raw images are redistributed in this package.",
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta


def build_topics() -> tuple[dict[str, np.ndarray], dict]:
    Y, U = prepare_topic_data(K=10, seed=SEED)
    arrays = {
        "Y": Y,
        "U": U,
        "doc_index": np.arange(len(Y), dtype=np.int32),
    }
    meta = {
        "task_id": "topics_20ng",
        "task_name": "20 Newsgroups topics",
        "subset": "real",
        "n_samples": int(Y.shape[0]),
        "simplex_dim": int(Y.shape[1]),
        "source_asset": "20 Newsgroups",
        "predictor": "TF-IDF to topic-mixture kNN regressor",
        "redistribution": "derived-only",
        "notes": "This release exposes only derived simplex arrays, not the raw text corpus.",
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta


def build_samson(root: Path) -> tuple[dict[str, np.ndarray], dict]:
    data = load_hyperspectral(str(root / "data/raw/hyperspectral"), "samson")
    U = unmix_nmf(data["pixels"], data["K"], seed=SEED)
    arrays = {
        "Y": data["abundances"],
        "U": U,
        "pixels": data["pixels"],
        "endmembers": data["endmembers"],
        "endmember_names": np.asarray(data["names"]),
        "image_shape": np.asarray(data["shape"], dtype=np.int32),
    }
    meta = {
        "task_id": "samson_unmixing",
        "task_name": "Samson hyperspectral unmixing",
        "subset": "real",
        "n_samples": int(arrays["Y"].shape[0]),
        "simplex_dim": int(arrays["Y"].shape[1]),
        "source_asset": "Samson ROI hyperspectral benchmark",
        "predictor": "NMF abundance estimator",
        "redistribution": "source-cited",
        "notes": "The public bundle did not ship an explicit license file; keep attribution with any downstream reuse.",
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta


def build_pbmc(root: Path) -> tuple[dict[str, np.ndarray], dict]:
    import scanpy as sc
    from sklearn.cluster import KMeans
    from sklearn.decomposition import PCA

    h5ad_path = root / "data/pbmc3k_raw.h5ad"
    adata = sc.read_h5ad(h5ad_path)
    expr = adata.X
    if hasattr(expr, "toarray"):
        expr = expr.toarray()
    expr = np.asarray(expr, dtype=np.float64)

    celltype_key = "cell_type"
    if celltype_key in adata.obs.columns:
        labels = adata.obs[celltype_key].values
    else:
        pca = PCA(n_components=30, random_state=42)
        X_pca = pca.fit_transform(expr)
        kmeans = KMeans(n_clusters=8, random_state=42, n_init=10)
        labels = np.asarray([f"ct_{v}" for v in kmeans.fit_predict(X_pca)])
    cell_type_names = sorted(np.unique(labels).tolist())
    gene_names = adata.var_names.to_numpy()

    pb = generate_pseudobulk(
        expr=expr,
        labels=labels,
        cell_type_names=cell_type_names,
        gene_names=gene_names.tolist(),
        n_samples=5000,
        cells_per_sample=200,
        concentration=1.0,
        noise_sd=0.1,
        seed=SEED,
    )
    U = nnls_deconv(pb.bulk, pb.signature)
    arrays = {
        "Y": pb.proportions,
        "U": U,
        "bulk_expression": pb.bulk,
        "signature": pb.signature,
        "cell_type_names": np.asarray(pb.cell_type_names),
        "gene_names": np.asarray(pb.gene_names),
    }
    meta = {
        "task_id": "pbmc3k_pseudobulk",
        "task_name": "PBMC3K pseudobulk deconvolution",
        "subset": "real",
        "n_samples": int(arrays["Y"].shape[0]),
        "simplex_dim": int(arrays["Y"].shape[1]),
        "source_asset": "PBMC3K single-cell reference",
        "predictor": "NNLS deconvolution from pseudobulk mixtures",
        "redistribution": "derived-only",
        "notes": "This package exports pseudobulk-derived arrays and signatures rather than the raw single-cell matrix.",
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta


def build_utkface(root: Path) -> tuple[dict[str, np.ndarray], dict]:
    data_dir = root / "data/raw/UTKFace"
    ages, Y, image_paths = load_utkface(str(data_dir), K=10, sigma=2.0)
    U = get_age_predictions(ages, Y, image_paths, K=10, method="image_knn", seed=SEED)
    X_feat = extract_image_features(image_paths, image_size=16, cache_name=f"utkface_imgfeat_{len(image_paths)}_s16.npz")
    arrays = {
        "Y": Y,
        "U": U,
        "age": ages.astype(np.int16),
        "image_features": X_feat,
    }
    meta = {
        "task_id": "utkface_age_ldl",
        "task_name": "UTKFace age label distributions",
        "subset": "real",
        "n_samples": int(arrays["Y"].shape[0]),
        "simplex_dim": int(arrays["Y"].shape[1]),
        "source_asset": "UTKFace aligned-and-cropped images",
        "predictor": "Thumbnail feature PCA+kNN age-distribution regressor",
        "redistribution": "derived-only",
        "notes": "The package omits raw face images and keeps only derived features and simplex arrays.",
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta


def build_affective(root: Path) -> tuple[dict[str, np.ndarray], dict]:
    data_dir = root / "data/raw/AffectiveText.Semeval.2007"
    cache_path = root / "data/processed/affective_text_predictions.jsonl"
    data = load_affective_text(data_dir)
    pred_raw, U = build_prediction_matrix(data["ids"], cache_path)
    arrays = {
        "Y": data["Y"],
        "U": U,
        "gold_raw_scores": data["raw_scores"],
        "pred_raw_scores": pred_raw,
        "instance_id": np.asarray(data["ids"]),
        "emotion_names": np.asarray(data["emotions"]),
    }
    meta = {
        "task_id": "affectivetext_emotions",
        "task_name": "SemEval AffectiveText emotions",
        "subset": "real",
        "n_samples": int(arrays["Y"].shape[0]),
        "simplex_dim": int(arrays["Y"].shape[1]),
        "source_asset": "SemEval-2007 Task 14 AffectiveText",
        "predictor": "Frozen zero-shot emotion scorer; open TF-IDF+SVD+kNN fallback script available",
        "redistribution": "derived-only",
        "notes": "The package omits raw headlines and raw API responses, keeps only derived score arrays, and provides an open fallback cache builder in scripts/cache_affective_text_open_predictions.py.",
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta


def build_synthetic_task(root: Path, task_id: str, config_rel: str, regime_label: str) -> tuple[dict[str, np.ndarray], dict, Path]:
    cfg_path = root / config_rel
    cfg = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
    dgp_cfg = dict(cfg["dgp"])
    dgp_name = dgp_cfg.pop("name")
    dgp = DGP_MAP[dgp_name](**dgp_cfg)

    data_cfg = cfg["data"]
    n_train = int(data_cfg.get("n_train", 0))
    n_scale = int(data_cfg.get("n_scale_est", 0))
    n_cal = int(data_cfg.get("n_cal", 0))
    n_test = int(data_cfg.get("n_test", 0))
    n_total = n_train + n_scale + n_cal + n_test

    rng = np.random.default_rng(SEED)
    sample = dgp.sample(n_total, rng)
    split = np.concatenate(
        [
            np.repeat("train", n_train),
            np.repeat("scale", n_scale),
            np.repeat("cal", n_cal),
            np.repeat("test", n_test),
        ]
    )
    arrays = {
        "X": sample.X,
        "Y": sample.Y,
        "U": sample.U,
        "R": sample.R,
        "sigma_true": sample.sigma_true if sample.sigma_true is not None else np.full(n_total, np.nan),
        "split": split.astype("U8"),
    }
    meta = {
        "task_id": task_id,
        "task_name": task_id.replace("_", " "),
        "subset": "synthetic",
        "n_samples": int(sample.Y.shape[0]),
        "simplex_dim": int(sample.Y.shape[1]),
        "source_asset": "Synthetic DGP",
        "predictor": "Oracle mean predictor from the configured DGP",
        "redistribution": "open",
        "regime_label": regime_label,
        "config_file": config_rel,
        "seed": SEED,
        "available_arrays": sorted(arrays.keys()),
    }
    return arrays, meta, cfg_path


def write_task(task_dir: Path, arrays: dict[str, np.ndarray], metadata: dict) -> None:
    save_npz(task_dir / "task.npz", arrays)
    write_json(task_dir / "metadata.json", metadata)


def build_package_readme(manifest: dict) -> str:
    real_lines = []
    for task in manifest["real_tasks"]:
        real_lines.append(f"| `{task['task_id']}` | {task['task_name']} | {task['n_samples']} | {task['simplex_dim']} | {task['predictor']} |")
    synthetic_lines = []
    for task in manifest["synthetic_tasks"]:
        synthetic_lines.append(f"| `{task['task_id']}` | {task['regime_label']} | {task['n_samples']} | {task['simplex_dim']} |")
    return f"""
---
pretty_name: {PACKAGE_NAME}
license: other
task_categories:
- other
tags:
- conformal-prediction
- uncertainty-estimation
- simplex
- benchmark
- task-collection
---

# {PACKAGE_NAME}

{PACKAGE_NAME} is the processed task collection behind the SimplexUQ benchmark. It packages 12 simplex-valued prediction tasks: 6 real tasks with fixed predictors and 6 synthetic regimes with canonical reference draws.

## What is inside

- Standardized `task.npz` files with at least `Y` and `U` for every task.
- Per-task `metadata.json` files with provenance, redistribution notes, and task-specific semantics.
- Canonical synthetic configs copied alongside the synthetic tasks.
- A `manifest.json` file that summarizes the full release.

## Real Tasks

| Task ID | Task | Samples | K | Predictor |
| --- | --- | ---: | ---: | --- |
{chr(10).join(real_lines)}

## Synthetic Tasks

| Task ID | Regime | Samples | K |
| --- | --- | ---: | ---: |
{chr(10).join(synthetic_lines)}

## Redistribution Notes

This package is release-oriented rather than raw-data-complete. Some tasks include only derived simplex arrays or derived features because the underlying source assets carry their own terms of use. In particular, raw UTKFace images, raw AffectiveText headlines, and the original CIFAR-10 image archive are not redistributed here.

## Loading Example

```python
from pathlib import Path
import numpy as np

root = Path("{PACKAGE_SLUG}")
task = np.load(root / "real/cifar10_softmax/task.npz")
Y = task["Y"]
U = task["U"]
```
"""


def build_license_notes() -> str:
    return """
# License and Usage Notes

SimplexTasks-12 is a processed task collection. It should not be interpreted as a license override for the underlying source assets.

- CIFAR-10: this release exposes derived simplex arrays only.
- 20 Newsgroups: this release exposes derived topic-mixture arrays only.
- AffectiveText: this release omits raw headlines and raw API responses.
- Samson: keep attribution with the source benchmark bundle.
- PBMC3K: this release exports derived pseudobulk and deconvolution artifacts.
- UTKFace: this release omits raw face images and keeps only derived features and simplex arrays.
"""


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default=f"release/{PACKAGE_SLUG}")
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    root = Path(__file__).resolve().parent.parent
    out_dir = root / args.output_dir
    if out_dir.exists():
        if not args.force:
            raise FileExistsError(f"{out_dir} already exists. Use --force to overwrite.")
        shutil.rmtree(out_dir)
    (out_dir / "real").mkdir(parents=True, exist_ok=True)
    (out_dir / "synthetic").mkdir(parents=True, exist_ok=True)

    real_builders = [
        ("cifar10_softmax", lambda: build_cifar(root)),
        ("topics_20ng", build_topics),
        ("samson_unmixing", lambda: build_samson(root)),
        ("pbmc3k_pseudobulk", lambda: build_pbmc(root)),
        ("utkface_age_ldl", lambda: build_utkface(root)),
        ("affectivetext_emotions", lambda: build_affective(root)),
    ]

    real_manifest = []
    for task_id, builder in real_builders:
        arrays, metadata = builder()
        task_dir = out_dir / "real" / task_id
        write_task(task_dir, arrays, metadata)
        real_manifest.append(metadata)

    synthetic_manifest = []
    for task_id, cfg_rel, regime_label in SYNTHETIC_SPECS:
        arrays, metadata, cfg_path = build_synthetic_task(root, task_id, cfg_rel, regime_label)
        task_dir = out_dir / "synthetic" / task_id
        write_task(task_dir, arrays, metadata)
        shutil.copy2(cfg_path, task_dir / "config.yaml")
        synthetic_manifest.append(metadata)

    manifest = {
        "name": PACKAGE_NAME,
        "slug": PACKAGE_SLUG,
        "version": VERSION,
        "seed": SEED,
        "task_count": len(real_manifest) + len(synthetic_manifest),
        "real_tasks": real_manifest,
        "synthetic_tasks": synthetic_manifest,
    }
    write_json(out_dir / "manifest.json", manifest)
    write_text(out_dir / "README.md", build_package_readme(manifest))
    write_text(out_dir / "LICENSE_NOTES.md", build_license_notes())


if __name__ == "__main__":
    main()