"""Reproducible GPU Batch01 V1.1 evidence runner for the ARM64 CUDA image.

The runner deliberately evaluates the CPU historical port and the CUDA engine
separately.  CUDA unavailability is an error: GPU results are never replaced by
CPU results.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import platform
import sys
import time
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import numpy as np
import torch

from gpu_feature_engine_v1 import evaluate_batch, psar_trace, supertrend_trace
from gpu_feature_parity_contract_v1_1 import (
    calibrate,
    canonical_bytes,
    compare_stateful_trace,
    corpus_manifest,
    cpu_psar_trace,
    cpu_supertrend_trace,
    deterministic_adversarial_ohlcv,
    freeze_contract,
    historical_cpu_oracle,
    nan_gap_semantics_manifest,
    validate_frozen_contract,
)

ARTIFACT_PREFIX = "gpu_batch01_v1_1"


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


def write_json(path: Path, payload: Mapping[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(canonical_bytes(dict(payload)) + b"\n")


def read_ohlcv(path: Path) -> dict[str, np.ndarray]:
    with path.open(newline="", encoding="utf-8") as handle:
        rows = list(csv.DictReader(handle))
    required = ("close", "high", "low", "volume")
    if not rows or not set(required).issubset(rows[0]):
        raise ValueError("CSV must contain non-empty close, high, low, volume columns")
    values = {
        name: np.asarray([float(row[name]) for row in rows], dtype=np.float64) for name in required
    }
    if not all(np.isfinite(value).all() for value in values.values()):
        raise ValueError(
            "historical NaN/gap semantics are undefined; non-finite OHLCV inputs are rejected"
        )
    return values


def write_ohlcv_csv(path: Path, ohlcv: Mapping[str, np.ndarray]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=["close", "high", "low", "volume"])
        writer.writeheader()
        writer.writerows(
            {name: float(ohlcv[name][index]) for name in writer.fieldnames}
            for index in range(len(ohlcv["close"]))
        )


def gpu_outputs(
    request: Mapping[str, Any], ohlcv: Mapping[str, np.ndarray], device: torch.device
) -> tuple[dict[str, np.ndarray], dict[str, float]]:
    torch.cuda.reset_peak_memory_stats(device)
    transfer_start = time.perf_counter()
    tensors = tuple(
        torch.as_tensor(ohlcv[name], dtype=torch.float64, device=device)
        for name in ("close", "high", "low", "volume")
    )
    torch.cuda.synchronize(device)
    transfer_seconds = time.perf_counter() - transfer_start
    compute_start = time.perf_counter()
    outputs = evaluate_batch(dict(request), *tensors)
    torch.cuda.synchronize(device)
    compute_seconds = time.perf_counter() - compute_start
    host_start = time.perf_counter()
    result = {key: value.detach().cpu().numpy() for key, value in outputs.items()}
    torch.cuda.synchronize(device)
    return result, {
        "host_to_device_seconds": transfer_seconds,
        "compute_seconds": compute_seconds,
        "device_to_host_seconds": time.perf_counter() - host_start,
        "max_gpu_memory_bytes": int(torch.cuda.max_memory_allocated(device)),
    }


def traces(
    requests: list[Mapping[str, Any]], ohlcv: Mapping[str, np.ndarray], device: torch.device
) -> dict[str, Any]:
    close, high, low = (ohlcv[name] for name in ("close", "high", "low"))
    gpu_values = tuple(
        torch.as_tensor(value, dtype=torch.float64, device=device) for value in (close, high, low)
    )
    records = []
    for item in requests:
        indicator = int(item["indicator_id"])
        if indicator == 19:
            expected = cpu_supertrend_trace(
                close, high, low, int(item["period"]), float(item["p1"])
            )
            actual = supertrend_trace(*gpu_values, int(item["period"]), float(item["p1"]))
            trace_type = "supertrend"
        elif indicator == 28:
            expected = cpu_psar_trace(close, high, low, float(item["p1"]))
            actual = psar_trace(*gpu_values, float(item["p1"]))
            trace_type = "psar"
        else:
            continue
        torch.cuda.synchronize(device)
        compared = compare_stateful_trace(
            expected,
            {key: value.detach().cpu().numpy() for key, value in actual.items()},
            trace_type=trace_type,
        )
        records.append({"request_id": str(item["request_id"]), "family": trace_type, **compared})
    return {
        "artifact": "GPU_BATCH01_STATEFUL_TRACE_COMPARISON_V1_1",
        "passed": all(item["exact"] for item in records),
        "records": records,
    }


def performance(
    request: Mapping[str, Any], ohlcv: Mapping[str, np.ndarray], device: torch.device
) -> dict[str, Any]:
    # Each family is measured through the same CUDA evaluate_batch entry point.
    families: dict[str, list[Mapping[str, Any]]] = {}
    names = {
        17: "bollinger",
        18: "bollinger",
        19: "supertrend",
        20: "donchian",
        21: "donchian",
        23: "keltner",
        24: "keltner",
        28: "psar",
    }
    for item in request["requests"]:
        families.setdefault(names[int(item["indicator_id"])], []).append(item)
    results = []
    for family, items in sorted(families.items()):
        family_request = {"requests": items}
        _, cold = gpu_outputs(family_request, ohlcv, device)
        _, warm = gpu_outputs(family_request, ohlcv, device)
        values = len(ohlcv["close"]) * len(items)
        results.append(
            {
                "family": family,
                "variants": len(items),
                "values": values,
                "cold": cold,
                "warm": warm,
                "cold_values_per_second": values / cold["compute_seconds"]
                if cold["compute_seconds"]
                else None,
                "warm_values_per_second": values / warm["compute_seconds"]
                if warm["compute_seconds"]
                else None,
            }
        )
    return {
        "artifact": "GPU_BATCH01_PERFORMANCE_METRICS_V1_1",
        "measurement": (
            "CUDA synchronize before each interval; cold is first family invocation, warm is second"
        ),
        "families": results,
    }


def role_semantics(path: Path | None) -> dict[str, Any]:
    if path is None:
        return {
            "status": "not_reconstructable",
            "reason": "Batch01 request records contain no strategy-role assignment; provide "
            "--role-usage-artifact generated from historical lineage.",
        }
    if not path.is_file():
        raise ValueError(f"role usage artifact does not exist: {path}")
    return {"status": "source_artifact_recorded", "path": str(path), "sha256": sha256_file(path)}


def implementation_hashes() -> dict[str, str]:
    root = Path(__file__).resolve().parent
    return {
        path.name: sha256_file(path)
        for path in (
            root / "gpu_batch01_v1_1_runner.py",
            root / "gpu_feature_engine_v1.py",
            root / "gpu_feature_parity_contract_v1_1.py",
        )
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--calibration-csv", type=Path, required=True)
    parser.add_argument("--holdout-csv", type=Path, required=True)
    parser.add_argument("--request-json", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--adversarial-length", type=int, default=256)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--role-usage-artifact", type=Path)
    parser.add_argument("--image-env", default=os.getenv("GPU_BATCH01_IMAGE_ENV"))
    args = parser.parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA GPU is required; no CPU fallback is available")
    args.output_dir.mkdir(parents=True, exist_ok=True)
    request = json.loads(args.request_json.read_text(encoding="utf-8"))
    requests = request.get("requests")
    if not isinstance(requests, list) or not requests:
        raise ValueError("request JSON must contain a non-empty requests list")
    device = torch.device("cuda")

    # Calibration is the only stage allowed to derive and freeze acceptance limits.
    calibration_ohlcv = read_ohlcv(args.calibration_csv)
    calibration_cpu = historical_cpu_oracle(args.calibration_csv, requests)
    calibration_gpu, calibration_timing = gpu_outputs(request, calibration_ohlcv, device)
    calibration_manifest = corpus_manifest(
        {"calibration": calibration_ohlcv},
        {
            "source_csv": str(args.calibration_csv),
            "source_sha256": sha256_file(args.calibration_csv),
        },
    )
    calibration_evidence = calibrate(
        calibration_gpu,
        calibration_cpu["outputs"],
        requests,
        calibration_ohlcv["close"],
        calibration_manifest,
    )
    contract = freeze_contract(calibration_evidence)
    write_json(args.output_dir / "gpu_feature_parity_calibration_v1_1.json", calibration_evidence)
    write_json(args.output_dir / "gpu_feature_parity_contract_v1_1.json", contract)
    np.savez_compressed(
        args.output_dir / f"{ARTIFACT_PREFIX}_calibration_cpu_oracle.npz",
        **calibration_cpu["outputs"],
    )
    np.savez_compressed(
        args.output_dir / f"{ARTIFACT_PREFIX}_calibration_gpu_outputs.npz", **calibration_gpu
    )

    adversarial = deterministic_adversarial_ohlcv(length=args.adversarial_length, seed=args.seed)
    adversarial_csv = args.output_dir / f"{ARTIFACT_PREFIX}_adversarial.csv"
    write_ohlcv_csv(adversarial_csv, adversarial)
    corpora = {
        "calibration": calibration_ohlcv,
        "holdout": read_ohlcv(args.holdout_csv),
        "adversarial": adversarial,
    }
    corpus = corpus_manifest(
        corpora,
        {
            "adversarial_generator": "deterministic_adversarial_ohlcv",
            "seed": args.seed,
            "role_semantics": role_semantics(args.role_usage_artifact),
        },
    )
    write_json(args.output_dir / "gpu_feature_parity_corpus_manifest_v1_1.json", corpus)
    write_json(
        args.output_dir / "gpu_feature_parity_nan_gap_semantics_v1_1.json",
        nan_gap_semantics_manifest(),
    )

    validations, trace_reports = [], []
    for name in ("calibration", "holdout", "adversarial"):
        if name == "calibration":
            csv_path, cpu, gpu, timing = (
                args.calibration_csv,
                calibration_cpu,
                calibration_gpu,
                calibration_timing,
            )
        else:
            csv_path = args.holdout_csv if name == "holdout" else adversarial_csv
            cpu = historical_cpu_oracle(csv_path, requests)
            gpu, timing = gpu_outputs(request, corpora[name], device)
        validation = validate_frozen_contract(contract, gpu, cpu["outputs"], corpora[name]["close"])
        validation.update(
            {
                "corpus": name,
                "cpu_oracle": {key: value for key, value in cpu.items() if key != "outputs"},
                "gpu_timing": timing,
            }
        )
        validations.append(validation)
        trace_reports.append({"corpus": name, **traces(requests, corpora[name], device)})
        if name != "calibration":
            np.savez_compressed(
                args.output_dir / f"{ARTIFACT_PREFIX}_{name}_cpu_oracle.npz", **cpu["outputs"]
            )
            np.savez_compressed(
                args.output_dir / f"{ARTIFACT_PREFIX}_{name}_gpu_outputs.npz", **gpu
            )
    validation_payload = {
        "artifact": "GPU_FEATURE_PARITY_VALIDATION_REPORT_V1_1",
        "contract_sha256": hashlib.sha256(canonical_bytes(contract)).hexdigest(),
        "passed": all(value["passed"] for value in validations),
        "corpora": validations,
    }
    trace_payload = {
        "artifact": "GPU_BATCH01_STATEFUL_TRACE_REPORT_V1_1",
        "passed": all(value["passed"] for value in trace_reports),
        "corpora": trace_reports,
    }
    write_json(args.output_dir / "gpu_feature_parity_validation_v1_1.json", validation_payload)
    write_json(args.output_dir / "gpu_batch01_stateful_trace_comparison_v1_1.json", trace_payload)
    write_json(
        args.output_dir / "gpu_batch01_performance_metrics_v1_1.json",
        performance(request, corpora["holdout"], device),
    )
    root = Path(__file__).resolve().parent
    manifest = {
        "artifact": "GPU_BATCH01_V1_1_EVIDENCE_MANIFEST",
        "schema_version": "1.1",
        "python": sys.version,
        "platform": platform.platform(),
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "gpu": torch.cuda.get_device_properties(device).name,
        "image_env": args.image_env,
        "source_hashes": {
            "historical_band_channel.py": sha256_file(
                root / "control_plane/trading_studio/indicators/historical_band_channel.py"
            )
        },
        "implementation_hashes": implementation_hashes(),
        "data_hashes": {
            "request_json": sha256_file(args.request_json),
            "calibration_csv": sha256_file(args.calibration_csv),
            "holdout_csv": sha256_file(args.holdout_csv),
            "adversarial_csv": sha256_file(adversarial_csv),
        },
        "calibration_gpu_timing": calibration_timing,
        "role_semantics": role_semantics(args.role_usage_artifact),
    }
    write_json(args.output_dir / "gpu_batch01_v1_1_evidence_manifest.json", manifest)
    print(
        json.dumps(
            {
                "output_dir": str(args.output_dir),
                "validation_passed": validation_payload["passed"],
                "trace_passed": trace_payload["passed"],
            },
            allow_nan=False,
        )
    )


if __name__ == "__main__":
    main()