#!/usr/bin/env python3
"""Download one validated LRF-IMU fold pair and generate standardized IMU windows."""

from __future__ import annotations

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

import numpy as np
import torch
from huggingface_hub import hf_hub_download

from lrf_imu.checkpoints import load_flow_checkpoint, load_vae_checkpoint
from lrf_imu.generation import sample_reverse_euler


DEFAULT_REPO_ID = "Aminrezaei/LRF-IMU"
ACTIVITIES = {
    0: {"name": "walking", "realdisp_code": 1},
    1: {"name": "running", "realdisp_code": 3},
    2: {"name": "jump_up", "realdisp_code": 4},
    3: {"name": "cycling", "realdisp_code": 33},
}
ACTIVITY_IDS = {details["name"]: class_id for class_id, details in ACTIVITIES.items()}
SENSOR_CHANNELS = {"six_channel": 6, "accelerometer_only": 3}
CONFIG_NAMES = {
    "six_channel": "six_channel_160_40.yaml",
    "accelerometer_only": "accelerometer_only_160_40.yaml",
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def download_file(
    *,
    repo_id: str,
    filename: str,
    revision: str,
    cache_dir: Path | None,
) -> Path:
    return Path(
        hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            revision=revision,
            cache_dir=None if cache_dir is None else str(cache_dir),
        )
    )


def read_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as stream:
        value = json.load(stream)
    if not isinstance(value, dict):
        raise ValueError(f"expected a JSON object in {path.name}")
    return value


def require_hash(path: Path, expected: str, label: str) -> str:
    actual = sha256_file(path)
    if actual.lower() != str(expected).lower():
        raise ValueError(f"{label} SHA-256 mismatch: expected {expected}, got {actual}")
    return actual


def resolve_device(requested: str) -> torch.device:
    if requested == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if requested == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but torch.cuda.is_available() is false")
    return torch.device(requested)


def selected_model(index: dict[str, Any], sensor: str, fold: int) -> dict[str, Any]:
    for entry in index.get("models", []):
        if (
            entry.get("sensor_configuration") == sensor
            and int(entry.get("fold", -1)) == fold
        ):
            return entry
    raise ValueError(f"no checkpoint pair for sensor={sensor!r}, fold={fold:02d}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Download one matched LRF-IMU fold pair, verify its hashes, and "
            "generate standardized VAE-decoder-space IMU windows."
        )
    )
    parser.add_argument("--repo-id", default=DEFAULT_REPO_ID)
    parser.add_argument("--revision", default="main")
    parser.add_argument(
        "--sensor",
        choices=tuple(SENSOR_CHANNELS),
        default="six_channel",
    )
    parser.add_argument(
        "--fold",
        type=int,
        choices=(1, 2, 3, 5, 8, 9, 10, 11, 12, 13, 14, 16),
        default=1,
        help="held-out REALDISP participant/fold",
    )
    parser.add_argument(
        "--activity",
        choices=tuple(ACTIVITY_IDS),
        default="walking",
    )
    parser.add_argument("--count", type=int, default=4)
    parser.add_argument("--steps", type=int, default=10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    parser.add_argument("--cache-dir", type=Path)
    parser.add_argument("--output", type=Path, default=Path("synthetic_imu.npz"))
    parser.add_argument("--overwrite", action="store_true")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    class_id = ACTIVITY_IDS[args.activity]
    if args.count < 1:
        raise ValueError("--count must be positive")
    if args.steps < 1:
        raise ValueError("--steps must be positive")

    output = args.output.expanduser().resolve()
    metadata_output = output.with_suffix(".metadata.json")
    if not args.overwrite:
        existing = [path for path in (output, metadata_output) if path.exists()]
        if existing:
            raise FileExistsError(
                "refusing to overwrite existing output: "
                + ", ".join(str(path) for path in existing)
            )

    cache_dir = (
        None if args.cache_dir is None else args.cache_dir.expanduser().resolve()
    )
    index_path = download_file(
        repo_id=args.repo_id,
        filename="model_index.json",
        revision=args.revision,
        cache_dir=cache_dir,
    )
    index = read_json(index_path)
    model = selected_model(index, args.sensor, args.fold)

    config_filename = str(model["config"])
    expected_config = f"configs/{CONFIG_NAMES[args.sensor]}"
    if config_filename != expected_config:
        raise ValueError(
            f"index config mismatch: expected {expected_config}, got {config_filename}"
        )

    selected_paths = {
        "config": config_filename,
        "manifest": str(model["manifest"]),
        "vae": str(model["vae"]["path"]),
        "flow": str(model["flow"]["path"]),
    }
    downloaded = {
        key: download_file(
            repo_id=args.repo_id,
            filename=filename,
            revision=args.revision,
            cache_dir=cache_dir,
        )
        for key, filename in selected_paths.items()
    }

    manifest = read_json(downloaded["manifest"])
    if manifest.get("sensor_configuration") != args.sensor:
        raise ValueError("manifest sensor does not match the requested sensor")
    if int(manifest.get("fold", -1)) != args.fold:
        raise ValueError("manifest fold does not match the requested fold")

    index_config = index["configs"][args.sensor]
    require_hash(downloaded["config"], index_config["sha256"], "config")
    vae_hash = require_hash(
        downloaded["vae"], manifest["files"]["vae"]["sha256"], "VAE checkpoint"
    )
    flow_hash = require_hash(
        downloaded["flow"], manifest["files"]["flow"]["sha256"], "Flow checkpoint"
    )
    if vae_hash != str(model["vae"]["sha256"]).lower():
        raise ValueError("VAE hash differs between model_index.json and manifest.json")
    if flow_hash != str(model["flow"]["sha256"]).lower():
        raise ValueError("Flow hash differs between model_index.json and manifest.json")

    device = resolve_device(args.device)
    channels = SENSOR_CHANNELS[args.sensor]
    vae, vae_info = load_vae_checkpoint(
        downloaded["vae"],
        channels=channels,
        latent_channels=48,
        down_levels=2,
        device=device,
    )
    flow, flow_info = load_flow_checkpoint(
        downloaded["flow"],
        channels=channels,
        latent_channels=48,
        num_classes=4,
        model_width=256,
        device=device,
    )
    vae.eval()
    flow.eval()

    if vae_info.latent_time_steps != 40 or flow_info.model_width != 256:
        raise ValueError("loaded checkpoint geometry differs from the published index")
    labels = torch.full(
        (args.count,),
        class_id,
        dtype=torch.long,
        device=device,
    )
    with torch.inference_mode():
        latents = sample_reverse_euler(
            flow,
            labels,
            (48, 40),
            num_steps=args.steps,
            seed=args.seed,
            device=device,
        )
        samples = vae.decode(latents).detach().cpu().numpy().astype(np.float32)
    label_values = labels.detach().cpu().numpy().astype(np.int64)

    expected_shape = (args.count, channels, 160)
    if samples.shape != expected_shape:
        raise ValueError(
            f"decoded shape mismatch: expected {expected_shape}, got {samples.shape}"
        )

    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, samples=samples, labels=label_values)
    output_hash = sha256_file(output)
    activity = ACTIVITIES[class_id]
    metadata = {
        "schema_version": "1.0",
        "model_id": args.repo_id,
        "revision": args.revision,
        "sensor_configuration": args.sensor,
        "fold": args.fold,
        "held_out_subject": args.fold,
        "activity": {
            "class_id": class_id,
            "name": activity["name"],
            "realdisp_code": activity["realdisp_code"],
        },
        "generation": {
            "solver": "explicit_reverse_euler",
            "steps": args.steps,
            "paper_default_steps": 10,
            "uses_paper_step_count": args.steps == 10,
            "seed": args.seed,
            "device": str(device),
        },
        "signal": {
            "shape": list(samples.shape),
            "dtype": str(samples.dtype),
            "orientation": "B,C,T",
            "sampling_frequency_hz": 50,
            "window_samples": 160,
            "coordinate_system": "standardized_vae_decoder_space",
            "physical_units_assigned": False,
        },
        "checkpoints": {
            "vae": {
                "filename": selected_paths["vae"],
                "sha256": vae_hash,
            },
            "flow": {
                "filename": selected_paths["flow"],
                "sha256": flow_hash,
            },
        },
        "output": {
            "filename": output.name,
            "sha256": output_hash,
        },
        "exact_paper_reproduction": False,
    }
    metadata_output.write_text(
        json.dumps(metadata, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    print(f"Wrote {output}")
    print(f"Wrote {metadata_output}")
    print(
        "Signal values are standardized VAE-decoder outputs; "
        "no physical units are assigned."
    )
    return 0


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