"""
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

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

USE_HARDWARE   = True   # ← set True to run on IonQ Forte-1
SHOTS          = 100      # ← 100 shots recommended for first hardware run
S3_BUCKET      = "amazon-braket-us-east-1-803694230509"
S3_PREFIX      = "quantum-control-flip"

# IonQ Forte-1 pricing
COST_PER_SHOT  = 0.08    # USD per shot on Forte-1
TASK_FEE       = 0.30    # USD per task submission

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

if USE_HARDWARE:
    from braket.aws import AwsDevice
    device       = AwsDevice("arn:aws:braket:us-east-1::device/qpu/ionq/Forte-1")
    DEVICE_LABEL = "ionq_forte1"
    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:      IonQ Forte-1 (us-east-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("╚" + "═" * 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_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}")
    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()