#!/usr/bin/env python3
"""
Quantum Context Frame Experiment - Real Hardware Version
ChronoAI Solutions - Enhanced for Hardware Deployment

This experiment demonstrates how quantum hardware effects (noise, decoherence)
can influence AI decision-making by changing which context frame is selected.

The key insight: Quantum noise isn't a bug - it's a feature that adds
natural randomness to AI context selection.
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime
import numpy as np

# ── Braket imports ────────────────────────────────────────────────────────────
try:
    from braket.aws import AwsDevice, AwsQuantumTask
    from braket.circuits import Circuit
    from braket.devices import LocalSimulator
    print("✅ Braket SDK loaded successfully")
except ImportError:
    print("ERROR: amazon-braket-sdk not installed.")
    print("Run: pip install amazon-braket-sdk")
    sys.exit(1)

# ── Anthropic import (optional) ───────────────────────────────────────────────
try:
    import anthropic
    ANTHROPIC_AVAILABLE = True
    print("✅ Anthropic SDK available")
except ImportError:
    ANTHROPIC_AVAILABLE = False
    print("ℹ️  Anthropic SDK not available (optional)")

# ── Configuration ─────────────────────────────────────────────────────────────
S3_BUCKET = "amazon-braket-quantum-control-flip-eu"  # Your working IQM bucket
S3_PREFIX = "quantum-frame-experiment"
AWS_REGION = "eu-north-1"  # Changed from us-east-1

# Device ARNs
DEVICES = {
    "iqm":   "arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet",  # Changed from ionq
    "sim":   "arn:aws:braket:::device/quantum-simulator/amazon/sv1",
    "local": "local",
}

SHOTS = 1000

# ── Context frames ────────────────────────────────────────────────────────────
CONTEXT_FRAMES = {
    "00": {
        "label": "SCRAPPY / NO BUDGET",
        "prompt": "You are helping a solo entrepreneur with no budget and high hustle."
    },
    "01": {
        "label": "ACTIVE SALES MODE", 
        "prompt": "You are helping someone in active sales mode with real clients on the line."
    },
    "10": {
        "label": "DEEP TECHNICAL BUILD",
        "prompt": "You are in deep technical build mode. Full code, no hand-holding."
    },
    "11": {
        "label": "STRATEGIC / BIG PICTURE",
        "prompt": "You are in strategic planning mode. Think long-term."
    },
}

USER_QUESTION = "Should I spend today building the quantum context layer or making cold calls?"

# ── Circuit builder ───────────────────────────────────────────────────────────
def build_frame_circuit(revenue_pressure: float, is_technical: float) -> Circuit:
    """
    Build the quantum context selection circuit.
    
    - Qubit 0: Revenue pressure (0.0 = no pressure → 1.0 = urgent)  
    - Qubit 1: Technical vs strategic (0.0 = strategic → 1.0 = technical)
    - CNOT: Entangles the two contexts so they influence each other
    """
    circuit = Circuit()
    circuit.ry(0, revenue_pressure * np.pi)
    circuit.ry(1, is_technical * np.pi)
    circuit.cnot(0, 1)
    return circuit

# ── Circuit runner ────────────────────────────────────────────────────────────
def run_circuit(circuit: Circuit, device_key: str) -> dict:
    """Run the quantum circuit and return measurement counts."""
    print(f"\n🔬 Running circuit on {device_key.upper()}...")
    
    if device_key == "local":
        device = LocalSimulator()
        result = device.run(circuit, shots=SHOTS).result()
        print(f"   ✅ Local simulation completed ({SHOTS} shots)")
        return result.measurement_counts
    
    elif device_key == "sim":
        print(f"   📤 Submitting to AWS SV1 simulator...")
        device = AwsDevice(DEVICES[device_key])
        s3_folder = (S3_BUCKET, S3_PREFIX)
        
        task = device.run(circuit, s3_folder, shots=SHOTS)
        print(f"   📋 Task ID: {task.id}")
        print(f"   ⏳ Waiting for completion...")
        
        # Wait for completion
        start_time = time.time()
        while True:
            state = task.state()
            elapsed = int(time.time() - start_time)
            print(f"   🔄 [{elapsed}s] State: {state}", end="\r")
            
            if state == "COMPLETED":
                print(f"\n   ✅ Task completed in {elapsed}s!")
                break
            elif state in ("FAILED", "CANCELLED"):
                print(f"\n   ❌ Task {state}")
                raise RuntimeError(f"Task {state}: {task.id}")
            
            time.sleep(5)
        
        result = task.result()
        return result.measurement_counts
    
    elif device_key == "iqm":
        print(f"   📤 Submitting to IQM Garnet...")
        device = AwsDevice(DEVICES[device_key])
        s3_folder = (S3_BUCKET, S3_PREFIX)
        
        # Check device status
        print(f"   📡 Device status: {device.status}")
        print(f"   📊 Queue depth: {device.queue_depth()}")
        
        task = device.run(circuit, s3_folder, shots=SHOTS)
        print(f"   📋 Task ID: {task.id}")
        print(f"   ⏳ Waiting for completion (this may take several minutes)...")
        
        # Wait for completion
        start_time = time.time()
        while True:
            state = task.state()
            elapsed = int(time.time() - start_time)
            
            if state == "COMPLETED":
                print(f"\n   ✅ Task completed in {elapsed}s!")
                break
            elif state in ("FAILED", "CANCELLED"):
                print(f"\n   ❌ Task {state}")
                raise RuntimeError(f"Task {state}: {task.id}")
            else:
                print(f"   🔄 [{elapsed}s] State: {state}", end="\r")
            
            time.sleep(10)
        
        result = task.result()
        return result.measurement_counts

# ── Results analyzer ──────────────────────────────────────────────────────────
def counts_to_frame(counts: dict) -> tuple[str, dict]:
    """Convert measurement counts to frame probabilities."""
    total = sum(counts.values())
    if total == 0:
        return "00", {"00": 1.0, "01": 0.0, "10": 0.0, "11": 0.0}
    
    # Normalize to probabilities
    probabilities = {}
    for state, count in counts.items():
        key = state.zfill(2)[-2:]  # Ensure 2-bit format
        probabilities[key] = probabilities.get(key, 0) + count / total
    
    # Ensure all states are represented
    for state in ["00", "01", "10", "11"]:
        if state not in probabilities:
            probabilities[state] = 0.0
    
    # Winner takes all
    winner = max(probabilities, key=probabilities.get)
    return winner, probabilities

# ── Main experiment runner ────────────────────────────────────────────────────
def run_scenario(label: str, revenue: float, technical: float, device_key: str):
    """Run one complete scenario."""
    
    print(f"\n{'─' * 70}")
    print(f"🧪 SCENARIO: {label}")
    print(f"   Revenue pressure: {revenue:.1f}  |  Technical focus: {technical:.1f}")
    print(f"   Device: {device_key.upper()}")
    print(f"{'─' * 70}")

    # Build circuit
    circuit = build_frame_circuit(revenue, technical)
    print(f"📐 Circuit: RY({revenue:.1f}π) → RY({technical:.1f}π) → CNOT → Measure")

    # Run circuit
    try:
        counts = run_circuit(circuit, device_key)
    except Exception as e:
        print(f"❌ Circuit execution failed: {e}")
        return None

    # Analyze results
    frame_id, probabilities = counts_to_frame(counts)
    frame = CONTEXT_FRAMES.get(frame_id, {"label": "UNKNOWN"})

    # Display results
    print(f"\n📊 Quantum Measurement Results:")
    for state in ["00", "01", "10", "11"]:
        prob = probabilities.get(state, 0)
        percentage = prob * 100
        bar_length = int(percentage / 3)
        bar = "█" * bar_length
        frame_label = CONTEXT_FRAMES.get(state, {}).get("label", "Unknown")
        winner_mark = " ◄ SELECTED" if state == frame_id else ""
        
        print(f"   {state} ({frame_label:<25}) {percentage:5.1f}%  {bar}{winner_mark}")

    print(f"\n🎯 Quantum Selected Frame: {frame_id} - {frame['label']}")
    
    return {
        "scenario": label,
        "revenue_pressure": revenue,
        "is_technical": technical,
        "device": device_key,
        "frame_selected": frame_id,
        "frame_label": frame["label"],
        "distribution": probabilities,
        "raw_counts": counts,
        "timestamp": datetime.now().isoformat(),
    }

# ── Cost estimation ───────────────────────────────────────────────────────────
def print_cost_estimate():
    """Print cost estimates."""
    print("\n💰 COST ESTIMATES (approximate)")
    print("─" * 55)
    print(f"  Local simulator  : FREE")
    print(f"  SV1 (managed)    : ~$0.075/min × ~2min = ~$0.15 per scenario")
    print(f"  IQM Garnet       : ~$0.30 base + $0.08/shot")
    print(f"                     1000 shots = ~$80.30 per scenario")
    print(f"")
    print(f"  📋 Recommended testing order:")
    print(f"    1. --device local      (FREE - verify circuit works)")
    print(f"    2. --device sim        (~$0.15 - test AWS integration)")
    print(f"    3. --device iqm        (~$80.30 - real quantum hardware)")
    print(f"")
    print(f"  💡 For learning: Start with local, then try sim")
    print(f"     For research: Compare sim vs iqm results")

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

# ── Main function ─────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Quantum Frame Experiment - Hardware Version")
    parser.add_argument("--device", choices=["ionq", "iqm", "sim", "local"], default="local",
                       help="Device to run on (default: local)")
    parser.add_argument("--cost", action="store_true",
                       help="Show cost estimate and exit")
    parser.add_argument("--shots", type=int, default=1000, 
                       help="Number of shots (default: 1000)")
    parser.add_argument("--scenario", type=int, choices=[0,1,2,3],
                       help="Run single scenario by index (0-3)")
    args = parser.parse_args()

    global SHOTS
    SHOTS = args.shots

    if args.cost:
        print_cost_estimate()
        return

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    print(f"\n{'=' * 70}")
    print(f"🚀 QUANTUM CONTEXT FRAME EXPERIMENT - HARDWARE VERSION")
    print(f"ChronoAI Solutions")
    print(f"Device: {args.device.upper()}")
    print(f"Shots:  {SHOTS}")
    print(f"Time:   {timestamp}")
    print(f"{'=' * 70}")
    print(f"\nQuestion: \"{USER_QUESTION}\"")
    print("Same question every scenario. Frame changes based on quantum measurement.")

    results_log = []

    if args.scenario is not None:
        # Run single scenario
        scenario = SCENARIOS[args.scenario]
        result = run_scenario(
            scenario["label"],
            scenario["revenue"],
            scenario["technical"],
            args.device
        )
        if result:
            results_log.append(result)
    else:
        # Run all scenarios
        for i, scenario in enumerate(SCENARIOS):
            print(f"\n🔄 Running scenario {i+1}/{len(SCENARIOS)}")
            result = run_scenario(
                scenario["label"],
                scenario["revenue"],
                scenario["technical"],
                args.device
            )
            if result:
                results_log.append(result)

    # Save results
    if results_log:
        output_file = f"quantum_results_{args.device}_{timestamp}.json"
        with open(output_file, "w") as f:
            json.dump({
                "experiment": "quantum_context_frame",
                "device": args.device,
                "shots": SHOTS,
                "question": USER_QUESTION,
                "timestamp": timestamp,
                "results": results_log,
            }, f, indent=2)

        print(f"\n{'=' * 70}")
        print(f"🎉 Results saved → {output_file}")
        print(f"📤 Upload this JSON to share verifiable results.")
        print(f"{'=' * 70}")

if __name__ == "__main__":
    main()