# %% [markdown]
# # Notebook 4: File Co-occurrence & Dependency Graph
# Build co-occurrence matrices and file clusters from sessions.

# %%
import sys, os, json
from itertools import combinations

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import pandas as pd
import numpy as np
from collections import Counter
from utils.connection import init, query_df
from utils import sessions
from IPython.display import display, Markdown

pd.set_option("display.max_columns", 30)
pd.set_option("display.max_rows", 100)
pd.set_option("display.max_colwidth", 120)

conn = init()
display(Markdown("**Connected.**"))

# %% [markdown]
# ## 1 — Collect all files touched per session (read + edit)

# %%
TOUCH_TOOLS = ("Edit", "Write", "MultiEdit", "MultiEditTool",
               "Read", "Grep", "Glob", "View", "ReadFile",
               "read_file", "file_search")

all_tool_files = query_df(conn, """
    SELECT tc.tool_name, tc.tool_input,
           m.session_id, s.user_email, s.repo_name, s.cwd
    FROM   tool_calls tc
    JOIN   messages m ON m.id = tc.message_id
    JOIN   sessions s ON s.id = m.session_id
    WHERE  tc.tool_name = ANY(%s)
""", (list(TOUCH_TOOLS),))

display(Markdown(f"**Total file-touching tool calls:** {len(all_tool_files)}"))

# %%
def extract_path(tool_input):
    if tool_input is None:
        return None
    if isinstance(tool_input, str):
        try:
            tool_input = json.loads(tool_input)
        except (json.JSONDecodeError, TypeError):
            return None
    if isinstance(tool_input, dict):
        return (tool_input.get("file_path") or tool_input.get("path")
                or tool_input.get("target_file") or tool_input.get("command"))
    return None

def make_relative(fp, cwd):
    if not fp or not cwd:
        return fp
    if fp.startswith(cwd):
        return fp[len(cwd):].lstrip("/") or fp
    return fp

all_tool_files["file_path"] = all_tool_files["tool_input"].apply(extract_path)
all_tool_files["rel_path"] = all_tool_files.apply(
    lambda r: make_relative(r["file_path"], r["cwd"]), axis=1
)
all_tool_files = all_tool_files[all_tool_files["rel_path"].notna()].copy()

# Filter to likely file paths (has extension or is in a directory)
all_tool_files = all_tool_files[
    all_tool_files["rel_path"].str.contains(r"\.\w{1,6}$", na=False)
].copy()

display(Markdown(f"**File paths extracted:** {len(all_tool_files)}"))

# %% [markdown]
# ## 2 — Per-session file sets

# %%
session_files = (
    all_tool_files.groupby("session_id")["rel_path"]
    .apply(lambda x: sorted(set(x)))
    .to_frame("files")
)
session_files["file_count"] = session_files["files"].apply(len)

display(Markdown("### Files per session distribution"))
display(session_files["file_count"].describe().round(1).to_frame())

# Add metadata
all_sess = sessions.list_all(conn, limit=500)
session_files = session_files.join(all_sess.set_index("id")[["user_email", "repo_name"]])

# %% [markdown]
# ## 3 — Co-occurrence matrix (file pairs)

# %%
pair_counter = Counter()
for _, row in session_files.iterrows():
    files = row["files"]
    if len(files) < 2:
        continue
    # Limit to avoid combinatorial explosion
    files_limited = files[:50]
    for a, b in combinations(files_limited, 2):
        pair = tuple(sorted([a, b]))
        pair_counter[pair] += 1

# Convert to DataFrame
pairs_df = pd.DataFrame(
    [(a, b, count) for (a, b), count in pair_counter.most_common(100)],
    columns=["file_a", "file_b", "co_occurrence_count"]
)

display(Markdown("### Top co-occurring file pairs (across all sessions)"))
display(pairs_df.head(30))

# %% [markdown]
# ## 4 — Strong pairs (appear together in >50% of sessions where either appears)

# %%
file_session_count = (
    all_tool_files.groupby("rel_path")["session_id"].nunique()
    .to_dict()
)

pairs_df["file_a_sessions"] = pairs_df["file_a"].map(file_session_count)
pairs_df["file_b_sessions"] = pairs_df["file_b"].map(file_session_count)
pairs_df["min_sessions"] = pairs_df[["file_a_sessions", "file_b_sessions"]].min(axis=1)
pairs_df["co_occurrence_pct"] = (pairs_df["co_occurrence_count"] / pairs_df["min_sessions"] * 100).round(1)

display(Markdown("### Strongly coupled file pairs (>50% co-occurrence)"))
strong = pairs_df[pairs_df["co_occurrence_pct"] > 50].sort_values("co_occurrence_pct", ascending=False)
display(strong.head(30))

# %% [markdown]
# ## 5 — File clusters (connected components from strong pairs)

# %%
# Simple union-find clustering
from collections import defaultdict

def find_clusters(pairs, threshold=2):
    """Group files into clusters based on co-occurrence."""
    parent = {}
    def find(x):
        if x not in parent:
            parent[x] = x
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    for _, row in pairs.iterrows():
        if row["co_occurrence_count"] >= threshold:
            union(row["file_a"], row["file_b"])

    clusters = defaultdict(set)
    for f in parent:
        clusters[find(f)].add(f)

    return [sorted(c) for c in clusters.values() if len(c) >= 2]

clusters = find_clusters(pairs_df, threshold=2)
clusters.sort(key=len, reverse=True)

display(Markdown(f"### File clusters found: {len(clusters)}"))
for i, cluster in enumerate(clusters[:10]):
    display(Markdown(f"**Cluster {i+1}** ({len(cluster)} files):"))
    for f in cluster:
        print(f"  {f}")

# %% [markdown]
# ## 6 — Per-repo dependency map

# %%
display(Markdown("### Per-repo: if you touch file A, you'll likely also touch..."))

for repo in all_tool_files["repo_name"].dropna().unique():
    repo_files = all_tool_files[all_tool_files["repo_name"] == repo]
    repo_session_files = repo_files.groupby("session_id")["rel_path"].apply(set)

    repo_pairs = Counter()
    for files in repo_session_files:
        for a, b in combinations(sorted(files)[:30], 2):
            repo_pairs[tuple(sorted([a, b]))] += 1

    top_repo_pairs = repo_pairs.most_common(15)
    if top_repo_pairs:
        display(Markdown(f"#### {repo}"))
        rp_df = pd.DataFrame(
            [(a, b, c) for (a, b), c in top_repo_pairs],
            columns=["file_a", "file_b", "sessions_together"]
        )
        display(rp_df)

# %% [markdown]
# ## 7 — Per-user file affinity

# %%
display(Markdown("### Per-user top files (by session count)"))
for email in all_tool_files["user_email"].dropna().unique():
    user_files = all_tool_files[all_tool_files["user_email"] == email]
    affinity = (
        user_files.groupby("rel_path")["session_id"].nunique()
        .sort_values(ascending=False)
        .head(15)
    )
    display(Markdown(f"#### {email}"))
    display(affinity.to_frame("sessions"))

# %% [markdown]
# ---
# *End of Notebook 4.*
