"""Evaluator-only timing and exact assignment; independent global B_cost.
This module is never imported by the local policy or its scheduler.
"""
import heapq,time
from common import CAP,IntegrityError,canonical
from host import neighbors
from legacy import package

def tree_paths(h,parent):
    root=next(t['site'] for t in h['terminals'] if t['kind']=='query'); response=next(t['site'] for t in h['terminals'] if t['kind']=='response')
    writes={t['a']:t['site'] for t in h['terminals'] if t['kind']=='write'}
    delay={tuple(sorted((e['u'],e['v']))):e['delay'] for e in h['bonds']}
    def path(u):
        out=[]
        while u is not None: out.append(u);u=parent[u]
        return list(reversed(out))
    paths={a:path(u) for a,u in writes.items()}
    def distances(p):
        d=[0]
        for u,v in zip(p,p[1:]):d.append(d[-1]+delay[tuple(sorted((u,v)))])
        return d
    dist={a:distances(p) for a,p in paths.items()}; rp=path(response); rd=distances(rp)[-1]
    eligible={a:[u for u in p if u not in (root,response) and (u not in writes.values() or u==writes[a])] for a,p in paths.items()}
    return dict(root=root,response=response,writes=writes,paths=paths,distances=dist,response_path=rp,response_distance=rd,eligible=eligible)

def exact_assignment(costs):
    """Rectangular Hungarian primal/dual assignment on integer costs, K<=8.
    Infeasible entries omitted by caller, represented by a finite dominating M.
    Returns row->column indexes; deterministic ascending column ties.
    """
    n=len(costs);m=len(costs[0]);assert n<=m
    u=[0]*(n+1);v=[0]*(m+1);p=[0]*(m+1);way=[0]*(m+1);ops=0
    for i in range(1,n+1):
        p[0]=i;j0=0;minv=[10**30]*(m+1);used=[False]*(m+1)
        while True:
            used[j0]=True;i0=p[j0];delta=10**30;j1=0
            for j in range(1,m+1):
                if not used[j]:
                    cur=costs[i0-1][j-1]-u[i0]-v[j];ops+=1
                    if cur<minv[j]:minv[j]=cur;way[j]=j0
                    if minv[j]<delta:delta=minv[j];j1=j
            for j in range(m+1):
                if used[j]:u[p[j]]+=delta;v[j]-=delta
                else:minv[j]-=delta
            j0=j1
            if p[j0]==0:break
        while True:
            j1=way[j0];p[j0]=p[j1];j0=j1
            if j0==0:break
    answer=[None]*n
    for j in range(1,m+1):
        if p[j]:answer[p[j]-1]=j-1
    return answer,ops

def optimum(h,parent,workload):
    start=time.process_time();t=tree_paths(h,parent);sites=sorted(set(u for row in t['eligible'].values() for u in row))
    costs=[];M=10**18
    for a in range(h['K']):
        by={}
        for u in t['eligible'][a]:
            i=t['paths'][a].index(u);d=t['distances'][a][i];D=t['distances'][a][-1]
            by[u]=workload['wi'][a]*(1+D-d)+workload['qi'][a]*(2+2*d+t['response_distance'])
        costs.append([by.get(u,M) for u in sites])
    answer,ops=exact_assignment(costs);storage={a:sites[i] for a,i in enumerate(answer)}
    total=sum(costs[a][i] for a,i in enumerate(answer))
    if any(costs[a][i]==M for a,i in enumerate(answer)):raise IntegrityError('exact assignment infeasible')
    return dict(storage=storage,J=total/workload['denominator'],integer_cost=total,denominator=workload['denominator'],work_units=ops,cpu_seconds=time.process_time()-start,sites=len(sites))

def latencies(a):
    cells=a['cells'];root=next(t['site'] for t in a['terminals'] if t['kind']=='query');response=next(t['site'] for t in a['terminals'] if t['kind']=='response')
    writes={t['a']:t['site'] for t in a['terminals'] if t['kind']=='write'};stores={c['address']:u for u,c in cells.items() if c['role']=='latch'}
    edge={tuple(sorted((e['u'],e['v']))):e['delay'] for e in a['edges']}
    def distance(u,v,key):
        d=0;seen=set()
        while u!=v:
            if u in seen:raise IntegrityError('latency loop')
            seen.add(u);n=cells[u]['routes'][key];d+=edge[tuple(sorted((u,n)))];u=n
        return d
    response_distance=distance(root,response,'A')
    return {'write':[1+distance(writes[x],stores[x],'W'+str(x)) for x in range(a['K'])],
            'query':[2+distance(root,stores[x],'Q'+str(x))+distance(stores[x],root,'V')+response_distance for x in range(a['K'])]}
def expected(a,w):
    lt=latencies(a);num=sum(w['wi'][i]*lt['write'][i]+w['qi'][i]*lt['query'][i] for i in range(a['K']))
    return dict(J=num/w['denominator'],integer_cost=num,denominator=w['denominator'],latencies=lt)

def b_cost(h,w):
    start=time.perf_counter();cpu=time.process_time();adj=neighbors(h);root=next(t['site'] for t in h['terminals'] if t['kind']=='query')
    distances={root:0};parent={root:None};heap=[(0,root)];visits=0;pops=0;pushes=1
    while heap:
        d,u=heapq.heappop(heap);pops+=1
        if d!=distances[u]:continue
        for v,cost in adj[u]:
            visits+=1;nd=d+cost
            if nd<distances.get(v,10**9):distances[v]=nd;parent[v]=u;heapq.heappush(heap,(nd,v));pushes+=1
    opt=optimum(h,parent,w);storage=opt['storage'];t=tree_paths(h,parent);active={root};planning=0
    for path in list(t['paths'].values())+[t['response_path']]:active.update(path)
    cells={u:dict(role='interface' if u==root else 'latch' if u in storage.values() else 'router',address=next((a for a,v in storage.items() if v==u),None),response=u==t['response'],routes={}) for u in sorted(active)}
    def route(path,key):
        nonlocal planning
        for u,v in zip(path,path[1:]):cells[u]['routes'][key]=v;planning+=1
    for a,p in t['paths'].items():
        i=p.index(storage[a]);route(list(reversed(p[i:])),'W'+str(a));route(p[:i+1],'Q'+str(a));route(list(reversed(p[:i+1])),'V')
    route(t['response_path'],'A')
    c=dict(occupied=len(cells),development_ticks=0,development_messages=0,development_bytes=0,peak_development_queue=0,peak_local_state_bytes=0,peak_total_development_state_bytes=0,
           rule_applications=0,global_builder_work=visits+pops+pushes+planning+opt['work_units'],neighbor_scans=visits,heap_pops=pops,heap_pushes=pushes,path_route_visits=planning,
           assignment_work_units=opt['work_units'],assignment_cpu_seconds=opt['cpu_seconds'],construction_wall_seconds=time.perf_counter()-start,construction_cpu_seconds=time.process_time()-cpu)
    return package(h,cells,c,[])