"""Changes to route costs and endpoint merging on an identical action graph."""

from dataclasses import replace
from .graph import Graph, Node
from .exact import solve


def zero_cost_policy(graph, log_rewards):
    """Compute the exact policy after setting every route cost to zero."""
    altered = Graph(
        graph.nodes,
        [replace(e, cost=0.0) for e in graph.edges],
        graph.root,
        graph.metadata,
    )
    return solve(altered, log_rewards, 1.0).forward


def duplicate_endpoint_policy(graph, log_rewards, temperature):
    """Assign a separate reward to each incoming terminal edge.

    Edge ordering and all nonterminal nodes are retained. The returned policy
    can be evaluated on the original graph to measure canonical endpoint bias.
    """
    nodes = [n for n in graph.nodes if n.outcome is None]
    edges = []
    rewards = {}
    for index, edge in enumerate(graph.edges):
        destination = graph.nodes[graph.node_index[edge.target]]
        if destination.outcome is None:
            edges.append(edge)
            continue
        outcome = f"copy_{index}"
        node_id = f"terminal_copy_{index}"
        nodes.append(Node(node_id, outcome, destination.features))
        edges.append(replace(edge, target=node_id))
        rewards[outcome] = log_rewards[destination.outcome]
    altered = Graph(nodes, edges, graph.root, graph.metadata)
    return solve(altered, rewards, temperature).forward