#!/usr/bin/env python3
"""OpenQASM 2.0, read into the gate lists the `quantum` module runs.

A compiler hands over its output as OpenQASM 2.0 text, and this reads it
without importing any compiler. `read(text)` gives a `Circuit`: its qubit count
and its gates in order, each on qubit indices counted across the registers in
the order they are declared. `write(circuit)` writes one back.

WHAT IT KEEPS EXACT. An angle is a float, computed as the text writes it, and
also, when the text makes it one, an exact rational multiple of pi plus an exact
rational: `pi/2` is exactly a half of pi, `0.3` exactly three tenths. A function
(`sin`, `cos`, `tan`, `exp`, `ln`, `sqrt`) or a power leaves only the float. So
a reader can tell `rz(pi/2)` — a quarter turn, a Clifford gate — from
`rz(1.5707963267948966)`, which is not one exactly.

WHAT A GATE IS. A gate the module runs keeps its name: `h x y z s t sx rx ry rz`,
`cnot` (written `cx` or `CX`), `cz`, `ecr` and `swap`, and `id`, which a device
may name. A gate the file defines with `gate` is replaced by its body, its
parameters substituted, since that is what the definition says it is. A gate the
module runs keeps its name even when the file also defines it — Qiskit's
exporter defines `ecr`, which `qelib1.inc` lacks — and the definition is kept in
`Circuit.defined`, so a test can hold the module's matrix to it. Any other name
is refused.

`measure` and `reset` are kept as gates of those names, so a file with them
reads, and a comparison of two circuits can refuse them. `barrier` is dropped.
`opaque`, `if` and `U` are refused.

Usage:
    python -m dke.quantum.qasm <file>     print the circuit, one gate a line
"""

from __future__ import annotations

import math
import re
import sys
from dataclasses import dataclass, field
from fractions import Fraction

NATIVE = {"h", "x", "y", "z", "s", "t", "sx", "rx", "ry", "rz", "cnot", "cz", "ecr", "swap", "id"}
RENAMED = {"cx": "cnot", "CX": "cnot"}
ARITY = {"cnot": 2, "cz": 2, "ecr": 2, "swap": 2}
ANGLES = {"rx": 1, "ry": 1, "rz": 1}
FUNCTIONS = {"sin": math.sin, "cos": math.cos, "tan": math.tan,
             "exp": math.exp, "ln": math.log, "sqrt": math.sqrt}


class QasmError(ValueError):
    pass


@dataclass(frozen=True)
class Angle:
    """An angle in radians: `value` as the text computes it in floating point,
    and `exact`, the pair (a, b) for a·pi + b with a and b rational, when the
    text makes it exactly that."""
    value: float
    exact: tuple[Fraction, Fraction] | None = None

    def pi_multiple(self) -> Fraction | None:
        """a, when the angle is exactly a·pi."""
        if self.exact is None or self.exact[1] != 0:
            return None
        return self.exact[0]


@dataclass(frozen=True)
class Gate:
    name: str
    qubits: tuple[int, ...]
    angles: tuple[Angle, ...] = ()


@dataclass
class Definition:
    params: list[str]
    args: list[str]
    body: list[tuple[str, list[str], list[str]]]


@dataclass
class Circuit:
    qubits: int
    gates: list[Gate]
    defined: dict[str, Definition] = field(default_factory=dict)


# ── Expressions ─────────────────────────────────────────────────────────────

_TOKEN = re.compile(r"\s*(?:(\d+\.\d*(?:[eE][-+]?\d+)?|\.\d+(?:[eE][-+]?\d+)?|\d+[eE][-+]?\d+|\d+)"
                    r"|([A-Za-z_][A-Za-z0-9_]*)|(.))")


def _tokens(text: str) -> list[tuple[str, str]]:
    out = []
    pos = 0
    while pos < len(text):
        m = _TOKEN.match(text, pos)
        if m is None or m.end() == pos:
            break
        pos = m.end()
        num, name, op = m.groups()
        if num is not None:
            out.append(("num", num))
        elif name is not None:
            out.append(("name", name))
        elif op is not None and not op.isspace():
            out.append(("op", op))
    return out


class _Expr:
    """A recursive-descent reader of one OpenQASM 2.0 expression. Binding, from
    loosest: `+ -`, then `* /`, then unary minus, then `^`, which groups to the
    right, as in Python."""

    def __init__(self, text: str, env: dict[str, Angle]) -> None:
        self.toks = _tokens(text)
        self.i = 0
        self.env = env
        self.text = text

    def parse(self) -> Angle:
        v = self.sum()
        if self.i != len(self.toks):
            raise QasmError(f"cannot read the expression {self.text!r}")
        return v

    def peek(self) -> tuple[str, str] | None:
        return self.toks[self.i] if self.i < len(self.toks) else None

    def take(self) -> tuple[str, str]:
        t = self.peek()
        if t is None:
            raise QasmError(f"the expression {self.text!r} ends early")
        self.i += 1
        return t

    def sum(self) -> Angle:
        v = self.product()
        while self.peek() in (("op", "+"), ("op", "-")):
            op = self.take()[1]
            w = self.product()
            if op == "+":
                v = Angle(v.value + w.value, _add(v.exact, w.exact, 1))
            else:
                v = Angle(v.value - w.value, _add(v.exact, w.exact, -1))
        return v

    def product(self) -> Angle:
        v = self.unary()
        while self.peek() in (("op", "*"), ("op", "/")):
            op = self.take()[1]
            w = self.unary()
            if op == "*":
                v = Angle(v.value * w.value, _mul(v.exact, w.exact))
            else:
                if w.value == 0.0:
                    raise QasmError(f"division by zero in {self.text!r}")
                v = Angle(v.value / w.value, _div(v.exact, w.exact))
        return v

    def unary(self) -> Angle:
        if self.peek() == ("op", "-"):
            self.take()
            v = self.unary()
            return Angle(-v.value, None if v.exact is None else (-v.exact[0], -v.exact[1]))
        if self.peek() == ("op", "+"):
            self.take()
            return self.unary()
        return self.power()

    def power(self) -> Angle:
        v = self.primary()
        if self.peek() == ("op", "^"):
            self.take()
            w = self.unary()
            return Angle(v.value ** w.value)
        return v

    def primary(self) -> Angle:
        kind, tok = self.take()
        if kind == "num":
            return Angle(float(tok), (Fraction(0), Fraction(tok)))
        if kind == "name":
            if tok == "pi":
                return Angle(math.pi, (Fraction(1), Fraction(0)))
            if tok in FUNCTIONS:
                if self.take() != ("op", "("):
                    raise QasmError(f"{tok} takes its argument in parentheses: {self.text!r}")
                v = self.sum()
                if self.take() != ("op", ")"):
                    raise QasmError(f"unclosed parenthesis in {self.text!r}")
                return Angle(FUNCTIONS[tok](v.value))
            if tok in self.env:
                return self.env[tok]
            raise QasmError(f"no parameter named {tok} in {self.text!r}")
        if tok == "(":
            v = self.sum()
            if self.take() != ("op", ")"):
                raise QasmError(f"unclosed parenthesis in {self.text!r}")
            return v
        raise QasmError(f"cannot read the expression {self.text!r}")


_Exact = tuple[Fraction, Fraction] | None


def _add(a: _Exact, b: _Exact, sign: int) -> _Exact:
    if a is None or b is None:
        return None
    return (a[0] + sign * b[0], a[1] + sign * b[1])


def _mul(a: _Exact, b: _Exact) -> _Exact:
    if a is None or b is None:
        return None
    if a[0] == 0:
        return (a[1] * b[0], a[1] * b[1])
    if b[0] == 0:
        return (a[0] * b[1], a[1] * b[1])
    return None


def _div(a: _Exact, b: _Exact) -> _Exact:
    if a is None or b is None or b[0] != 0 or b[1] == 0:
        return None
    return (a[0] / b[1], a[1] / b[1])


def angle(text: str, env: dict[str, Angle] | None = None) -> Angle:
    return _Expr(text, env or {}).parse()


# ── Statements ──────────────────────────────────────────────────────────────

def _split_top(text: str) -> list[str]:
    """`text` split at commas outside parentheses."""
    out, depth, cur = [], 0, []
    for ch in text:
        if ch == "(":
            depth += 1
        elif ch == ")":
            depth -= 1
        if ch == "," and depth == 0:
            out.append("".join(cur).strip())
            cur = []
        else:
            cur.append(ch)
    if "".join(cur).strip():
        out.append("".join(cur).strip())
    return out


_CALL = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*)\s*(?:\((.*)\))?\s+(.+)$", re.S)
_REG = re.compile(r"^(qreg|creg)\s+([A-Za-z_][A-Za-z0-9_]*)\s*\[\s*(\d+)\s*\]$")
_ARG = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*)\s*(?:\[\s*(\d+)\s*\])?$")
_GATE = re.compile(r"^gate\s+([A-Za-z_][A-Za-z0-9_]*)\s*(?:\(([^)]*)\))?\s+([^{]+)\{(.*)\}$", re.S)


def _statements(text: str) -> list[str]:
    """The statements of the file, comments removed: a `gate` definition is one
    statement, its body included."""
    text = re.sub(r"//[^\n]*", "", text)
    out, cur, depth = [], [], 0
    for ch in text:
        cur.append(ch)
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                out.append("".join(cur).strip())
                cur = []
        elif ch == ";" and depth == 0:
            out.append("".join(cur)[:-1].strip())
            cur = []
    if "".join(cur).strip():
        raise QasmError(f"the file ends inside a statement: {''.join(cur).strip()[:60]!r}")
    return [s for s in out if s]


def read(text: str) -> Circuit:
    stmts = _statements(text)
    if not stmts or re.sub(r"\s+", " ", stmts[0]) != "OPENQASM 2.0":
        raise QasmError("not OpenQASM 2.0: the file must open with `OPENQASM 2.0;`")
    qregs: dict[str, tuple[int, int]] = {}
    cregs: set[str] = set()
    width = 0
    gates: list[Gate] = []
    defined: dict[str, Definition] = {}

    def qubits_of(arg: str) -> list[int]:
        m = _ARG.match(arg)
        if m is None or m.group(1) not in qregs:
            raise QasmError(f"no quantum register for {arg!r}")
        off, size = qregs[m.group(1)]
        if m.group(2) is None:
            return list(range(off, off + size))
        i = int(m.group(2))
        if i >= size:
            raise QasmError(f"{arg!r} is past the end of a register of {size}")
        return [off + i]

    def emit(name: str, params: list[Angle], qs: list[int], depth: int) -> None:
        name = RENAMED.get(name, name)
        if name in NATIVE:
            want = ARITY.get(name, 1)
            if len(qs) != want or len(set(qs)) != want:
                raise QasmError(f"{name} takes {want} distinct qubit(s), not {qs}")
            if len(params) != ANGLES.get(name, 0):
                raise QasmError(f"{name} takes {ANGLES.get(name, 0)} parameter(s), not {len(params)}")
            gates.append(Gate(name, tuple(qs), tuple(params)))
            return
        if name not in defined:
            raise QasmError(f"no gate named {name}: neither one the module runs nor defined in the file")
        if depth > 64:
            raise QasmError(f"definitions nest past 64 deep at {name}")
        d = defined[name]
        if len(params) != len(d.params) or len(qs) != len(d.args):
            raise QasmError(f"{name} takes {len(d.params)} parameter(s) and {len(d.args)} qubit(s)")
        env = dict(zip(d.params, params))
        where = dict(zip(d.args, qs))
        for bname, bparams, bargs in d.body:
            if bname == "barrier":
                continue
            emit(bname, [angle(p, env) for p in bparams], [where[a] for a in bargs], depth + 1)

    for st in stmts[1:]:
        flat = re.sub(r"\s+", " ", st).strip()
        if flat.startswith("include "):
            if flat != 'include "qelib1.inc"':
                raise QasmError(f"only qelib1.inc may be included, not {flat[8:]}")
            continue
        m = _REG.match(flat)
        if m:
            kind, name, size = m.group(1), m.group(2), int(m.group(3))
            if name in qregs or name in cregs:
                raise QasmError(f"register {name} is declared twice")
            if kind == "qreg":
                qregs[name] = (width, size)
                width += size
            else:
                cregs.add(name)
            continue
        if flat.startswith("gate "):
            g = _GATE.match(st.strip())
            if g is None:
                raise QasmError(f"cannot read the definition {flat[:60]!r}")
            body = []
            for b in _statements(g.group(4)):
                c = _CALL.match(re.sub(r"\s+", " ", b).strip())
                if c is None:
                    raise QasmError(f"cannot read {b!r} in the definition of {g.group(1)}")
                body.append((c.group(1), _split_top(c.group(2) or ""), [a.strip() for a in c.group(3).split(",")]))
            defined[g.group(1)] = Definition(_split_top(g.group(2) or ""),
                                             [a.strip() for a in g.group(3).split(",")], body)
            continue
        if flat.startswith("opaque ") or re.match(r"(if|U)\s*\(", flat):
            raise QasmError(f"not read here: {flat[:40]!r}")
        if flat.startswith("barrier"):
            continue
        if flat.startswith("measure "):
            src, _, dst = flat[8:].partition("->")
            for q in qubits_of(src.strip()):
                gates.append(Gate("measure", (q,)))
            continue
        if flat.startswith("reset "):
            for q in qubits_of(flat[6:].strip()):
                gates.append(Gate("reset", (q,)))
            continue
        c = _CALL.match(flat)
        if c is None:
            raise QasmError(f"cannot read the statement {flat[:60]!r}")
        params = [angle(p) for p in _split_top(c.group(2) or "")]
        args = [qubits_of(a.strip()) for a in c.group(3).split(",")]
        sizes = {len(a) for a in args if len(a) > 1}
        if len(sizes) > 1:
            raise QasmError(f"registers of different sizes in {flat[:60]!r}")
        reps = sizes.pop() if sizes else 1
        for r in range(reps):
            emit(c.group(1), params, [a[r] if len(a) > 1 else a[0] for a in args], 0)
    return Circuit(width, gates, defined)


def _angle_text(a: Angle) -> str:
    # An angle that is exactly a rational multiple of pi is written as one,
    # so reading the text back keeps it exact; any other as its float.
    m = a.pi_multiple()
    if m is None:
        return repr(a.value)
    if m == 0:
        return "0"
    num = "pi" if abs(m.numerator) == 1 else f"{abs(m.numerator)}*pi"
    text = num if m.denominator == 1 else f"{num}/{m.denominator}"
    return "-" + text if m < 0 else text


def write(circuit: Circuit) -> str:
    """The circuit as OpenQASM 2.0 on one register `q`, with `cnot` written
    `cx`. `ecr`, which `qelib1.inc` lacks, comes with the definition Qiskit's
    exporter writes for it, which equals the module's `ecr` up to a global
    phase. A circuit holding `measure` or `reset` is refused: there is no
    classical register to write it to."""
    lines = ["OPENQASM 2.0;", 'include "qelib1.inc";']
    if any(g.name == "ecr" for g in circuit.gates):
        lines.append("gate ecr q0,q1 { s q0; sx q1; cx q0,q1; x q0; }")
    lines.append(f"qreg q[{circuit.qubits}];")
    for g in circuit.gates:
        if g.name in ("measure", "reset"):
            raise QasmError(f"{g.name} is not written: there is no classical register")
        name = "cx" if g.name == "cnot" else g.name
        params = f"({', '.join(_angle_text(a) for a in g.angles)})" if g.angles else ""
        lines.append(f"{name}{params} {','.join(f'q[{q}]' for q in g.qubits)};")
    return "\n".join(lines) + "\n"


def main(argv: list[str] | None = None) -> int:
    args = sys.argv[1:] if argv is None else argv
    if len(args) != 1:
        print(__doc__.split("Usage:")[1], file=sys.stderr)
        return 2
    with open(args[0], encoding="utf-8") as f:
        c = read(f.read())
    print(f"{c.qubits} qubits")
    for g in c.gates:
        print(g.name, *g.qubits, *(a.value for a in g.angles))
    return 0


if __name__ == "__main__":
    sys.exit(main())