"""Independent formulas and exhaustive small cases, not substitutes for Lean."""
from fractions import Fraction as F
from itertools import product
import random
import unittest

from example import (Branch, DegradationStar, FiniteRelation, JoinNode, PassiveTail,
                     PriceProfile, dot, gluing_examples, paper_profiles, solve)


class ScientificChecks(unittest.TestCase):
    def test_strict_price_alternative_and_scalar_loss(self):
        n1, n2 = paper_profiles()
        self.assertEqual(n1.decide(2)["kind"], "feasible")
        self.assertFalse(n1.contains(2, (1, 2)))  # zero margins at equality are not strict
        self.assertEqual(n1.compose(n1).decide(4)["kind"], "infeasible")
        result = n2.compose(n1).decide(4)
        self.assertEqual(result["kind"], "feasible")
        self.assertEqual(result["residual"], (0, 0))
        for a, b in product(range(1, 6), repeat=2):
            model = PriceProfile(((1, 0), (0, 1)), ((0, b), (a, 0)))
            for q in (F(i, 2) for i in range(1, 12)):
                result = model.decide(q)
                # Analytic two-reaction threshold sqrt(a*b).
                self.assertEqual(result["kind"] == "feasible", q*q <= a*b)
                if result["kind"] == "feasible":
                    self.assertEqual(sum(result["flux"]), 1)
                    self.assertTrue(all(x >= 0 for x in result["residual"]))
                else:
                    self.assertTrue(model.contains(q, result["prices"]))
        # The paper's weak-price false positive: A -> 2A, B -> B, q=1.
        weak = PriceProfile(((1, 0), (0, 1)), ((2, 0), (0, 1)))
        self.assertFalse(weak.contains(1, (0, 1)))
        self.assertEqual(weak.decide(1)["kind"], "feasible")

    def test_price_intersection_retains_every_halfspace(self):
        a, b = paper_profiles()
        for q, pa, pb in product((F(1), F(2), F(3), F(4), F(5)), range(6), range(6)):
            self.assertEqual(a.compose(b).contains(q, (pa, pb)),
                             a.contains(q, (pa, pb)) and b.contains(q, (pa, pb)))

    def test_passive_tail_reconstruction_and_literal_solve(self):
        rng = random.Random(11)
        for length in range(1, 9):
            for _ in range(12):
                c = tuple(F(rng.randint(1, 8), rng.randint(1, 4)) for _ in range(length))
                b = tuple(F(rng.randint(1, 8), rng.randint(1, 4)) for _ in range(length))
                d = tuple(F(rng.randint(0, 4), 3) for _ in range(length-1))
                tail = PassiveTail(c, b, d)
                x, y = F(3, 2), F(2, 7)
                result = tail.audit(x, y)
                size = length-1
                matrix = [[F(0)] * size for _ in range(size)]
                rhs = [F(0)] * size
                for j in range(size):
                    matrix[j][j] = b[j]+c[j+1]+d[j]
                    if j:
                        matrix[j][j-1] = -c[j]
                    else:
                        rhs[j] += c[0]*x
                    if j+1 < size:
                        matrix[j][j+1] = -b[j+1]
                    else:
                        rhs[j] += b[-1]*y
                self.assertEqual(result["states"], (x,)+solve(matrix, rhs)+(y,))
                jl, jr = result["boundary_currents"]
                self.assertEqual(jl-jr, result["internal_loss"])
                self.assertTrue(all(v > 0 for v in result["states"]))
        with self.assertRaises(ValueError):
            PassiveTail((0,), (1,), ())

    def test_tail_substitution_inside_boundary_context(self):
        # The endpoints are now unknown: forcing and endpoint degradation provide
        # an outside environment. Solve the full chain and the compressed network.
        tail = PassiveTail((2, 3, 1, 4), (1, 2, 2, 1), (F(1,5), F(1,2), F(1,10)))
        port = tail.port()
        for dl, dr, fl, fr in product((F(1, 3), F(2)), repeat=4):
            endpoints = solve(((dl+port.forward+port.left_leak, -port.backward),
                               (-port.forward, dr+port.backward+port.right_leak)), (fl, fr))
            n = len(tail.forward)+1
            matrix = [[F(0)] * n for _ in range(n)]
            for i, (c, b) in enumerate(zip(tail.forward, tail.backward)):
                matrix[i][i] += c
                matrix[i][i+1] -= b
                matrix[i+1][i] -= c
                matrix[i+1][i+1] += b
            for i, d in enumerate((dl,)+tail.degradation+(dr,)):
                matrix[i][i] += d
            literal = solve(matrix, (fl,)+(F(0),)*(n-2)+(fr,))
            self.assertEqual((literal[0], literal[-1]), endpoints)
            self.assertEqual(literal, tail.reconstruct(*endpoints))

    def test_star_load_signs_and_criticality(self):
        b = Branch(((-15,),), (3,), (3,))
        self.assertEqual(b.load(), F(3,5))
        self.assertEqual(DegradationStar(-1, (b,)).certificate()["state"], "extinction")
        self.assertEqual(DegradationStar(-1, (b,b)).certificate()["state"], "growth")
        self.assertEqual(DegradationStar(F(-6,5), (b,b)).certificate()["state"], "critical")
        general = Branch(((-4, 1), (1, -5)), (2, 3), (1, 2))
        for branch in (b, general):
            for count in range(1, 5):
                for shift in (F(-1), F(0), F(1)):
                    star = DegradationStar(-count*branch.load()+shift, (branch,)*count)
                    certificate = star.certificate()
                    self.assertEqual(certificate["schur_load"], shift)
                    self.assertTrue(all(x > 0 for x in certificate["positive_vector"]))
                    for row in star.matrix():
                        value = dot(row, certificate["positive_vector"])
                        self.assertEqual((value > 0)-(value < 0), (shift > 0)-(shift < 0))
        with self.assertRaises(ValueError):
            Branch(((1,),), (3,), (3,))

    def test_join_trees_exhaustive_relations_and_failed_gluing(self):
        good, bad, repaired = gluing_examples()
        self.assertEqual(good.message((0,1)), good.global_trace((0,1)))
        with self.assertRaisesRegex(ValueError, "Running intersection"):
            bad.message((0,1))
        self.assertEqual(bad.global_trace((0,1)).rows, frozenset())
        self.assertEqual(repaired.message((0,1)).rows, frozenset())
        truth = FiniteRelation((), frozenset(((),)))
        universe = tuple(product((0,1), repeat=2))
        subsets = [frozenset(r for i,r in enumerate(universe) if mask & (1<<i)) for mask in range(16)]
        for left, right in product(subsets, repeat=2):
            tree = JoinNode("root", ("x","a","b"), ("a","b"), truth,
                            (JoinNode("left", ("x","a"), ("x","a"), FiniteRelation(("x","a"), left)),
                             JoinNode("right", ("x","b"), ("x","b"), FiniteRelation(("x","b"), right))))
            self.assertEqual(tree.message((0,1)), tree.global_trace((0,1)))
            self.assertEqual(tree.width(), 2)
        # A local relation cannot read an undeclared/out-of-bag variable.
        with self.assertRaises(ValueError):
            JoinNode("bad_scope", (), (), FiniteRelation(("x",), frozenset(((0,),))))

    def test_trace_equality_contexts_and_cardinality_loss(self):
        traces = [FiniteRelation(("x",), frozenset((i,) for i in range(3) if mask & (1<<i))) for mask in range(8)]
        for a, b in product(traces, repeat=2):
            indistinguishable = all(a.compatible(c) == b.compatible(c) for c in traces)
            self.assertEqual(a == b, indistinguishable)
            separator = a.separating_context(b)
            if a == b:
                self.assertIsNone(separator)
            else:
                self.assertNotEqual(a.compatible(separator), b.compatible(separator))
        self.assertEqual(len(traces[1].rows), len(traces[2].rows))
        self.assertNotEqual(traces[1], traces[2])


if __name__ == "__main__":
    unittest.main()