"""
ChronoAI Quantum Experiment – Control Qubit Flip (Hardware Ready)
Triangulation experiment: tests variable dominance theory by flipping
which qubit holds CNOT control position.

Hypothesis: The Scenario 3 anomaly moves to Scenario 2 when control flips,
proving the circuit detects variable dominance, not individual signal strength.

Cost guards:
  - Prints exact cost estimate before any hardware execution
  - Requires typed confirmation before submitting to hardware
  - Single DEVICE_LABEL / SHOTS block – no override bug
  - 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
import boto3

# ╔══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION – edit only this block
# ╚══════════════════════════════════════════════════════════════════════════════

USE_HARDWARE   = True   # ← set True to run on IQM Garnet
SHOTS          = 100    # ← 100 shots recommended for first hardware run
S3_BUCKET      = "amazon-braket-quantum-control-flip-eu"  # ← New bucket for eu-north-1
S3_PREFIX      = "quantum-control-flip"

# IQM Garnet pricing (update these with current rates)
COST_PER_SHOT  = 0.08    # USD per shot on IQM Garnet (check current pricing)
TASK_FEE       = 0.30    # USD per task submission

# ╔══════════════════════════════════════════════════════════════════════════════
# DEVICE SETUP – nothing to edit below this line
# ╚══════════════════════════════════════════════════════════════════════════════

def create_s3_bucket_if_needed():
    """Create S3 bucket for IQM Garnet results if it doesn't exist"""
    try:
        s3_client = boto3.client('s3', region_name='eu-north-1')
        s3_client.create_bucket(
            Bucket=S3_BUCKET,
            CreateBucketConfiguration={'LocationConstraint': 'eu-north-1'}
        )
        print(f"Created S3 bucket: {S3_BUCKET}")
    except Exception as e:
        if "BucketAlreadyOwnedByYou" in str(e) or "BucketAlreadyExists" in str(e):
            print(f"S3 bucket {S3_BUCKET} already exists")
        else:
            print(f"S3 bucket creation note: {e}")

if USE_HARDWARE:
    from braket.aws import AwsDevice
    create_s3_bucket_if_needed()  # Create bucket before device setup
    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",
}

# What the math predicts for each circuit type
PREDICTIONS = {
    "original": ["11", "10", "01", "00"],
    "flipped":  ["10", "01", "11", "00"],
}

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

def estimate_cost() -> float:
    """
    2 circuits per scenario (original + flipped)
    × 4 scenarios
    × SHOTS per circuit
    + task fee per circuit submission
    """
    n_circuits   = len(scenarios) * 2
    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.08 = ${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("║  Note: IQM Garnet uses different gate set than IonQ        ║")
    print("╚" + "═" * 60 + "╝")

    # Check device status
    try:
        print(f"\nDevice Status: {device.status}")
        print(f"Queue Depth: {device.queue_depth()}")
    except Exception as e:
        print(f"Could not check device status: {e}")

    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_original(revenue: float, technical: float) -> Circuit:
    """Revenue qubit (q0) controls technical qubit (q1) – original design."""
    c = Circuit()
    c.ry(0, revenue   * np.pi)
    c.ry(1, technical * np.pi)
    c.cnot(0, 1)
    return c

def build_flipped(revenue: float, technical: float) -> Circuit:
    """Technical qubit (q1) controls revenue qubit (q0) – triangulation test."""
    c = Circuit()
    c.ry(0, revenue   * np.pi)
    c.ry(1, technical * np.pi)
    c.cnot(1, 0)              # ← flipped control
    return c

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

def run_circuit(circuit: Circuit, label: str) -> dict:
    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)
        import time
        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 – Control Qubit Flip Triangulation Experiment")
    print(f"Device: {DEVICE_LABEL.upper()} | Shots: {SHOTS}")
    if USE_HARDWARE:
        print("Target: IQM Garnet (20-qubit superconducting, eu-north-1)")
    print("=" * 65)

    if USE_HARDWARE:
        confirm_hardware_run()

    print("Predictions if variable-dominance theory holds:")
    for i, sc in enumerate(scenarios):
        o = frame_map[PREDICTIONS["original"][i]]
        f = frame_map[PREDICTIONS["flipped"][i]]
        tag = " ← SHOULD CHANGE" if PREDICTIONS["original"][i] != PREDICTIONS["flipped"][i] else ""
        print(f"  S{i+1} ({sc['revenue']}, {sc['technical']}): {o} → {f}{tag}")
    print()

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results = {
        "experiment":  "control_flip_triangulation",
        "device":      DEVICE_LABEL,
        "shots":       SHOTS,
        "timestamp":   timestamp,
        "hypothesis":  (
            "Flipping CNOT control from q0 (revenue) to q1 (technical) moves "
            "the Scenario 3 anomaly to Scenario 2, confirming the circuit detects "
            "variable dominance relationships, not individual signal strength."
        ),
        "results": [],
    }

    anomaly_moved   = False
    anomaly_gone    = False
    confirmations   = 0

    for i, sc in enumerate(scenarios):
        print(f"{'─' * 65}")
        print(f"S{i+1}: {sc['name']}")
        print(f"  Revenue: {sc['revenue']} | Technical: {sc['technical']}")

        orig = run_circuit(build_original(sc["revenue"], sc["technical"]), "original (q0→q1)")
        flip = run_circuit(build_flipped( sc["revenue"], sc["technical"]), "flipped  (q1→q0)")

        po = PREDICTIONS["original"][i]
        pf = PREDICTIONS["flipped"][i]
        om = orig["frame_selected"] == po
        fm = flip["frame_selected"] == pf
        if om and fm:
            confirmations += 1

        o_pct = orig["distribution"].get(orig["frame_selected"], 0) * 100
        f_pct = flip["distribution"].get(flip["frame_selected"], 0) * 100

        print(f"  ORIGINAL → {orig['frame_selected']} {orig['frame_label']:<26} "
              f"@ {o_pct:.0f}%  pred:{po} {'✓' if om else '✗'}")
        print(f"  FLIPPED  → {flip['frame_selected']} {flip['frame_label']:<26} "
              f"@ {f_pct:.0f}%  pred:{pf} {'✓' if fm else '✗'}")

        if i == 1 and flip["frame_selected"] == "01":
            print("  ★ NEW ANOMALY at S2 confirmed")
            anomaly_moved = True
        if i == 2 and flip["frame_selected"] != "01":
            print("  ★ S3 anomaly disappeared")
            anomaly_gone = True

        results["results"].append({
            "scenario_index":   i + 1,
            "scenario":         sc["name"],
            "revenue_pressure": sc["revenue"],
            "is_technical":     sc["technical"],
            "original": {**orig, "circuit_type": "CNOT(q0→q1)", "predicted": po, "hit": om},
            "flipped":  {**flip, "circuit_type": "CNOT(q1→q0)", "predicted": pf, "hit": fm},
            "timestamp": datetime.now().isoformat(),
        })

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

    print("\n" + "=" * 65)
    print("SUMMARY")
    print(f"  Predictions confirmed : {confirmations} / {len(scenarios)}")
    print(f"  New anomaly at S2     : {'YES' if anomaly_moved else 'NO'}")
    print(f"  S3 anomaly gone       : {'YES' if anomaly_gone  else 'NO'}")
    print()
    if confirmations == 4 and anomaly_moved and anomaly_gone:
        print("  STRONG CONFIRMATION")
        print("  The circuit detects variable dominance relationships.")
        print("  The Scenario 3 finding is not a happy accident.")
    elif confirmations >= 3:
        print("  PARTIAL CONFIRMATION – trending toward dominance theory.")
        print("  Increase shots or rerun for stronger signal.")
    else:
        print("  INCONCLUSIVE – hardware noise may be overwhelming signal.")
        print("  Try 100 shots and rerun.")
    print(f"\n  Saved: {filename}")
    print("=" * 65)

if __name__ == "__main__":
    main()
