import importlib.util
from pathlib import Path
import pandas as pd

def _has_pkg(name):
    return importlib.util.find_spec(name) is not None


class SUNTVisualizer:
    def __init__(self, df):
        self.df = df

    def build_od_graph(
        self,
        origin_col="stop_id",
        trip_col="trip_id",
        sequence_col="pt_sequence",
        time_col="stop_time",
        weight_col="n-boardings",
    ):
        if not _has_pkg("networkx"):
            raise ImportError("NetworkX not found. Install with: pip install networkx")

        import networkx as nx

        required_cols = {origin_col, trip_col}
        if sequence_col:
            required_cols.add(sequence_col)
        if time_col:
            required_cols.add(time_col)
        if weight_col:
            required_cols.add(weight_col)

        missing = [col for col in required_cols if col not in self.df.columns]
        if missing:
            raise ValueError(f"Columns not found for generating OD graph: {missing}")

        df = self.df.copy()
        if time_col and time_col in df.columns:
            df[time_col] = pd.to_datetime(df[time_col], errors="coerce")

        sort_cols = [trip_col]
        if sequence_col and sequence_col in df.columns:
            sort_cols.append(sequence_col)
        elif time_col and time_col in df.columns:
            sort_cols.append(time_col)
        else:
            raise ValueError("Info sequence_col or time_col invalid for order.")

        df = df.sort_values(sort_cols)
        df["next_stop_id"] = df.groupby(trip_col)[origin_col].shift(-1)
        edges_df = df.dropna(subset=[origin_col, "next_stop_id"])
        edges_df = edges_df[edges_df[origin_col] != edges_df["next_stop_id"]]

        if edges_df.empty:
            return nx.DiGraph()

        if weight_col and weight_col in edges_df.columns:
            edges_df[weight_col] = pd.to_numeric(edges_df[weight_col], errors="coerce").fillna(0)
            grouped = (
                edges_df.groupby([origin_col, "next_stop_id"], as_index=False)
                .agg(
                    trips=(trip_col, "nunique"),
                    passengers=(weight_col, "sum"),
                )
            )
        else:
            grouped = (
                edges_df.groupby([origin_col, "next_stop_id"], as_index=False)
                .agg(trips=(trip_col, "nunique"))
            )
            grouped["passengers"] = grouped["trips"]

        graph = nx.DiGraph()
        for row in grouped.itertuples(index=False):
            origin = str(getattr(row, origin_col))
            destination = str(row.next_stop_id)
            trips = float(row.trips)
            passengers = float(row.passengers)
            graph.add_edge(
                origin,
                destination,
                trips=trips,
                passengers=passengers,
                weight=passengers,
            )

        return graph

    def export_od_graph(self, output_path, fmt="graphml", **build_kwargs):
        if not _has_pkg("networkx"):
            raise ImportError("NetworkX not found. Install with: pip install networkx")

        import networkx as nx

        graph = self.build_od_graph(**build_kwargs)
        output = Path(output_path)
        output.parent.mkdir(parents=True, exist_ok=True)

        fmt_normalized = fmt.strip().lower()
        if fmt_normalized == "graphml":
            nx.write_graphml(graph, output)
        elif fmt_normalized == "gexf":
            nx.write_gexf(graph, output)
        else:
            raise ValueError("Invalid format. Use 'graphml' or 'gexf'.")

        return output