fithia2/apps/tools/wikimedia_low_attention_probe.py

"""Research probe for Wikimedia pageviews as a low-attention axis.

This is a standalone analysis tool. It does not modify the backtest engine.

The probe answers one narrow question:
Does filtering for low Wikimedia attention improve the quality of idle-alpha
events, relative to the same candidate set without the filter?

It reports:
- pageview spike distribution
- forward continuation statistics for low-attention vs high-attention buckets
- a simple threshold sweep over low-attention cutoffs

The script is intentionally lightweight and uses only Wikimedia pageviews plus
existing snapshot outcome columns.
"""

from __future__ import annotations

import argparse
import asyncio
import datetime as dt
import json
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd
import requests
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

from libs.common.config import get_settings

USER_AGENT = "codex-fithia2-wikimedia-low-attention-probe/1.0"
WIKI_SEARCH_URL = "https://en.wikipedia.org/w/api.php"
WIKI_PAGEVIEWS_URL = (
    "https://wikimedia.org/api/rest_v1/metrics/pageviews/per-article/"
    "en.wikipedia.org/all-access/all-agents/{article}/daily/{start}/{end}"
)

STOPWORDS = {
    "inc",
    "incorporated",
    "corp",
    "corporation",
    "ltd",
    "holdings",
    "group",
    "co",
    "company",
    "plc",
    "nv",
}

MANUAL_WIKI_TITLE = {
    "AMERICAN AIRLINES GROUP INC.": "American Airlines Group",
    "ASTRONICS CORPORATION": "Astronics",
    "CENTURY ALUMINUM COMPANY": "Century Aluminum",
    "CLEANSPARK, INC.": "CleanSpark",
    "CCC INTELLIGENT SOLUTIONS HOLDINGS INC.": "CCC Intelligent Solutions",
    "FLUENCE ENERGY, INC.": "Fluence Energy",
    "IMMUNITYBIO, INC.": "ImmunityBio",
    "LYFT, INC.": "Lyft",
    "MIRION TECHNOLOGIES, INC.": "Mirion Technologies",
    "MOSAIC CO": "The Mosaic Company",
    "NORWEGIAN CRUISE LINE HOLDINGS LTD.": "Norwegian Cruise Line Holdings",
    "PAR PACIFIC HOLDINGS, INC.": "Par Pacific Holdings",
    "PATTERSON-UTI ENERGY, INC.": "Patterson-UTI Energy",
    "RITHM CAPITAL CORP.": "Rithm Capital",
    "SOUNDHOUND AI, INC.": "SoundHound AI",
    "TANGO THERAPEUTICS, INC.": "Tango Therapeutics",
    "TERNS PHARMACEUTICALS, INC.": "Terns Pharmaceuticals",
    "UNITY SOFTWARE INC.": "Unity Technologies",
}


@dataclass(frozen=True)
class EventRow:
    event_id: str
    ticker: str
    issuer_name: str
    event_date: dt.date
    reaction_day_return: float
    fwd_return_3d: float
    fwd_return_5d: float


def _session() -> requests.Session:
    session = requests.Session()
    session.headers.update({"User-Agent": USER_AGENT})
    return session


def _clean_tokens(text_value: str) -> list[str]:
    tokens = re.findall(r"[A-Za-z0-9]+", text_value.lower().replace("\xa0", " "))
    return [token for token in tokens if token not in STOPWORDS]


def _candidate_names(name: str) -> list[str]:
    clean = name.replace("\xa0", " ").strip()
    manual = MANUAL_WIKI_TITLE.get(clean.upper())
    candidates = [candidate for candidate in [manual, clean] if candidate]
    token_name = " ".join(_clean_tokens(clean))
    if token_name:
        candidates.append(token_name)
    seen: set[str] = set()
    result: list[str] = []
    for candidate in candidates:
        stripped = candidate.strip(" ,.")
        if stripped and stripped not in seen:
            result.append(stripped)
            seen.add(stripped)
    return result


def _title_match_score(name: str, title: str) -> float:
    name_tokens = set(_clean_tokens(name))
    title_tokens = set(_clean_tokens(title))
    if not name_tokens or not title_tokens:
        return 0.0
    overlap = len(name_tokens & title_tokens)
    score = overlap / max(1, len(name_tokens))
    first_word = name.split()[0].lower() if name.split() else ""
    if first_word and title.lower().startswith(first_word):
        score += 0.1
    return score


def resolve_wikipedia_title(session: requests.Session, issuer_name: str) -> tuple[str | None, float]:
    best_score = 0.0
    best_title: str | None = None
    for candidate in _candidate_names(issuer_name):
        response = None
        for attempt in range(4):
            response = session.get(
                WIKI_SEARCH_URL,
                params={
                    "action": "query",
                    "list": "search",
                    "srsearch": candidate,
                    "format": "json",
                    "srlimit": 5,
                },
                timeout=20,
            )
            if response.status_code != 429:
                break
            time.sleep(1.5 * (attempt + 1))
        if response is None or response.status_code == 429:
            continue
        response.raise_for_status()
        hits = response.json().get("query", {}).get("search", [])
        for hit in hits:
            title = hit["title"]
            score = _title_match_score(candidate, title)
            if score > best_score:
                best_score = score
                best_title = title
    if best_score >= 0.55:
        return best_title, best_score
    return None, best_score


def fetch_pageview_spike(
    session: requests.Session,
    article_title: str,
    event_date: dt.date,
) -> dict[str, float] | None:
    start = (event_date - dt.timedelta(days=20)).strftime("%Y%m%d")
    end = (event_date + dt.timedelta(days=2)).strftime("%Y%m%d")
    response = None
    for attempt in range(4):
        response = session.get(
            WIKI_PAGEVIEWS_URL.format(
                article=article_title.replace(" ", "_"),
                start=start,
                end=end,
            ),
            timeout=20,
        )
        if response.status_code != 429:
            break
        time.sleep(1.5 * (attempt + 1))
    if response is None or response.status_code == 429:
        return None
    if response.status_code != 200:
        return None
    items = response.json().get("items", [])
    if len(items) < 8:
        return None
    views = pd.DataFrame(
        [(pd.to_datetime(item["timestamp"][:8]), item["views"]) for item in items],
        columns=["date", "views"],
    ).sort_values("date")
    event_ts = pd.Timestamp(event_date)
    pre_event = views.loc[views["date"] < event_ts, "views"]
    event_views = views.loc[views["date"] == event_ts, "views"]
    if pre_event.empty or event_views.empty:
        return None
    baseline = float(pre_event.tail(10).median())
    if baseline <= 0:
        return None
    event_value = float(event_views.iloc[0])
    return {
        "event_views": event_value,
        "baseline_views": baseline,
        "pageview_spike": event_value / baseline,
    }


async def load_event_rows(snapshot_path: Path, limit: int) -> list[EventRow]:
    snapshot = pd.read_parquet(snapshot_path)
    snapshot = snapshot[
        ["event_id", "event_date", "reaction_day_return", "fwd_return_3d", "fwd_return_5d"]
    ].copy()
    snapshot["event_date"] = pd.to_datetime(snapshot["event_date"]).dt.date

    engine = create_async_engine(get_settings().postgres_dsn)
    try:
        async with engine.connect() as conn:
            result = await conn.execute(
                text(
                    """
                    select e.event_id, e.event_type, sm.ticker, i.issuer_name
                    from events e
                    left join issuer_master i on i.issuer_id = e.issuer_id
                    left join symbol_master sm on sm.symbol_id = e.symbol_id
                    where e.event_id = any(:ids)
                    """
                ),
                {"ids": snapshot["event_id"].tolist()},
            )
            meta = pd.DataFrame(result.fetchall(), columns=result.keys())
    finally:
        await engine.dispose()

    merged = snapshot.merge(meta, on="event_id", how="left")
    merged = merged[merged["event_type"] == "earnings_release"].copy()
    generic = (
        (merged["issuer_name"].fillna("") == merged["ticker"].fillna("") + " Corporation")
        | (merged["issuer_name"].fillna("") == merged["ticker"].fillna("") + " Inc.")
        | (merged["issuer_name"].fillna("") == merged["ticker"].fillna("") + " Ltd.")
    )
    filtered = merged.loc[~generic & merged["issuer_name"].notna()].sort_values("event_date")
    if limit > 0:
        filtered = filtered.head(limit)
    rows: list[EventRow] = []
    for _, row in filtered.iterrows():
        rows.append(
            EventRow(
                event_id=str(row["event_id"]),
                ticker=str(row["ticker"]),
                issuer_name=str(row["issuer_name"]),
                event_date=row["event_date"],
                reaction_day_return=float(row["reaction_day_return"]),
                fwd_return_3d=float(row["fwd_return_3d"]),
                fwd_return_5d=float(row["fwd_return_5d"]),
            )
        )
    return rows


def _summarize(df: pd.DataFrame) -> dict[str, Any]:
    if df.empty:
        return {"count": 0}
    low = df[df["pageview_spike"] <= df["pageview_spike"].median()]
    high = df[df["pageview_spike"] > df["pageview_spike"].median()]
    return {
        "count": int(len(df)),
        "median_spike": round(float(df["pageview_spike"].median()), 4),
        "mean_3d_low": round(float(low["signed_cont_3d"].mean()), 4) if not low.empty else None,
        "mean_3d_high": round(float(high["signed_cont_3d"].mean()), 4) if not high.empty else None,
        "mean_5d_low": round(float(low["signed_cont_5d"].mean()), 4) if not low.empty else None,
        "mean_5d_high": round(float(high["signed_cont_5d"].mean()), 4) if not high.empty else None,
        "corr_3d": round(float(df["pageview_spike"].corr(df["signed_cont_3d"])), 4),
        "corr_5d": round(float(df["pageview_spike"].corr(df["signed_cont_5d"])), 4),
    }


def _threshold_sweep(df: pd.DataFrame, thresholds: list[float]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for thr in thresholds:
        subset = df[df["pageview_spike"] <= thr]
        if subset.empty:
            continue
        rows.append(
            {
                "threshold": thr,
                "count": int(len(subset)),
                "mean_3d": round(float(subset["signed_cont_3d"].mean()), 4),
                "mean_5d": round(float(subset["signed_cont_5d"].mean()), 4),
                "median_3d": round(float(subset["signed_cont_3d"].median()), 4),
                "median_5d": round(float(subset["signed_cont_5d"].median()), 4),
                "mean_spike": round(float(subset["pageview_spike"].mean()), 4),
            }
        )
    return rows


async def main() -> None:
    parser = argparse.ArgumentParser(description="Probe Wikimedia low-attention axis for idle-alpha candidates.")
    parser.add_argument("--snapshot", default="data/datasets/snapshots/midcap-filtered/test.parquet")
    parser.add_argument("--limit", type=int, default=120, help="Number of events to sample before joining pageviews.")
    parser.add_argument("--thresholds", default="1.0,1.2,1.5,2.0,3.0", help="Comma-separated low-attention thresholds.")
    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
    args = parser.parse_args()

    rows = await load_event_rows(Path(args.snapshot), args.limit)
    session = _session()
    title_cache: dict[str, tuple[str | None, float]] = {}
    pageview_cache: dict[tuple[str, dt.date], dict[str, float] | None] = {}
    resolved: list[dict[str, Any]] = []
    for row in rows:
        if row.issuer_name not in title_cache:
            title_cache[row.issuer_name] = resolve_wikipedia_title(session, row.issuer_name)
        title, match_score = title_cache[row.issuer_name]
        if not title:
            continue
        cache_key = (title, row.event_date)
        if cache_key not in pageview_cache:
            pageview_cache[cache_key] = fetch_pageview_spike(session, title, row.event_date)
        pageviews = pageview_cache[cache_key]
        if not pageviews:
            continue
        signed_cont_3d = (1.0 if row.reaction_day_return >= 0 else -1.0) * row.fwd_return_3d
        signed_cont_5d = (1.0 if row.reaction_day_return >= 0 else -1.0) * row.fwd_return_5d
        resolved.append(
            {
                "ticker": row.ticker,
                "issuer_name": row.issuer_name,
                "article_title": title,
                "match_score": round(match_score, 3),
                "event_date": row.event_date.isoformat(),
                "reaction_day_return": row.reaction_day_return,
                "fwd_return_3d": row.fwd_return_3d,
                "fwd_return_5d": row.fwd_return_5d,
                "signed_cont_3d": signed_cont_3d,
                "signed_cont_5d": signed_cont_5d,
                **pageviews,
            }
        )

    df = pd.DataFrame(resolved)
    if df.empty:
        payload = {"resolved_rows": 0}
        if args.json:
            print(json.dumps(payload, ensure_ascii=False, indent=2))
        else:
            print("resolved_rows=0")
        return

    summary = _summarize(df)
    thresholds = [float(item) for item in args.thresholds.split(",") if item.strip()]
    sweep = _threshold_sweep(df, thresholds)
    result = {
        "resolved_rows": int(len(df)),
        "summary": summary,
        "threshold_sweep": sweep,
        "sample_rows": resolved[:10],
    }

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return

    print(f"resolved_rows={result['resolved_rows']}")
    print(f"median_pageview_spike={summary['median_spike']:.3f}")
    print(f"corr(pageview_spike,signed_cont_3d)={summary['corr_3d']:.4f}")
    print(f"corr(pageview_spike,signed_cont_5d)={summary['corr_5d']:.4f}")
    print("threshold_sweep:")
    for row in sweep:
        print(
            f"  <= {row['threshold']:.2f}: n={row['count']} mean3d={row['mean_3d']:.4f} "
            f"mean5d={row['mean_5d']:.4f} median3d={row['median_3d']:.4f} "
            f"median5d={row['median_5d']:.4f}"
        )


if __name__ == "__main__":
    asyncio.run(main())