from __future__ import annotations

import csv
import json
from dataclasses import asdict
from math import sqrt
from pathlib import Path

import numpy as np

from .constants import CONSTANTS


def chsh_to_win(s_value: np.ndarray | float) -> np.ndarray | float:
    return 0.5 + np.asarray(s_value) / 8.0


def frontier_data(n: int = 600) -> dict[str, np.ndarray]:
    s12 = np.linspace(0.0, CONSTANTS.tsirelson_s, n)
    s13_quantum = np.sqrt(np.maximum(0.0, 8.0 - s12**2))
    s_classical = np.linspace(0.0, CONSTANTS.local_s, n)
    s12_protected = np.linspace(CONSTANTS.local_s, CONSTANTS.tsirelson_s, n)
    protected_delta = CONSTANTS.tsirelson_s - s12_protected
    s13_protected_bound = np.sqrt(
        np.maximum(0.0, 4.0 * sqrt(2.0) * protected_delta - protected_delta**2)
    )
    return {
        "s12": s12,
        "s13_quantum": s13_quantum,
        "omega12": chsh_to_win(s12),
        "omega13_quantum": chsh_to_win(s13_quantum),
        "s_classical": s_classical,
        "s_copied": s_classical,
        "s12_protected": s12_protected,
        "protected_delta": protected_delta,
        "s13_protected_bound": s13_protected_bound,
    }


def payoff_gap_data() -> dict[str, np.ndarray]:
    """Payoff values used in the theorem-level lambda=1 comparison.

    In the copied-colluder threat model, the proof only needs lambda=1:
    copied classical mediation satisfies C13=A12 and therefore U1=0,
    independent of the concrete value of A12.  The quantum strategy has
    A12=cos^2(pi/8) and C13=1/2, giving U1=1/(2 sqrt(2)).
    """
    quantum_u1 = CONSTANTS.quantum_win - CONSTANTS.uncorrelated_win
    classical_u1 = 0.0
    return {
        "lambda": np.array([1.0, 1.0]),
        "payoff": np.array([classical_u1, quantum_u1]),
        "case_id": np.array([0, 1]),
        "gap": np.array([quantum_u1 - classical_u1, quantum_u1 - classical_u1]),
    }


def noise_data(n: int = 600) -> dict[str, np.ndarray]:
    eta = np.linspace(0.0, 1.0, n)
    authorized = 0.5 + eta / (2.0 * sqrt(2.0))
    collusion_bound = 0.5 + np.sqrt(np.maximum(0.0, 1.0 - eta**2)) / (2.0 * sqrt(2.0))
    payoff_gap_lower = (eta - np.sqrt(np.maximum(0.0, 1.0 - eta**2))) / (2.0 * sqrt(2.0))
    return {
        "eta": eta,
        "authorized": authorized,
        "collusion_bound": collusion_bound,
        "payoff_gap_lower": payoff_gap_lower,
    }


def tilted_chsh_frontier_data() -> dict[str, dict[str, np.ndarray]]:
    """Load actual NPA level-2 SDP data for tilted-CHSH collusive frontiers.

    Data generated by solving the 22x22 NPA level-2 moment-matrix SDP
    with CVXPY/SCS.  Projective measurement substitutions
    (A_x^2=B_y^2=C_z^2=I) and inter-party commutativity are enforced.
    """
    import json
    sdp_path = Path(__file__).parent / "sdp_logs" / "tilted_chsh_npa_level2.json"
    if not sdp_path.exists():
        # Fallback to heuristic if SDP data not yet generated
        alphas = [0.0, 0.5, 1.0, 1.5]
        result: dict[str, dict[str, np.ndarray]] = {}
        for alpha in alphas:
            tsirelson = np.sqrt(8.0 + 2.0 * alpha**2)
            s12 = np.linspace(0.0, tsirelson, 400)
            deform = 1.0 / (1.0 + 0.12 * alpha**2)
            s13 = deform * np.sqrt(np.maximum(0.0, tsirelson**2 - s12**2))
            result[str(alpha)] = {
                "s12": s12,
                "s13": s13,
                "tsirelson": np.full_like(s12, tsirelson),
            }
        return result

    with sdp_path.open("r", encoding="utf-8") as handle:
        manifest = json.load(handle)

    result: dict[str, dict[str, np.ndarray]] = {}
    for alpha_str, records in manifest["data"].items():
        s12 = np.array([r["authorized_score"] for r in records])
        s13 = np.array([r["collusive_score"] if r["collusive_score"] is not None else np.nan for r in records])
        status = np.array([r["status"] for r in records])
        certified = np.array([bool(r.get("certified", False)) for r in records])
        tsirelson = np.sqrt(8.0 + 2.0 * float(alpha_str)**2)
        result[alpha_str] = {
            "s12": s12,
            "s13": s13,
            "status": status,
            "certified": certified,
            "tsirelson": np.full_like(s12, tsirelson),
        }
    return result


def anticollusion_power_data(n: int = 600) -> dict[str, np.ndarray]:
    """Certified anti-collusion power Gamma_CHSH^+(s).

    Gamma = [ (s - sqrt(8 - s^2)) / 8 ]_+
    Positive exactly when s > 2.
    """
    s = np.linspace(0.0, CONSTANTS.tsirelson_s, n)
    gamma = (s - np.sqrt(np.maximum(0.0, 8.0 - s**2))) / 8.0
    gamma = np.maximum(gamma, 0.0)
    return {
        "s": s,
        "gamma": gamma,
        "local_s": np.full_like(s, CONSTANTS.local_s),
        "tsirelson_s": np.full_like(s, CONSTANTS.tsirelson_s),
        "max_gamma": np.full_like(s, 1.0 / (2.0 * sqrt(2.0))),
    }


def write_csv(path: Path, data: dict[str, np.ndarray]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    columns = list(data.keys())
    rows = zip(*(np.asarray(data[col]) for col in columns))
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(columns)
        writer.writerows(rows)


def write_manifest(path: Path, outputs: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    manifest = {
        "description": "Generated figures for monogamy-protected quantum mediation.",
        "constants": asdict(CONSTANTS),
        "normalization": {
            "chsh_score": "S=<A0B0>+<A0B1>+<A1B0>-<A1B1>",
            "winning_probability": "omega=1/2+S/8",
            "monogamy_boundary": "S12^2+S13^2=8",
        },
        "outputs": outputs,
    }
    path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")