"""Original LearnAI reference. Standard-library-only bounded A* and HMM filtering."""
import argparse
import heapq
import json
import math


def finite(x):
    try:
        return isinstance(x, (int, float)) and not isinstance(x, bool) and math.isfinite(x)
    except OverflowError:
        return False


def astar(graph, heuristic, start, goal, budget=64):
    """Nonnegative directed costs; strict improvements reopen states. Goal selected at cap is valid."""
    if not isinstance(graph, dict) or not 1 <= len(graph) <= 32 or start not in graph or goal not in graph:
        raise ValueError('Use 1–32 named states including start and goal')
    if any(not isinstance(n, str) or not 1 <= len(n) <= 32 for n in graph):
        raise ValueError('Bounded string state names required')
    if type(budget) is not int or not 1 <= budget <= 256 or not isinstance(heuristic, dict) or set(heuristic) != set(graph):
        raise ValueError('Invalid budget or heuristic keys')
    if any(not finite(v) or v < 0 for v in heuristic.values()) or heuristic[goal] != 0:
        raise ValueError('Nonnegative finite heuristics with zero goal estimate required')
    if any(not isinstance(row, list) for row in graph.values()) or sum(map(len, graph.values())) > 128:
        raise ValueError('At most 128 edge records')
    for row in graph.values():
        for edge in row:
            if not isinstance(edge, (tuple, list)) or len(edge) != 2 or edge[0] not in graph or not finite(edge[1]) or edge[1] < 0:
                raise ValueError('Invalid directed edge')
    order = {n: i for i, n in enumerate(graph)}
    best, parents, closed = {start: 0.0}, {}, set()
    frontier = [(heuristic[start], order[start], 0.0, start)]
    expansions = edge_reads = reopens = 0
    trace = []
    while frontier:
        f, _, g, node = heapq.heappop(frontier)
        if g != best.get(node) or node in closed:
            continue
        if node == goal:
            path = [node]
            while path[-1] != start:
                path.append(parents[path[-1]])
            return dict(status='found', route=path[::-1], cost=g, expansions=expansions, edgeReads=edge_reads, reopens=reopens, trace=trace)
        if expansions == budget:
            return dict(status='budget', route=[], cost=None, expansions=expansions, edgeReads=edge_reads, reopens=reopens, trace=trace)
        expansions += 1
        closed.add(node)
        trace.append(dict(node=node, g=g, h=heuristic[node], f=f))
        for to, cost in graph[node]:
            edge_reads += 1
            candidate = g + cost
            if not finite(candidate) or not finite(candidate + heuristic[to]):
                raise ValueError('Accumulated priority exceeds numeric range')
            if candidate >= best.get(to, math.inf):
                continue
            if to in closed:
                closed.remove(to)
                reopens += 1
            best[to], parents[to] = candidate, node
            heapq.heappush(frontier, (candidate + heuristic[to], order[to], candidate, to))
    return dict(status='exhausted', route=[], cost=None, expansions=expansions, edgeReads=edge_reads, reopens=reopens, trace=trace)


def distribution(row, n):
    if not isinstance(row, (list, tuple)) or len(row) != n or any(not finite(p) or p < 0 or p > 1 for p in row) or abs(sum(row)-1) > 1e-12:
        raise ValueError('Invalid distribution')


def filter_belief(initial, transition, likelihoods):
    """One movement before each supplied likelihood row; all-ones row means missing evidence."""
    if not isinstance(initial, (list, tuple)) or not 2 <= len(initial) <= 8:
        raise ValueError('Use 2–8 states')
    n = len(initial)
    distribution(initial, n)
    if not isinstance(transition, (list, tuple)) or len(transition) != n:
        raise ValueError('Invalid transition matrix')
    for row in transition:
        distribution(row, n)
    if not isinstance(likelihoods, (list, tuple)) or len(likelihoods) > 16:
        raise ValueError('At most 16 observations')
    for row in likelihoods:
        if not isinstance(row, (list, tuple)) or len(row) != n or any(not finite(p) or p < 0 or p > 1 for p in row):
            raise ValueError('Invalid observation likelihood')
    belief, steps, log_evidence = list(initial), [], 0.0
    for time, likelihood in enumerate(likelihoods, 1):
        predicted = [sum(belief[i] * transition[i][j] for i in range(n)) for j in range(n)]
        weights = [predicted[i] * likelihood[i] for i in range(n)]
        evidence = sum(weights)
        posterior = [w/evidence for w in weights] if evidence else None
        steps.append(dict(time=time, predicted=predicted, weights=weights, evidence=evidence, posterior=posterior))
        if posterior is None:
            return dict(status='impossible', steps=steps, logEvidence=None)
        belief = posterior
        log_evidence += math.log(evidence)
    return dict(status='complete', steps=steps, logEvidence=log_evidence)


GRAPH = {'S':[('A',2),('B',1),('X',1)], 'A':[('C',2),('G',8)], 'B':[('D',2)], 'C':[('G',2)], 'D':[('C',.5),('G',8)], 'X':[], 'G':[]}
ESTIMATES = {'uniform':[0,0,0,0,0,0,0], 'consistent':[5,4,4,2,2,6,0], 'inconsistent':[0,0,4.5,0,2.5,0,0], 'misleading':[0,0,9,0,9,0,0]}
TRANSITIONS = {'slow':[[.8,.2,0],[.1,.8,.1],[0,.2,.8]], 'fixed':[[1,0,0],[0,1,0],[0,0,1]]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='command', required=True)
    search = sub.add_parser('search')
    search.add_argument('output'); search.add_argument('--method', choices=ESTIMATES, default='uniform'); search.add_argument('--budget', type=int, choices=range(1,33), default=16)
    belief = sub.add_parser('belief')
    belief.add_argument('output'); belief.add_argument('--motion', choices=TRANSITIONS, default='slow'); belief.add_argument('--steps', type=int, choices=range(1,7), default=3)
    args = parser.parse_args()
    if args.command == 'search':
        result = astar(GRAPH, dict(zip(GRAPH, ESTIMATES[args.method])), 'S', 'G', args.budget)
        result.update(scenario='open', method=args.method, budget=args.budget)
        artifact = dict(format='informed-search-practice-v1', run=result)
    else:
        readings = [0,0,1,2,2,1][:args.steps]
        result = filter_belief([1,0,0], TRANSITIONS[args.motion], [[.8 if p == r else (1-.8)/2 for p in range(3)] for r in readings])
        artifact = dict(format='hmm-filter-practice-v1', initial=[1,0,0], motion=args.motion, accuracy=.8, readings=readings, consumed=len(result['steps']), pending=False, run=result)
    artifact['provenance'] = 'Supplied LearnAI Python reference, not independently authored learner code. Browser recomputes results from supported settings.'
    # Exclusive create: never replace earlier evidence.
    with open(args.output, 'x', encoding='utf8') as output:
        json.dump(artifact, output, indent=2, allow_nan=False)
        output.write('\n')
    print('Wrote', args.output)


if __name__ == '__main__':
    main()