"""
ChronoAI Quantum Experiment — Entanglement Proof
Runs identical scenarios WITH and WITHOUT the CNOT gate on real IonQ hardware.
Direct response to the "no entanglement needed" claim.
Results saved to JSON matching existing file naming convention.
"""

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

# ── Swap to real IonQ when ready ──────────────────────────────────────────────
from braket.aws import AwsDevice
device = AwsDevice("arn:aws:braket:us-east-1::device/qpu/ionq/Forte-1")  # Updated to Forte-1!
DEVICE_LABEL = "ionq_entanglement"   # Unique label for this experiment
SHOTS = 50  # Start with 50 shots (~$2-3) to test

# from braket.aws import AwsDevice
# device = AwsDevice("arn:aws:braket:us-east-1::device/qpu/ionq/Aria-1")
#device = LocalSimulator()
DEVICE_LABEL = "ionq"   # change to "local" if testing locally first
SHOTS = 100

# ── 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 ─────────────────────────────────────────────────────────────────
frame_map = {
    "00": "SCRAPPY / NO BUDGET",
    "01": "ACTIVE SALES MODE",
    "10": "DEEP TECHNICAL BUILD",
    "11": "STRATEGIC / BIG PICTURE",
}

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

def build_entangled_circuit(revenue: float, technical: float) -> Circuit:
    """Our original circuit — RY encoding + CNOT entanglement."""
    c = Circuit()
    c.ry(0, revenue * np.pi)
    c.ry(1, technical * np.pi)
    c.cnot(0, 1)   # signals interact — this is the key gate
    return c


def build_independent_circuit(revenue: float, technical: float) -> Circuit:
    """Grok's approach — RY encoding, NO entanglement. Signals are independent."""
    c = Circuit()
    c.ry(0, revenue * np.pi)
    c.ry(1, technical * np.pi)
    # No CNOT — each qubit resolves independently
    return c


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

def run_scenario(circuit: Circuit, shots: int) -> dict:
    task   = device.run(circuit, shots=shots)
    counts = task.result().measurement_counts
    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():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    question  = "Should I spend today building the quantum context layer or making cold calls?"

    results = {
        "experiment":  "entanglement_proof",
        "device":      DEVICE_LABEL,
        "shots":       SHOTS,
        "question":    question,
        "timestamp":   timestamp,
        "hypothesis":  "CNOT entanglement produces non-obvious frame collapse that independent qubits cannot replicate.",
        "results":     [],
    }

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

    for sc in scenarios:
        print(f"\nScenario: {sc['name']}")
        print(f"  Revenue pressure: {sc['revenue']} | Is technical: {sc['technical']}")

        # Run both circuits
        entangled   = run_scenario(build_entangled_circuit(sc["revenue"], sc["technical"]),   SHOTS)
        independent = run_scenario(build_independent_circuit(sc["revenue"], sc["technical"]), SHOTS)

        # Flag where the winner differs — that's the interesting result
        same_winner = entangled["frame_selected"] == independent["frame_selected"]

        print(f"  ENTANGLED  → {entangled['frame_selected']} ({entangled['frame_label']}) "
              f"@ {entangled['distribution'].get(entangled['frame_selected'], 0)*100:.0f}%")
        print(f"  INDEPENDENT→ {independent['frame_selected']} ({independent['frame_label']}) "
              f"@ {independent['distribution'].get(independent['frame_selected'], 0)*100:.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(),
        })

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

    # ── Summary ───────────────────────────────────────────────────────────────
    changed = [r for r in results["results"] if not r["same_winner"]]
    print("\n" + "=" * 65)
    print(f"SUMMARY: {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"\nResults saved to: {filename}")
    print("=" * 65)


if __name__ == "__main__":
    main()