#
#   Imandra Inc.
#
#   speclogician/modeling/trace_reachability_compute.py
#

from __future__ import annotations

from collections import Counter
from math import inf
from typing import Iterable
from dataclasses import dataclass

from speclogician.data.container import ArtifactContainer
from speclogician.data.mapping import ArtifactMap
from speclogician.modeling.reachability import ScenarioReachabilityGraph
from speclogician.modeling.spec import Spec
from speclogician.modeling.trace_reachability import (
    ObservedScenarioHop,
    TraceScenarioAlignment,
    TraceReachabilityAnalysis,
    TraceReachabilitySummary,
)

@dataclass(frozen=True)
class _DPParams:
    reward_edge: float = 2.0
    penalty_no_edge: float = -4.0
    penalty_unmatched_hop: float = -3.0
    penalty_choose_none: float = -0.2


def _dp_choose_step_best(
    *,
    step_matches: list[list[str]],
    reach: ScenarioReachabilityGraph,
    params: _DPParams = _DPParams(),
    allow_none: bool = True,
) -> list[str | None]:
    """
    Choose one label per step to maximize supported hops across the trace.
    Deterministic tie-break: among equal-score solutions, choose the
    lexicographically smallest *entire* sequence.

    We implement tie-break by tracking a per-state "path key" for the best
    solution ending in that state.
    """
    n = len(step_matches)
    if n == 0:
        return []

    # Deterministic state ordering per step
    states: list[list[str | None]] = []
    for cands in step_matches:
        uniq = sorted(set(cands))
        states.append(uniq + ([None] if allow_none else []))

    for i in range(n):
        if not states[i]:
            states[i] = [None]

    def _node_penalty(x: str | None) -> float:
        return params.penalty_choose_none if x is None else 0.0

    def _trans_score(a: str | None, b: str | None) -> float:
        if a is None or b is None:
            return params.penalty_unmatched_hop
        return params.reward_edge if (b in reach.successors(a)) else params.penalty_no_edge

    def _tok(x: str | None) -> tuple[int, str]:
        # None sorts *after* strings
        return (1, "") if x is None else (0, x)

    # dp_score[i][j] = best score up to step i ending in states[i][j]
    # dp_key[i][j]   = lexicographic key for the chosen sequence up to i
    # back[i][j]     = best predecessor index
    dp_score: list[list[float]] = []
    dp_key: list[list[tuple[tuple[int, str], ...]]] = []
    back: list[list[int]] = []

    # init (step 0)
    s0 = []
    k0 = []
    b0 = []
    for x in states[0]:
        s0.append(_node_penalty(x))
        k0.append((_tok(x),))
        b0.append(-1)
    dp_score.append(s0)
    dp_key.append(k0)
    back.append(b0)

    # transitions
    for i in range(1, n):
        cur_scores: list[float] = []
        cur_keys: list[tuple[tuple[int, str], ...]] = []
        cur_back: list[int] = []

        prev_states = states[i - 1]
        cur_states = states[i]

        for j, cur in enumerate(cur_states):
            best_score = -inf
            best_key: tuple[tuple[int, str], ...] | None = None
            best_k = 0

            for k, prev in enumerate(prev_states):
                score = dp_score[i - 1][k] + _trans_score(prev, cur) + _node_penalty(cur)
                key = dp_key[i - 1][k] + (_tok(cur),)

                if score > best_score:
                    best_score = score
                    best_key = key
                    best_k = k
                elif score == best_score:
                    assert best_key is not None
                    # global lexicographic tie-break on the full sequence key
                    if key < best_key:
                        best_key = key
                        best_k = k

            cur_scores.append(best_score)
            cur_keys.append(best_key if best_key is not None else ())
            cur_back.append(best_k)

        dp_score.append(cur_scores)
        dp_key.append(cur_keys)
        back.append(cur_back)

    # pick best final state by (score, key)
    last_scores = dp_score[-1]
    last_keys = dp_key[-1]

    best_j = 0
    for j in range(1, len(states[-1])):
        if last_scores[j] > last_scores[best_j]:
            best_j = j
        elif last_scores[j] == last_scores[best_j]:
            if last_keys[j] < last_keys[best_j]:
                best_j = j

    # backtrack
    out: list[str | None] = [None] * n
    j = best_j
    for i in range(n - 1, -1, -1):
        out[i] = states[i][j]
        j = back[i][j]
        if j < 0:
            break

    return out

def compute_trace_reachability_analysis(
    *,
    spec: Spec,
    art_cont: ArtifactContainer,
    art_map: ArtifactMap,
    reach: ScenarioReachabilityGraph,
    include_traces: Iterable[str] | None = None,
    cap_step_matches: int | None = None,
    include_shortest_path: bool = True,
) -> TraceReachabilityAnalysis:
    """
    Compute TraceReachabilityAnalysis based on:
      - reachability graph between scenarios (reach)
      - step-level mapping in art_map using trace.step_ref(i)
 synchronously, purely in-memory.

    Assumptions:
      - Scenario IDs used in the graph are scenario.comp_id.
      - Step-level mapping uses trace.step_ref(i) -> [scenario.comp_id, ...]
    """

    scenario_ids = [sc.comp_id for sc in spec.scenarios]
    scenario_set = set(scenario_ids)

    # Choose traces
    traces = art_cont.all_traces()
    if include_traces is not None:
        want = set(include_traces)
        traces = [t for t in traces if t.art_id in want]

    per_trace: dict[str, TraceScenarioAlignment] = {}

    # Summary accumulators
    total_steps = 0
    total_hops = 0
    supported_hops = 0
    unsupported_hops = 0

    scenarios_hit: set[str] = set()
    unsupported_edge_counts: Counter[tuple[str, str]] = Counter()
    frontier_end_counts: Counter[str] = Counter()

    for tr in traces:
        if tr.art_id is None:
            continue

        n = len(tr.steps)
        total_steps += n
        total_hops += max(0, n - 1)

        step_matches: list[list[str]] = []
        step_best: list[str | None] = []

        # ---- step matches ----
        step_matches: list[list[str]] = []

        for i in range(n):
            ref = tr.step_ref(i)
            cands = list(art_map.get_comp_for_artifact(ref))
            cands = [c for c in cands if c in scenario_set]   # filter stale
            if cap_step_matches is not None:
                cands = cands[:cap_step_matches]
            step_matches.append(cands)
            for c in cands:
                scenarios_hit.add(c)

        # NEW: choose globally-optimal step_best
        step_best = _dp_choose_step_best(step_matches=step_matches, reach=reach)

        # ---- hops ----
        observed_hops: list[ObservedScenarioHop] = []
        supported_local = 0
        unsupported_local = 0

        for i in range(n - 1):
            src = step_best[i]
            dst = step_best[i + 1]

            hop = ObservedScenarioHop(
                trace_id=tr.art_id,
                step_from=i,
                step_to=i + 1,
                src=src or "",
                dst=dst or "",
                supported_by_graph=False,
                shortest_path=None,
                reason=None,
            )

            if src is None or dst is None:
                hop.supported_by_graph = False
                hop.reason = "unmatched_step"
                unsupported_local += 1
            else:
                if dst in reach.successors(src):
                    hop.supported_by_graph = True
                    supported_local += 1
                else:
                    hop.supported_by_graph = False
                    hop.reason = "no_graph_edge"
                    unsupported_local += 1
                    unsupported_edge_counts[(src, dst)] += 1

                    if include_shortest_path:
                        path = reach.shortest_path(src, dst)
                        hop.shortest_path = path
                        if path is None:
                            hop.reason = "no_path"

            observed_hops.append(hop)

        # “frontier ending” = last matched label (useful for UI triage)
        for j in range(n - 1, -1, -1):
            if step_best[j] is not None:
                frontier_end_counts[step_best[j]] += 1
                break

        supported_hops += supported_local
        unsupported_hops += unsupported_local

        per_trace[tr.art_id] = TraceScenarioAlignment(
            trace_id=tr.art_id,
            step_matches=step_matches,
            step_best=step_best,
            observed_hops=observed_hops,
            num_steps=n,
            num_matched_steps=sum(1 for x in step_best if x is not None),
            num_supported_hops=supported_local,
            num_unsupported_hops=unsupported_local,
        )

    scenarios_unhit = sorted(list(scenario_set - scenarios_hit))
    scenarios_hit_list = sorted(list(scenarios_hit))

    top_unsupported_edges = [
        (src, dst, cnt) for (src, dst), cnt in unsupported_edge_counts.most_common(25)
    ]
    top_frontier_trace_endings = [
        (sc, cnt) for sc, cnt in frontier_end_counts.most_common(25)
    ]

    summary = TraceReachabilitySummary(
        total_traces=len(per_trace),
        total_steps=total_steps,
        total_hops=total_hops,
        supported_hops=supported_hops,
        unsupported_hops=unsupported_hops,
        scenarios_hit=scenarios_hit_list,
        scenarios_unhit=scenarios_unhit,
        top_unsupported_edges=top_unsupported_edges,
        top_frontier_trace_endings=top_frontier_trace_endings,
    )

    return TraceReachabilityAnalysis(
        per_trace=per_trace,
        summary=summary,
        # fingerprints: fill if you have stable hashing utilities
        spec_fp=None,
        artifacts_fp=None,
        mapping_fp=None,
    )