import gradio as gr
import pandas as pd
import networkx as nx
import plotly.graph_objects as go
from huggingface_hub import hf_hub_download
import numpy as np

REPO_ID = "overthelex/ua-court-citation-graph"

def load_parquet(filename: str) -> pd.DataFrame:
    path = hf_hub_download(repo_id=REPO_ID, filename=filename, repo_type="dataset")
    return pd.read_parquet(path)

print("Loading dataset files...")
edges_df = load_parquet("cocitation_edges.parquet")
centrality_df = load_parquet("article_centrality.parquet")
stats_df = load_parquet("article_citation_stats.parquet")
communities_df = load_parquet("community_labels.parquet")

edges_df["node_a"] = edges_df["law_a"] + " ст. " + edges_df["article_a"]
edges_df["node_b"] = edges_df["law_b"] + " ст. " + edges_df["article_b"]

ALL_LAWS = sorted(edges_df["law_a"].unique().tolist())

DOMAIN_COLORS = {
    "Кримінальний кодекс України": "#ef4444",
    "Кримінальний процесуальний кодекс України": "#f97316",
    "Цивільний кодекс України": "#3b82f6",
    "Цивільний процесуальний кодекс України": "#60a5fa",
    "Господарський кодекс України": "#22c55e",
    "Господарський процесуальний кодекс України": "#4ade80",
    "Кодекс адміністративного судочинства України": "#eab308",
    "КУпАП": "#f59e0b",
    "Податковий кодекс України": "#a855f7",
    "Сімейний кодекс України": "#14b8a6",
    "Земельний кодекс України": "#f97316",
    "Конституція України": "#fbbf24",
}
DEFAULT_COLOR = "#94a3b8"

print(f"Loaded {len(edges_df):,} edges, {len(centrality_df)} centrality nodes, {len(stats_df):,} citation stats")


def get_color(law_name: str) -> str:
    for key, color in DOMAIN_COLORS.items():
        if key in law_name:
            return color
    return DEFAULT_COLOR


SHORT_NAMES = [
    ("Кримінальний процесуальний кодекс України", "КПК"),
    ("Кримінальний кодекс України", "ККУ"),
    ("Цивільний процесуальний кодекс України", "ЦПК"),
    ("Цивільний кодекс України", "ЦКУ"),
    ("Господарський процесуальний кодекс України", "ГПК"),
    ("Господарський кодекс України", "ГКУ"),
    ("Кодекс адміністративного судочинства України", "КАСУ"),
    ("Податковий кодекс України", "ПКУ"),
    ("Сімейний кодекс України", "СКУ"),
    ("Земельний кодекс України", "ЗКУ"),
    ("Конституція України", "КУ"),
]


def shorten(name: str) -> str:
    for long, short in SHORT_NAMES:
        if name == long:
            return short
    if len(name) > 20:
        return name[:18] + "..."
    return name


def build_graph(law_filter: list[str], min_weight: int, max_edges: int, layout: str):
    filtered = edges_df[edges_df["weight"] >= min_weight]

    if law_filter:
        mask = filtered["law_a"].isin(law_filter) | filtered["law_b"].isin(law_filter)
        filtered = filtered[mask]

    filtered = filtered.nlargest(max_edges, "weight")

    if filtered.empty:
        fig = go.Figure()
        fig.update_layout(
            paper_bgcolor="#0f172a", plot_bgcolor="#0f172a",
            annotations=[dict(text="Немає ребер. Зменшіть мінімальну вагу.",
                            showarrow=False, font=dict(size=18, color="#94a3b8"),
                            xref="paper", yref="paper", x=0.5, y=0.5)],
            xaxis=dict(visible=False), yaxis=dict(visible=False),
            height=750,
        )
        return fig

    G = nx.Graph()
    for _, row in filtered.iterrows():
        G.add_edge(row["node_a"], row["node_b"], weight=int(row["weight"]))

    degree_dict = dict(G.degree(weight="weight"))
    pr = nx.pagerank(G, weight="weight")

    if layout == "Kamada-Kawai":
        pos = nx.kamada_kawai_layout(G, weight="weight")
    else:
        pos = nx.spring_layout(G, k=2.5 / np.sqrt(G.number_of_nodes()),
                                iterations=100, weight="weight", seed=42)

    max_w = filtered["weight"].max()
    min_w = filtered["weight"].min()
    w_range = max_w - min_w if max_w > min_w else 1

    edge_x, edge_y, edge_hover = [], [], []
    for u, v, d in G.edges(data=True):
        x0, y0 = pos[u]
        x1, y1 = pos[v]
        edge_x.extend([x0, x1, None])
        edge_y.extend([y0, y1, None])

    edge_trace = go.Scatter(
        x=edge_x, y=edge_y,
        mode="lines",
        line=dict(width=0.5, color="rgba(148,163,184,0.2)"),
        hoverinfo="none",
        showlegend=False,
    )

    thick_edge_traces = []
    sorted_edges = sorted(G.edges(data=True), key=lambda e: e[2]["weight"], reverse=True)
    top_n = min(50, len(sorted_edges))
    for u, v, d in sorted_edges[:top_n]:
        x0, y0 = pos[u]
        x1, y1 = pos[v]
        w = d["weight"]
        width = 1 + 5 * ((w - min_w) / w_range)
        opacity = 0.3 + 0.5 * ((w - min_w) / w_range)
        thick_edge_traces.append(go.Scatter(
            x=[x0, x1, None], y=[y0, y1, None],
            mode="lines",
            line=dict(width=width, color=f"rgba(96,165,250,{opacity:.2f})"),
            hoverinfo="text",
            hovertext=f"{u} ↔ {v}<br>Вага: {w:,} рішень",
            showlegend=False,
        ))

    max_pr = max(pr.values()) if pr else 1
    node_groups: dict[str, dict] = {}
    for node in G.nodes():
        parts = node.split(" ст. ", 1)
        law = parts[0] if len(parts) == 2 else "Інше"
        article = parts[1] if len(parts) == 2 else node
        color = get_color(law)

        if law not in node_groups:
            node_groups[law] = {"x": [], "y": [], "text": [], "hover": [],
                                "size": [], "color": color}

        x, y = pos[node]
        size = 8 + 40 * (pr.get(node, 0) / max_pr)
        hover = (
            f"<b>{node}</b><br>"
            f"Зважений ступінь: {degree_dict.get(node, 0):,}<br>"
            f"PageRank: {pr.get(node, 0):.6f}<br>"
            f"Зв'язків: {G.degree(node)}"
        )
        node_groups[law]["x"].append(x)
        node_groups[law]["y"].append(y)
        node_groups[law]["text"].append(f"ст. {article}")
        node_groups[law]["hover"].append(hover)
        node_groups[law]["size"].append(size)

    node_traces = []
    for group_name in sorted(node_groups.keys()):
        data = node_groups[group_name]
        node_traces.append(go.Scatter(
            x=data["x"], y=data["y"],
            mode="markers+text",
            marker=dict(
                size=data["size"],
                color=data["color"],
                line=dict(width=1.5, color="rgba(255,255,255,0.4)"),
            ),
            text=data["text"],
            textposition="top center",
            textfont=dict(size=9, color="#cbd5e1"),
            hoverinfo="text",
            hovertext=data["hover"],
            name=shorten(group_name),
        ))

    fig = go.Figure(data=[edge_trace] + thick_edge_traces + node_traces)

    fig.update_layout(
        paper_bgcolor="#0f172a",
        plot_bgcolor="#0f172a",
        height=750,
        margin=dict(l=5, r=5, t=45, b=5),
        title=dict(
            text=(f"Вузлів: {G.number_of_nodes()} | "
                  f"Ребер: {G.number_of_edges()} | "
                  f"Мін. вага: {filtered['weight'].min():,} | "
                  f"Макс. вага: {filtered['weight'].max():,}"),
            font=dict(size=13, color="#94a3b8"),
            x=0.5,
        ),
        xaxis=dict(visible=False, showgrid=False, zeroline=False),
        yaxis=dict(visible=False, showgrid=False, zeroline=False),
        legend=dict(
            font=dict(color="#e2e8f0", size=11),
            bgcolor="rgba(15,23,42,0.9)",
            bordercolor="rgba(71,85,105,0.5)",
            borderwidth=1,
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="center",
            x=0.5,
        ),
        hoverlabel=dict(
            bgcolor="#1e293b",
            font_size=13,
            font_color="#e2e8f0",
            bordercolor="#475569",
        ),
        dragmode="pan",
    )

    fig.update_layout(
        modebar=dict(
            bgcolor="rgba(15,23,42,0.8)",
            color="#94a3b8",
            activecolor="#60a5fa",
        )
    )

    return fig


def get_top_articles(n: int = 30):
    top = stats_df.nlargest(n, "total_citations")[
        ["law_number", "law_article", "citation_type", "total_citations", "unique_decisions"]
    ].copy()
    top.columns = ["Закон", "Стаття", "Тип", "Цитувань", "Унікальних рішень"]
    top["Цитувань"] = top["Цитувань"].apply(lambda x: f"{x:,}")
    top["Унікальних рішень"] = top["Унікальних рішень"].apply(lambda x: f"{x:,}")
    return top


def get_top_centrality():
    top = centrality_df.nlargest(20, "pagerank")[
        ["law_number", "law_article", "degree", "pagerank", "authority"]
    ].copy()
    top.columns = ["Закон", "Стаття", "Ступінь", "PageRank", "Authority"]
    top["Ступінь"] = top["Ступінь"].apply(lambda x: f"{x:,}")
    top["PageRank"] = top["PageRank"].apply(lambda x: f"{x:.6f}")
    top["Authority"] = top["Authority"].apply(lambda x: f"{x:.6f}")
    return top


def get_communities():
    df = communities_df.copy()
    df.columns = ["Період", "Ранг", "Розмір", "Домінуючий закон", "Частка", "Цитувань"]
    df["Частка"] = df["Частка"].apply(lambda x: f"{x:.1%}")
    df["Цитувань"] = df["Цитувань"].apply(lambda x: f"{x:,}")
    return df


with gr.Blocks(
    title="Ukrainian Court Citation Graph",
    theme=gr.themes.Base(
        primary_hue="blue",
        neutral_hue="slate",
    ),
    css="footer { display: none !important; }",
) as demo:
    gr.Markdown(
        """
        # Ukrainian Court Citation Graph Explorer
        Інтерактивна візуалізація мережі со-цитувань з **99.5 млн** судових рішень ЄДРСР (2003-2026).
        Два нормативні положення з'єднані, якщо вони цитуються разом у одному рішенні.
        Вага ребра -- кількість рішень, що со-цитують пару.
        """
    )

    with gr.Tab("Граф со-цитувань"):
        with gr.Row():
            with gr.Column(scale=1):
                law_filter = gr.Dropdown(
                    choices=ALL_LAWS,
                    multiselect=True,
                    label="Фільтр за законом",
                    info="Залиште порожнім для всіх",
                )
                min_weight = gr.Slider(
                    minimum=10, maximum=500000, value=50000, step=1000,
                    label="Мінімальна вага ребра",
                    info="Кількість рішень, що со-цитують пару",
                )
                max_edges = gr.Slider(
                    minimum=50, maximum=2000, value=300, step=50,
                    label="Макс. ребер",
                    info="Топ-N найважчих ребер",
                )
                layout = gr.Radio(
                    choices=["Spring", "Kamada-Kawai"],
                    value="Spring",
                    label="Алгоритм розташування",
                )
                btn = gr.Button("Побудувати граф", variant="primary", size="lg")

            with gr.Column(scale=3):
                graph_output = gr.Plot()

        btn.click(
            fn=build_graph,
            inputs=[law_filter, min_weight, max_edges, layout],
            outputs=graph_output,
        )

    with gr.Tab("Топ статей за цитуваннями"):
        gr.Markdown("### Найбільш цитовані нормативні положення")
        gr.Dataframe(value=get_top_articles(30), interactive=False)

    with gr.Tab("Centrality (PageRank, HITS)"):
        gr.Markdown("### Топ-20 за PageRank у графі со-цитувань")
        gr.Dataframe(value=get_top_centrality(), interactive=False)

    with gr.Tab("Спільноти (Louvain)"):
        gr.Markdown("### Виявлені спільноти за періодами (Louvain community detection)")
        gr.Dataframe(value=get_communities(), interactive=False)

    gr.Markdown(
        """
        ---
        **Джерело**: [overthelex/ua-court-citation-graph](https://huggingface.co/datasets/overthelex/ua-court-citation-graph)
        | **Платформа**: [legal.org.ua](https://legal.org.ua)
        | **Автор**: Volodymyr Ovcharov, LEX AI LLC
        """
    )

demo.launch()