from __future__ import annotations

import argparse
import itertools
import json
import random
from dataclasses import dataclass
from typing import Iterable, Iterator

from sage.all import GF, Graph, Matrix, matroids
from sage.matroids.constructor import Matroid


def all_pairings(elements: tuple[int, ...]) -> Iterator[tuple[tuple[int, int], ...]]:
    if not elements:
        yield ()
        return
    first = elements[0]
    for i in range(1, len(elements)):
        pair = (first, elements[i])
        rest = elements[1:i] + elements[i + 1 :]
        for tail in all_pairings(rest):
            yield (pair,) + tail


def pairing_to_masks(pairing: tuple[tuple[int, int], ...]) -> tuple[int, ...]:
    masks = []
    for choices in itertools.product([0, 1], repeat=len(pairing)):
        mask = 0
        for bit, pair in zip(choices, pairing):
            mask |= 1 << pair[bit]
        masks.append(mask)
    return tuple(masks)


def min_cover_size(rainbow_masks: tuple[int, ...], full_mask: int) -> int | None:
    for k in range(1, len(rainbow_masks) + 1):
        for combo in itertools.combinations(rainbow_masks, k):
            union = 0
            for mask in combo:
                union |= mask
            if union == full_mask:
                return k
    return None


def has_double_cover(rainbow_masks: tuple[int, ...], n: int) -> bool:
    target = tuple(2 for _ in range(n))
    for combo in itertools.combinations_with_replacement(rainbow_masks, 4):
        counts = [0] * n
        for mask in combo:
            bitmask = mask
            idx = 0
            while bitmask:
                if bitmask & 1:
                    counts[idx] += 1
                bitmask >>= 1
                idx += 1
        if tuple(counts) == target:
            return True
    return False


def subset_to_mask(subset: Iterable[int]) -> int:
    mask = 0
    for e in subset:
        mask |= 1 << int(e)
    return mask


@dataclass(frozen=True)
class PairingResult:
    pairing: tuple[tuple[int, int], ...]
    rainbow_count: int
    min_cover: int | None
    has_double_cover: bool


@dataclass(frozen=True)
class MatroidResult:
    name: str
    rank: int
    num_bases: int
    complementary_basis_pairs: int
    worst_cover: int | None
    no_three_cover_pairings: tuple[PairingResult, ...]
    no_double_cover_pairings: tuple[PairingResult, ...]


def analyze_base_mask_set(
    *,
    name: str,
    n: int,
    r: int,
    base_masks: tuple[int, ...],
    pairing_data,
) -> MatroidResult | None:
    full_mask = (1 << n) - 1
    base_mask_set = set(base_masks)
    complementary_basis_pairs = sum(1 for mask in base_masks if (full_mask ^ mask) in base_mask_set) // 2
    if complementary_basis_pairs == 0:
        return None

    no_three_cover = []
    no_double_cover = []
    worst_cover = 0
    for pairing, transversal_masks in pairing_data:
        rainbow_masks = tuple(mask for mask in transversal_masks if mask in base_mask_set)
        if not rainbow_masks:
            result = PairingResult(pairing, 0, None, False)
            no_three_cover.append(result)
            no_double_cover.append(result)
            worst_cover = None
            continue

        cover = min_cover_size(rainbow_masks, full_mask)
        double_cover = has_double_cover(rainbow_masks, n)
        result = PairingResult(pairing, len(rainbow_masks), cover, double_cover)
        if cover is None or cover > 3:
            no_three_cover.append(result)
            worst_cover = None if cover is None else max(worst_cover, cover)
        else:
            worst_cover = max(worst_cover, cover)
        if not double_cover:
            no_double_cover.append(result)

    return MatroidResult(
        name=name,
        rank=r,
        num_bases=len(base_masks),
        complementary_basis_pairs=complementary_basis_pairs,
        worst_cover=worst_cover,
        no_three_cover_pairings=tuple(no_three_cover),
        no_double_cover_pairings=tuple(no_double_cover),
    )


def analyze_matroid(M, pairing_data) -> MatroidResult | None:
    ground = tuple(sorted(int(e) for e in M.groundset()))
    n = len(ground)
    r = M.rank()
    assert n == 2 * r
    base_masks = tuple(sorted(subset_to_mask(B) for B in M.bases()))
    return analyze_base_mask_set(
        name=str(M),
        n=n,
        r=r,
        base_masks=base_masks,
        pairing_data=pairing_data,
    )


def is_spanning_tree(mask: int, chosen_edges: tuple[tuple[int, int], ...], num_vertices: int) -> bool:
    parent = list(range(num_vertices))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    edges_used = 0
    for label, (u, v) in enumerate(chosen_edges):
        if not (mask >> label) & 1:
            continue
        edges_used += 1
        ru = find(u)
        rv = find(v)
        if ru == rv:
            return False
        parent[ru] = rv
    if edges_used != num_vertices - 1:
        return False
    root = find(0)
    return all(find(v) == root for v in range(num_vertices))


def exhaustive_search(ranks: list[int], limit_per_rank: int | None = None) -> dict:
    summary = {"mode": "exhaustive", "ranks": []}
    for r in ranks:
        n = 2 * r
        pairings = tuple(all_pairings(tuple(range(n))))
        pairing_data = tuple((pairing, pairing_to_masks(pairing)) for pairing in pairings)
        checked = 0
        qualifying = 0
        max_cover = 0
        counterexamples = []
        double_cover_failures = []
        for idx, M in enumerate(matroids.AllMatroids(n, r)):
            if limit_per_rank is not None and idx >= limit_per_rank:
                break
            checked += 1
            result = analyze_matroid(M, pairing_data)
            if result is None:
                continue
            qualifying += 1
            if result.worst_cover is None:
                max_cover = None
            elif max_cover is not None:
                max_cover = max(max_cover, result.worst_cover)
            if result.no_three_cover_pairings:
                counterexamples.append(
                    {
                        "matroid": result.name,
                        "pairings": [
                            {
                                "pairing": pr.pairing,
                                "rainbow_count": pr.rainbow_count,
                                "min_cover": pr.min_cover,
                                "has_double_cover": pr.has_double_cover,
                            }
                            for pr in result.no_three_cover_pairings
                        ],
                    }
                )
            if result.no_double_cover_pairings:
                double_cover_failures.append(
                    {
                        "matroid": result.name,
                        "pairings": [
                            {
                                "pairing": pr.pairing,
                                "rainbow_count": pr.rainbow_count,
                                "min_cover": pr.min_cover,
                            }
                            for pr in result.no_double_cover_pairings
                        ],
                    }
                )

        summary["ranks"].append(
            {
                "rank": r,
                "elements": n,
                "pairings": len(pairings),
                "matroids_checked": checked,
                "qualifying_matroids": qualifying,
                "max_min_cover": max_cover,
                "three_cover_counterexamples": counterexamples,
                "double_cover_failures": double_cover_failures,
            }
        )
    return summary


def random_full_rank_matrix(field_size: int, rows: int, cols: int, rng: random.Random):
    while True:
        entries = [rng.randrange(field_size) for _ in range(rows * cols)]
        A = Matrix(GF(field_size), rows, cols, entries)
        if A.rank() == rows:
            return A


def random_linear_search(rank: int, field_size: int, samples: int, seed: int) -> dict:
    n = 2 * rank
    rng = random.Random(seed)
    pairings = tuple(all_pairings(tuple(range(n))))
    pairing_data = tuple((pairing, pairing_to_masks(pairing)) for pairing in pairings)

    seen = set()
    qualifying = 0
    max_cover = 0
    counterexamples = []
    double_cover_failures = []

    for _ in range(samples):
        A = random_full_rank_matrix(field_size, rank, n, rng)
        M = Matroid(matrix=A)
        key = tuple(sorted(tuple(sorted(int(e) for e in B)) for B in M.bases()))
        if key in seen:
            continue
        seen.add(key)
        result = analyze_matroid(M, pairing_data)
        if result is None:
            continue
        qualifying += 1
        if result.worst_cover is None:
            max_cover = None
        elif max_cover is not None:
            max_cover = max(max_cover, result.worst_cover)
        if result.no_three_cover_pairings:
            counterexamples.append(
                {
                    "matroid": result.name,
                    "matrix": str(A),
                    "pairings": [
                        {
                            "pairing": pr.pairing,
                            "rainbow_count": pr.rainbow_count,
                            "min_cover": pr.min_cover,
                            "has_double_cover": pr.has_double_cover,
                        }
                        for pr in result.no_three_cover_pairings
                    ],
                }
            )
        if result.no_double_cover_pairings:
            double_cover_failures.append(
                {
                    "matroid": result.name,
                    "matrix": str(A),
                    "pairings": [
                        {
                            "pairing": pr.pairing,
                            "rainbow_count": pr.rainbow_count,
                            "min_cover": pr.min_cover,
                        }
                        for pr in result.no_double_cover_pairings
                    ],
                }
            )

    return {
        "mode": "random_linear",
        "rank": rank,
        "elements": n,
        "field_size": field_size,
        "samples_requested": samples,
        "distinct_matroids_checked": len(seen),
        "qualifying_matroids": qualifying,
        "max_min_cover": max_cover,
        "three_cover_counterexamples": counterexamples,
        "double_cover_failures": double_cover_failures,
    }


def simple_graphic_search(num_vertices: int, num_edges: int) -> dict:
    complete_edges = tuple(itertools.combinations(range(num_vertices), 2))
    pairings = tuple(all_pairings(tuple(range(num_edges))))
    pairing_data = tuple((pairing, pairing_to_masks(pairing)) for pairing in pairings)

    checked = 0
    connected = 0
    qualifying = 0
    max_cover = 0
    counterexamples = []
    double_cover_failures = []

    for chosen_edges in itertools.combinations(complete_edges, num_edges):
        checked += 1
        G = Graph()
        G.add_vertices(range(num_vertices))
        for label, (u, v) in enumerate(chosen_edges):
            G.add_edge(u, v, label)
        if not G.is_connected():
            continue
        connected += 1
        base_masks = []
        for labels in itertools.combinations(range(num_edges), num_vertices - 1):
            mask = 0
            for label in labels:
                mask |= 1 << label
            if is_spanning_tree(mask, chosen_edges, num_vertices):
                base_masks.append(mask)
        result = analyze_base_mask_set(
            name=str(chosen_edges),
            n=num_edges,
            r=num_vertices - 1,
            base_masks=tuple(sorted(base_masks)),
            pairing_data=pairing_data,
        )
        if result is None:
            continue
        qualifying += 1
        if result.worst_cover is None:
            max_cover = None
        elif max_cover is not None:
            max_cover = max(max_cover, result.worst_cover)
        if result.no_three_cover_pairings:
            counterexamples.append(
                {
                    "graph_edges": chosen_edges,
                    "matroid": result.name,
                    "pairings": [
                        {
                            "pairing": pr.pairing,
                            "rainbow_count": pr.rainbow_count,
                            "min_cover": pr.min_cover,
                            "has_double_cover": pr.has_double_cover,
                        }
                        for pr in result.no_three_cover_pairings
                    ],
                }
            )
        if result.no_double_cover_pairings:
            double_cover_failures.append(
                {
                    "graph_edges": chosen_edges,
                    "matroid": result.name,
                    "pairings": [
                        {
                            "pairing": pr.pairing,
                            "rainbow_count": pr.rainbow_count,
                            "min_cover": pr.min_cover,
                        }
                        for pr in result.no_double_cover_pairings
                    ],
                }
            )

    return {
        "mode": "simple_graphic",
        "vertices": num_vertices,
        "edges": num_edges,
        "graphs_checked": checked,
        "connected_graphs": connected,
        "qualifying_matroids": qualifying,
        "pairings": len(pairings),
        "max_min_cover": max_cover,
        "three_cover_counterexamples": counterexamples,
        "double_cover_failures": double_cover_failures,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="mode", required=True)

    exhaustive = subparsers.add_parser("exhaustive")
    exhaustive.add_argument("--ranks", nargs="+", type=int, required=True)
    exhaustive.add_argument("--limit-per-rank", type=int, default=None)

    random_linear = subparsers.add_parser("random-linear")
    random_linear.add_argument("--rank", type=int, required=True)
    random_linear.add_argument("--field-size", type=int, required=True)
    random_linear.add_argument("--samples", type=int, default=1000)
    random_linear.add_argument("--seed", type=int, default=0)

    simple_graphic = subparsers.add_parser("simple-graphic")
    simple_graphic.add_argument("--vertices", type=int, required=True)
    simple_graphic.add_argument("--edges", type=int, required=True)

    args = parser.parse_args()
    if args.mode == "exhaustive":
        result = exhaustive_search(args.ranks, args.limit_per_rank)
    elif args.mode == "random-linear":
        result = random_linear_search(args.rank, args.field_size, args.samples, args.seed)
    else:
        result = simple_graphic_search(args.vertices, args.edges)
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()