
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron AI Quantum Assistant
#
#  No internet or API key needed.
#  Uses a built-in keyword engine to:
#    1. Convert plain English → QuantumCircuit
#    2. Explain any circuit step-by-step in simple language
# ─────────────────────────────────────────────────────────────────

# Keywords the user might type → list of (gate, qubit, ...) tuples
_RECIPES = {
    "bell":          [("h",0), ("cx",0,1)],
    "entangle":      [("h",0), ("cx",0,1)],
    "superposition": [("h",0)],
    "hadamard":      [("h",0)],
    "flip":          [("x",0)],
    "not gate":      [("x",0)],
    "phase flip":    [("z",0)],
    "phase":         [("z",0)],
    "teleport":      [("h",0), ("cx",0,1), ("cx",1,2), ("h",1)],
    "grover":        [("h",0), ("h",1), ("x",0), ("x",1),
                      ("cx",0,1), ("x",0), ("x",1), ("h",0), ("h",1)],
    "swap":          [("swap",0,1)],
    "toffoli":       [("ccx",0,1,2)],
    "rotate x":      [("rx",0,1.5708)],
    "rotate y":      [("ry",0,1.5708)],
    "rotate z":      [("rz",0,1.5708)],
    "qft":           [("h",0), ("rz",1,1.5708), ("cx",0,1), ("h",1)],
    "measure":       [("measure",0)],
}

# Gate → plain English sentence
_EXPLAIN = {
    "h":    "H (Hadamard) on qubit {a} → puts qubit into superposition: 50% |0⟩, 50% |1⟩. Like a coin spinning in the air.",
    "x":    "X (NOT) on qubit {a} → flips the qubit. |0⟩ becomes |1⟩ and vice versa. Like a light switch.",
    "y":    "Y on qubit {a} → flips the qubit AND adds an imaginary phase. A flip with a twist.",
    "z":    "Z on qubit {a} → leaves |0⟩ alone, flips the sign of |1⟩. Invisible unless in superposition.",
    "s":    "S on qubit {a} → quarter-turn (90°) phase rotation on |1⟩.",
    "t":    "T on qubit {a} → eighth-turn (45°) phase rotation. Key gate for universal quantum computing.",
    "cx":   "CNOT: qubit {a} controls qubit {b} → flips {b} only if {a} is |1⟩. Creates entanglement.",
    "cz":   "CZ: qubit {a} and {b} → flips the phase only when both are |1⟩.",
    "swap": "SWAP: qubit {a} and {b} exchange their full quantum states.",
    "ccx":  "Toffoli: qubits {a} and {b} control qubit {c} → flips {c} only if BOTH controls are |1⟩.",
    "rx":   "Rx({angle:.3f} rad) on qubit {a} → rotates around X-axis of the Bloch sphere.",
    "ry":   "Ry({angle:.3f} rad) on qubit {a} → rotates around Y-axis of the Bloch sphere.",
    "rz":   "Rz({angle:.3f} rad) on qubit {a} → rotates around Z-axis (phase rotation).",
    "measure": "Measure qubit {a} → collapses superposition to a definite 0 or 1.",
}


def generate_circuit(prompt: str):
    """
    Convert a plain English prompt into a PKTron QuantumCircuit.

    Examples:
        qc = generate_circuit("create a bell state")
        qc = generate_circuit("put qubit in superposition")
        qc = generate_circuit("make a toffoli gate")
    """
    from pktron import QuantumCircuit
    import pktron.gates as g

    prompt_lower = prompt.lower()
    ops = []

    for keyword, recipe in _RECIPES.items():
        if keyword in prompt_lower:
            ops.extend(recipe)
            break   # use first match to keep circuit clean

    if not ops:
        print(f'[AI Assistant] Could not understand: "{prompt}"')
        print("  → Defaulting to: Hadamard on qubit 0")
        ops = [("h", 0)]

    # Work out how many qubits are needed
    max_q = 0
    for op in ops:
        for item in op[1:]:
            if isinstance(item, int):
                max_q = max(max_q, item)
    n = max_q + 1

    qc = QuantumCircuit(n)
    qc._ai_ops = []   # store ops so explain_circuit can read them

    _GATE_FN = {
        "h": lambda q: g.h(qc, q),
        "x": lambda q: g.x(qc, q),
        "y": lambda q: g.y(qc, q),
        "z": lambda q: g.z(qc, q),
        "s": lambda q: g.s(qc, q),
        "t": lambda q: g.t(qc, q),
        "cx":   lambda a, b: g.cnot(qc, a, b),
        "cz":   lambda a, b: g.cz(qc, a, b),
        "swap": lambda a, b: g.swap(qc, a, b),
        "ccx":  lambda a, b, c: g.toffoli(qc, a, b, c),
        "rx":   lambda a, angle: g.rx(qc, a, angle),
        "ry":   lambda a, angle: g.ry(qc, a, angle),
        "rz":   lambda a, angle: g.rz(qc, a, angle),
    }

    for op in ops:
        gate = op[0]
        args = op[1:]
        try:
            if gate != "measure":
                _GATE_FN[gate](*args)
            qc._ai_ops.append(op)
        except Exception as e:
            print(f"  [AI Assistant] Skipped {gate}: {e}")

    print(f'[AI Assistant] Built circuit: {n} qubit(s), {len(qc._ai_ops)} gate(s)')
    print(f'  Prompt: "{prompt}"')
    return qc


def explain_circuit(circuit) -> str:
    """
    Explain every gate in a circuit in plain English.

    Works on any PKTron QuantumCircuit.

    Example:
        qc = generate_circuit("bell state")
        print(explain_circuit(qc))
    """
    # Try to get ops from AI assistant first, then fall back to ops_list
    ops = getattr(circuit, "_ai_ops",
          getattr(circuit, "ops_list",
          getattr(circuit, "gate_history", [])))

    if not ops:
        return "This circuit has no recorded gates to explain."

    lines = [
        "=" * 58,
        "  PKTron AI Assistant — Circuit Explanation",
        "=" * 58,
        "",
    ]

    for i, op in enumerate(ops):
        gate = op[0].lower()
        tmpl = _EXPLAIN.get(gate, f"Gate: {gate} (no explanation available)")
        try:
            if gate in ("h","x","y","z","s","t","measure"):
                msg = tmpl.format(a=op[1])
            elif gate in ("cx","cz","swap"):
                msg = tmpl.format(a=op[1], b=op[2])
            elif gate in ("rx","ry","rz"):
                msg = tmpl.format(a=op[1], angle=op[2])
            elif gate == "ccx":
                msg = tmpl.format(a=op[1], b=op[2], c=op[3])
            else:
                msg = tmpl
        except Exception:
            msg = f"{gate} gate"

        lines.append(f"  Step {i+1}: {msg}")

    lines += ["", "=" * 58]
    return "\n".join(lines)
