# /// script
# requires-python = ">=3.12"
# dependencies = ["numpy==2.3.3", "pandas==2.3.2", "scipy==1.16.2", "statsmodels==0.14.5"]
# ///
"""Reproduce Report 002: CrUX ad metrics joined to crawler inventory signals.

    uv run scripts/research/report-002.py            # aggregate from the committed snapshot
    uv run scripts/research/report-002.py --refresh  # rebuild the snapshot from local source profiles first

Snapshot (internal): data/research/report-002-observations.json, one row per origin.
Outputs: apps/waitlist/src/content/reports/report-002.data.json (chart and text numbers),
apps/waitlist/public/research/report-002-summary.json (same aggregates, public) and
apps/waitlist/public/research/report-002-crux.json (CrUX fields only, public).
Every number in the article comes from report-002.data.json. Random draws use a fixed seed.
"""
import argparse, json, os
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from scipy import stats
from statsmodels.stats.multitest import multipletests

ROOT = Path(__file__).resolve().parents[2]
SP = ROOT / 'data/source-profiles'
SNAPSHOT = ROOT / 'data/research/report-002-observations.json'
DATA_OUT = ROOT / 'apps/waitlist/src/content/reports/report-002.data.json'
PUBLIC = ROOT / 'apps/waitlist/public/research'
SEED = 20261022
B = 1000
AD = {'experimental_ad_cpu': 'ad_cpu', 'experimental_ad_kilobytes': 'ad_kb', 'experimental_ad_count': 'ad_count', 'experimental_ad_density': 'ad_density'}
CWV = {'interaction_to_next_paint': 'inp', 'largest_contentful_paint': 'lcp', 'cumulative_layout_shift': 'cls'}
CRUX_PUBLIC = ['domain', 'origin', 'crux_last', *AD.values(), *CWV.values(), 'ff_phone']


def write(path, value, compact=False):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=None if compact else 2, separators=(',', ':') if compact else None, allow_nan=False) + '\n')


def load(path):
    return json.loads(path.read_text()) if path.exists() else None


def refresh():
    pt = load(ROOT / 'data/publisher-types/batches/20260925T163436Z/publisher-types.json')['byDomain']
    jev = {}
    for line in open(ROOT / 'data/classifications/jev-2026-09-20-v1/results.ndjson'):
        r = json.loads(line); jev[r['domain']] = r['answer']['choice']
    rows = []
    for name in sorted(os.listdir(SP)):
        if not name.startswith('crux-data-model-'):
            continue
        src = load(SP / name)
        recs = [ep['data']['record'] for ep in src['endpoints'] if ep.get('id') == 'queryRecord' and ep.get('status') == 200 and (ep.get('data') or {}).get('record')]
        if len(recs) != 1:
            if len(recs) > 1: raise ValueError('Multiple CrUX records for ' + src['domain'])
            continue
        rec, dom = recs[0], src['domain']
        m = rec['metrics']
        p75 = lambda k: (lambda v: None if v is None else float(v))((m.get(k) or {}).get('percentiles', {}).get('p75'))
        row = {'domain': dom, 'origin': rec['key'].get('origin'), 'crux_last': '%(year)04d-%(month)02d-%(day)02d' % rec['collectionPeriod']['lastDate'],
               **{v: p75(k) for k, v in {**AD, **CWV}.items()},
               'ff_phone': (m.get('form_factors') or {}).get('fractions', {}).get('phone')}
        pub = None
        prof = load(SP / f'data-model-{dom}.json')
        if prof:
            pub = next((ep['data'] for ep in prof['endpoints'] if ep.get('id') == 'publishers' and ep.get('status') == 200 and isinstance(ep.get('data'), dict)), None)
        row['crawler_match'] = pub is not None
        if pub:
            dl = pub.get('device_level_metrics') or {}
            ts = pub.get('tech_stack') or {}
            row.update({'traffic_rank': pub.get('traffic_rank'), 'acr': pub.get('avg_ads_to_content_ratio'),
                        'acr_mobile': (dl.get('mobile') or {}).get('avg_ads_to_content_ratio'), 'acr_desktop': (dl.get('desktop') or {}).get('avg_ads_to_content_ratio'),
                        'ads_in_view': pub.get('avg_ads_in_view'), 'ad_units': pub.get('total_unique_gpids'),
                        'refresh_slots_desktop': (dl.get('desktop') or {}).get('percentage_of_ad_slots_with_refresh'),
                        'refresh_slots_mobile': (dl.get('mobile') or {}).get('percentage_of_ad_slots_with_refresh'),
                        'supply_paths': pub.get('total_supply_paths'), 'resellers': pub.get('reseller_count'),
                        'wrappers': len(ts.get('wrappers') or []), 'identity_providers': len(ts.get('identity_providers') or []),
                        'video_players': len(ts.get('video_players') or []), 'ad_servers': len(ts.get('ad_servers') or [])})
        t = pt.get(dom) or {}
        row['pub_type'] = t.get('type') or 'unclassified'
        row['ad_monetized_publisher'] = jev.get(dom)
        aer = load(SP / f'ad-experience-data-model-{dom}.json')
        if aer:
            s = next((ep['data'].get('mobileSummary') for ep in aer['endpoints'] if ep.get('status') == 200 and isinstance(ep.get('data'), dict)), None)
            if s is not None:
                row['aer_mobile'] = s.get('betterAdsStatus') or 'NOT_REVIEWED'
                row['aer_mobile_changed'] = s.get('lastChangeTime')
        rows.append(row)
    if len(rows) != len({r['domain'] for r in rows}):
        raise ValueError('Duplicate domains')
    write(SNAPSHOT, {'schemaVersion': 1, 'builtAt': datetime.now().astimezone().isoformat(timespec='seconds'),
                     'cruxSelection': 'Origin-level queryRecord, all form factors, HTTP 200 with record present.',
                     'crawlerJoin': 'Exact domain match to the crawler publisher profile; HTTP 200 only.',
                     'observations': rows}, compact=True)
    print(f'Snapshot: {len(rows)} CrUX origins')


def r(x, nd=3):
    return None if x is None or (isinstance(x, float) and not np.isfinite(x)) else round(float(x), nd)


def spearman_ci(x, y, rng):
    x, y = np.asarray(x), np.asarray(y)
    rho = stats.spearmanr(x, y)[0]
    rx, ry = stats.rankdata(x), stats.rankdata(y)
    idx = rng.integers(0, len(x), (B, len(x)))
    boots = [stats.spearmanr(rx[i], ry[i])[0] for i in idx]
    lo, hi = np.percentile(boots, [2.5, 97.5])
    return rho, lo, hi, stats.spearmanr(x, y)[1]


def rank_beta(df, y, x, controls, fe=True):
    cols = [y, x, *controls]
    t = df[cols + (['pub_type'] if fe else [])].dropna().copy()
    for c in cols:
        rk = t[c].rank(); t[c] = (rk - rk.mean()) / rk.std()
    f = f'{y} ~ {x}' + ''.join(f' + {c}' for c in controls) + (' + C(pub_type)' if fe else '')
    m = smf.ols(f, t).fit(cov_type='HC3')
    lo, hi = m.conf_int().loc[x]
    return {'n': len(t), 'beta': r(m.params[x]), 'ci': [r(lo), r(hi)]}


def aggregate():
    rng = np.random.default_rng(SEED)
    snap = load(SNAPSHOT)
    d = pd.DataFrame(snap['observations'])
    a = d[d.ad_cpu.notna()].copy()
    s = a[a.crawler_match].copy()
    s['log_rank'] = np.log10(s.traffic_rank)
    out = {'reportId': 'report-002', 'seed': SEED, 'bootstrap': B, 'snapshotBuiltAt': snap['builtAt']}

    out['population'] = {
        'cruxOrigins': len(d), 'adMetricOrigins': len(a), 'adMetricShare': r(len(a) / len(d)),
        'perMetricN': {k: int(a[k].notna().sum()) for k in AD.values()},
        'crawlerMatched': len(s), 'collectionWindows': {k: int(v) for k, v in a.crux_last.value_counts().items()},
        'adMonetizedPublisherShareOfMatched': r((s.ad_monetized_publisher == 'yes').mean()),
        'pubTypeOfMatched': {k: int(v) for k, v in s.pub_type.value_counts().items()},
        'matchedVsUnmatchedMedian': {k: {'matched': r(s[k].median(), 1), 'unmatched': r(a[~a.crawler_match][k].median(), 1)} for k in [*AD.values(), 'inp']}}

    q = [.1, .25, .5, .75, .9]
    out['distribution'] = {k: {'n': int(a[k].notna().sum()), **{f'p{int(p*100)}': r(a[k].quantile(p), 2) for p in q}, 'max': r(a[k].max(), 2)} for k in AD.values()}
    hist = {}
    for k, edges in {'ad_cpu': np.geomspace(100, 100000, 31), 'ad_kb': np.geomspace(10, 100000, 31),
                     'ad_density': np.arange(0, 101, 2), 'ad_count': np.arange(0, 5.01, 0.25)}.items():
        v = a[k].dropna(); c, e = np.histogram(v.clip(edges[0], edges[-1]), bins=edges)
        hist[k] = {'edges': [r(x, 2) for x in e], 'counts': c.tolist()}
    out['histograms'] = hist
    out['densityAtLeast30'] = {'n': int((a.ad_density >= 30).sum()), 'of': int(a.ad_density.notna().sum()), 'share': r((a.ad_density >= 30).mean())}

    keys = [*AD.values(), *CWV.values()]
    mat = a[keys].corr('spearman')
    out['metricCorrelation'] = {'keys': keys, 'rho': [[r(mat.loc[i, j], 2) for j in keys] for i in keys]}

    stack = ['wrappers', 'identity_providers', 'ad_units', 'refresh_slots_desktop', 'refresh_slots_mobile', 'video_players',
             'ads_in_view', 'acr', 'supply_paths', 'resellers', 'log_rank']
    corr = []
    for x in AD.values():
        for y in stack:
            t = s[[x, y]].dropna()
            rho, lo, hi, p = spearman_ci(t[x], t[y], rng)
            ctr = [] if y == 'log_rank' else ['log_rank']
            corr.append({'metric': x, 'signal': y, 'n': len(t), 'rho': r(rho), 'ci': [r(lo), r(hi)], 'p': float(p),
                         'controlled': rank_beta(s, y, x, ctr)})
    qv = multipletests([c['p'] for c in corr], method='fdr_bh')[1]
    for c, v in zip(corr, qv):
        c['q'] = float(v); c['p'] = float(c['p'])
    out['signalCorrelations'] = corr

    out['byWrapperCount'] = []
    for k, g in s.groupby(s.wrappers.clip(upper=3)):
        out['byWrapperCount'].append({'wrappers': int(k), 'n': len(g), **{m: {'median': r(g[m].median(), 2), 'p25': r(g[m].quantile(.25), 2), 'p75': r(g[m].quantile(.75), 2)} for m in AD.values()}})

    cols = ['wrappers', 'identity_providers', 'ad_units', 'refresh_slots_desktop', 'video_players', 'ads_in_view', 'log_rank']
    t = s[['ad_cpu', *cols, 'pub_type']].dropna().copy()
    for c in ['ad_cpu', *cols]:
        rk = t[c].rank(); t[c] = (rk - rk.mean()) / rk.std()
    m = smf.ols('ad_cpu ~ ' + ' + '.join(cols) + ' + C(pub_type)', t).fit(cov_type='HC3')
    ci = m.conf_int()
    out['cpuModel'] = {'n': len(t), 'r2': r(m.rsquared), 'terms': [{'term': c, 'beta': r(m.params[c]), 'ci': [r(ci.loc[c, 0]), r(ci.loc[c, 1])]} for c in cols]}

    t = s[['ad_density', 'acr']].dropna().copy()
    t['bin'] = pd.qcut(t.acr, 10)
    out['densityVsAcr'] = {'n': len(t), 'deciles': [{'acrMedian': r(g.acr.median()), 'n': len(g), 'densityMedian': r(g.ad_density.median(), 1),
                                                    'p25': r(g.ad_density.quantile(.25), 1), 'p75': r(g.ad_density.quantile(.75), 1)} for _, g in t.groupby('bin', observed=True)]}
    for dev in ['acr_mobile', 'acr_desktop']:
        u = s[['ad_density', dev]].dropna(); rho, lo, hi, _ = spearman_ci(u.ad_density, u[dev], rng)
        out['densityVsAcr'][dev] = {'n': len(u), 'rho': r(rho), 'ci': [r(lo), r(hi)]}
    u = a[['ad_density', 'ff_phone']].dropna(); rho, lo, hi, _ = spearman_ci(u.ad_density, u.ff_phone, rng)
    out['densityVsPhoneShare'] = {'n': len(u), 'rho': r(rho), 'ci': [r(lo), r(hi)]}

    def profile(col):
        res = []
        s['_q'] = pd.qcut(s[col], 4, labels=['Q1', 'Q2', 'Q3', 'Q4'])
        for qn, g in s.groupby('_q', observed=True):
            res.append({'quartile': qn, 'n': len(g), 'range': [r(g[col].min(), 2), r(g[col].max(), 2)],
                        'wrapperShare': r((g.wrappers > 0).mean()), 'identityShare': r((g.identity_providers > 0).mean()),
                        'refreshShare': r((g.refresh_slots_desktop > 0).mean()), 'adUnitsMedian': r(g.ad_units.median(), 1),
                        'acrMedian': r(g.acr.median()), 'resellersMedian': r(g.resellers.median(), 1), 'rankMedian': r(g.traffic_rank.median(), 0),
                        'adMonetizedShare': r((g.ad_monetized_publisher == 'yes').mean())})
        return res
    out['quartileProfiles'] = {'ad_cpu': profile('ad_cpu'), 'ad_density': profile('ad_density')}

    a['_dec'] = pd.qcut(a.ad_cpu, 10)
    out['pageExperienceByCpuDecile'] = [{'cpuMedian': r(g.ad_cpu.median(), 0), 'n': len(g), 'inpMedian': r(g.inp.median(), 0),
                                         'inpPoorShare': r((g.inp > 200).mean()), 'lcpMedian': r(g.lcp.median(), 0), 'lcpPoorShare': r((g.lcp > 2500).mean()),
                                         'clsMedian': r(g.cls.median(), 3), 'clsPoorShare': r((g.cls > 0.1).mean())} for _, g in a.groupby('_dec', observed=True)]
    a['_q'] = pd.qcut(a.ad_cpu, 4, labels=['Q1', 'Q2', 'Q3', 'Q4'])
    out['pageExperienceByCpuQuartile'] = [{'quartile': qn, 'n': len(g), 'inpOver200Share': r((g.inp > 200).mean()), 'lcpOver2500Share': r((g.lcp > 2500).mean()),
                                           'clsOver01Share': r((g.cls > 0.1).mean())} for qn, g in a.groupby('_q', observed=True)]
    sub = a[['inp', 'ad_cpu', 'ff_phone', 'ad_kb', 'lcp', 'cls']].dropna()
    out['inpControl'] = {'subsetN': len(sub), 'raw': rank_beta(sub, 'inp', 'ad_cpu', [], fe=False), 'phoneControlled': rank_beta(sub, 'inp', 'ad_cpu', ['ff_phone'], fe=False)}
    rho, lo, hi, _ = spearman_ci(a.dropna(subset=['inp']).ad_cpu, a.dropna(subset=['inp']).inp, rng)
    out['inpControl']['fullSpearman'] = {'n': int(a.inp.notna().sum()), 'rho': r(rho), 'ci': [r(lo), r(hi)]}
    out['inpControl']['crawlerControlled'] = rank_beta(s, 'inp', 'ad_cpu', ['log_rank'])

    s['heavy'] = s.ad_cpu > s.ad_cpu.median(); s['dense'] = s.ad_density > s.ad_density.median()
    names = {(False, False): 'Light, sparse', (False, True): 'Light, dense', (True, False): 'Heavy, sparse', (True, True): 'Heavy, dense'}
    out['archetypeSplit'] = {'cpuMedian': r(s.ad_cpu.median(), 0), 'densityMedian': r(s.ad_density.median(), 0)}
    out['archetypes'] = [{'name': names[k], 'n': len(g), 'wrapperShare': r((g.wrappers > 0).mean()), 'identityShare': r((g.identity_providers > 0).mean()),
                          'refreshShare': r((g.refresh_slots_desktop > 0).mean()), 'acrMedian': r(g.acr.median()), 'resellersMedian': r(g.resellers.median(), 1),
                          'rankMedian': r(g.traffic_rank.median(), 0), 'inpMedian': r(g.inp.median(), 0), 'adMonetizedShare': r((g.ad_monetized_publisher == 'yes').mean())}
                         for k, g in s.groupby(['heavy', 'dense'])]
    u = s[['resellers', 'heavy', 'dense']].dropna()
    hd = u[u.heavy & u.dense].resellers; rest = u[~(u.heavy & u.dense)].resellers
    out['archetypeResellerTest'] = {'heavyDenseN': len(hd), 'otherN': len(rest), 'mannWhitneyP': float(stats.mannwhitneyu(hd, rest).pvalue)}

    rv = a[a.aer_mobile.notna()]
    rated = rv[rv.aer_mobile.isin(['PASSING', 'WARNING', 'FAILING'])]
    years = pd.to_datetime(rated.aer_mobile_changed, errors='coerce', utc=True).dt.year
    out['betterAds'] = {'checked': len(rv), 'status': {k: int(v) for k, v in rv.aer_mobile.value_counts().items()}, 'rated': len(rated),
                        'ratedDensityAtLeast30': int((rated.ad_density >= 30).sum()), 'ratedDensityAtLeast30Share': r((rated.ad_density >= 30).mean()),
                        'failingDensities': sorted(int(x) for x in rated[rated.aer_mobile == 'FAILING'].ad_density.dropna()),
                        'lastChangeYearMedian': int(years.median()), 'lastChangeBefore2020Share': r((years < 2020).mean())}

    rob = []
    subsets = {'All matched': s, 'Ad-monetized publishers only': s[s.ad_monetized_publisher == 'yes'], 'Window ending 2026-09-27': s[s.crux_last == '2026-09-27'],
               'Rank above median': s[s.traffic_rank <= s.traffic_rank.median()], 'Rank below median': s[s.traffic_rank > s.traffic_rank.median()]}
    pairs = [('ad_cpu', 'wrappers'), ('ad_cpu', 'refresh_slots_desktop'), ('ad_density', 'acr'), ('ad_density', 'wrappers'), ('ad_cpu', 'inp')]
    for name, sub in subsets.items():
        row = {'subset': name, 'n': len(sub)}
        for x, y in pairs:
            t = sub[[x, y]].dropna(); row[f'{x}~{y}'] = r(stats.spearmanr(t[x], t[y])[0], 2)
        rob.append(row)
    out['robustness'] = rob

    write(DATA_OUT, out)
    write(PUBLIC / 'report-002-summary.json', out)
    write(PUBLIC / 'report-002-crux.json', {'schemaVersion': 1, 'source': 'Chrome UX Report API queryRecord, origin level, all form factors.',
                                            'observations': json.loads(a[CRUX_PUBLIC].to_json(orient='records'))}, compact=True)
    print(f'Aggregated {len(a)} ad-metric origins, {len(s)} crawler-matched.')


if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--refresh', action='store_true')
    args = p.parse_args()
    if args.refresh:
        refresh()
    aggregate()