"""Exact fixed-margin cycle acceleration with a general QF_NRA backend.

This implements Algorithm 1's control flow. It does NOT implement the paper's
projected three-variable CAD backend, positive-infinitesimal run, or FPT bound.
Strict decision is a separate full-system nonlinear-real-arithmetic query.
"""
from dataclasses import dataclass
from fractions import Fraction as F
import z3


class BudgetUnknown(RuntimeError):pass


def q(x):return z3.RealVal(str(F(x)))
def exact(x):return z3.simplify(x)


def truth(formula):
    value=z3.simplify(formula)
    if z3.is_true(value):return True
    if z3.is_false(value):return False
    raise BudgetUnknown('ground algebraic comparison did not simplify')


def numeric(value):
    value=exact(value)
    if z3.is_rational_value(value):return value.numerator_as_long()/value.denominator_as_long()
    return numeric(value.approx(25))


def encode(value):
    value=exact(value)
    return dict(exact=value.sexpr(),approximation=numeric(value))


def status(solver):
    result=solver.check()
    if result==z3.unknown:raise BudgetUnknown(solver.reason_unknown())
    return result


def full_system(problem,margin=None,timeout_ms=30000):
    solver=z3.SolverFor('QF_NRA');solver.set(timeout=timeout_ms)
    values={v:z3.Real(f'x_{i}') for i,v in enumerate(problem.boxes)}
    for v,b in problem.boxes.items():solver.add(values[v]>=q(b.lower),values[v]<=q(b.upper))
    for e in problem.cores:
        x,y=values[e.tail],values[e.head];lo=e.lower(x);hi=e.upper(x)
        if margin is None:solver.add(y>lo,y<hi)
        else:solver.add(y>=lo+q(margin*e.lower_weight),y+q(margin*e.upper_weight)<=hi)
    return solver,values


@dataclass(frozen=True)
class Implication:
    index: int
    core: object
    forward: bool
    margin: F

    @property
    def tail(self):return self.core.tail if self.forward else self.core.head
    @property
    def head(self):return self.core.head if self.forward else self.core.tail

    def apply(self,value):
        e=self.core
        if self.forward:return exact(e.lower(value)+q(self.margin*e.lower_weight))
        r=q(e.ratio_lower);target=value+q(self.margin*e.upper_weight)
        return exact((-r+z3.Sqrt(r*r+4*(r+1)*target))/2)

    def relation(self,x,y):
        e=self.core
        if self.forward:return y==e.lower(x)+q(self.margin*e.lower_weight)
        return e.upper(y)==x+q(self.margin*e.upper_weight)


@dataclass(frozen=True)
class Label:
    seed: str
    vertices: tuple
    arcs: tuple
    value: object


class CycleAccelerator:
    def __init__(self,problem,margin,timeout_ms=30000,max_jumps=200,max_root_candidates=1000):
        self.problem=problem;self.margin=F(margin);self.timeout_ms=timeout_ms;self.max_jumps=max_jumps;self.max_root_candidates=max_root_candidates
        if self.margin<0 or timeout_ms<=0 or max_jumps<0 or max_root_candidates<1:raise ValueError('invalid margin or budgets')
        self.arcs=tuple(Implication(2*i+j,e,j==0,self.margin) for i,e in enumerate(problem.cores) for j in range(2))

    def least_cycle_root(self,cycle,current,upper):
        solver=z3.SolverFor('QF_NRA');solver.set(timeout=self.timeout_ms)
        values=z3.Reals(' '.join(f'cycle_{i}' for i in range(len(cycle))))
        solver.add(values[0]>current,values[0]<=q(upper))
        for i,arc in enumerate(cycle):
            x,y=values[i],values[(i+1)%len(cycle)];solver.add(x>=0,y>=0,arc.relation(x,y))
        # Intermediate transport values are not constrained by species boxes.
        # The monotone inverse is selected by nonnegativity, excluding spurious roots.
        candidate=None
        for _ in range(self.max_root_candidates):
            result=status(solver)
            if result==z3.unsat:return candidate
            candidate=exact(solver.model().eval(values[0],model_completion=True))
            solver.add(values[0]<candidate)
        raise BudgetUnknown('cycle-root candidate budget exhausted')

    def solve(self):
        trace=[];anchors={v:q(b.lower) for v,b in self.problem.boxes.items()};seeds={v:dict(kind='lower_endpoint',value=encode(a)) for v,a in anchors.items()};jumps=0;rounds=0;max_path=0
        try:
            while True:
                labels={v:Label(v,(v,),(),a) for v,a in anchors.items()}
                for phase_round in range(len(anchors)+1):
                    rounds+=1
                    if any(truth(label.value>q(self.problem.boxes[v].upper)) for v,label in labels.items()):
                        return dict(status='infeasible',reason='forced lower bound exceeds upper box',trace=trace,jumps=jumps)
                    candidates={};detection=None
                    for arc in self.arcs:
                        incoming=labels[arc.tail];value=arc.apply(incoming.value)
                        if not truth(value>labels[arc.head].value):continue
                        if arc.head in incoming.vertices:
                            start=incoming.vertices.index(arc.head);cycle=incoming.arcs[start:]+(arc,);detection=(arc.head,cycle);break
                        candidate=Label(incoming.seed,incoming.vertices+(arc.head,),incoming.arcs+(arc,),value)
                        old=candidates.get(arc.head)
                        if old is None or truth(value>old.value):candidates[arc.head]=candidate
                    if detection is not None:
                        root_vertex,cycle=detection;current=labels[root_vertex].value;image=current
                        for arc in cycle:image=arc.apply(image)
                        assert truth(image>current)
                        if jumps>=self.max_jumps:raise BudgetUnknown('cycle-jump budget exhausted')
                        root=self.least_cycle_root(cycle,current,self.problem.boxes[root_vertex].upper)
                        item=dict(root_vertex=root_vertex,cycle_arcs=[a.index for a in cycle],before=encode(current),cycle_value_before=encode(image),phase_round=phase_round+1)
                        if root is None:
                            item['root']=None;trace.append(item)
                            return dict(status='infeasible',reason='violated cycle has no admissible fixed point above the current bound',trace=trace,jumps=jumps)
                        assert truth(root>current) and truth(root>anchors[root_vertex])
                        image=root
                        for arc in cycle:image=arc.apply(image)
                        assert truth(image==root)
                        anchors[root_vertex]=root;seeds[root_vertex]=dict(kind='cycle_root',cycle_arcs=[a.index for a in cycle],value=encode(root));item['root']=encode(root);trace.append(item);jumps+=1
                        break
                    if not candidates:
                        state={v:l.value for v,l in labels.items()}
                        return dict(status='feasible',state=state,trace=trace,jumps=jumps,rounds=rounds,max_simple_transport=max_path,labels={v:dict(seed=l.seed,seed_description=seeds[l.seed],path=list(l.vertices),arcs=[a.index for a in l.arcs]) for v,l in labels.items()})
                    labels.update(candidates);max_path=max(max_path,max(len(l.arcs) for l in labels.values()))
                else:raise AssertionError('simple-label phase exceeded its finite vertex bound')
        except BudgetUnknown as error:
            return dict(status='unknown',reason=str(error),trace=trace,jumps=jumps)

    def audit_least(self,state):
        """Independent full-system checks of feasibility and no smaller feasible coordinate."""
        for v,b in self.problem.boxes.items():assert truth(state[v]>=q(b.lower)) and truth(state[v]<=q(b.upper))
        for e in self.problem.cores:
            a,b=e.margins(state[e.tail],state[e.head]);assert truth(a>=q(self.margin)) and truth(b>=q(self.margin))
        s,z=full_system(self.problem,self.margin,self.timeout_ms);s.add(z3.Or([z[v]<value for v,value in state.items()]))
        result=status(s)
        if result!=z3.unsat:raise AssertionError('a smaller feasible coordinate exists')
        return True


def round_down(problem,state,margin):
    """Exact one-sided dyadic rounding, preserving boxes including singleton boxes."""
    margin=F(margin)
    if margin<=0:raise ValueError('strictly positive margin required for rational reconstruction')
    minimum_weight=min((min(e.lower_weight,e.upper_weight) for e in problem.cores),default=F(1))
    N=1
    while F(1,N)>margin*minimum_weight/4:N*=2
    result={}
    for v,value in state.items():
        low,high=0,N+1
        while high-low>1:
            mid=(low+high)//2
            if truth(q(F(mid,N))<=value):low=mid
            else:high=mid
        result[v]=max(problem.boxes[v].lower,F(low,N))
    assert problem.check_rational(result,margin/2)
    return result


def strict_decision(problem,timeout_ms=30000):
    """General exact algebraic backend, not the paper's infinitesimal/FPT algorithm."""
    try:
        s,z=full_system(problem,None,timeout_ms);answer=status(s)
        if answer==z3.unsat:return dict(status='infeasible',method='full-system QF_NRA strict constraints')
        state={v:exact(s.model().eval(value,model_completion=True)) for v,value in z.items()}
        if not problem.cores:return dict(status='feasible',witness={v:b.lower for v,b in problem.boxes.items()},margin=None)
        margins=[exact(m) for e in problem.cores for m in e.margins(state[e.tail],state[e.head])]
        t=F(1,2)
        while not all(truth(m>=q(t)) for m in margins):t/=2
        witness=round_down(problem,state,t)
        assert problem.check_rational(witness,strict=True)
        return dict(status='feasible',witness=witness,margin=t/2,method='full-system QF_NRA plus exact downward rational reconstruction')
    except BudgetUnknown as error:return dict(status='unknown',reason=str(error))


def capacity_bracket(problem,bits=12,timeout_ms=30000):
    """Bracket a nonnegative capacity; negative capacities are not estimated here."""
    if type(bits) is not int or bits<1:raise ValueError('positive bit count required')
    if not problem.cores:return dict(status='unbounded',scope='No core constraints; every nonnegative margin is feasible.')
    try:
        s,_=full_system(problem,F(0),timeout_ms)
        if status(s)==z3.unsat:return dict(status='negative',scope='even zero margin infeasible; negative optimum not calculated')
        low,high=F(0),max(F(1),max(F(1)/(e.lower_weight+e.upper_weight) for e in problem.cores))
        for _ in range(bits):
            mid=(low+high)/2;s,_=full_system(problem,mid,timeout_ms)
            if status(s)==z3.sat:low=mid
            else:high=mid
        return dict(status='bracketed',lower=low,upper=high,precision=high-low)
    except BudgetUnknown as error:return dict(status='unknown',reason=str(error))
