#!/usr/bin/env python3
"""Exact checks for the JT fixed-fan ambiguity note.

The proofs in the companion note are analytic.  This script checks the finite
algebraic identities used as controls: the chamber vertices and adjacent edge
lengths of the deformed generalized permutohedra, and the all-orders-in-k
string/dilaton mechanism through six labelled defects.
"""

from fractions import Fraction
from itertools import combinations, permutations
from math import comb, factorial
import json


def p_poly(m, x, y):
    return sum(
        Fraction(comb(m, j), factorial(j) * factorial(m - j))
        * x**j
        * y ** (m - j)
        for j in range(m + 1)
    )


def contact(q, p, eta):
    return eta * q * q * p * p * (1 - q - p) ** 2


def weight(q, p, c, eta):
    return c * q * p + contact(q, p, eta)


def is_forest(k, edges):
    parent = list(range(k))

    def root(a):
        while parent[a] != a:
            parent[a] = parent[parent[a]]
            a = parent[a]
        return a

    for a, b in edges:
        ra, rb = root(a), root(b)
        if ra == rb:
            return False
        parent[ra] = rb
    return True


def forests(k):
    edge_list = list(combinations(range(k), 2))
    for mask in range(1 << len(edge_list)):
        edges = [edge_list[j] for j in range(len(edge_list)) if mask >> j & 1]
        if is_forest(k, edges):
            yield edges


def forest_volume(k, qs, x, y, c, eta):
    total = Fraction(0)
    for edges in forests(k):
        term = p_poly(k - 1 - len(edges), x, y)
        for a, b in edges:
            term *= weight(qs[a], qs[b], c, eta)
        total += term
    return total


def soft_derivative(k, old_qs, x, y, c, eta):
    """Derivative in q_k at zero, evaluated from the forest expansion."""
    total = Fraction(0)
    for old_edges in forests(k - 1):
        old_term = p_poly(k - 2 - len(old_edges), x, y)
        for a, b in old_edges:
            old_term *= weight(old_qs[a], old_qs[b], c, eta)
        total += old_term * c * sum(old_qs)
    return total


def chamber_vertex(order, weights, x, y):
    k = len(order)
    v = [Fraction(0) for _ in range(k)]
    v[order[0]] += x
    v[order[-1]] -= y
    rank = {label: place for place, label in enumerate(order)}
    for (a, b), w in weights.items():
        v[a if rank[a] < rank[b] else b] += w
    return tuple(v)


def vector_sub(a, b):
    return tuple(x - y for x, y in zip(a, b))


def root_vector(k, a, b, scale):
    result = [Fraction(0) for _ in range(k)]
    result[a] = scale
    result[b] = -scale
    return tuple(result)


def check_fan(k, qs, x, y, c, eta):
    ws = {
        (a, b): weight(qs[a], qs[b], c, eta)
        for a, b in combinations(range(k), 2)
    }
    vertices = {}
    adjacent_checks = 0
    for order in permutations(range(k)):
        vertex = chamber_vertex(order, ws, x, y)
        vertices[order] = vertex
        for r in range(k - 1):
            swapped = list(order)
            a, b = swapped[r], swapped[r + 1]
            swapped[r], swapped[r + 1] = swapped[r + 1], swapped[r]
            swapped = tuple(swapped)
            other = chamber_vertex(swapped, ws, x, y)
            length = ws[tuple(sorted((a, b)))]
            if r == 0:
                length += x
            if r == k - 2:
                length += y
            assert length > 0
            assert vector_sub(vertex, other) == root_vector(k, a, b, length)
            adjacent_checks += 1
    assert len(set(vertices.values())) == factorial(k)
    return {
        "chambers": factorial(k),
        "distinct_vertices": len(set(vertices.values())),
        "oriented_adjacent_checks": adjacent_checks,
    }


def main():
    x = Fraction(2, 3)
    y = Fraction(3, 5)
    c = Fraction(4)
    eta = Fraction(7)

    fan_checks = {}
    soft_checks = {}
    for k in range(2, 7):
        qs = [Fraction(i + 1, 12 * k) for i in range(k)]
        assert sum(qs) < 1
        fan_checks[str(k)] = check_fan(k, qs, x, y, c, eta)

        old_qs = qs[:-1]
        left_string = forest_volume(k, old_qs + [Fraction(0)], x, y, c, eta)
        right_string = sum(
            p_poly(k - 1 - len(edges), x, y)
            * product(weight(old_qs[a], old_qs[b], c, eta) for a, b in edges)
            for edges in forests(k - 1)
        )
        assert left_string == right_string

        left_dilaton = soft_derivative(k, old_qs, x, y, c, eta)
        right_dilaton = c * sum(old_qs) * forest_volume(
            k - 1, old_qs, x, y, c, eta
        )
        assert left_dilaton == right_dilaton
        soft_checks[str(k)] = {
            "string_value": str(left_string),
            "dilaton_value": str(left_dilaton),
        }

    q, p = Fraction(1, 4), Fraction(1, 5)
    undeformed = forest_volume(2, [q, p], x, y, c, Fraction(0))
    deformed = forest_volume(2, [q, p], x, y, c, eta)
    shift = contact(q, p, eta)
    assert deformed - undeformed == shift and shift != 0

    # h(q,p)=eta*q^2*p^2*(1-q-p)^2 has zero value and first p-derivative
    # on p=1-q.  The derivative is expanded directly here.
    wall_p = 1 - q
    wall_value = contact(q, wall_p, eta)
    wall_derivative = eta * q * q * (
        2 * wall_p * (1 - q - wall_p) ** 2
        - 2 * wall_p * wall_p * (1 - q - wall_p)
    )
    assert wall_value == 0 and wall_derivative == 0

    result = {
        "status": "pass",
        "arithmetic": "exact rational",
        "parameters": {
            "x": str(x),
            "y": str(y),
            "c": str(c),
            "eta": str(eta),
        },
        "fan_checks": fan_checks,
        "soft_identity_checks": soft_checks,
        "two_defect_control": {
            "q": str(q),
            "p": str(p),
            "undeformed_volume": str(undeformed),
            "deformed_volume": str(deformed),
            "nonzero_shift": str(shift),
            "collision_wall_value": str(wall_value),
            "collision_wall_first_derivative": str(wall_derivative),
        },
        "scope": (
            "Finite exact controls only. The fixed-fan theorem and relative-"
            "cohomology obstruction are proved analytically in the packet."
        ),
    }
    print(json.dumps(result, indent=2, sort_keys=True))


def product(values):
    answer = Fraction(1)
    for value in values:
        answer *= value
    return answer


if __name__ == "__main__":
    main()