#!/usr/bin/env python3
"""Generate and verify the exact instances in Section 7 of the paper."""
from __future__ import annotations

import argparse
from fractions import Fraction
from itertools import combinations
import json
from pathlib import Path

from exact import check_matching, incidence, maximum_matching, patterns, product, rank

ROOT = Path(__file__).resolve().parents[1]
VALUES = [-34, -28, -20, 7, 10, 13, 14, 17, 21]
COUNTS = [2, 2, 2, 1, 3, 1, 3, 3, 1]
TRIPLES = [(-20, 10, 10), (-20, 7, 13), (-28, 14, 14),
           (-28, 7, 21), (-34, 17, 17), (-34, 13, 21)]
ALPHA = list(map(Fraction, ['3/2', '1/2', '3/2', '1/2', '3/2', '1/2']))
FIVE = [(-20, 10, 10), (-20, 7, 13), (-28, 14, 14),
        (-28, 10, 14), (-34, 17, 17)]


def base_columns():
    return [tuple(e.count(v) for v in VALUES) for e in TRIPLES]


def build_family(blocks: int) -> dict:
    if blocks < 1:
        raise ValueError("The number of copies must be positive")
    values = [100**j * v for j in range(blocks) for v in VALUES]
    counts = COUNTS * blocks
    denom = 1000 * 100**(blocks - 1)
    columns = []
    for j in range(blocks):
        for p in base_columns():
            column = [0] * (9 * blocks)
            column[9*j:9*j+9] = p
            columns.append(column)
    matching = [[100**j * v for v in e] for j in range(blocks) for e in FIVE]
    matching += [[-34*100**j, 17*100**j, 21*100**(j-1)] for j in range(1, blocks)]
    return {
        "copies": blocks, "values": values, "multiplicities": counts,
        "sizes": [str(Fraction(1, 3) + Fraction(v, denom)) for v in values],
        "patterns": columns, "fractional_coefficients": [str(x) for x in ALPHA*blocks],
        "matching": matching,
        "expected": {"items": 18*blocks, "lp": 6*blocks, "opt": 6*blocks+1,
                     "dimension": 6*blocks-1, "zero_matching": 4*blocks,
                     "nonpositive_matching": 6*blocks-1},
    }


def verify_family(data: dict, compute_matching: bool = False) -> dict:
    b = data['copies']; values = data['values']; counts = data['multiplicities']
    columns = data['patterns']; P = incidence(columns)
    alpha = list(map(Fraction, data['fractional_coefficients']))
    sizes = list(map(Fraction, data['sizes']))
    assert len(values) == len(set(values)) == 9*b
    assert sum(counts) == 18*b and sum(x*n for x, n in zip(values, counts)) == 0
    assert all(Fraction(1, 4) < x < Fraction(1, 2) for x in sizes)
    assert sum(x*n for x, n in zip(sizes, counts)) == 6*b
    assert set(patterns(values, counts)) == set(map(tuple, columns))
    assert rank(P) == 6*b
    assert product(P, alpha) == counts and sum(alpha) == 6*b
    assert all(sum(p) == 3 and sum(x*y for x, y in zip(p, values)) == 0 for p in columns)
    assert product(P, [1]*(6*b)) == [2]*(9*b)
    sums, remainder = check_matching(values, counts, data['matching'])
    assert len(sums) == 6*b-1 and sum(remainder.values()) == 3

    # The four disjoint triangles on positive items give the upper bound 4b.
    # Check their exact pair list on one copy and exhibit a matching meeting it.
    positives = [x for x, n in zip(VALUES, COUNTS) if x > 0 for _ in range(n)]
    pair_edges = {e for e in combinations(range(12), 2) if sum(positives[i] for i in e) in {20, 28, 34}}
    groups = [[i for i, x in enumerate(positives) if x == a] for a in [10, 14, 17]]
    groups.append([i for i, x in enumerate(positives) if x in {7, 13, 21}])
    assert all(len(g) == 3 for g in groups)
    assert pair_edges == {e for g in groups for e in combinations(g, 2)}
    check_matching(VALUES, COUNTS, [TRIPLES[i] for i in [0, 2, 4, 1]])

    floor = [int(x) for x in alpha]
    residual = [n-q for n, q in zip(counts, product(P, floor))]
    assert residual == [1]*(9*b)
    zero_residual = patterns(values, residual)
    assert len(zero_residual) == 3*b
    for j in range(b):
        local = [p for p in zero_residual if any(p[9*j:9*j+9])]
        assert len(local) == 3
        assert all(sum(min(x, y) for x, y in zip(p, q)) >= 1 for p, q in combinations(local, 2))
    if b >= 2:
        assert b < 3*b-3
    report = {"copies": b, "verified": True, "patterns": len(columns),
              "rank": rank(P), "items": sum(counts), "matching_constructed": len(sums),
              "remainder_zero_matching": b}
    if compute_matching:
        labelled = [v for v, n in zip(values, counts) for _ in range(n)]
        report['zero_matching'] = maximum_matching(labelled, 'zero')
        report['nonpositive_matching'] = maximum_matching(labelled, 'nonpositive')
        assert report['zero_matching']['size'] == 4
        assert report['nonpositive_matching']['size'] == 5
    return report


def main() -> None:
    if not __debug__:
        raise RuntimeError("Run without -O: assertion checks are required")
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--write', action='store_true', help='write the two canonical data files')
    parser.add_argument('--output', type=Path, help='write the verification report')
    args = parser.parse_args()
    directory = ROOT / 'data' / 'examples'
    directory.mkdir(parents=True, exist_ok=True)
    result = []
    for b, filename in [(1, 'eighteen_items.json'), (2, 'twelve_pattern_remainder.json')]:
        canonical = build_family(b)
        path = directory / filename
        if args.write:
            path.write_text(json.dumps(canonical, indent=2) + '\n')
        data = json.loads(path.read_text())
        assert data == canonical, f"Unexpected instance: {filename}"
        result.append(verify_family(data, compute_matching=(b == 1)))
    # Additional finite checks of the formulas; the general proof is in the paper.
    result.append(verify_family(build_family(3)))
    text = json.dumps({"all_examples_verified": True, "instances": result}, indent=2) + '\n'
    if args.output:
        args.output.write_text(text)
    print(text, end='')


if __name__ == '__main__':
    main()