#
#   Imandra Inc.
#
#   speclogician/state/state_stats.py
#

from __future__ import annotations

from dataclasses import dataclass

from .inst import StateInstance


@dataclass(frozen=True)
class StatsCalculator:
    """
    Compute time-series stats over a collection of State objects.

    Assumptions:
    - Each State has `instances` ordered newest-first (index 0 = latest).
    - We return lists oldest-first for plotting.
    """

    @staticmethod
    def stats(instances: list[StateInstance]) -> dict[str, list[int | str]]:
        """ 
        Note: we reverse the order of StateInstances here so the 0s element (latest state instance) is all the way on the right
        """
        out: dict[str, list[int | str]] = {
            # identity / time
            "created_at": [],
            "num_changes": [],

            # ---- domain model counts (from DomainModel) ----
            "num_state_preds_total": [],
            "num_state_preds_valid_logic": [],
            "num_state_preds_matched": [],

            "num_action_preds_total": [],
            "num_action_preds_valid_logic": [],
            "num_action_preds_matched": [],

            "num_preds_total": [],
            "num_preds_valid_logic": [],
            "num_preds_matched": [],

            "num_trans_total": [],
            "num_trans_valid_logic": [],
            "num_trans_matched": [],

            # ---- scenario stats (from Spec) ----
            "num_sc_total": [],
            "num_sc_missing": [],
            "num_sc_matched": [],
            "num_sc_inconsistent": [],
            "num_sc_conflicted": [],

            # ---- artifact stats (from ArtifactContainer) ----
            "num_test_traces_total": [],
            "num_test_traces_matched": [],
            "num_test_traces_logic_good": [],

            "num_log_traces_total": [],
            "num_log_traces_matched": [],
            "num_log_traces_logic_good": [],

            "num_doc_arts_total": [],
            "num_doc_arts_matched": [],
            "num_src_code_arts_total": [],
            "num_src_code_arts_matched": [],

            # ---- optional derived metrics (handy for “progress”) ----
            "num_traces_total": [],
            "num_traces_logic_good": [],
            "num_traces_matched": [],
        }

        for inst in instances[::-1]:
            # identity / time
            out["created_at"].append(inst.created_at.isoformat())
            out["num_changes"].append(len(inst.changes or []))

            # domain model
            dm = inst.spec.domain_model
            out["num_state_preds_total"].append(dm.num_state_preds_total)
            out["num_state_preds_valid_logic"].append(dm.num_state_preds_valid_logic)
            out["num_state_preds_matched"].append(dm.num_state_preds_matched)

            out["num_action_preds_total"].append(dm.num_action_preds_total)
            out["num_action_preds_valid_logic"].append(dm.num_action_preds_valid_logic)
            out["num_action_preds_matched"].append(dm.num_action_preds_matched)

            out["num_preds_total"].append(dm.num_preds_total)
            out["num_preds_valid_logic"].append(dm.num_preds_valid_logic)
            out["num_preds_matched"].append(dm.num_preds_matched)

            out["num_trans_total"].append(dm.num_trans_total)
            out["num_trans_valid_logic"].append(dm.num_trans_valid_logic)
            out["num_trans_matched"].append(dm.num_trans_matched)

            # spec / scenarios
            sp = inst.spec
            out["num_sc_total"].append(sp.num_sc_total)
            out["num_sc_missing"].append(sp.num_sc_missing)
            out["num_sc_matched"].append(sp.num_sc_matched)
            out["num_sc_inconsistent"].append(sp.num_sc_inconsistent)
            out["num_sc_conflicted"].append(sp.num_sc_conflicted)

            # artifacts
            ac = inst.art_container
            out["num_test_traces_total"].append(ac.num_test_traces_total)
            out["num_test_traces_matched"].append(ac.num_test_traces_matched)
            out["num_test_traces_logic_good"].append(ac.num_test_traces_logic_good)

            out["num_log_traces_total"].append(ac.num_log_traces_total)
            out["num_log_traces_matched"].append(ac.num_log_traces_matched)
            out["num_log_traces_logic_good"].append(ac.num_log_traces_logic_good)

            out["num_doc_arts_total"].append(ac.num_doc_arts_total)
            out["num_doc_arts_matched"].append(ac.num_doc_arts_matched)
            out["num_src_code_arts_total"].append(ac.num_src_code_arts_total)
            out["num_src_code_arts_matched"].append(ac.num_src_code_arts_matched)

            # derived
            traces_total = ac.num_test_traces_total + ac.num_log_traces_total
            traces_good = ac.num_test_traces_logic_good + ac.num_log_traces_logic_good
            traces_matched = ac.num_test_traces_matched + ac.num_log_traces_matched

            out["num_traces_total"].append(traces_total)
            out["num_traces_logic_good"].append(traces_good)
            out["num_traces_matched"].append(traces_matched)

        return out