import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
import marimo as mo
import numpy as np
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
import os

import torch
from gl.datasets import graph, summary
from gl.nb import json_cache
from torch_geometric.data import Data, HeteroData
from torch_geometric.utils import (
    from_networkx,
    to_undirected,
    is_undirected,
)
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

SEED = 0
rng = np.random.default_rng(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)  # several open notebooks/kernels share the CPU; more threads = contention
CACHE_DIR = mo.notebook_dir() / "cache"
CACHE_DIR.mkdir(exist_ok=True)
plt.rcParams.update(
    {
        "figure.dpi": 100,
        "axes.spines.top": False,
        "axes.spines.right": False,
    }
)

karate = nx.karate_club_graph()
club = np.array(
    [karate.nodes[v]["club"] == "Officer" for v in karate.nodes], dtype=int
)  # 0 = Mr. Hi, 1 = Officer
degrees = np.array([d for _, d in karate.degree()])

{
    "nodes N": karate.number_of_nodes(),
    "edges |E|": karate.number_of_edges(),
    "directed?": karate.is_directed(),
    "average degree 2|E|/N": round(
        2 * karate.number_of_edges() / karate.number_of_nodes(), 2
    ),
    "max degree (node)": (int(degrees.max()), int(degrees.argmax())),
    "density |E| / (N choose 2)": round(nx.density(karate), 3),
    "Officer members": int(club.sum()),
}

karate_pos = nx.spring_layout(karate, seed=SEED)

def draw_karate(
    highlight=None,
    ax=None,
    title="Karate Club — colour = faction after the split",
):
    ax = ax or plt.subplots(figsize=(6, 4))[1]
    colors = ["#4C78A8" if c == 0 else "#F58518" for c in club]
    nx.draw_networkx_edges(karate, karate_pos, ax=ax, alpha=0.3)
    nx.draw_networkx_nodes(
        karate,
        karate_pos,
        ax=ax,
        node_color=colors,
        node_size=[60 + 25 * d for d in degrees],
        edgecolors="k",
        linewidths=0.5,
    )
    if highlight is not None:
        nx.draw_networkx_nodes(
            karate,
            karate_pos,
            nodelist=list(highlight),
            ax=ax,
            node_color="none",
            edgecolors="crimson",
            linewidths=2.5,
            node_size=500,
        )
    nx.draw_networkx_labels(karate, karate_pos, ax=ax, font_size=7)
    ax.set_title(title, fontsize=10)
    ax.axis("off")
    return ax

draw_karate()

_edges = torch.tensor(
    list(karate.edges()), dtype=torch.long
).t()  # 2 x 78, one direction only
karate_edge_index = to_undirected(_edges)  # 2 x 156, both directions
karate_data = Data(
    x=torch.tensor(degrees, dtype=torch.float).unsqueeze(
        1
    ),  # one feature per node: its degree
    edge_index=karate_edge_index,
    y=torch.tensor(club),
)
print(karate_data)
print("first 8 columns of edge_index:\n", karate_edge_index[:, :8])
print(
    "undirected?",
    is_undirected(karate_edge_index),
    "| stored edges:",
    karate_edge_index.size(1),
)

_v = node_picker.value
_from_nx = sorted(karate.neighbors(_v))
_from_ei = sorted(
    karate_edge_index[1, karate_edge_index[0] == _v].tolist()
)  # targets of edges leaving v
_A_row = np.flatnonzero(
    nx.to_numpy_array(karate, nodelist=range(34), weight=None)[_v]
).tolist()
mo.hstack(
    [
        draw_karate(
            highlight=_from_nx,
            title=f"neighbours of node {_v} (red rings)",
        ),
        mo.md(f"""
| source | neighbours of {_v} |
|---|---|
| `karate.neighbors(v)` | {_from_nx} |
| `edge_index[1, edge_index[0]==v]` | {_from_ei} |
| nonzeros of row $A_{{v,:}}$ | {_A_row} |

degree = **{len(_from_nx)}**, all three agree: **{_from_nx == _from_ei == _A_row}**
"""),
    ],
    widths=[3, 2],
    align="center",
)

A_karate = torch.tensor(
    nx.to_numpy_array(karate, nodelist=range(34), weight=None),
    dtype=torch.float,
)
x_degree = A_karate.sum(1, keepdim=True)  # node feature = degree

torch.manual_seed(SEED)
flat_mlp = torch.nn.Sequential(
    torch.nn.Linear(34 * 34, 64), torch.nn.ReLU(), torch.nn.Linear(64, 34)
)

def mlp_on_adjacency(A, x):
    return flat_mlp(A.reshape(1, -1)).squeeze(
        0
    )  # one output per *position* 0..33

def neighbour_sum(A, x):
    return (A @ x).squeeze(1)  # h_v = sum of neighbours' features

perm_seed = mo.ui.slider(
    0, 50, value=1, label="permutation seed", show_value=True
)
perm_seed

def equivariance_gap(model, seed):
    """Relabel nodes with permutation `perm` (new position i holds old node perm[i]); compare outputs per original node."""
    perm = torch.from_numpy(np.random.default_rng(seed).permutation(34))
    A_perm = A_karate[perm][:, perm]  # P A P^T
    x_perm = x_degree[perm]
    with torch.no_grad():
        out_orig = model(A_karate, x_degree)
        out_perm = model(A_perm, x_perm)
    out_back = torch.empty_like(out_perm)
    out_back[perm] = out_perm  # undo the relabelling
    return out_orig, out_back

_rows = []
for _name, _model in [
    ("MLP on vec(A)", mlp_on_adjacency),
    ("neighbour sum", neighbour_sum),
]:
    _o, _b = equivariance_gap(_model, perm_seed.value)
    _rows.append(
        {
            "model": _name,
            "node 0 (orig)": round(_o[0].item(), 3),
            "node 0 (relabelled)": round(_b[0].item(), 3),
            "node 33 (orig)": round(_o[33].item(), 3),
            "node 33 (relabelled)": round(_b[33].item(), 3),
            "max |difference| over nodes": round(
                (_o - _b).abs().max().item(), 4
            ),
        }
    )
pd.DataFrame(_rows)

def build_views():
    camp_of = dict(zip(users_df.user, users_df.camp))
    views = {}
    g = nx.DiGraph()
    g.add_nodes_from(users_df.user)
    g.add_edges_from(follows_df.itertuples(index=False))
    views["follows (directed)"] = g
    rt = retweets_df.merge(tweets_df, on="tweet")  # who retweeted whom
    g = nx.DiGraph()
    g.add_nodes_from(users_df.user)
    for (u, a), w in rt.groupby(["user", "author"]).size().items():
        g.add_edge(u, a, weight=int(w))
    views["retweets (directed, weighted)"] = g
    ut = tweets_df.groupby("author").hashtag.apply(set)
    g = nx.Graph()
    g.add_nodes_from(users_df.user)
    for i, u in enumerate(ut.index):
        for v in ut.index[i + 1 :]:
            shared = len(ut[u] & ut[v])
            if shared:
                g.add_edge(u, v, weight=shared)
    views["co-hashtag (folded, undirected)"] = g
    g = nx.Graph()
    g.add_nodes_from(users_df.user, kind="user")
    g.add_nodes_from(tweets_df.tweet, kind="tweet")
    g.add_nodes_from(tweets_df.hashtag.unique(), kind="hashtag")
    g.add_edges_from(zip(tweets_df.author, tweets_df.tweet), rel="posts")
    g.add_edges_from(
        zip(retweets_df.user, retweets_df.tweet), rel="retweets"
    )
    g.add_edges_from(
        zip(tweets_df.tweet, tweets_df.hashtag), rel="has_tag"
    )
    views["heterogeneous (users, tweets, hashtags)"] = g
    return views, camp_of

social_views, camp_of = build_views()
view_picker = mo.ui.dropdown(
    list(social_views), value="follows (directed)", label="representation"
)
view_picker

def to_heterodata():
    uid = {u: i for i, u in enumerate(users_df.user)}
    tid = {t: i for i, t in enumerate(tweets_df.tweet)}
    tags = sorted(tweets_df.hashtag.unique())
    hid = {h: i for i, h in enumerate(tags)}
    data = HeteroData()
    data["user"].num_nodes = len(uid)
    data["user"].y = torch.tensor(
        users_df.camp.values
    )  # a node label to predict later
    data["tweet"].num_nodes = len(tid)
    data["hashtag"].num_nodes = len(hid)
    pairs = lambda a, b: torch.tensor([list(a), list(b)], dtype=torch.long)
    data["user", "follows", "user"].edge_index = pairs(
        follows_df.follower.map(uid), follows_df.followee.map(uid)
    )
    data["user", "posts", "tweet"].edge_index = pairs(
        tweets_df.author.map(uid), tweets_df.tweet.map(tid)
    )
    data["user", "retweets", "tweet"].edge_index = pairs(
        retweets_df.user.map(uid), retweets_df.tweet.map(tid)
    )
    data["tweet", "has_tag", "hashtag"].edge_index = pairs(
        tweets_df.tweet.map(tid), tweets_df.hashtag.map(hid)
    )
    return data

social_hetero = to_heterodata()
print(social_hetero)
print("metadata (node types, relation types):", social_hetero.metadata())

movie_users = ["Ali", "Sara", "Reza"]
movie_names = ["m1", "m2", "m3", "m4"]
B_movies = np.array([[1, 1, 0, 0], [1, 1, 1, 0], [0, 0, 1, 1]])

add_blockbuster = mo.ui.checkbox(
    label="add a blockbuster column that everyone rated"
)
downweight = mo.ui.checkbox(
    label="down-weight shared items by 1 / log(1 + item degree)"
)
mo.vstack(
    [
        mo.md("**Try it:** toggle a hub item and the usual fix."),
        add_blockbuster,
        downweight,
    ]
)

def fold(B, item_weights=None):
    """Return (user-user, item-item) folded matrices; optional per-item weights for the user side."""
    w = np.ones(B.shape[1]) if item_weights is None else item_weights
    return (B * w) @ B.T, B.T @ B

def heatmap(M, rows, cols, title, ax):
    ax.imshow(M, cmap="Blues", aspect="auto")
    ax.set_xticks(range(len(cols)), cols, rotation=45, fontsize=8)
    ax.set_yticks(range(len(rows)), rows, fontsize=8)
    for (i, j), val in np.ndenumerate(M):
        ax.text(
            j,
            i,
            f"{val:.2f}".rstrip("0").rstrip("."),
            ha="center",
            va="center",
            fontsize=8,
            color="white" if val > 0.6 * M.max() else "black",
        )
    ax.set_title(title, fontsize=9)

_B = (
    np.hstack([B_movies, np.ones((3, 1), dtype=int)])
    if add_blockbuster.value
    else B_movies
)
_items = movie_names + (["blockbuster"] if add_blockbuster.value else [])
_w = 1 / np.log1p(_B.sum(0)) if downweight.value else None
_UU, _VV = fold(_B, _w)

_fig, _axs = plt.subplots(
    1, 3, figsize=(9, 2.8), gridspec_kw={"width_ratios": [1.3, 1, 1.3]}
)
heatmap(_B, movie_users, _items, "B (users × movies)", _axs[0])
heatmap(
    _UU,
    movie_users,
    movie_users,
    "user–user  B·diag(w)·Bᵀ" if downweight.value else "user–user  BBᵀ",
    _axs[1],
)
heatmap(_VV, _items, _items, "movie–movie  BᵀB", _axs[2])
plt.tight_layout()
_fig

_dist0 = nx.single_source_shortest_path_length(karate, 0)
_dist33 = nx.single_source_shortest_path_length(karate, 33)
_pagerank = nx.pagerank(karate, weight=None)
_clust = nx.clustering(karate)
node_features = pd.DataFrame(
    {
        "degree": degrees,
        "clustering": [_clust[v] for v in karate],
        "pagerank": [_pagerank[v] for v in karate],
        "dist_to_0 - dist_to_33": [_dist0[v] - _dist33[v] for v in karate],
    }
)

def cv_accuracy(cols, repeats=5):
    """Mean accuracy of standardised logistic regression over `repeats` shuffled, stratified 5-fold splits."""
    clf = make_pipeline(
        StandardScaler(), LogisticRegression(max_iter=1000)
    )
    return np.mean(
        [
            cross_val_score(
                clf,
                node_features[cols].values,
                club,
                cv=StratifiedKFold(5, shuffle=True, random_state=r),
            ).mean()
            for r in range(repeats)
        ]
    )

pd.DataFrame(
    {
        "features": [
            "degree only",
            "degree + clustering + pagerank",
            "+ relative distance to the two leaders",
        ],
        "accuracy (5×5-fold CV)": [
            cv_accuracy(["degree"]),
            cv_accuracy(["degree", "clustering", "pagerank"]),
            cv_accuracy(list(node_features.columns)),
        ],
    }
).round(3)

hide_frac = mo.ui.slider(
    0.05,
    0.5,
    step=0.05,
    value=0.15,
    label="fraction of edges hidden",
    show_value=True,
)
hide_frac

def link_prediction_demo(frac, seed=SEED):
    r = np.random.default_rng(seed)
    edges = list(karate.edges())
    hidden_idx = r.choice(
        len(edges), int(frac * len(edges)), replace=False
    )
    hidden = [edges[i] for i in hidden_idx]
    g_obs = karate.copy()
    g_obs.remove_edges_from(hidden)
    non_edges = list(
        nx.non_edges(karate)
    )  # truly absent pairs = negatives
    candidates = hidden + non_edges
    labels = np.r_[np.ones(len(hidden)), np.zeros(len(non_edges))]
    scorers = {
        "common neighbours": lambda u, v: len(
            list(nx.common_neighbors(g_obs, u, v))
        ),
        "Jaccard": lambda u, v: next(
            nx.jaccard_coefficient(g_obs, [(u, v)])
        )[2],
        "Adamic–Adar": lambda u, v: next(
            nx.adamic_adar_index(g_obs, [(u, v)])
        )[2],
        "random guess": lambda u, v: r.random(),
    }
    return {
        name: roc_auc_score(labels, [f(u, v) for u, v in candidates])
        for name, f in scorers.items()
    }, len(hidden)

_aucs, _n_hidden = link_prediction_demo(hide_frac.value)
mo.md(
    f"Hidden **{_n_hidden}** edges, scored against all {len(list(nx.non_edges(karate)))} non-edges. ROC AUC:\n\n"
    + "\n".join(f"- **{k}**: {v:.3f}" for k, v in _aucs.items())
)

REAL_NETS = ["karate", "dolphins", "football", "polblogs", "email-eu-core", "ego-facebook", "higgs-retweet"]

def describe(name):
    g = graph(name)
    st = summary(g)
    und = g.to_undirected() if g.is_directed() else g
    deg = np.array([d for _, d in und.degree()])
    vals, counts = np.unique(deg[deg > 0], return_counts=True)
    ccdf = 1 - np.cumsum(counts) / counts.sum() + counts / counts.sum()  # P(degree >= k)
    st["degree"] = vals.tolist()
    st["ccdf"] = ccdf.tolist()
    return st

real_run = mo.ui.run_button(label="Load 7 real networks (downloads ≈ 3 MB the first time; then cached)")
real_run

_cache = CACHE_DIR / "L01_real_networks.json"
mo.stop(
    not (real_run.value or os.environ.get("GL_RUN_EXPERIMENTS") == "1") and not _cache.exists(),
    mo.md("Press the button. Results are cached in `notebooks/cs224w/cache/L01_real_networks.json`."),
)
real_stats = {n: json_cache(_cache, {"network": n, "v": 1}, lambda n=n: describe(n)) for n in REAL_NETS}
_cols = ["nodes", "edges", "directed", "mean degree", "max degree / mean", "avg clustering", "avg path (sampled)",
         "reciprocity", "giant component"]
real_table = pd.DataFrame({n: {c: v[c] for c in _cols} for n, v in real_stats.items()}).T
# an Erdős–Rényi graph with the same density has expected clustering = density
real_table["clustering ÷ ER"] = real_table["avg clustering"] / (
    real_table["mean degree"] / (real_table["nodes"] - 1)
)
_fig, _ax = plt.subplots(figsize=(5.2, 3.3))
for _n, _v in real_stats.items():
    _ax.loglog(_v["degree"], _v["ccdf"], ".-", ms=3, lw=0.8, label=_n)
_ax.set_xlabel("degree k")
_ax.set_ylabel("P(degree ≥ k)")
_ax.set_title("Degree distributions (log-log)", fontsize=9)
_ax.legend(fontsize=6.5)
plt.tight_layout()
mo.vstack([_fig, real_table.round(3)])

ex1_answer = None
ex2_answer = None

_one_way = torch.tensor(list(karate.edges()), dtype=torch.long).t()
_in_deg_one_way = torch.bincount(_one_way[1], minlength=34)
_in_deg_fixed = torch.bincount(to_undirected(_one_way)[1], minlength=34)
pd.DataFrame(
    {
        "": ["one-way edge_index", "to_undirected(edge_index)"],
        "is_undirected": [
            is_undirected(_one_way),
            is_undirected(to_undirected(_one_way)),
        ],
        "columns": [_one_way.size(1), to_undirected(_one_way).size(1)],
        "in-degree of node 0": [
            int(_in_deg_one_way[0]),
            int(_in_deg_fixed[0]),
        ],
        "in-degree of node 33": [
            int(_in_deg_one_way[33]),
            int(_in_deg_fixed[33]),
        ],
    }
)