"""
ChronoAI Solutions - Quantum Threshold Sweep
Series 2: Where Does the Anomaly Begin and End?

Circuit: RY(revenue×π) → RY(0.9×π) → CNOT(q0→q1) → Measure
Revenue: swept 0.0 to 1.0 in 0.1 increments
Technical: fixed at 0.9 (Scenario 3 conditions)
Expected: find the exact crossover point where routing shifts

Known: revenue=0.1 → ACTIVE SALES (01) anomaly at 83-96%
Question: at what revenue value does this anomaly appear/disappear?

Cost estimate (IonQ Forte-1):
  11 tasks × $0.30 task fee = $3.30
  11 × 100 shots × $0.08   = $88.00
  Total                    = ~$91.30
"""

import boto3
import numpy as np
import json
import time
from datetime import datetime
from braket.aws import AwsDevice, AwsQuantumTask
from braket.circuits import Circuit
from braket.devices import LocalSimulator

# ─── CONFIG ──────────────────────────────────────────────────
# Set USE_SIMULATOR = True to test locally before spending QPU time
# Set USE_SIMULATOR = False to run on real hardware

USE_SIMULATOR   = True   # ← flip to False when ready for QPU
SHOTS           = 100
TECHNICAL_FIXED = 0.9
REVENUE_VALUES  = [round(x * 0.1, 1) for x in range(11)]  # 0.0 → 1.0

# AWS config (only used when USE_SIMULATOR = False)
S3_BUCKET   = "amazon-braket-chronoai-results"   # your S3 bucket
S3_PREFIX   = "threshold-sweep"
AWS_REGION  = "us-east-1"

# Device ARNs
IONQ_FORTE  = "arn:aws:braket:us-east-1::device/qpu/ionq/Forte-1"
IQM_GARNET  = "arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet"

DEVICE_ARN  = IONQ_FORTE   # ← swap to IQM_GARNET for cross-validation

# ─── FRAME LABELS ────────────────────────────────────────────
FRAMES = {
    "00": "Scrappy / No Budget",
    "01": "Active Sales Mode",
    "10": "Deep Technical Build",
    "11": "Strategic / Big Picture"
}

# ─── BUILD CIRCUIT ────────────────────────────────────────────
def build_circuit(revenue, technical):
    """
    RY(revenue×π) on q0 - encodes revenue signal
    RY(technical×π) on q1 - encodes technical signal
    CNOT(q0→q1) - entangles, allows q0 to influence q1 routing
    Measure both
    """
    circ = Circuit()
    circ.ry(0, revenue * np.pi)
    circ.ry(1, technical * np.pi)
    circ.cnot(0, 1)
    return circ

# ─── GET DEVICE ───────────────────────────────────────────────
def get_device():
    if USE_SIMULATOR:
        print("Using local simulator")
        return LocalSimulator()
    else:
        print(f"Using QPU: {DEVICE_ARN}")
        return AwsDevice(DEVICE_ARN)

# ─── RUN ONE COMBINATION ─────────────────────────────────────
def run_combination(device, revenue, technical, shot_count):
    circ = build_circuit(revenue, technical)

    if USE_SIMULATOR:
        task = device.run(circ, shots=shot_count)
        result = task.result()
    else:
        task = device.run(
            circ,
            s3_destination_folder=(S3_BUCKET, S3_PREFIX),
            shots=shot_count
        )
        print(f"  Task ID: {task.id} - waiting...")
        result = task.result()

    # FIXED: use measurement_counts (not result.values[0])
    counts = result.measurement_counts

    total = sum(counts.values())
    probs = {}
    for bitstring, count in counts.items():
        # Normalize bitstring to 2 chars
        key = bitstring.zfill(2) if len(bitstring) <= 2 else bitstring[-2:]
        probs[key] = probs.get(key, 0) + (count / total)

    dominant = max(probs, key=probs.get)
    dominant_pct = probs[dominant] * 100

    return {
        "probabilities": {k: round(v, 4) for k, v in sorted(probs.items())},
        "dominant": dominant,
        "dominant_label": FRAMES.get(dominant, "Unknown"),
        "dominant_pct": round(dominant_pct, 1),
        "counts": dict(counts)
    }

# ─── MAIN SWEEP ───────────────────────────────────────────────
def run_threshold_sweep():
    print("\n" + "="*60)
    print("CHRONOAI - QUANTUM THRESHOLD SWEEP")
    print(f"Technical fixed at: {TECHNICAL_FIXED}")
    print(f"Revenue sweep: {REVENUE_VALUES}")
    print(f"Shots per combination: {SHOTS}")
    print(f"Mode: {'SIMULATOR' if USE_SIMULATOR else 'QPU - ' + DEVICE_ARN}")
    print("="*60 + "\n")

    device = get_device()
    results = {}
    run_time = datetime.now().isoformat()

    for revenue in REVENUE_VALUES:
        label = f"R={revenue:.1f} T={TECHNICAL_FIXED}"
        print(f"Running {label}...")

        try:
            data = run_combination(device, revenue, TECHNICAL_FIXED, SHOTS)
            results[str(revenue)] = data

            anomaly_flag = " ← ANOMALY" if data["dominant"] == "01" and revenue < 0.5 else ""
            crossover_flag = " ← CROSSOVER?" if data["dominant"] in ["10", "11"] and revenue < 0.4 else ""

            print(f"  Result: {data['dominant']} ({data['dominant_label']}) "
                  f"at {data['dominant_pct']}%{anomaly_flag}{crossover_flag}")
            print(f"  All probs: {data['probabilities']}")

            # Small delay between QPU tasks to be respectful
            if not USE_SIMULATOR:
                time.sleep(2)

        except Exception as e:
            print(f"  ERROR on revenue={revenue}: {e}")
            results[str(revenue)] = {"error": str(e)}

    # ─── ANALYSIS ────────────────────────────────────────────
    print("\n" + "="*60)
    print("THRESHOLD SWEEP SUMMARY")
    print("="*60)
    print(f"{'Revenue':<10} {'Technical':<10} {'Frame':<6} {'Label':<25} {'Confidence'}")
    print("-"*60)

    crossover_point = None
    prev_dominant = None

    for revenue in REVENUE_VALUES:
        r = str(revenue)
        if r not in results or "error" in results[r]:
            print(f"{revenue:<10.1f} {'ERROR'}")
            continue
        d = results[r]
        marker = ""
        if prev_dominant and d["dominant"] != prev_dominant:
            crossover_point = revenue
            marker = " ← CROSSOVER"
        print(f"{revenue:<10.1f} {TECHNICAL_FIXED:<10} {d['dominant']:<6} "
              f"{d['dominant_label']:<25} {d['dominant_pct']}%{marker}")
        prev_dominant = d["dominant"]

    print("-"*60)
    if crossover_point is not None:
        print(f"\nCROSSOVER DETECTED at revenue ≈ {crossover_point}")
        print("Below this point: Anomaly active (Active Sales dominates)")
        print("Above this point: Expected routing resumes")
    else:
        print("\nNo clean crossover detected - anomaly may span full range")
        print("or may not appear at all under these conditions.")

    # ─── SAVE RESULTS ────────────────────────────────────────
    output = {
        "experiment": "Quantum Threshold Sweep - S2",
        "run_time": run_time,
        "config": {
            "technical_fixed": TECHNICAL_FIXED,
            "revenue_values": REVENUE_VALUES,
            "shots": SHOTS,
            "device": "LocalSimulator" if USE_SIMULATOR else DEVICE_ARN
        },
        "crossover_point": crossover_point,
        "results": results
    }

    filename = f"threshold_sweep_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(filename, "w") as f:
        json.dump(output, f, indent=2)

    print(f"\nResults saved to: {filename}")
    print("="*60 + "\n")

    return output

# ─── RUN ─────────────────────────────────────────────────────
if __name__ == "__main__":
    run_threshold_sweep()
