"""
ChronoAI Quantum Experiment — Entanglement Proof (Hardware Ready)
Runs identical scenarios WITH and WITHOUT the CNOT gate on real hardware.
Direct response to the "no entanglement needed" claim.

Cost guards:
  - Prints exact cost estimate before any hardware execution
  - Requires typed confirmation before submitting
  - Defaults to local simulator (free) until you explicitly switch
"""

import json
import sys
import numpy as np
from datetime import datetime
from braket.circuits import Circuit
from braket.devices import LocalSimulator

# ══════════════════════════════════════════════════════════════════════
# CONFIGURATION — edit only this block
# ══════════════════════════════════════════════════════════════════════

USE_HARDWARE   = True
SHOTS          = 100
S3_BUCKET = "amazon-braket-quantum-control-flip-eu"
S3_PREFIX      = "quantum-entanglement-proof"

COST_PER_SHOT  = 0.00145   # IQM Garnet per shot
TASK_FEE       = 0.30      # per task submission

# ══════════════════════════════════════════════════════════════════════
# DEVICE SETUP — nothing to edit below this line
# ══════════════════════════════════════════════════════════════════════

if USE_HARDWARE:
    from braket.aws import AwsDevice
    device       = AwsDevice("arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet")
    DEVICE_LABEL = "iqm_garnet"
    s3_folder    = (S3_BUCKET, S3_PREFIX)
else:
    device       = LocalSimulator()
    DEVICE_LABEL = "local"
    s3_folder    = None

# ── Scenarios ─────────────────────────────────────────────────────────
scenarios = [
    {"name": "High money pressure, strategic task",  "revenue": 0.9, "technical": 0.2},
    {"name": "High money pressure, technical task",  "revenue": 0.9, "technical": 0.9},
    {"name": "Low pressure, deep technical build",   "revenue": 0.1, "technical": 0.9},
    {"name": "Low pressure, big picture thinking",   "revenue": 0.1, "technical": 0.1},
]

frame_map = {
    "00": "SCRAPPY / NO BUDGET",
    "01": "ACTIVE SALES MODE",
    "10": "DEEP TECHNICAL BUILD",
    "11": "STRATEGIC / BIG PICTURE",
}

# ── Cost estimator ────────────────────────────────────────────────────

def estimate_cost():
    n_circuits = len(scenarios) * 2   # entangled + independent per scenario
    shot_cost  = n_circuits * SHOTS * COST_PER_SHOT
    task_cost  = n_circuits * TASK_FEE
    return shot_cost, task_cost, shot_cost + task_cost

def confirm_hardware_run():
    shot_cost, task_cost, total = estimate_cost()
    n_circuits = len(scenarios) * 2

    print("\n" + "╔" + "═" * 60 + "╗")
    print("║  HARDWARE RUN — COST ESTIMATE                              ║")
    print("╠" + "═" * 60 + "╣")
    print(f"║  Device:      IQM Garnet (eu-north-1)                     ║")
    print(f"║  Circuits:    {n_circuits} ({len(scenarios)} scenarios × 2 circuit types)           ║")
    print(f"║  Shots each:  {SHOTS}                                            ║")
    print(f"║  Shot cost:   {n_circuits} × {SHOTS} × $0.00145 = ${shot_cost:.2f}               ║")
    print(f"║  Task fees:   {n_circuits} × $0.30     = ${task_cost:.2f}                   ║")
    print(f"║  TOTAL EST:   ${total:.2f}                                    ║")
    print("╠" + "═" * 60 + "╣")
    print("║  This will use AWS Braket credits.                         ║")
    print("╚" + "═" * 60 + "╝")

    answer = input("\n  Type 'run' to proceed, anything else to cancel: ").strip().lower()
    if answer != "run":
        print("\n  Cancelled. No charges incurred.")
        sys.exit(0)
    print()

# ── Circuit builders ──────────────────────────────────────────────────

def build_entangled_circuit(revenue: float, technical: float) -> Circuit:
    """Original circuit — RY encoding + CNOT entanglement."""
    c = Circuit()
    c.ry(0, revenue * np.pi)
    c.ry(1, technical * np.pi)
    c.cnot(0, 1)
    return c

def build_independent_circuit(revenue: float, technical: float) -> Circuit:
    """Control circuit — RY encoding, NO entanglement. Signals resolve independently."""
    c = Circuit()
    c.ry(0, revenue * np.pi)
    c.ry(1, technical * np.pi)
    # No CNOT — this is what the critics said should be sufficient
    return c

# ── Runner ────────────────────────────────────────────────────────────

def run_circuit(circuit: Circuit, label: str) -> dict:
    import time
    print(f"    Submitting {label}...", end=" ", flush=True)

    if USE_HARDWARE:
        task = device.run(circuit, s3_folder, shots=SHOTS)
        print(f"Task ID: {task.id}")
        print(f"    Waiting", end="", flush=True)
        while task.state() not in ("COMPLETED", "FAILED", "CANCELLED"):
            print(".", end="", flush=True)
            time.sleep(10)
        print()
        if task.state() != "COMPLETED":
            raise RuntimeError(f"Task {task.state()}: {task.id}")
        counts = task.result().measurement_counts
    else:
        counts = device.run(circuit, shots=SHOTS).result().measurement_counts
        print("done")

    total  = sum(counts.values())
    dist   = {k: round(v / total, 4) for k, v in counts.items()}
    winner = max(dist, key=dist.get)
    return {
        "frame_selected": winner,
        "frame_label":    frame_map.get(winner, "UNKNOWN"),
        "distribution":   dist,
        "raw_counts":     dict(counts),
    }

# ── Main ──────────────────────────────────────────────────────────────

def main():
    print("=" * 65)
    print("ChronoAI — Entanglement Proof Experiment")
    print(f"Device: {DEVICE_LABEL.upper()} | Shots: {SHOTS}")
    print("=" * 65)

    if USE_HARDWARE:
        confirm_hardware_run()

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results = {
        "experiment":  "entanglement_proof",
        "device":      DEVICE_LABEL,
        "shots":       SHOTS,
        "timestamp":   timestamp,
        "hypothesis":  "CNOT entanglement produces non-obvious frame collapse that independent qubits cannot replicate.",
        "results":     [],
    }

    for sc in scenarios:
        print(f"\n{'─' * 65}")
        print(f"Scenario: {sc['name']}")
        print(f"  Revenue: {sc['revenue']} | Technical: {sc['technical']}")

        entangled   = run_circuit(build_entangled_circuit(sc["revenue"], sc["technical"]),   "entangled  (RY+RY+CNOT)")
        independent = run_circuit(build_independent_circuit(sc["revenue"], sc["technical"]), "independent (RY+RY only)")

        same_winner = entangled["frame_selected"] == independent["frame_selected"]

        e_pct = entangled["distribution"].get(entangled["frame_selected"], 0) * 100
        i_pct = independent["distribution"].get(independent["frame_selected"], 0) * 100

        print(f"  ENTANGLED   → {entangled['frame_selected']} {entangled['frame_label']:<26} @ {e_pct:.0f}%")
        print(f"  INDEPENDENT → {independent['frame_selected']} {independent['frame_label']:<26} @ {i_pct:.0f}%")
        print(f"  Same winner : {'YES' if same_winner else 'NO  ← entanglement changed the outcome'}")

        results["results"].append({
            "scenario":         sc["name"],
            "revenue_pressure": sc["revenue"],
            "is_technical":     sc["technical"],
            "entangled":   {**entangled,   "circuit_type": "RY + RY + CNOT"},
            "independent": {**independent, "circuit_type": "RY + RY (no CNOT)"},
            "same_winner": same_winner,
            "timestamp":   datetime.now().isoformat(),
        })

    filename = f"quantum_results_entanglement_{DEVICE_LABEL}_{timestamp}.json"
    with open(filename, "w") as f:
        json.dump(results, f, indent=2)

    changed = [r for r in results["results"] if not r["same_winner"]]
    print("\n" + "=" * 65)
    print("SUMMARY")
    print(f"  {len(changed)} of {len(scenarios)} scenarios changed winner with entanglement.")
    if changed:
        print("  Scenarios where CNOT made a difference:")
        for r in changed:
            print(f"    • {r['scenario']}")
            print(f"      Entangled:   {r['entangled']['frame_label']}")
            print(f"      Independent: {r['independent']['frame_label']}")
    print(f"\n  Saved: {filename}")
    print("=" * 65)

if __name__ == "__main__":
    main()