#!/usr/bin/env python3
"""
Algorithmic Doeblin witness search on a finite host instance.

RU is the universality engine. This tool searches for an allowed word (sequence of edges)
whose product matrix is strictly positive.

Inputs:
- --host_instance path to a standardized host instance JSON
- --Lmax maximum word length

Outputs:
- found word (edge ids and template symbols) and matrix summary (delta, U)
"""
from __future__ import annotations
import json, argparse
from collections import defaultdict, deque
from typing import Dict, Any, List

def load(p: str) -> Dict[str, Any]:
    with open(p,"r",encoding="utf-8") as f:
        return json.load(f)

def mm(X: List[List[float]], Y: List[List[float]]) -> List[List[float]]:
    return [[sum(X[i][k]*Y[k][j] for k in range(len(Y))) for j in range(len(Y[0]))] for i in range(len(X))]

def is_strict_pos(M: List[List[float]]) -> bool:
    return all(M[i][j] > 0 for i in range(len(M)) for j in range(len(M[0])))

def min_entry(M: List[List[float]]) -> float:
    return float(min(min(r) for r in M))

def max_entry(M: List[List[float]]) -> float:
    return float(max(max(r) for r in M))

def identity(n: int) -> List[List[float]]:
    return [[1.0 if i==j else 0.0 for j in range(n)] for i in range(n)]

def main()->int:
    ap=argparse.ArgumentParser()
    ap.add_argument("--host_instance", required=True)
    ap.add_argument("--Lmax", type=int, default=8)
    args=ap.parse_args()

    inst=load(args.host_instance)
    n=int(inst["cone_dimension"])
    templates=inst["templates"]
    edges=inst["constraint_graph"]["edges"]
    states=inst["constraint_graph"]["states"]

    out_edges=defaultdict(list)
    for e in edges:
        out_edges[e["from"]].append(e)

    q=deque()
    for s in states:
        q.append((s,0,identity(n),[],[]))

    seen=set()
    while q:
        cur,L,M,ids,syms=q.popleft()
        if L>=args.Lmax:
            continue
        for e in out_edges[cur]:
            T=templates[e["template"]]["matrix"]
            M2=mm(T,M)
            ids2=ids+[e["id"]]
            syms2=syms+[e["template"]]
            cur2=e["to"]
            if is_strict_pos(M2):
                print("FOUND")
                print("L:", L+1)
                print("edge_ids:", ",".join(ids2))
                print("templates:", ",".join(syms2))
                print("delta:", min_entry(M2))
                print("U:", max_entry(M2))
                return 0
            pat="".join("1" if M2[i][j]>0 else "0" for i in range(n) for j in range(n))
            sig=(cur2,L+1,pat)
            if sig in seen:
                continue
            seen.add(sig)
            q.append((cur2,L+1,M2,ids2,syms2))

    print("NOT_FOUND up to Lmax", args.Lmax)
    return 1

if __name__=="__main__":
    raise SystemExit(main())
