"""
Supplementary Material: Collocation Analysis around Target Terms
----------------------------------------------------------------
This script computes sentence-bounded, window-based collocations around
user-specified target lemmas *with explicit target POS* (e.g., office/PROPN,
agency/NOUN, department/NOUN, secretary/PROPN). It supports POS filtering for
collocates (ADJ, VERB, NOUN/PROPN) and an adjustable symmetric window.

For each (target, collocate) pair it reports:
- Frequency (co-occurrence count within the window)
- Pointwise Mutual Information (PMI; base 2)
- T-Score ((O - E) / sqrt(E))
- Z-Score (approx. under independence, from 2x2 variance)
- Log-Likelihood G^2 (Dunning, 1993)
- Phi coefficient (effect size for 2x2 table)

Selection for figures uses *adaptive LL gates*:
  Try LL ≥ 10.83 (p < .001); if a target has < k items, relax to 6.63 (.01),
  then 3.84 (.05), always ranking by T-score. We drop self-pairs and (optionally)
  keep only attractive associations (T > 0 and Phi > 0).

References:
- Dunning, T. (1993). Accurate Methods for the Statistics of Surprise and Coincidence.
- Church & Hanks (1990). Word association norms, mutual information, and lexicography.
- Evert (2009). The Statistics of Word Cooccurrences.
"""

from __future__ import annotations
import math
from collections import Counter, defaultdict
from typing import Iterable, List, Tuple, Dict, Set

import pandas as pd
import numpy as np
import spacy


# ----------------------------- Configuration -----------------------------

# Robust English model; PROPN kept distinct (no PROPN->NOUN mapping).
NLP = spacy.load("en_core_web_md", disable=["ner", "parser", "textcat"])
NLP.add_pipe("sentencizer")
NLP.max_length = 5_000_000

STOPWORDS = NLP.Defaults.stop_words  # optional; not applied unless requested


# ----------------------------- Preprocessing -----------------------------

def _keep_token(tok: spacy.tokens.Token) -> bool:
    """Include alphabetic tokens only (drop numbers, punctuation)."""
    return tok.is_alpha


def _token_repr(tok: spacy.tokens.Token) -> Tuple[str, str]:
    """Return lowercased lemma and UD POS tag (PROPN kept as PROPN)."""
    return tok.lemma_.lower(), tok.pos_


# ----------------------------- Counting Pairs -----------------------------

def collect_cooccurrences(
    texts: Iterable[str],
    target_pos_by_lemma: Dict[str, Set[str]],
    collocate_pos_filter: Set[str] = frozenset({"ADJ"}),
    window: int = 7,
    batch_size: int = 200,
    n_process: int = 1,
    drop_stopwords_collocates: bool = False,
    min_lemma_len: int = 1,
) -> Tuple[Counter, Counter, Counter, int]:
    """
    Iterate through texts and collect symmetric window-based co-occurrence events
    within sentence boundaries for (target, collocate) pairs.

    Parameters
    ----------
    texts : iterable of str
        Corpus texts (e.g., paragraphs). Windows do not cross sentence boundaries.
    target_pos_by_lemma : dict[str, set[str]]
        Allowed POS tags per target lemma, e.g. {"office": {"PROPN"}, "agency": {"NOUN"}}.
    collocate_pos_filter : set[str]
        Allowed POS for collocates (e.g., {"ADJ"} or {"VERB"} or {"NOUN","PROPN"}).
    window : int
        Symmetric window size around the target (±window).
    drop_stopwords_collocates : bool
        If True, remove collocate lemmas that are stopwords.
    min_lemma_len : int
        Minimum length of a collocate lemma (filters very short items).

    Returns
    -------
    pair_counts : Counter keyed by (t_lemma, t_pos, c_lemma, c_pos)
    target_marginals : Counter keyed by t_lemma with sum over its collocates
    colloc_marginals : Counter keyed by c_lemma with sum over its targets
    N_events : int total co-occurrence events across all pairs
    """
    pair_counts: Counter = Counter()
    target_marginals: Counter = Counter()
    colloc_marginals: Counter = Counter()

    def valid_collocate(lemma: str, pos: str) -> bool:
        if len(lemma) < min_lemma_len:
            return False
        if pos not in collocate_pos_filter:
            return False
        if drop_stopwords_collocates and lemma in STOPWORDS:
            return False
        return True

    search_lemmas = set(target_pos_by_lemma.keys())

    for doc in NLP.pipe(texts, batch_size=batch_size, n_process=n_process):
        for sent in doc.sents:
            sent_tokens = []
            for tok in sent:
                if not _keep_token(tok):
                    continue
                lemma, pos = _token_repr(tok)
                sent_tokens.append((lemma, pos))

            for i, (lem_i, pos_i) in enumerate(sent_tokens):
                if lem_i not in search_lemmas:
                    continue
                if pos_i not in target_pos_by_lemma.get(lem_i, set()):
                    continue

                start = max(0, i - window)
                end = min(len(sent_tokens), i + window + 1)
                for j in range(start, end):
                    if j == i:
                        continue
                    lem_j, pos_j = sent_tokens[j]
                    if not valid_collocate(lem_j, pos_j):
                        continue

                    key = (lem_i, pos_i, lem_j, pos_j)
                    pair_counts[key] += 1
                    target_marginals[lem_i] += 1
                    colloc_marginals[lem_j] += 1

    N_events = sum(pair_counts.values())
    return pair_counts, target_marginals, colloc_marginals, N_events


# ----------------------------- Statistics -----------------------------

def compute_stats_for_pairs(
    pair_counts: Counter,
    target_marginals: Counter,
    colloc_marginals: Counter,
    N_events: int,
    min_freq: int = 3
) -> pd.DataFrame:
    """
    Compute PMI, T-Score, Z-Score, Log-Likelihood (G^2), and Phi for each pair.
    A minimum frequency filter (default >= 3) reduces noise.
    """
    rows = []
    eps = 1e-12

    for (t_lem, t_pos, c_lem, c_pos), k11 in pair_counts.items():
        if k11 < min_freq:
            continue

        k1_ = target_marginals[t_lem]
        k_1 = colloc_marginals[c_lem]
        k12 = max(k1_ - k11, 0)
        k21 = max(k_1 - k11, 0)
        k22 = max(N_events - k11 - k12 - k21, 0)

        # Expected under independence
        E11 = (k1_ * k_1) / max(N_events, eps)

        # PMI (base 2)
        p_xy = k11 / max(N_events, eps)
        p_x  = k1_ / max(N_events, eps)
        p_y  = k_1 / max(N_events, eps)
        PMI = math.log2(p_xy / max(p_x * p_y, eps)) if p_xy > 0 and p_x > 0 and p_y > 0 else 0.0

        # T-Score
        T_score = (k11 - E11) / math.sqrt(E11) if E11 > 0 else 0.0

        # Z-Score (approximate variance of 2x2 table)
        k2_ = max(N_events - k1_, 0)
        k_2 = max(N_events - k_1, 0)
        var = (k1_ * k_1 * k2_ * k_2) / max((N_events**2) * (N_events - 1), eps)
        Z_score = (k11 - E11) / math.sqrt(var) if var > 0 else 0.0

        # Log-Likelihood G^2 (Dunning, 1993)
        def term(cell, exp):
            return 0.0 if cell == 0 or exp == 0 else cell * (math.log(cell) - math.log(exp))
        E12 = (k1_ * (N_events - k_1)) / max(N_events, eps)
        E21 = ((N_events - k1_) * k_1) / max(N_events, eps)
        E22 = ((N_events - k1_) * (N_events - k_1)) / max(N_events, eps)
        G2 = 2.0 * (term(k11, E11) + term(k12, E12) + term(k21, E21) + term(k22, E22))

        # Phi
        denom = math.sqrt((k11 + k12) * (k11 + k21) * (k12 + k22) * (k21 + k22))
        Phi = ((k11 * k22) - (k12 * k21)) / denom if denom > 0 else 0.0

        rows.append({
            "Collocation": f"{t_lem}_{t_pos} :: {c_lem}_{c_pos}",
            "Frequency": int(k11),
            "PMI": PMI,
            "T-Score": T_score,
            "Z-Score": Z_score,
            "Log-Likelihood": G2,
            "Phi": Phi,
            "Word_1": t_lem,
            "POS_1": t_pos,
            "Word_2": c_lem,
            "POS_2": c_pos
        })

    df = pd.DataFrame(rows)
    # Sort by LL then freq for stable display; ranking for figures will use T-Score later.
    if not df.empty:
        df.sort_values(["Log-Likelihood", "Frequency"], ascending=[False, False], inplace=True)
    return df


# ----------------------- Adaptive top-k selection ------------------------

def adaptive_top_k_per_target(
    df: pd.DataFrame,
    k: int = 10,
    ll_bins: Tuple[float, ...] = (10.83, 6.63, 3.84),
    enforce_attractive: bool = True,
    drop_self_pairs: bool = True
) -> pd.DataFrame:
    """
    For each target (Word_1, POS_1), return up to k collocates ranked by T-Score,
    using adaptive LL gates: try the highest LL threshold; if <k items, relax.

    - drop_self_pairs: remove items with identical lemmas (regardless of POS)
    - enforce_attractive: keep only T-Score > 0 and Phi > 0
    """
    if df.empty:
        return df.copy()

    out_rows = []

    # Optional cleaning
    work = df.copy()
    if drop_self_pairs:
        work = work[work["Word_1"] != work["Word_2"]]
    if enforce_attractive:
        work = work[(work["T-Score"] > 0) & (work["Phi"] > 0)]

    # For reporting LL bins later in plots (optional)
    def ll_bin(v: float) -> str:
        if v >= 10.83:
            return "LL≥10.83 (p<.001)"
        elif v >= 6.63:
            return "LL≥6.63 (p<.01)"
        elif v >= 3.84:
            return "LL≥3.84 (p<.05)"
        else:
            return "LL<3.84"

    work["LL_bin"] = work["Log-Likelihood"].apply(ll_bin)

    # Process each target independently
    for (t_lem, t_pos), sub in work.groupby(["Word_1", "POS_1"]):
        selected = pd.DataFrame(columns=sub.columns)
        for thr in ll_bins + (0.0,):  # final 0.0 as last-resort (if you want to allow)
            cand = sub[sub["Log-Likelihood"] >= thr]
            cand = cand.sort_values(["T-Score", "Frequency"], ascending=[False, False])
            if len(cand) >= k:
                selected = cand.head(k)
                break
            # keep the best-so-far and continue relaxing
            selected = cand if len(cand) > len(selected) else selected

        out_rows.append(selected)

    out = pd.concat(out_rows, axis=0) if out_rows else work.iloc[0:0].copy()

    # Order nicely for tables/figures
    out = out.sort_values(
        ["Word_1", "POS_1", "T-Score", "Log-Likelihood", "Frequency"],
        ascending=[True, True, False, False, False]
    )
    return out


# ----------------------------- High-level API -----------------------------

def collocation_table_from_csv(
    csv_path: str,
    text_col: str,
    targets_with_pos: List[Tuple[str, str]],
    window: int = 7,
    collocate_pos_filter: Iterable[str] = ("ADJ",),
    n_process: int = 1,
    batch_size: int = 200,
    min_freq: int = 3,
    drop_stopwords_collocates: bool = False,
) -> pd.DataFrame:
    """
    Build a collocation table for given target (lemma, POS) pairs.
    """
    df_in = pd.read_csv(csv_path, usecols=[text_col])
    texts = df_in[text_col].fillna("").astype(str).tolist()

    # Map: lemma -> allowed POS set (keeps PROPN distinct)
    target_pos_by_lemma: Dict[str, Set[str]] = defaultdict(set)
    for lem, pos in targets_with_pos:
        target_pos_by_lemma[lem.lower()].add(pos)

    pair_counts, targ_marg, coll_marg, N = collect_cooccurrences(
        texts=texts,
        target_pos_by_lemma=target_pos_by_lemma,
        collocate_pos_filter=set(collocate_pos_filter),
        window=window,
        batch_size=batch_size,
        n_process=n_process,
        drop_stopwords_collocates=drop_stopwords_collocates,
    )

    table = compute_stats_for_pairs(
        pair_counts=pair_counts,
        target_marginals=targ_marg,
        colloc_marginals=coll_marg,
        N_events=N,
        min_freq=min_freq
    )
    return table


def print_diagnostics(df: pd.DataFrame, label: str):
    """Compact diagnostics for LL distribution and target coverage."""
    print(f"\n=== DIAGNOSTICS {label} ===")
    if df.empty:
        print("No pairs found.")
        return
    total = len(df)
    n_001 = (df["Log-Likelihood"] >= 10.83).sum()
    n_01  = (df["Log-Likelihood"] >= 6.63).sum()
    n_05  = (df["Log-Likelihood"] >= 3.84).sum()
    by_target = df.groupby(["Word_1", "POS_1"]).size()
    print(f"Total pairs: {total}")
    print("Pairs by LL bins:", n_001, n_01, n_05)
    print("Pairs by target:\n", by_target)


# ----------------------------- Example Usage -----------------------------
if __name__ == "__main__":
    # Targets for RQ2 (as discussed): keep PROPN distinct
    nodes = [
        ("administration", "PROPN"),
        ("agency", "NOUN"),
        ("policy", "NOUN"),
        ("program", "NOUN"),
    ]

    corpus_csv = "raw_data/Project2025.csv"   # one paragraph per row
    text_col = "text"
    nproc = 1  # macOS-safe; increase if you wrap in `if __name__ == '__main__'`

    # -------- ADJECTIVE collocates --------
    coll_adj = collocation_table_from_csv(
        csv_path=corpus_csv,
        text_col=text_col,
        targets_with_pos=nodes,
        window=7,
        collocate_pos_filter=("ADJ",),
        n_process=nproc,
        batch_size=200,
        min_freq=3,
        drop_stopwords_collocates=False,
    )
    print_diagnostics(coll_adj, label="(ADJ)")

    top_adj = adaptive_top_k_per_target(
        coll_adj, k=10,
        ll_bins=(10.83, 6.63, 3.84),
        enforce_attractive=True,
        drop_self_pairs=True
    )
    top_adj.to_csv("P2025_colloc_ADJ_w7_top10_per_node.csv", index=False)

    # -------- VERB collocates --------
    coll_verb = collocation_table_from_csv(
        csv_path=corpus_csv,
        text_col=text_col,
        targets_with_pos=nodes,
        window=7,
        collocate_pos_filter=("VERB",),
        n_process=nproc,
        batch_size=200,
        min_freq=3,
        drop_stopwords_collocates=False,
    )
    print_diagnostics(coll_verb, label="(VERB)")

    top_verb = adaptive_top_k_per_target(
        coll_verb, k=10,
        ll_bins=(10.83, 6.63, 3.84),
        enforce_attractive=True,
        drop_self_pairs=True
    )
    top_verb.to_csv("P2025_colloc_VERB_w7_top10_per_node.csv", index=False)

    # -------- NOUN/PROPN collocates (optional for supplement) --------
    coll_noun = collocation_table_from_csv(
        csv_path=corpus_csv,
        text_col=text_col,
        targets_with_pos=nodes,
        window=7,
        collocate_pos_filter=("NOUN", "PROPN"),
        n_process=nproc,
        batch_size=200,
        min_freq=3,
        drop_stopwords_collocates=False,
    )
    print_diagnostics(coll_noun, label="(NOUN/PROPN)")

    top_noun = adaptive_top_k_per_target(
        coll_noun, k=10,
        ll_bins=(10.83, 6.63, 3.84),
        enforce_attractive=True,
        drop_self_pairs=True
    )
    top_noun.to_csv("P2025_colloc_NOUN_w7_top10_per_node.csv", index=False)

    # Console peek for sanity
    if not top_adj.empty:
        print("\nTop ADJ per node:\n", top_adj.head(20).to_string(index=False))
    if not top_verb.empty:
        print("\nTop VERB per node:\n", top_verb.head(20).to_string(index=False))
    if not top_noun.empty:
        print("\nTop NOUN/PROPN per node:\n", top_noun.head(20).to_string(index=False))

# ----------------------------- Example sentences -----------------------------

from pathlib import Path

def iter_sentence_tokens(texts, nlp=NLP):
    """
    Yields (doc_id, sent_id, tokens), where tokens is a list of dicts:
      {"text": surface, "lemma": lemma_lower, "pos": UD_POS}
    Sentences are bounded by spaCy's sentencizer.
    """
    for di, doc in enumerate(nlp.pipe(texts, batch_size=200, n_process=1)):
        si = 0
        for sent in doc.sents:
            toks = []
            for tok in sent:
                if tok.is_alpha:
                    toks.append({
                        "text": tok.text,
                        "lemma": tok.lemma_.lower(),
                        "pos": tok.pos_
                    })
            yield di, si, toks
            si += 1


def find_pair_hits_in_sentence(tokens, target_lemma, target_pos, colloc_lemma, colloc_pos, window=7):
    """
    Returns list of (t_idx, c_idx, distance) for occurrences of target/collocate
    in the same sentence within ±window (token indices within sentence tokens).
    """
    t_idx = [i for i, t in enumerate(tokens) if t["lemma"] == target_lemma and (target_pos is None or t["pos"] == target_pos)]
    c_idx = [i for i, t in enumerate(tokens) if t["lemma"] == colloc_lemma and (colloc_pos is None or t["pos"] == colloc_pos)]
    hits = []
    if not t_idx or not c_idx:
        return hits
    for i in t_idx:
        for j in c_idx:
            dist = abs(i - j)
            if dist <= window and dist > 0:
                hits.append((i, j, dist))
    return hits


def reconstruct_sentence(tokens):
    """Rebuild a simple sentence string from surface tokens."""
    return " ".join(tok["text"] for tok in tokens)


def highlight_sentence(tokens, t_idx, c_idx):
    """Return a simple markdown-highlighted sentence string."""
    out = []
    for k, tok in enumerate(tokens):
        surf = tok["text"]
        if k == t_idx:
            surf = f"**{surf}**"         # target: bold
        if k == c_idx:
            surf = f"__{surf}__"         # collocate: italics
        out.append(surf)
    return " ".join(out)


def extract_examples_for_pairs(
    texts,
    pairs_df,
    window=7,
    max_per_pair=20,
    require_pos=True
):
    """
    pairs_df: must have columns Word_1, POS_1, Word_2, POS_2 (wie deine colloc-CSV)
    Returns a DataFrame with example sentences for each pair.
    """
    rows = []
    # Preload all sentences once
    sentences = list(iter_sentence_tokens(texts, NLP))

    for _, row in pairs_df.iterrows():
        t_lem = str(row["Word_1"]).lower()
        t_pos = str(row["POS_1"]) if require_pos and pd.notnull(row["POS_1"]) else None
        c_lem = str(row["Word_2"]).lower()
        c_pos = str(row["POS_2"]) if require_pos and pd.notnull(row["POS_2"]) else None

        collected = 0
        # iterate sentences, keep closest hit per sentence
        for di, si, toks in sentences:
            hits = find_pair_hits_in_sentence(toks, t_lem, t_pos, c_lem, c_pos, window=window)
            if not hits:
                continue
            # keep the closest target–collocate pair in this sentence
            t_idx, c_idx, dist = sorted(hits, key=lambda x: x[2])[0]
            sent_plain = reconstruct_sentence(toks)
            sent_hl    = highlight_sentence(toks, t_idx, c_idx)
            rows.append({
                "Word_1": row["Word_1"], "POS_1": row["POS_1"],
                "Word_2": row["Word_2"], "POS_2": row["POS_2"],
                "doc_id": di, "sent_id": si, "distance": dist,
                "sentence": sent_plain,
                "sentence_markdown": sent_hl
            })
            collected += 1
            if collected >= max_per_pair:
                break

    if not rows:
        return pd.DataFrame(columns=[
            "Word_1","POS_1","Word_2","POS_2","doc_id","sent_id","distance","sentence","sentence_markdown"
        ])
    out = pd.DataFrame(rows)
    # order: tighter distance first, then by target/collocate
    out = out.sort_values(["Word_1","POS_1","Word_2","POS_2","distance","doc_id","sent_id"])
    return out


# ----------------------------- Glue it in main -----------------------------
if __name__ == "__main__":
    RAW_CSV  = "raw_data/Project2025.csv"   # same as above
    TEXT_COL = "text"

    df_raw = pd.read_csv(RAW_CSV, usecols=[TEXT_COL])
    texts  = df_raw[TEXT_COL].fillna("").astype(str).tolist()

    # … (hier laufen bereits deine top_adj / top_verb Selektionen)

    # Ordner für Supplement-Tabellen
    Path("tables").mkdir(exist_ok=True)

    # Beispiel-Extraktion für ADJ-Top10
    if not top_adj.empty:
        ex_adj = extract_examples_for_pairs(
            texts,
            pairs_df = top_adj[["Word_1","POS_1","Word_2","POS_2"]].drop_duplicates(),
            window = 7,
            max_per_pair = 10,    # z.B. 10 Beispiele pro Paar
            require_pos = True
        )
        ex_adj.to_csv("tables/Supp_Collocation_Examples_ADJ.csv", index=False)

    # Beispiel-Extraktion für VERB-Top10
    if not top_verb.empty:
        ex_verb = extract_examples_for_pairs(
            texts,
            pairs_df = top_verb[["Word_1","POS_1","Word_2","POS_2"]].drop_duplicates(),
            window = 7,
            max_per_pair = 10,
            require_pos = True
        )
        ex_verb.to_csv("tables/Supp_Collocation_Examples_VERB.csv", index=False)