#!/usr/bin/env python3
"""A device's calibration snapshot, written as facts for the `quantum` module.

A snapshot is a directory named for the device, holding `conf_<name>.json` and
`props_<name>.json`, the two files Qiskit's fake backends ship for each device:
its configuration and its calibration properties. `read` turns one into a
`Device`, and `programs` turns that into DKE Python programs, each sized for one
run, through `dke.factwriter`. Nothing here imports Qiskit; it reads the files.

WHAT IS WRITTEN, by kind:

- `Device.<name>`: `qubits`; `snapshot_at`, the snapshot's own date; and
  `digest`, the snapshot's digest (`provenance.snapshot_digest`), which a
  result computed for the device records as the calibration it read.
- `PhysicalQubit.<name>_q<i>`: `device`, `index` (the qubit's number as text,
  which is how a compiled circuit's gates name it), `t1` and `t2` in
  microseconds, `readout_error`, and `measured_at`, the latest date among those
  three.
- `Native.<name>_<gate>`: `device`, `op`, `a`, `b` for a pair, `error`, `length`
  in nanoseconds, and `measured_at`, the date the error was measured. A pair is
  ordered as the snapshot lists it, because a device may run a two-qubit
  operation one way round only.

An operation keeps the device's name for it, except `cx`, which the module
calls `cnot`. `reset`, `measure` and `delay` act on no state the module models
and are not written. A value the snapshot lacks is not written, never as a
placeholder, so a rule that asks whether a cell is empty is answered truly.

A store refuses a program larger than a size it is set up with, so the size to
split at is yours to pass, as `FactLoad` describes.

Usage:
    dke-quantum-device <snapshot-dir> --max-bytes N [--module M]
                                       print the programs, separated by a line
                                       holding only `# ---`
"""

from __future__ import annotations

import argparse
import json
import pathlib
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone

from dke.factwriter import FactLoad, sanitize
from dke.quantum.provenance import snapshot_digest

MODULE = "quantum"
SOURCE = "ibm-calibration"
RENAMED = {"cx": "cnot"}
NOT_OPERATIONS = {"reset", "measure", "delay"}
UNITS = {"T1": "us", "T2": "us", "readout_error": "", "gate_error": "", "gate_length": "ns"}


@dataclass
class Qubit:
    subject: str
    index: str
    t1: float | None = None
    t2: float | None = None
    readout_error: float | None = None
    measured_at: str | None = None


@dataclass
class Native:
    subject: str
    op: str
    a: str
    b: str | None
    error: float | None = None
    length: float | None = None
    measured_at: str | None = None


@dataclass
class Device:
    name: str
    qubits: int
    snapshot_at: str
    digest: str
    physical: list[Qubit] = field(default_factory=list)
    natives: list[Native] = field(default_factory=list)
    skipped: dict[str, int] = field(default_factory=dict)


def utc(stamp: str) -> str:
    """A snapshot date as RFC 3339 UTC, which is how DKE Python writes one.
    The snapshots carry an offset, and an instant is kept, not a wall clock."""
    t = datetime.fromisoformat(stamp).astimezone(timezone.utc)
    frac = f".{t.microsecond:06d}".rstrip("0") if t.microsecond else ""
    return t.strftime("%Y-%m-%dT%H:%M:%S") + frac + "Z"


def _param(params: list[dict], name: str) -> dict | None:
    for p in params:
        if p.get("name") == name:
            unit = p.get("unit", "")
            if unit != UNITS[name]:
                raise ValueError(f"{name} is in {unit!r}, and this reader reads {UNITS[name]!r}")
            return p
    return None


def read(directory: pathlib.Path) -> Device:
    """One snapshot directory, as a `Device`."""
    name = directory.name
    conf = json.loads((directory / f"conf_{name}.json").read_text(encoding="utf-8"))
    props = json.loads((directory / f"props_{name}.json").read_text(encoding="utf-8"))
    dev = Device(name, int(conf["n_qubits"]), utc(props["last_update_date"]), snapshot_digest(directory))
    for i, params in enumerate(props["qubits"]):
        q = Qubit(sanitize(f"{name}_q{i}"), str(i))
        dates = []
        for key, attr in (("T1", "t1"), ("T2", "t2"), ("readout_error", "readout_error")):
            p = _param(params, key)
            if p is not None:
                setattr(q, attr, float(p["value"]))
                dates.append(utc(p["date"]))
        q.measured_at = max(dates) if dates else None
        dev.physical.append(q)
    for g in props["gates"]:
        op = g["gate"]
        if op in NOT_OPERATIONS:
            dev.skipped[op] = dev.skipped.get(op, 0) + 1
            continue
        qs = [str(x) for x in g["qubits"]]
        if len(qs) not in (1, 2):
            raise ValueError(f"{g['name']} acts on {len(qs)} qubits")
        n = Native(sanitize(f"{name}_{g['name']}"), RENAMED.get(op, op), qs[0],
                   qs[1] if len(qs) == 2 else None)
        err = _param(g["parameters"], "gate_error")
        length = _param(g["parameters"], "gate_length")
        if err is not None:
            n.error = float(err["value"])
            n.measured_at = utc(err["date"])
        if length is not None:
            n.length = float(length["value"])
        dev.natives.append(n)
    subjects = [q.subject for q in dev.physical] + [n.subject for n in dev.natives]
    if len(set(subjects)) != len(subjects):
        raise ValueError(f"{name}: two facts would share a subject")
    return dev


def programs(dev: Device, *, max_bytes: int, module: str = MODULE) -> list[str]:
    """The programs that write `dev`, in the order they run, into a store
    holding the module as `module`, none larger than `max_bytes`. Each carries
    the load marker `Load.quantum_device_<name>_<n>`, so a store can tell a load
    that finished from one that stopped part-way."""
    load = FactLoad(module, SOURCE, max_bytes=max_bytes)
    q = load.q
    load.comment(f"Device {dev.name}: {dev.qubits} qubits, snapshot of {dev.snapshot_at}.")
    load.comment("Written by dke.quantum.device.")
    stamp = int(datetime.fromisoformat(dev.snapshot_at.replace("Z", "+00:00")).timestamp())
    # The subject a `device` field holds, written from one place: a join reads
    # the field as the name of a `Device`, so the two must agree exactly.
    dname = sanitize(dev.name)
    load.load_marker(f"quantum_device_{dname}", stamp)

    def when(t: str) -> str:
        return f"datetime({q(t)})"

    def column(table: str, attr: str, rows: list[str]) -> None:
        load.table(table, rows, [load.put(attr, "r[1]")])

    # A table's name becomes a variable of the program, and a path segment that
    # names a variable resolves to its value. So no table may be named as a
    # field is: a table called `device` turns `PhysicalQubit.s.device` into a
    # write to an attribute named by the table. Every name here is `rows_…`,
    # and no field of the module begins that way.

    load.table("rows_device", [f"({q(dname)}, {dev.qubits}, {when(dev.snapshot_at)}, {q(dev.digest)})"],
               [load.put("Device.s.qubits", "r[1]"), load.put("Device.s.snapshot_at", "r[2]"),
                load.put("Device.s.digest", "r[3]")])
    load.table("rows_qubit", [f"({q(x.subject)}, {q(dname)}, {q(x.index)})" for x in dev.physical],
               [load.put("PhysicalQubit.s.device", "r[1]"), load.put("PhysicalQubit.s.index", "r[2]")])
    for attr in ("t1", "t2", "readout_error"):
        column(f"rows_qubit_{attr}", f"PhysicalQubit.s.{attr}",
               [f"({q(x.subject)}, {getattr(x, attr)!r})" for x in dev.physical
                if getattr(x, attr) is not None])
    column("rows_qubit_measured", "PhysicalQubit.s.measured_at",
           [f"({q(x.subject)}, {when(x.measured_at)})" for x in dev.physical if x.measured_at])
    load.table("rows_native", [f"({q(n.subject)}, {q(dname)}, {q(n.op)}, {q(n.a)})" for n in dev.natives],
               [load.put("Native.s.device", "r[1]"), load.put("Native.s.op", "r[2]"),
                load.put("Native.s.a", "r[3]")])
    column("rows_native_b", "Native.s.b", [f"({q(n.subject)}, {q(n.b)})" for n in dev.natives if n.b])
    column("rows_native_error", "Native.s.error",
           [f"({q(n.subject)}, {n.error!r})" for n in dev.natives if n.error is not None])
    column("rows_native_length", "Native.s.length",
           [f"({q(n.subject)}, {n.length!r})" for n in dev.natives if n.length is not None])
    column("rows_native_measured", "Native.s.measured_at",
           [f"({q(n.subject)}, {when(n.measured_at)})" for n in dev.natives if n.measured_at])
    return load.render()


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(prog="dke-quantum-device",
                                description="A calibration snapshot, written as facts for the quantum module.")
    p.add_argument("snapshot", type=pathlib.Path, help="the snapshot's directory, named for the device")
    p.add_argument("--max-bytes", type=int, required=True,
                   help="the largest program your store accepts, from its documentation or its operator")
    p.add_argument("--module", default=MODULE, help="the name the store holds the module under")
    a = p.parse_args(argv)
    print("\n# ---\n".join(programs(read(a.snapshot), max_bytes=a.max_bytes, module=a.module)))
    return 0


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