"""SEC Filing Text Signals - FinBERT sentiment and embedding-based alpha factors."""

import warnings

import matplotlib.pyplot as plt
import numpy as np
import polars as pl
import torch

from utils.paths import get_chapter_dir
from utils.reproducibility import set_global_seeds

warnings.filterwarnings("ignore")

# Production defaults (0 = all symbols; full dataset has 477 companies, 6770 filings)
# FinBERT scoring is sequential per-filing, so runtime scales linearly.
# 50 symbols (~750 filings) runs in ~5 minutes on GPU; set to 0 for full dataset.
SEED = 42
MAX_SYMBOLS = 50
MAX_FILINGS = 0
BATCH_SIZE = 8
MAX_TOKENS = 512
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
SENTIMENT_MODEL = "yiyanghkust/finbert-tone"

OUTPUT_DIR = get_chapter_dir(10) / "output" / "filing_signals"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Reproducibility — set_global_seeds covers Python random / NumPy / Torch.
set_global_seeds(SEED)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
if device.type == "cuda":
    print(f"GPU: {torch.cuda.get_device_name()}")

from data import load_sp500_10q_mda

filings = load_sp500_10q_mda()

print(f"Loaded {len(filings):,} MD&A sections from {filings['symbol'].n_unique()} companies")
print(f"Date range: {filings['filing_date'].min()} to {filings['filing_date'].max()}")

if MAX_SYMBOLS > 0:
    # Break ties on filing count by symbol so the selected universe is stable
    # across runs — group_by returns rows in a nondeterministic order, so a
    # sort on `len` alone would pick different symbols at the 50-symbol cutoff
    # each run, shifting every downstream statistic.
    top_symbols = (
        filings.group_by("symbol")
        .len()
        .sort(["len", "symbol"], descending=[True, False])
        .head(MAX_SYMBOLS)["symbol"]
        .to_list()
    )
    filings = filings.filter(pl.col("symbol").is_in(top_symbols))
    print(f"Filtered to {MAX_SYMBOLS} symbols: {len(filings):,} filings")

if MAX_FILINGS > 0 and len(filings) > MAX_FILINGS:
    filings = filings.sort(["filing_date", "symbol"]).head(MAX_FILINGS)
    print(f"Reduced to first {MAX_FILINGS} filings for test run")

# One MD&A signal per (symbol, filing_date). A few firms file multiple quarters'
# 10-Qs on the same date (late/catch-up filers, e.g. KHC on 2017-11-07 filed three).
# The investable event is the filing date, so keep the current quarter (latest
# period_end). Without this, both the sentiment and narrative-change tables carry
# duplicate keys and the section-4 join fans out on both sides, double-counting
# those firms in the reported information coefficients.
n_before = len(filings)
filings = filings.sort(["symbol", "filing_date", "period_end"]).unique(
    subset=["symbol", "filing_date"], keep="last", maintain_order=True
)
if len(filings) < n_before:
    print(
        f"Collapsed {n_before - len(filings)} same-date multi-quarter filings -> {len(filings):,}"
    )

# Compute MD&A word count from the canonical `text` column.
filings = filings.with_columns(pl.col("text").str.split(" ").list.len().alias("word_count"))

filings.head(5).select(["symbol", "filing_date", "period_end", "word_count"])

# Word count distribution
print("MD&A word count statistics:")
print(filings["word_count"].describe())

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

axes[0].hist(filings["word_count"].to_numpy(), bins=50, edgecolor="white")
axes[0].set_xlabel("Word Count")
axes[0].set_ylabel("Frequency")
axes[0].set_title("MD&A Length Distribution")
axes[0].axvline(filings["word_count"].median(), color="red", linestyle="--", label="Median")
axes[0].legend()

# Filings per quarter
quarterly = (
    filings.with_columns(
        quarter=pl.col("filing_date").dt.year().cast(pl.String)
        + "-Q"
        + pl.col("filing_date").dt.quarter().cast(pl.String)
    )
    .group_by("quarter")
    .len()
    .sort("quarter")
)
axes[1].bar(range(len(quarterly)), quarterly["len"].to_numpy())
axes[1].set_xticks(range(0, len(quarterly), 4))
axes[1].set_xticklabels(quarterly["quarter"].to_list()[::4], rotation=45)
axes[1].set_ylabel("Filings")
axes[1].set_title("Filings per Quarter")

fig.tight_layout()
fig.show()

from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline, set_seed

set_seed(SEED)
print(f"Loading FinBERT: {SENTIMENT_MODEL}")
tokenizer = AutoTokenizer.from_pretrained(SENTIMENT_MODEL)
model = AutoModelForSequenceClassification.from_pretrained(SENTIMENT_MODEL)
model = model.to(device)
model.eval()

sentiment_pipeline = pipeline(
    "sentiment-analysis",
    model=model,
    tokenizer=tokenizer,
    device=device,
    truncation=True,
    max_length=MAX_TOKENS,
    batch_size=BATCH_SIZE,
)
print("FinBERT loaded")

def chunk_text(text: str, max_chars: int = 1500) -> list[str]:
    """Split text into chunks suitable for FinBERT (roughly 512 tokens each)."""
    paragraphs = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 50]
    if not paragraphs:
        # Fall back to sentence splitting
        paragraphs = [s.strip() + "." for s in text.split(".") if len(s.strip()) > 30]

    chunks = []
    current = ""
    for para in paragraphs:
        if len(current) + len(para) > max_chars and current:
            chunks.append(current.strip())
            current = para
        else:
            current = current + "\n\n" + para if current else para

    if current.strip():
        chunks.append(current.strip())

    return chunks if chunks else [text[:max_chars]]

def score_document_sentiment(text: str) -> dict:
    """Score full MD&A document by aggregating chunk-level FinBERT predictions."""
    chunks = chunk_text(text)
    if not chunks:
        return {"sentiment_mean": 0.0, "sentiment_std": 0.0, "n_chunks": 0}

    # Score all chunks
    results = sentiment_pipeline(chunks)

    # Convert labels to numeric: positive=+1, neutral=0, negative=-1
    label_map = {"Positive": 1.0, "Neutral": 0.0, "Negative": -1.0}
    scores = []
    for r in results:
        label = r["label"]
        confidence = r["score"]
        numeric = label_map.get(label, 0.0) * confidence
        scores.append(numeric)

    scores_arr = np.array(scores)
    return {
        "sentiment_mean": float(scores_arr.mean()),
        "sentiment_std": float(scores_arr.std()) if len(scores_arr) > 1 else 0.0,
        "sentiment_pos_pct": float((scores_arr > 0).mean()),
        "sentiment_neg_pct": float((scores_arr < 0).mean()),
        "n_chunks": len(chunks),
    }

# Score all filings
print(f"Scoring {len(filings):,} MD&A sections with FinBERT...")
print("(This may take several minutes depending on GPU/CPU)")

sentiment_records = []
for i, row in enumerate(filings.iter_rows(named=True)):
    scores = score_document_sentiment(row["text"])
    scores["symbol"] = row["symbol"]
    scores["filing_date"] = row["filing_date"]
    sentiment_records.append(scores)

    if (i + 1) % 100 == 0 or (i + 1) == len(filings):
        print(f"  Scored {i + 1:,}/{len(filings):,} filings")

sentiment_df = pl.DataFrame(sentiment_records)
print(f"\nSentiment scoring complete: {len(sentiment_df):,} filings scored")
sentiment_df.head(5)

# Sentiment distribution
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

axes[0].hist(sentiment_df["sentiment_mean"].to_numpy(), bins=50, edgecolor="white")
axes[0].set_xlabel("Mean Sentiment Score")
axes[0].set_ylabel("Frequency")
axes[0].set_title("MD&A Sentiment Distribution")
axes[0].axvline(0, color="red", linestyle="--", alpha=0.5)

axes[1].hist(sentiment_df["sentiment_std"].to_numpy(), bins=50, edgecolor="white")
axes[1].set_xlabel("Sentiment Std Dev")
axes[1].set_title("Within-Filing Sentiment Dispersion")

axes[2].hist(
    sentiment_df["sentiment_pos_pct"].to_numpy(),
    bins=30,
    edgecolor="white",
    alpha=0.7,
    label="Positive %",
)
axes[2].hist(
    sentiment_df["sentiment_neg_pct"].to_numpy(),
    bins=30,
    edgecolor="white",
    alpha=0.7,
    label="Negative %",
)
axes[2].set_xlabel("Fraction of Chunks")
axes[2].set_title("Positive vs Negative Chunk Fractions")
axes[2].legend()

fig.tight_layout()
fig.show()

from sentence_transformers import SentenceTransformer

print(f"Loading embedding model: {EMBEDDING_MODEL}")
embed_model = SentenceTransformer(EMBEDDING_MODEL, device=str(device))
print(f"Embedding dimension: {embed_model.get_sentence_embedding_dimension()}")

def embed_document(text: str, model: SentenceTransformer) -> np.ndarray:
    """Compute document embedding by mean-pooling chunk embeddings."""
    chunks = chunk_text(text, max_chars=1200)
    if not chunks:
        return np.zeros(model.get_sentence_embedding_dimension())

    chunk_embeddings = model.encode(chunks, show_progress_bar=False, batch_size=32)
    return chunk_embeddings.mean(axis=0)

# Compute embeddings for all filings
print(f"Computing embeddings for {len(filings):,} filings...")

embeddings = []
for i, row in enumerate(filings.iter_rows(named=True)):
    emb = embed_document(row["text"], embed_model)
    embeddings.append(emb)

    if (i + 1) % 100 == 0 or (i + 1) == len(filings):
        print(f"  Embedded {i + 1:,}/{len(filings):,} filings")

embeddings_array = np.stack(embeddings)
print(f"Embedding matrix: {embeddings_array.shape}")

# Compute quarter-over-quarter narrative change (cosine distance)
# For each filing, compare its embedding to the previous quarter's filing for the same company

# Sort by symbol and filing date
filing_order = (
    filings.select(["symbol", "filing_date"]).with_row_index("idx").sort(["symbol", "filing_date"])
)

narrative_changes = []
prev_emb_by_symbol = {}

for row in filing_order.iter_rows(named=True):
    idx = row["idx"]
    symbol = row["symbol"]
    emb = embeddings_array[idx]

    if symbol in prev_emb_by_symbol:
        prev_emb = prev_emb_by_symbol[symbol]
        # Cosine distance (0 = identical, 2 = opposite)
        cos_sim = np.dot(emb, prev_emb) / (np.linalg.norm(emb) * np.linalg.norm(prev_emb) + 1e-8)
        cos_dist = 1.0 - cos_sim
    else:
        cos_dist = None  # No previous quarter

    narrative_changes.append(
        {
            "symbol": symbol,
            "filing_date": row["filing_date"],
            "narrative_change": cos_dist,
        }
    )
    prev_emb_by_symbol[symbol] = emb

narrative_df = pl.DataFrame(narrative_changes)
print(
    f"Narrative change computed for {narrative_df.drop_nulls('narrative_change').height:,} filings"
)
print("  (first filing per company has no prior quarter for comparison)")

narrative_df.drop_nulls("narrative_change")["narrative_change"].describe()

# Distribution of narrative change
valid_changes = narrative_df.drop_nulls("narrative_change")["narrative_change"].to_numpy()

fig, ax = plt.subplots(figsize=(8, 4))
ax.hist(valid_changes, bins=50, edgecolor="white")
ax.set_xlabel("Cosine Distance (Quarter-over-Quarter)")
ax.set_ylabel("Frequency")
ax.set_title("MD&A Narrative Change Distribution")
ax.axvline(
    np.median(valid_changes),
    color="red",
    linestyle="--",
    label=f"Median: {np.median(valid_changes):.3f}",
)
ax.legend()
fig.tight_layout()
fig.show()

# Merge sentiment and narrative change
signals = sentiment_df.join(
    narrative_df.select(["symbol", "filing_date", "narrative_change"]),
    on=["symbol", "filing_date"],
    how="left",
)
print(f"Combined signals: {signals.shape}")
signals.head(5)

# Load S&P 500 daily prices
from data import load_sp500_daily_bars

prices = load_sp500_daily_bars()
print(f"Loaded {len(prices):,} price observations for {prices['symbol'].n_unique()} symbols")

# Compute forward returns: return from day t to day t+N
price_returns = (
    prices.sort(["symbol", "timestamp"])
    .with_columns(
        fwd_1d=(pl.col("close").shift(-1) / pl.col("close") - 1).over("symbol"),
        fwd_5d=(pl.col("close").shift(-5) / pl.col("close") - 1).over("symbol"),
        fwd_20d=(pl.col("close").shift(-20) / pl.col("close") - 1).over("symbol"),
    )
    .select(["symbol", "timestamp", "fwd_1d", "fwd_5d", "fwd_20d"])
)

# Join signals to prices using asof join
# Match each filing_date to the next trading day on or after that date
# This is PIT-correct: signal is available when the filing is accepted

# Add trade_date column to prices so we can track which day was matched
prices_with_trade_date = price_returns.with_columns(trade_date=pl.col("timestamp")).sort(
    ["symbol", "timestamp"]
)

# Prepare signals: rename filing_date -> timestamp for the asof join key
signals_for_join = signals.rename({"filing_date": "timestamp"}).sort(["symbol", "timestamp"])

# Asof join: for each signal date, find the nearest price date >= signal date
# strategy="forward" means: match the next trading day on or after the signal date
eval_df = (
    signals_for_join.join_asof(
        prices_with_trade_date,
        on="timestamp",
        by="symbol",
        strategy="forward",
    )
    .rename({"timestamp": "filing_date"})
    .drop_nulls(["fwd_5d"])
)

print(f"Evaluation dataset: {len(eval_df):,} observations ({eval_df['symbol'].n_unique()} symbols)")
print(f"Date range: {eval_df['filing_date'].min()} to {eval_df['filing_date'].max()}")

from scipy.stats import spearmanr

signal_cols = [
    "sentiment_mean",
    "sentiment_std",
    "sentiment_pos_pct",
    "sentiment_neg_pct",
    "narrative_change",
]
return_cols = ["fwd_1d", "fwd_5d", "fwd_20d"]

# Compute pooled ICs and cluster-bootstrap inference (cluster = symbol).
print("Signal Evaluation: Pooled ICs with cluster bootstrap (cluster=symbol)")
print("=" * 70)

N_BOOT = 1000
# Single RNG shared across all (signal, horizon) iterations. Reproducibility
# of the table therefore depends on the iteration order of `signal_cols` ×
# `return_cols` — reorder either list and every subsequent pair gets a
# different bootstrap draw. If fewer than MIN_VALID_BOOT of N_BOOT replicates
# remain valid after NaN-drop, the percentile/p-value is suppressed (NaN) so
# unstable estimates from degenerate draws (e.g., a draw where one cluster
# dominates) do not leak into the table.
MIN_VALID_BOOT = 200
_boot_rng = np.random.default_rng(SEED)


def _pooled_spearman(sig_vals: np.ndarray, ret_vals: np.ndarray) -> float:
    """Pooled Spearman correlation, robust to constant-signal slices."""
    if len(sig_vals) < 5:
        return np.nan
    if np.std(sig_vals) == 0 or np.std(ret_vals) == 0:
        return np.nan
    return float(spearmanr(sig_vals, ret_vals)[0])


ic_results = []
for sig in signal_cols:
    for ret in return_cols:
        valid = eval_df.select(["symbol", sig, ret]).drop_nulls()
        if valid.height < 20:
            continue
        symbols = valid["symbol"].to_numpy()
        sig_vals = valid[sig].to_numpy()
        ret_vals = valid[ret].to_numpy()
        n = len(sig_vals)

        ic_point = _pooled_spearman(sig_vals, ret_vals)

        # Cluster bootstrap: resample whole symbols with replacement.
        unique_symbols = np.unique(symbols)
        idx_by_symbol = {s: np.where(symbols == s)[0] for s in unique_symbols}
        boot_ics = np.empty(N_BOOT)
        for b in range(N_BOOT):
            drawn = _boot_rng.choice(unique_symbols, size=len(unique_symbols), replace=True)
            indices = np.concatenate([idx_by_symbol[s] for s in drawn])
            boot_ics[b] = _pooled_spearman(sig_vals[indices], ret_vals[indices])
        boot_ics = boot_ics[~np.isnan(boot_ics)]

        if boot_ics.size >= MIN_VALID_BOOT:
            ci_lo = float(np.percentile(boot_ics, 2.5))
            ci_hi = float(np.percentile(boot_ics, 97.5))
            # Two-sided bootstrap p-value via percentile-method test inversion:
            # the smallest α such that the (1-α) bootstrap CI excludes zero.
            # Not a recentered-around-zero reflection test — for that, swap
            # in `mean(|boot - boot.mean()| >= |ic_point|)`.
            p_boot = 2.0 * min(
                float(np.mean(boot_ics <= 0.0)),
                float(np.mean(boot_ics >= 0.0)),
            )
        else:
            ci_lo = ci_hi = p_boot = np.nan

        ic_results.append(
            {
                "signal": sig,
                "horizon": ret,
                "ic": round(ic_point, 4) if not np.isnan(ic_point) else np.nan,
                "ci95_lo": round(ci_lo, 4) if not np.isnan(ci_lo) else np.nan,
                "ci95_hi": round(ci_hi, 4) if not np.isnan(ci_hi) else np.nan,
                "p_cluster_boot": round(p_boot, 4) if not np.isnan(p_boot) else np.nan,
                "n_obs": n,
                "n_symbols": int(len(unique_symbols)),
            }
        )

if ic_results:
    ic_summary = pl.DataFrame(ic_results).sort(["signal", "horizon"])
    print(ic_summary)
    print(
        "\nInference: cluster bootstrap with N_BOOT=1000 replicates; cluster = symbol.\n"
        "Bootstrap CIs and p-values supersede the i.i.d. t-stat formula because\n"
        "the pooled sample contains multiple filings per firm and overlapping\n"
        "forward-return windows. Treat ICs as screening; HAC on cross-sectional\n"
        "IC series (NB07, NB08) is the chapter's headline inference framework."
    )
else:
    print("Insufficient data for IC computation (need more symbols/filings)")

if len(ic_summary) > 0:
    signals = ic_summary["signal"].unique(maintain_order=True).to_list()
    horizons = ic_summary["horizon"].unique(maintain_order=True).to_list()
    x = np.arange(len(signals))
    width = 0.8 / max(len(horizons), 1)
    fig, ax = plt.subplots(figsize=(max(8.0, 1.6 * len(signals)), 5))
    for i, h in enumerate(horizons):
        by_sig = {
            r["signal"]: r for r in ic_summary.filter(pl.col("horizon") == h).iter_rows(named=True)
        }
        ic = np.array([by_sig.get(s, {}).get("ic", np.nan) for s in signals], dtype=float)
        lo = np.array([by_sig.get(s, {}).get("ci95_lo", np.nan) for s in signals], dtype=float)
        hi = np.array([by_sig.get(s, {}).get("ci95_hi", np.nan) for s in signals], dtype=float)
        yerr = np.nan_to_num(np.vstack([ic - lo, hi - ic]), nan=0.0)
        ax.bar(x + i * width, np.nan_to_num(ic), width, yerr=yerr, capsize=3, label=h)
    ax.axhline(0, color="black", linewidth=0.8)
    ax.set_xticks(x + width * (len(horizons) - 1) / 2)
    ax.set_xticklabels(signals, rotation=30, ha="right")
    ax.set_ylabel("Rank IC (pooled Spearman)")
    ax.set_title("Filing-text signal ICs with 95% cluster-bootstrap CIs")
    ax.legend(title="Horizon")
    fig.tight_layout()
    fig.show()
    ic_summary = pl.DataFrame()

def quintile_analysis(df: pl.DataFrame, signal_col: str, return_col: str) -> pl.DataFrame:
    """Sort into quintiles by signal, compute average return per quintile."""
    valid = df.drop_nulls([signal_col, return_col])
    if valid.height < 25:  # Need at least 5 per quintile
        return pl.DataFrame()

    # Assign quintiles using qcut on the signal column
    result = valid.with_columns(
        quintile=pl.col(signal_col).qcut(5, labels=["Q1", "Q2", "Q3", "Q4", "Q5"])
    )

    # Average return per quintile
    summary = (
        result.group_by("quintile")
        .agg(
            avg_return=pl.col(return_col).mean(),
            std_return=pl.col(return_col).std(),
            n_obs=pl.col(return_col).len(),
        )
        .sort("quintile")
    )

    return summary

# Quintile analysis for key signals
key_signals = ["sentiment_mean", "narrative_change"]
fig, axes = plt.subplots(1, len(key_signals), figsize=(6 * len(key_signals), 5))
if len(key_signals) == 1:
    axes = [axes]

for i, sig in enumerate(key_signals):
    ax = axes[i]
    q_df = quintile_analysis(eval_df, sig, "fwd_20d")

    if q_df.height > 0:
        quintiles = q_df["quintile"].to_list()
        returns = q_df["avg_return"].to_numpy()
        colors = ["#d32f2f" if r < 0 else "#2e7d32" for r in returns]

        ax.bar(range(len(quintiles)), returns * 100, color=colors, edgecolor="white")
        ax.set_xticks(range(len(quintiles)))
        ax.set_xticklabels(quintiles)
        ax.set_ylabel("Average 20-day Return (%)")
        ax.set_title(f"{sig}: Quintile Returns")
        ax.axhline(0, color="black", linewidth=0.5)

        # Long-short spread
        if len(returns) >= 2:
            spread = returns[-1] - returns[0]
            ax.text(
                0.95,
                0.95,
                f"Q5-Q1: {spread * 100:.2f}%",
                transform=ax.transAxes,
                ha="right",
                va="top",
                fontsize=10,
                bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
            )
    else:
        ax.text(0.5, 0.5, "Insufficient data", transform=ax.transAxes, ha="center")
        ax.set_title(f"{sig}: Quintile Returns")

fig.tight_layout()
fig.show()

# Save evaluation dataset
eval_df.write_parquet(OUTPUT_DIR / "filing_signals.parquet")
print(f"Saved {len(eval_df):,} signal observations to {OUTPUT_DIR / 'filing_signals.parquet'}")

# Save IC summary
if len(ic_summary) > 0:
    ic_summary.write_parquet(OUTPUT_DIR / "ic_summary.parquet")
    print(f"Saved IC summary to {OUTPUT_DIR / 'ic_summary.parquet'}")

print("\n" + "=" * 70)
print("NOTEBOOK COMPLETE: SEC Filing Text Signals")
print("=" * 70)
print(f"""
Signals computed:
  - sentiment_mean: FinBERT paragraph-level sentiment (mean across chunks)
  - sentiment_std: Within-filing sentiment dispersion
  - sentiment_pos_pct / sentiment_neg_pct: Fraction of positive/negative chunks
  - narrative_change: Cosine distance to prior quarter's MD&A embedding

Evaluation dataset: {len(eval_df):,} filing-date observations
Symbols: {eval_df["symbol"].n_unique()}
Date range: {eval_df["filing_date"].min()} to {eval_df["filing_date"].max()}
Output: {OUTPUT_DIR}
""")