import itertools
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd

from analyze import ols

ROOT = Path(__file__).parent
SEED = 20261017
rng = np.random.default_rng(SEED)
d = pd.read_json(ROOT / 'analysis_data.json')
categories = {
    '技术作品/工具效果展示': [7, 9, 20, 24, 25, 27],
    'Vibe Coding指南': [11, 15, 16],
    '模型动态/观点': [5, 6, 8, 17, 18],
    '安全/隐私事件': [14, 19, 21],
    '其他': [10, 12, 13, 22, 23, 26, 28],
}
mapping = {row: group for group, rows in categories.items() for row in rows}
d['category'] = d.source_row.map(mapping)
d['october'] = d['首次发布时间'].str.contains('10月').astype(int)
assert d.category.notna().all() and len(d) == 24


def holm(ps):
    ps = np.asarray(ps)
    order = np.argsort(ps)
    corrected = np.empty_like(ps)
    running = 0.0
    for i, index in enumerate(order):
        running = max(running, min(1.0, (len(ps) - i) * ps[index]))
        corrected[index] = running
    return corrected.tolist()


def rank_test(a, b):
    a, b = np.asarray(a), np.asarray(b)
    na, nb = len(a), len(b)
    rank2 = np.rint(pd.Series(np.r_[a, b]).rank().to_numpy() * 2).astype(int)
    observed = int(rank2[:na].sum())
    center = na * (na + nb + 1)
    distance = abs(observed - center)
    hits, total = 0, 0
    # Every allocation of the observed ranks is equally likely under exchangeability.
    for selection in itertools.combinations(rank2.tolist(), na):
        hits += abs(sum(selection) - center) >= distance
        total += 1
    u = observed / 2 - na * (na + 1) / 2
    superiority = u / (na * nb)
    return {'n_a': na, 'n_b': nb, 'U': u, 'p_exact': hits / total,
            'allocations': total, 'superiority_probability': superiority,
            'cliff_delta': 2 * superiority - 1,
            'median_a': float(np.median(a)), 'median_b': float(np.median(b))}


def boot_effect(a, b, draws=50000):
    a, b = np.asarray(a), np.asarray(b)
    aa = a[rng.integers(0, len(a), size=(draws, len(a)))]
    bb = b[rng.integers(0, len(b), size=(draws, len(b)))]
    med_ratio = np.median(aa, axis=1) / np.median(bb, axis=1)
    delta = (aa[:, :, None] > bb[:, None, :]).mean(axis=(1, 2)) - (aa[:, :, None] < bb[:, None, :]).mean(axis=(1, 2))
    return {'median_ratio': float(np.median(a) / np.median(b)),
            'median_ratio_ci95_percentile': np.quantile(med_ratio, [.025, .975]).tolist(),
            'cliff_delta_ci95_percentile': np.quantile(delta, [.025, .975]).tolist(),
            'bootstrap_draws': draws}


def kw_permutation(sub, outcome, draws=200000):
    groups = [g for _, g in sub.groupby('category', sort=False)]
    sizes = np.array([len(g) for g in groups])
    vals = np.concatenate([g[outcome].to_numpy() for g in groups])
    ranks = pd.Series(vals).rank().to_numpy()
    n = len(ranks)
    _, counts = np.unique(vals, return_counts=True)
    correction = 1 - np.sum(counts**3 - counts) / (n**3 - n)
    starts = np.r_[0, np.cumsum(sizes)[:-1]]

    def statistic(r):
        sums = np.add.reduceat(r, starts, axis=-1)
        return (12 / (n * (n + 1)) * np.sum(sums**2 / sizes, axis=-1) - 3 * (n + 1)) / correction

    h = float(statistic(ranks))
    hits = 0
    for begin in range(0, draws, 5000):
        batch = min(5000, draws - begin)
        shuffled = ranks[np.argsort(rng.random((batch, n)), axis=1)]
        hits += int(np.sum(statistic(shuffled) >= h - 1e-10))
    p = (hits + 1) / (draws + 1)
    return {'n': n, 'H': h, 'p_permutation': p, 'draws': draws,
            'monte_carlo_se': math.sqrt(p * (1 - p) / (draws + 1)),
            'epsilon_squared': (h - len(groups) + 1) / (n - len(groups))}


def loo_r2(y, x):
    y, x = np.asarray(y), np.asarray(x)
    X = np.column_stack([np.ones(len(y)), x])
    beta = np.linalg.lstsq(X, y, rcond=None)[0]
    e = y - X @ beta
    h = np.sum((X @ np.linalg.inv(X.T @ X)) * X, axis=1)
    press = np.sum((e / (1 - h))**2)
    return float(1 - press / np.sum((y - y.mean())**2))


result = {'seed': SEED, 'n': len(d), 'genre': {}, 'topic_omnibus': {},
          'topic_pairwise_views': [], 'regressions': {}, 'topic_summary': []}
genre_tests = []
for outcome in ['观看量', '曝光', '封面点击率']:
    test = rank_test(d[d.video == 1][outcome], d[d.video == 0][outcome])
    result['genre'][outcome] = test
    genre_tests.append(test['p_exact'])
for outcome, adjusted in zip(result['genre'], holm(genre_tests)):
    result['genre'][outcome]['p_holm_3'] = adjusted
result['genre']['观看量'].update(boot_effect(d[d.video == 1]['观看量'], d[d.video == 0]['观看量']))

for subset_name, sub in [('all', d), ('without_top', d[d.source_row != 9]), ('september', d[d.october == 0])]:
    if subset_name != 'all':
        result['genre'][subset_name + '_views'] = rank_test(sub[sub.video == 1]['观看量'], sub[sub.video == 0]['观看量'])
    for outcome in ['观看量', '曝光']:
        covariates = ['video', 'october'] if subset_name != 'september' else ['video']
        fit = ols(np.log(sub[outcome]), sub[covariates], covariates)
        fit['loo_r2_log_target'] = loo_r2(np.log(sub[outcome]), sub[covariates])
        result['regressions'][f'{subset_name}_{outcome}_format_cohort'] = fit

for outcome in ['观看量', '曝光', '封面点击率']:
    result['topic_omnibus'][outcome] = kw_permutation(d, outcome)
for outcome, adjusted in zip(result['topic_omnibus'], holm([x['p_permutation'] for x in result['topic_omnibus'].values()])):
    result['topic_omnibus'][outcome]['p_holm_3'] = adjusted
result['topic_omnibus']['without_top_views'] = kw_permutation(d[d.source_row != 9], '观看量')

for (ka, ra), (kb, rb) in itertools.combinations(categories.items(), 2):
    aa, bb = d[d.source_row.isin(ra)]['观看量'], d[d.source_row.isin(rb)]['观看量']
    test = rank_test(aa, bb)
    test.update({'a': ka, 'b': kb})
    result['topic_pairwise_views'].append(test)
for test, adjusted in zip(result['topic_pairwise_views'], holm([x['p_exact'] for x in result['topic_pairwise_views']])):
    test['p_holm_10'] = adjusted

trim_pairs = []
for (ka, ra), (kb, rb) in itertools.combinations(categories.items(), 2):
    aa = d[d.source_row.isin(ra) & (d.source_row != 9)]['观看量']
    bb = d[d.source_row.isin(rb) & (d.source_row != 9)]['观看量']
    test = rank_test(aa, bb)
    test.update({'a': ka, 'b': kb})
    trim_pairs.append(test)
for test, adjusted in zip(trim_pairs, holm([x['p_exact'] for x in trim_pairs])):
    test['p_holm_10'] = adjusted
result['topic_pairwise_views_without_top'] = trim_pairs
result['within_video_demo_vs_other'] = rank_test(
    d[(d.video == 1) & (d.category == '技术作品/工具效果展示')]['观看量'],
    d[(d.video == 1) & (d.category != '技术作品/工具效果展示')]['观看量'])

for name, rows in categories.items():
    sub = d[d.source_row.isin(rows)]
    result['topic_summary'].append({'group': name, 'n': len(sub), 'video_n': int(sub.video.sum()),
        'median_views': float(sub['观看量'].median()), 'median_impressions': float(sub['曝光'].median()),
        'median_ctr': float(sub['封面点击率'].median())})

# Conditional exact permutation within month cohorts, using the adjusted regression coefficient.
y = np.log(d['观看量'].to_numpy())
video = d.video.to_numpy()
month = d.october.to_numpy()
indices = [np.flatnonzero(month == value) for value in [0, 1]]
counts = [int(video[ix].sum()) for ix in indices]
centered_y = y.copy()
for ix in indices:
    centered_y[ix] -= y[ix].mean()
observed = float(centered_y @ video)
hits, total = 0, 0
for a in itertools.combinations(indices[0].tolist(), counts[0]):
    sa = centered_y[list(a)].sum()
    for b in itertools.combinations(indices[1].tolist(), counts[1]):
        value = sa + centered_y[list(b)].sum()
        hits += abs(value) >= abs(observed) - 1e-12
        total += 1
result['genre']['views_month_stratified_permutation'] = {'p_exact': hits / total, 'allocations': total}

(ROOT / 'content_statistics_results.json').write_text(json.dumps(result, ensure_ascii=False, indent=2))
print('CONTENT_STATISTICS', json.dumps(result, ensure_ascii=False, indent=2))