from __future__ import annotations

import math

from ..ir.types import Instruction


def _is_small(angle: float, eps: float) -> bool:
    """Normalize angle to (-pi, pi] and compare against eps."""
    a = (angle + math.pi) % (2 * math.pi) - math.pi
    return abs(a) < eps


def optimize_1q(circuit, eps: float = 1e-9):
    """
    Peephole optimizations on single-qubit gates (local, order-preserving w.r.t. other qubits):
      - Merge consecutive RZ on same qubit: RZ(a); RZ(b) -> RZ(a+b); drop if ~0.
      - Collapse SX·SX -> X on same qubit, exactly.
      - Cancel X·X -> I (drop both) on same qubit.
    Scope boundaries: a CX or MEASURE on a qubit ends the local 1q block
    for that qubit (we do not commute across multi-qubit ops).
    Returns a NEW circuit.
    """
    from ..circuit.quantum_circuit import QuantumCircuit  # avoid circular import

    n = circuit.n_qubits
    out = QuantumCircuit(n, getattr(circuit, "n_clbits", 0))
    out_insts: list[Instruction] = out._instructions

    # Track, per qubit, index in out_insts of the last 1q gate affecting that qubit
    last_idx: list[int | None] = [None] * n

    def bump(q: int, new_pos: int | None):
        last_idx[q] = new_pos

    for inst in circuit.instructions:
        name, qs, ps = inst.name, inst.qubits, inst.params

        # Multi-qubit gate or measure: flush context for involved qubits and copy through
        if name == "cx":
            out_insts.append(inst)
            c, t = qs
            bump(c, None)
            bump(t, None)
            continue

        if name == "measure":
            q = qs[0]
            out_insts.append(inst)
            bump(q, None)
            continue

        # Single-qubit gates we know: x, sx, rz
        if name not in ("x", "sx", "rz"):
            # Unknown 1q op: copy and reset its qubit context for safety
            out_insts.append(inst)
            bump(qs[0], None)
            continue

        q = qs[0]
        li = last_idx[q]

        if name == "rz":
            theta = float(ps[0])
            if li is not None and out_insts[li].name == "rz":
                # Merge angles at last RZ
                prev = out_insts[li]
                new_theta = float(prev.params[0]) + theta
                if _is_small(new_theta, eps):
                    # Remove previous; current consumed (drop both)
                    out_insts.pop(li)
                    bump(q, None)
                else:
                    out_insts[li] = Instruction("rz", (q,), (new_theta,), None)
                # Do not append current
            else:
                out_insts.append(inst)
                bump(q, len(out_insts) - 1)
            continue

        if name == "sx":
            if li is not None and out_insts[li].name == "sx":
                # SX · SX -> X (replace previous SX with X; current consumed)
                out_insts[li] = Instruction("x", (q,), (), None)
                # Keep context at this X (may cancel with next X)
            else:
                out_insts.append(inst)
                bump(q, len(out_insts) - 1)
            continue

        if name == "x":
            if li is not None and out_insts[li].name == "x":
                # X · X -> I (drop previous X and consume current)
                out_insts.pop(li)
                bump(q, None)
            else:
                out_insts.append(inst)
                bump(q, len(out_insts) - 1)
            continue

    return out
