#!/usr/bin/env python3
"""Verify every CPM-CSS construction in this sharing bundle."""

from __future__ import annotations

import hashlib
import json
from collections import deque
from pathlib import Path


ROOT = Path(__file__).resolve().parent


def digest(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def read_cpm(path: Path):
    tokens = path.read_text(encoding="utf-8").split()
    assert tokens[0] == "CPM_CSS_V1", path
    name = tokens[1]
    p, width, jx, jz = map(int, tokens[2:6])
    values = list(map(int, tokens[6:]))
    assert len(values) == (jx + jz) * width, path
    ex = [values[i * width:(i + 1) * width] for i in range(jx)]
    offset = jx * width
    dz = [
        values[offset + i * width:offset + (i + 1) * width]
        for i in range(jz)
    ]
    return name, p, width, ex, dz


def supports(exponents, p, width):
    return [
        [column * p + ((output - exponent) % p) for column, exponent in enumerate(row)]
        for row in exponents
        for output in range(p)
    ]


def bit_rows(rows):
    output = []
    for row in rows:
        value = 0
        for qubit in row:
            value ^= 1 << qubit
        output.append(value)
    return output


def rank(rows):
    basis = {}
    for value in rows:
        while value:
            pivot = value.bit_length() - 1
            if pivot in basis:
                value ^= basis[pivot]
            else:
                basis[pivot] = value
                break
    return len(basis)


def girth(rows, n):
    adjacency = [[] for _ in range(n + len(rows))]
    for index, row in enumerate(rows):
        check = n + index
        for qubit in row:
            adjacency[check].append(qubit)
            adjacency[qubit].append(check)
    best = None
    for start in range(len(adjacency)):
        distance = [-1] * len(adjacency)
        parent = [-1] * len(adjacency)
        distance[start] = 0
        queue = deque([start])
        while queue:
            node = queue.popleft()
            if best is not None and 2 * distance[node] + 1 >= best:
                continue
            for neighbor in adjacency[node]:
                if distance[neighbor] < 0:
                    distance[neighbor] = distance[node] + 1
                    parent[neighbor] = node
                    queue.append(neighbor)
                elif parent[node] != neighbor and parent[neighbor] != node:
                    cycle = distance[node] + distance[neighbor] + 1
                    if best is None or cycle < best:
                        best = cycle
        if best == 4:
            break
    return best


def strict_pair_partition(ex, dz, p, width):
    for x_row in ex:
        for z_row in dz:
            groups = {}
            for column in range(width):
                value = (z_row[column] - x_row[column]) % p
                groups.setdefault(value, []).append(column)
            if len(groups) != width // 2 or any(len(group) != 2 for group in groups.values()):
                return False
    return True


def main():
    catalogue = json.loads((ROOT / "catalog.json").read_text(encoding="utf-8"))
    results = []
    for record in catalogue["codes"]:
        path = ROOT / record["file"]
        assert digest(path) == record["bundle_sha256"], record["metric_id"]
        name, p, width, ex, dz = read_cpm(path)
        assert name == record["name"]
        assert p == record["P"] and width == record["L"]
        assert len(ex) == record["Jx"] and len(dz) == record["Jz"]
        assert ex == record["E"] and dz == record["D"]
        n = p * width
        x_supports = supports(ex, p, width)
        z_supports = supports(dz, p, width)
        x_rows = bit_rows(x_supports)
        z_rows = bit_rows(z_supports)
        rank_x = rank(x_rows)
        rank_z = rank(z_rows)
        commute = all(
            bin(x_row & z_row).count("1") % 2 == 0
            for x_row in x_rows
            for z_row in z_rows
        )
        degree_x = [0] * n
        degree_z = [0] * n
        for row in x_supports:
            for qubit in row:
                degree_x[qubit] += 1
        for row in z_supports:
            for qubit in row:
                degree_z[qubit] += 1
        observed = {
            "N": n,
            "K": n - rank_x - rank_z,
            "rank_X": rank_x,
            "rank_Z": rank_z,
            "girth_X": girth(x_supports, n),
            "girth_Z": girth(z_supports, n),
            "check_weight_X": sorted(set(map(len, x_supports))),
            "check_weight_Z": sorted(set(map(len, z_supports))),
            "degree_X": sorted(set(degree_x)),
            "degree_Z": sorted(set(degree_z)),
            "css_orthogonal": commute,
            "strict_pair_partition": strict_pair_partition(ex, dz, p, width),
        }
        expected = {
            "N": record["N"],
            "K": record["K"],
            "rank_X": record["rank_X"],
            "rank_Z": record["rank_Z"],
            "girth_X": record["girth_X"],
            "girth_Z": record["girth_Z"],
            "check_weight_X": [record["check_weight"]],
            "check_weight_Z": [record["check_weight"]],
            "degree_X": [record["degree_X"]],
            "degree_Z": [record["degree_Z"]],
            "css_orthogonal": True,
            "strict_pair_partition": True,
        }
        assert observed == expected, (record["metric_id"], observed, expected)
        results.append({
            "metric_id": record["metric_id"],
            "parameters": record["display_parameters"],
            "sha256": record["bundle_sha256"],
            "status": "PASS",
        })
    assert len(results) == 19
    assert sum(code["distance"]["status"] == "exact" for code in catalogue["codes"]) == 18
    print(json.dumps({
        "schema": "pair-partition-sharing-bundle-verification-v1",
        "status": "PASS",
        "codes_verified": len(results),
        "results": results,
    }, indent=2))


if __name__ == "__main__":
    main()