from __future__ import annotations

import networkx as nx
import pandas as pd

from outhad_edge.edgelist import Edgelist
from outhad_edge.eventstream.types import EventstreamType
from outhad_edge.nodelist import Nodelist
from outhad_edge.tooling.typing.transition_graph import NormType


class _TransitionMatrix:
    def __init__(self, eventstream: EventstreamType) -> None:
        self.__eventstream = eventstream
        self.__nodelist = Nodelist(
            weight_cols=[eventstream.schema.event_name, *eventstream.schema.custom_cols],
            time_col=eventstream.schema.event_timestamp,
            event_col=eventstream.schema.event_name,
        )
        self.__nodelist.calculate_nodelist(self.__eventstream.to_dataframe())
        self.__edgelist = Edgelist(eventstream=eventstream)

    def _values(self, weight_col: str | None = None, norm_type: NormType = None) -> pd.DataFrame:
        weight_col = weight_col if weight_col else self.__eventstream.schema.user_id
        self.__edgelist.calculate_edgelist(norm_type=norm_type, weight_cols=[weight_col])
        edgelist: pd.DataFrame = self.__edgelist.edgelist_df
        graph = nx.DiGraph()
        graph.add_weighted_edges_from(edgelist.values)

        return nx.to_pandas_adjacency(G=graph)


__all__ = ("_TransitionMatrix",)
