#
#   Imandra Inc.
#
#   speclogician/tui/stats.py
#

from typing import Any
from textual.screen import Screen
from textual.app import ComposeResult
from textual.widgets import Header, Footer
from textual.containers import Container
from textual_plotext import PlotextPlot

from ..state.state_stats import StatsCalculator


class StatsScreen(Screen[Any]):
    CSS = """
    #body { height: 1fr; }

    /* 2x2 grid */
    #grid {
        height: 1fr;
        layout: grid;
        grid-size: 2 2;
        grid-gutter: 1 1;
    }

    PlotextPlot {
        height: 1fr;
        width: 1fr;
    }
    """

    def compose(self) -> ComposeResult:
        yield Header()
        with Container(id="body"):
            with Container(id="grid"):
                self.plot_model = PlotextPlot()
                self.plot_scenarios = PlotextPlot()
                self.plot_artifacts = PlotextPlot()
                self.plot_ratios = PlotextPlot()

                yield self.plot_model
                yield self.plot_scenarios
                yield self.plot_artifacts
                yield self.plot_ratios
        yield Footer()

    def on_mount(self) -> None:
        # If you want live updating while running:
        # self.set_interval(0.5, self.refresh_plot)
        self.refresh_plot()

    @staticmethod
    def _safe_ratio_series(num: list[int], den: list[int]) -> list[float]:
        out: list[float] = []
        for n, d in zip(num, den):
            out.append(n / (d if d else 1))
        return out

    def refresh_plot(self) -> None:

        v = StatsCalculator.stats(self.app.state.instances)

        # oldest -> newest already
        xs = list(range(len(v["created_at"])))

        # -------------------------
        # 1) Model growth & usage
        # -------------------------
        plt = self.plot_model.plt
        plt.clear_data()
        plt.title("Model growth & usage")
        plt.xlabel("Instance (oldest → newest)")
        plt.ylabel("count")
        plt.plot(xs, v["num_preds_total"])
        plt.plot(xs, v["num_preds_matched"])
        plt.plot(xs, v["num_trans_total"])
        plt.scatter(xs, v["num_preds_matched"])
        self.plot_model.refresh()

        # -------------------------
        # 2) Scenario health
        # -------------------------
        plt = self.plot_scenarios.plt
        plt.clear_data()
        plt.title("Scenario health")
        plt.xlabel("Instance (oldest → newest)")
        plt.ylabel("count")
        plt.plot(xs, v["num_sc_total"])
        plt.plot(xs, v["num_sc_matched"])
        plt.plot(xs, v["num_sc_missing"])
        plt.plot(xs, v["num_sc_inconsistent"])
        plt.plot(xs, v["num_sc_conflicted"])
        self.plot_scenarios.refresh()

        # -------------------------
        # 3) Artifact pipeline quality
        # -------------------------
        plt = self.plot_artifacts.plt
        plt.clear_data()
        plt.title("Artifacts: total → logic-good → matched")
        plt.xlabel("Instance (oldest → newest)")
        plt.ylabel("count")
        plt.plot(xs, v["num_traces_total"])
        plt.plot(xs, v["num_traces_logic_good"])
        plt.plot(xs, v["num_traces_matched"])
        plt.scatter(xs, v["num_traces_matched"])
        self.plot_artifacts.refresh()

        # -------------------------
        # 4) Coverage ratios
        # -------------------------
        plt = self.plot_ratios.plt
        plt.clear_data()
        plt.title("Coverage ratios")
        plt.xlabel("Instance (oldest → newest)")
        plt.ylabel("ratio (0..1)")
        preds_ratio = self._safe_ratio_series(v["num_preds_matched"], v["num_preds_total"])
        traces_ratio = self._safe_ratio_series(v["num_traces_matched"], v["num_traces_total"])
        plt.plot(xs, preds_ratio)
        plt.plot(xs, traces_ratio)
        plt.scatter(xs, preds_ratio)
        self.plot_ratios.refresh()