"""GPU_FEATURE_PARITY_CONTRACT_V1_2 state-first validation contract."""
from __future__ import annotations

import hashlib
import json
from typing import Any, Mapping

import numpy as np

ARTIFACT = "GPU_FEATURE_PARITY_CONTRACT_V1_2"
CALIBRATION_ARTIFACT = "GPU_FEATURE_PARITY_CALIBRATION_V1_2"
VALIDATION_ARTIFACT = "GPU_FEATURE_PARITY_VALIDATION_V1_2"
ROLE_SURFACE_ARTIFACT = "GPU_FEATURE_PARITY_ROLE_SURFACE_V1_2"
REVIEW_TEMPLATE_ARTIFACT = "GPU_FEATURE_PARITY_REVIEW_TEMPLATE_V1_2"
ROLE_A_EXACT = "A_EXACT"
ROLE_B_BOUNDED = "B_BOUNDED"
ROLE_C_UNVERIFIABLE = "C_UNVERIFIABLE"
UNVERIFIABLE = "UNVERIFIABLE_FROM_RECOVERED_SOURCE"
SUPERTREND_TRACE_KEYS = frozenset(("output", "true_range", "atr", "basic_upper", "basic_lower", "upper", "lower", "direction", "initialized", "prior_close_le_upper", "prior_close_ge_lower", "active_long", "close_ge_lower", "close_le_upper", "output_uses_lower", "direction_transition"))
EXACT_TRACE_KEYS = SUPERTREND_TRACE_KEYS - {"true_range", "atr", "basic_upper", "basic_lower"}
PREDICATE_KEYS = EXACT_TRACE_KEYS - {"output", "upper", "lower", "direction", "initialized", "direction_transition"}


class FrozenContractError(ValueError):
    pass


def canonical_bytes(value: object) -> bytes:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()


def _digest(value: object) -> str:
    return hashlib.sha256(canonical_bytes(value)).hexdigest()


def _same(expected: np.ndarray, actual: np.ndarray) -> bool:
    return expected.shape == actual.shape and np.array_equal(expected, actual, equal_nan=True)


def compare_supertrend_trace(expected: Mapping[str, np.ndarray], actual: Mapping[str, np.ndarray]) -> dict[str, Any]:
    """Check all discrete states and branch predicates exactly; ATR is excluded."""
    required = EXACT_TRACE_KEYS | {"atr"}
    missing = sorted(required - set(expected) | required - set(actual))
    mismatched = sorted(key for key in EXACT_TRACE_KEYS if key in expected and key in actual and not _same(np.asarray(expected[key]), np.asarray(actual[key])))
    atr = exact_output(np.asarray(expected["atr"]), np.asarray(actual["atr"])) if "atr" in expected and "atr" in actual else None
    atr_structure_exact = atr is not None and atr["shape_exact"] and atr["nan_mask_exact"]
    state_exact = not missing and not mismatched
    return {"state_exact": state_exact, "required_keys": sorted(required), "mismatched_exact_keys": mismatched, "missing_keys": missing, "atr": atr, "atr_structure_exact": atr_structure_exact, "branch_predicates_exact": state_exact and all(key not in mismatched for key in PREDICATE_KEYS), "transition_checks_exact": state_exact}


def atr_metrics(expected: np.ndarray, actual: np.ndarray) -> dict[str, float]:
    finite = np.isfinite(expected) & np.isfinite(actual)
    delta = np.abs(np.asarray(actual)[finite] - np.asarray(expected)[finite])
    return {"max_absolute_error": float(delta.max()) if delta.size else 0.0, "mae": float(delta.mean()) if delta.size else 0.0}


def output_metrics(expected: np.ndarray, actual: np.ndarray) -> dict[str, float]:
    """Stable, JSON-safe continuous-output measurements."""
    finite = np.isfinite(expected) & np.isfinite(actual)
    delta = np.abs(np.asarray(actual)[finite] - np.asarray(expected)[finite])
    denominator = np.abs(np.asarray(expected)[finite])
    relative = delta[denominator != 0] / denominator[denominator != 0]
    return {
        "max_absolute_error": float(delta.max()) if delta.size else 0.0,
        "max_relative_error": float(relative.max()) if relative.size else 0.0,
        "mae": float(delta.mean()) if delta.size else 0.0,
    }


def exact_output(expected: np.ndarray, actual: np.ndarray) -> dict[str, Any]:
    return {
        "shape_exact": expected.shape == actual.shape,
        "dtype_exact": expected.dtype == actual.dtype,
        "nan_mask_exact": expected.shape == actual.shape
        and np.array_equal(np.isnan(expected), np.isnan(actual)),
        "values_exact": _same(expected, actual),
    }


def family_for_indicator(indicator_id: int) -> str:
    return {17: "bollinger", 18: "bollinger", 19: "supertrend", 20: "donchian", 21: "donchian", 23: "keltner", 24: "keltner", 28: "psar"}[indicator_id]


def cpu_supertrend_trace(close: np.ndarray, high: np.ndarray, low: np.ndarray, period: int, multiplier: float) -> dict[str, np.ndarray]:
    """Historical Supertrend semantics with explicit V1.2 branch evidence."""
    from control_plane.trading_studio.indicators.historical_band_channel import _rma, _true_range
    n = len(close)
    true_range = _true_range(high, low, close)
    atr = _rma(true_range, period)
    output = np.full(n, np.nan)
    upper, lower, basic_upper, basic_lower = output.copy(), output.copy(), output.copy(), output.copy()
    direction = np.zeros(n, dtype=np.int8)
    initialized = np.zeros(n, dtype=bool)
    predicates = {name: np.zeros(n, dtype=bool) for name in PREDICATE_KEYS}
    direction_transition = np.zeros(n, dtype=bool)
    for index in range(period, n):
        midpoint = (high[index] + low[index]) / 2.0
        basic_upper[index], basic_lower[index] = midpoint + multiplier * atr[index], midpoint - multiplier * atr[index]
        if index == period:
            upper[index], lower[index] = basic_upper[index], basic_lower[index]
            predicates["close_ge_lower"][index] = close[index] > lower[index]
            predicates["output_uses_lower"][index] = predicates["close_ge_lower"][index]
            output[index] = lower[index] if predicates["output_uses_lower"][index] else upper[index]
        else:
            predicates["prior_close_le_upper"][index] = close[index - 1] <= upper[index - 1]
            predicates["prior_close_ge_lower"][index] = close[index - 1] >= lower[index - 1]
            upper[index] = min(basic_upper[index], upper[index - 1]) if predicates["prior_close_le_upper"][index] else basic_upper[index]
            lower[index] = max(basic_lower[index], lower[index - 1]) if predicates["prior_close_ge_lower"][index] else basic_lower[index]
            predicates["active_long"][index] = direction[index - 1] == 1
            predicates["close_ge_lower"][index] = close[index] >= lower[index]
            predicates["close_le_upper"][index] = close[index] <= upper[index]
            predicates["output_uses_lower"][index] = predicates["active_long"][index] and predicates["close_ge_lower"][index] or not predicates["active_long"][index] and not predicates["close_le_upper"][index]
            output[index] = lower[index] if predicates["output_uses_lower"][index] else upper[index]
        direction[index] = 1 if close[index] > output[index] else -1
        initialized[index] = True
        if index > period:
            direction_transition[index] = direction[index] != direction[index - 1]
    return {
        "output": output,
        "true_range": true_range,
        "atr": atr,
        "basic_upper": basic_upper,
        "basic_lower": basic_lower,
        "upper": upper,
        "lower": lower,
        "direction": direction,
        "initialized": initialized,
        "direction_transition": direction_transition,
        **predicates,
    }


def calibrate_supertrend(request_id: str, expected: Mapping[str, np.ndarray], actual: Mapping[str, np.ndarray], *, atr_limits: Mapping[str, float] | None = None) -> dict[str, Any]:
    trace = compare_supertrend_trace(expected, actual)
    measured = atr_metrics(np.asarray(expected["atr"]), np.asarray(actual["atr"])) if trace["state_exact"] and trace["atr_structure_exact"] else None
    return {"request_id": request_id, "family": "supertrend", "role": ROLE_A_EXACT, "role_classification": {"discrete_state_and_transitions": ROLE_A_EXACT, "atr": ROLE_B_BOUNDED}, "trace": trace, "atr": measured, "atr_limits": dict(atr_limits) if atr_limits is not None else measured}


def calibrate_output(request_id: str, indicator_id: int, expected: np.ndarray, actual: np.ndarray) -> dict[str, Any]:
    """Calibrate a non-Supertrend Batch01 output under its family evidence class."""
    family = family_for_indicator(indicator_id)
    exact = family in {"donchian", "psar"}
    result = exact_output(expected, actual)
    return {
        "request_id": request_id,
        "indicator_id": indicator_id,
        "family": family,
        "role": ROLE_A_EXACT if exact else ROLE_B_BOUNDED,
        "output": result,
        # Bounds are deliberately emitted only for continuous families.
        "limits": None if exact else output_metrics(expected, actual),
    }


def freeze_contract(calibration: Mapping[str, Any]) -> dict[str, Any]:
    if calibration.get("artifact") != CALIBRATION_ARTIFACT:
        raise FrozenContractError("only V1.2 calibration evidence can be frozen")
    limits = {}
    for record in calibration.get("records", []):
        if record.get("family") == "supertrend":
            if not record["trace"]["state_exact"] or not record["trace"]["atr_structure_exact"] or record["atr"] is None:
                raise FrozenContractError("cannot freeze a state-inexact calibration")
            limits[record["request_id"]] = {"family": "supertrend", "role": ROLE_A_EXACT, "atr_limits": dict(record["atr_limits"])}
        elif record.get("family") in {"donchian", "psar"}:
            if not all(record["output"].values()):
                raise FrozenContractError("cannot freeze an inexact discrete/state calibration")
            limits[record["request_id"]] = {"family": record["family"], "role": ROLE_A_EXACT}
        elif record.get("family") in {"bollinger", "keltner"}:
            if not record["output"]["shape_exact"] or not record["output"]["nan_mask_exact"]:
                raise FrozenContractError("cannot freeze a structurally inexact continuous calibration")
            limits[record["request_id"]] = {"family": record["family"], "role": ROLE_B_BOUNDED, "output_limits": dict(record["limits"])}
        else:
            raise FrozenContractError("unknown Batch01 family in calibration")
    return {"artifact": ARTIFACT, "schema_version": "1.2", "status": "frozen", "role_classification": {"discrete_state_and_transitions": ROLE_A_EXACT, "atr": ROLE_B_BOUNDED, "unrecovered_batch01_roles": ROLE_C_UNVERIFIABLE}, "calibration_sha256": _digest(calibration), "feature_limits": limits}


def validate_all_frozen_contract(contract: Mapping[str, Any], expected_outputs: Mapping[str, np.ndarray], actual_outputs: Mapping[str, np.ndarray], traces: Mapping[str, tuple[Mapping[str, np.ndarray], Mapping[str, np.ndarray]]]) -> dict[str, Any]:
    """Validate every frozen Batch01 feature without deriving new limits."""
    if contract.get("artifact") != ARTIFACT or contract.get("status") != "frozen":
        raise FrozenContractError("provided contract is not frozen GPU_FEATURE_PARITY_CONTRACT_V1_2")
    limits = contract.get("feature_limits", {})
    if set(limits) != set(expected_outputs) or set(limits) != set(actual_outputs):
        raise FrozenContractError("contract and outputs must name exactly the same features")
    records = []
    for request_id, limit in limits.items():
        family = limit["family"]
        expected, actual = np.asarray(expected_outputs[request_id]), np.asarray(actual_outputs[request_id])
        if family == "supertrend":
            if request_id not in traces:
                raise FrozenContractError("Supertrend trace missing")
            trace = compare_supertrend_trace(*traces[request_id])
            atr = atr_metrics(np.asarray(traces[request_id][0]["atr"]), np.asarray(traces[request_id][1]["atr"])) if trace["state_exact"] and trace["atr_structure_exact"] else None
            bounded = atr is not None and all(atr[key] <= float(limit["atr_limits"][key]) for key in atr)
            passed = trace["state_exact"] and trace["atr_structure_exact"] and bounded
            records.append({"request_id": request_id, "family": family, "role": ROLE_A_EXACT, "trace": trace, "atr": atr, "atr_bounded": bounded, "passed": passed})
        elif family in {"donchian", "psar"}:
            output = exact_output(expected, actual)
            records.append({"request_id": request_id, "family": family, "role": ROLE_A_EXACT, "output": output, "passed": all(output.values())})
        else:
            output = exact_output(expected, actual)
            metrics = output_metrics(expected, actual)
            bounded = all(metrics[key] <= float(limit["output_limits"][key]) for key in metrics)
            passed = output["shape_exact"] and output["nan_mask_exact"] and bounded
            records.append({"request_id": request_id, "family": family, "role": ROLE_B_BOUNDED, "output": output, "metrics": metrics, "bounded": bounded, "passed": passed})
    return {"artifact": VALIDATION_ARTIFACT, "schema_version": "1.2", "contract_sha256": _digest(contract), "passed": all(record["passed"] for record in records), "records": records}


def validate_frozen_contract(contract: Mapping[str, Any], traces: Mapping[str, tuple[Mapping[str, np.ndarray], Mapping[str, np.ndarray]]]) -> dict[str, Any]:
    """Mechanical validation only: it never derives or expands frozen limits."""
    if contract.get("artifact") != ARTIFACT or contract.get("schema_version") != "1.2" or contract.get("status") != "frozen":
        raise FrozenContractError("provided contract is not frozen GPU_FEATURE_PARITY_CONTRACT_V1_2")
    limits = contract.get("feature_limits")
    if not isinstance(limits, dict) or set(limits) != set(traces):
        raise FrozenContractError("contract and traces must name exactly the same features")
    records = []
    for request_id, (expected, actual) in traces.items():
        trace = compare_supertrend_trace(expected, actual)
        metric = atr_metrics(np.asarray(expected["atr"]), np.asarray(actual["atr"])) if trace["state_exact"] and trace["atr_structure_exact"] else None
        bounded = metric is not None and all(metric[name] <= float(limits[request_id]["atr_limits"][name]) for name in metric)
        records.append({"request_id": request_id, "role": ROLE_A_EXACT, "role_classification": {"discrete_state_and_transitions": ROLE_A_EXACT, "atr": ROLE_B_BOUNDED}, "trace": trace, "atr": metric, "atr_bounded": bounded, "passed": trace["state_exact"] and trace["atr_structure_exact"] and bounded})
    return {"artifact": VALIDATION_ARTIFACT, "schema_version": "1.2", "contract_sha256": _digest(contract), "passed": all(record["passed"] for record in records), "records": records}


def role_surface(requests: list[Mapping[str, Any]]) -> dict[str, Any]:
    """Classify only recovered Cohort001 role assignments; unknown Batch01 roles do not block parity."""
    try:
        from control_plane.trading_studio.indicators.registry import _HS22_REQUIRED_TRIPLES
    except ModuleNotFoundError:
        recovered = {}
        unavailable_reason = "control_plane.trading_studio.indicators.registry is unavailable"
    else:
        recovered = {(indicator, period, p1): slot for indicator, period, p1, slot in _HS22_REQUIRED_TRIPLES}
        unavailable_reason = None
    records = []
    for item in requests:
        key = (int(item["indicator_id"]), int(item["period"]), float(item["p1"]))
        slot = recovered.get(key)
        records.append({"request_id": str(item["request_id"]), "indicator_id": key[0], "period": key[1], "p1": key[2], "cohort001_role": slot, "classification": ROLE_A_EXACT if key[0] == 19 and slot else ROLE_C_UNVERIFIABLE, "role_status": "RECOVERED_COHORT001" if slot else UNVERIFIABLE, "reason": unavailable_reason if not slot else None, "parity_blocker": False if not slot else True})
    return {"artifact": ROLE_SURFACE_ARTIFACT, "schema_version": "1.2", "source": "Cohort001 recovered role surface", "records": records}


def review_template() -> dict[str, Any]:
    return {"artifact": REVIEW_TEMPLATE_ARTIFACT, "schema_version": "1.2", "status": "review_required_not_validated", "required_review": ["calibration evidence is separate from holdout and adversarial evidence", "exact state, transition, and branch-predicate checks pass", "ATR bounded checks were evaluated only after exact state validation", "C-unverifiable roles are recorded and excluded as parity blockers"], "decision": None}