#!/usr/bin/env python3
"""Step 3 of the ai-traffic-2026 post: from the website-traffic rows (out/ai-traffic/raw.json), keep the first 1,000
Tranco domains that have a Similarweb traffic estimate and compute every figure the post uses.

Writes:
  out/ai-traffic/sample.json   one compact row per sample domain (the page and the tables are built from this)
  out/ai-traffic/results.json  every aggregate quoted in the post
  out/ai-traffic/tables.md     the same figures as Markdown tables, for checking by eye

Definitions (all from Similarweb's estimates for the latest published month, as returned by the Actor):
  share of a channel     trafficSources.<channel>, a fraction of the site's visits
  AI chatbot visits      monthlyVisits x trafficSources.aiChatbots (same basis as every other channel)
  AI panel visits        aiChatbotVisits, Similarweb's separate AI-referral figure (a cross-check only)
  top chatbot            the first entry of topAiChatbots (Similarweb ranks the chatbots; it gives no per-chatbot counts)
  pooled share           sum of channel visits / sum of all visits, so big sites weigh more
Run: python3 distribution/analysis/ai-traffic-2026-analyse.py
"""
import json
import statistics as st
from collections import Counter, defaultdict
from pathlib import Path

HERE = Path(__file__).resolve().parent
OUT = HERE / 'out' / 'ai-traffic'
SAMPLE = 1000
MIN_VISITS_FOR_SHARE_RANKING = 1_000_000  # rankings by share skip tiny device/API hosts, whose shares are noise
MIN_SITES_PER_CATEGORY = 15
CHANNELS = ['direct', 'organicSearch', 'paidSearch', 'referrals', 'socialOrganic', 'socialPaid', 'email', 'displayAds', 'affiliates', 'aiChatbots']
LABELS = {'direct': 'Direct', 'organicSearch': 'Organic search', 'paidSearch': 'Paid search', 'referrals': 'Referrals (links)',
          'socialOrganic': 'Social', 'socialPaid': 'Paid social', 'email': 'Email', 'displayAds': 'Display ads',
          'affiliates': 'Affiliates', 'aiChatbots': 'AI chatbots'}
BOTS = {'chatgpt.com': 'ChatGPT', 'gemini.google.com': 'Gemini', 'perplexity.ai': 'Perplexity', 'claude.ai': 'Claude',
        'copilot.microsoft.com': 'Copilot', 'grok.com': 'Grok', 'chat.deepseek.com': 'DeepSeek'}

cand = json.loads((OUT / 'tranco-candidates.json').read_text())
raw = json.loads((OUT / 'raw.json').read_text())
by_domain = {r['domain']: r for r in raw['items']}

rows, no_estimate, missing = [], [], []
for c in cand['domains']:
    r = by_domain.get(c['domain'])
    if r is None:
        missing.append(c['domain'])
        continue
    if r.get('trafficAvailable') is not True:
        no_estimate.append(c['domain'])
        continue
    if len(rows) < SAMPLE:
        ts = r.get('trafficSources') or {}
        rows.append({
            'rank': c['rank'], 'domain': r['domain'], 'visits': r['monthlyVisits'],
            'category': r.get('category'), 'top': (r.get('category') or 'Unknown').split(' > ')[0],
            'sources': {k: ts.get(k) for k in CHANNELS} if ts else None,
            'aiShare': ts.get('aiChatbots') or 0,
            'aiPanelVisits': r.get('aiChatbotVisits'), 'aiPanelShare': r.get('aiChatbotShare'),
            'bots': r.get('topAiChatbots') or [],
        })
assert len(rows) == SAMPLE, len(rows)
last_rank = rows[-1]['rank']
no_estimate_in_range = [d for d in no_estimate if next(c['rank'] for c in cand['domains'] if c['domain'] == d) <= last_rank]
for r in rows:
    r['aiVisits'] = round(r['visits'] * r['aiShare'])

pct = lambda a, b: round(100 * a / b, 2) if b else None  # noqa: E731
total_visits = sum(r['visits'] for r in rows)
with_split = [r for r in rows if r['sources']]
split_visits = sum(r['visits'] for r in with_split)
channel_visits = {k: sum(r['visits'] * (r['sources'][k] or 0) for r in with_split) for k in CHANNELS}
channels = sorted(({'channel': k, 'label': LABELS[k], 'visits': round(v), 'pooledPct': pct(v, split_visits),
                    'medianPct': round(100 * st.median([(r['sources'][k] or 0) for r in with_split]), 2)} for k, v in channel_visits.items()),
                  key=lambda x: -x['visits'])
ai_total = channel_visits['aiChatbots']
with_ai = [r for r in rows if r['aiShare'] > 0]
with_bots = [r for r in rows if r['bots']]
big = [r for r in rows if r['visits'] >= MIN_VISITS_FOR_SHARE_RANKING and r['sources']]

def site(r):
    return {'rank': r['rank'], 'domain': r['domain'], 'visits': r['visits'], 'category': r['category'],
            'aiSharePct': round(100 * r['aiShare'], 2), 'aiVisits': r['aiVisits'],
            'organicPct': round(100 * (r['sources']['organicSearch'] or 0), 2) if r['sources'] else None,
            'topBot': BOTS.get(r['bots'][0], r['bots'][0]) if r['bots'] else None}

# AI companies' and AI-branded sites are listed on their own: a chatbot linking to its maker is not the question here.
AI_OWN = {'openai.com', 'chatgpt.com', 'anthropic.com', 'claude.ai', 'perplexity.ai', 'deepseek.com', 'huggingface.co', 'cursor.sh', 'chatgpt.site'}
is_ai_own = lambda r: r['domain'] in AI_OWN or r['top'] == 'Ai Chatbots and Tools'  # noqa: E731
top_by_share_all = sorted(big, key=lambda r: (-r['aiShare'], r['rank']))[:20]
top_by_share = sorted((r for r in big if not is_ai_own(r)), key=lambda r: (-r['aiShare'], r['rank']))[:20]
top_by_visits = sorted(rows, key=lambda r: (-r['aiVisits'], r['rank']))[:20]

# Which chatbot Similarweb ranks first, overall and by category; and how often each appears at all.
first = Counter(r['bots'][0] for r in with_bots)
present = Counter(b for r in with_bots for b in r['bots'])
first_weighted = Counter()
for r in with_bots:
    first_weighted[r['bots'][0]] += r['aiVisits']
bots = [{'bot': BOTS.get(b, b), 'host': b, 'firstSites': first[b], 'firstPct': pct(first[b], len(with_bots)),
         'listedSites': present[b], 'listedPct': pct(present[b], len(with_bots)),
         'aiVisitsAtSitesWhereFirstPct': pct(first_weighted[b], sum(first_weighted.values()))}
        for b in sorted(present, key=lambda b: (-first[b], -present[b]))]
# Where ChatGPT is not first: which chatbot is, and on which sites (largest first)
not_chatgpt = sorted((r for r in with_bots if r['bots'][0] != 'chatgpt.com'), key=lambda r: -r['visits'])
# Rank of Claude vs Perplexity vs Gemini when both listed
def ahead(a, b):
    both = [r for r in with_bots if a in r['bots'] and b in r['bots']]
    win = sum(1 for r in both if r['bots'].index(a) < r['bots'].index(b))
    return {'a': BOTS[a], 'b': BOTS[b], 'both': len(both), 'aAhead': win, 'aAheadPct': pct(win, len(both))}
pairs = [ahead('gemini.google.com', 'perplexity.ai'), ahead('claude.ai', 'perplexity.ai'), ahead('gemini.google.com', 'claude.ai'),
         ahead('chatgpt.com', 'gemini.google.com')]

cats = defaultdict(list)
for r in with_split:
    cats[r['top']].append(r)
overall_ai = ai_total / split_visits
categories = []
for name, rs in cats.items():
    if len(rs) < MIN_SITES_PER_CATEGORY or name == 'Unknown':
        continue
    v = sum(r['visits'] for r in rs)
    ai = sum(r['visits'] * (r['sources']['aiChatbots'] or 0) for r in rs)
    org = sum(r['visits'] * (r['sources']['organicSearch'] or 0) for r in rs)
    soc = sum(r['visits'] * (r['sources']['socialOrganic'] or 0) for r in rs)
    rb = [r for r in rs if r['bots']]
    fc = Counter(r['bots'][0] for r in rb)
    categories.append({
        'category': name, 'sites': len(rs), 'visits': v, 'aiPooledPct': pct(ai, v),
        'aiMedianPct': round(100 * st.median([r['sources']['aiChatbots'] or 0 for r in rs]), 2),
        'index': round((ai / v) / overall_ai, 2), 'organicPooledPct': pct(org, v), 'socialPooledPct': pct(soc, v),
        'aiPerOrganic': round(ai / org, 4) if org else None,
        'sitesWithBots': len(rb), 'topBot': {BOTS.get(b, b): n for b, n in fc.most_common()},
        'chatgptFirstPct': pct(fc['chatgpt.com'], len(rb)),
    })
categories.sort(key=lambda c: -c['aiPooledPct'])

# AI vs organic search and social, site by site
ai_beats_social = sum(1 for r in with_split if (r['sources']['aiChatbots'] or 0) > (r['sources']['socialOrganic'] or 0))
ai_beats_organic = [r for r in with_split if (r['sources']['aiChatbots'] or 0) > (r['sources']['organicSearch'] or 0)]
ai_beats_email = sum(1 for r in with_split if (r['sources']['aiChatbots'] or 0) > (r['sources']['email'] or 0))
ai_beats_paid = sum(1 for r in with_split if (r['sources']['aiChatbots'] or 0) > (r['sources']['paidSearch'] or 0))
ratio_org = [(r['sources']['organicSearch'] or 0) / r['aiShare'] for r in with_split if r['aiShare'] > 0]
dist = Counter()
for r in with_split:
    s = 100 * r['aiShare']
    band = '0' if s == 0 else '<0.5' if s < 0.5 else '0.5-1' if s < 1 else '1-2' if s < 2 else '2-5' if s < 5 else '5+'
    dist[band] += 1

# Cross-check against Similarweb's separate AI-referral panel
panel = [r for r in rows if r['aiPanelVisits']]
panel_total = sum(r['aiPanelVisits'] for r in panel)
panel_vs_split = st.median([r['aiPanelVisits'] / r['aiVisits'] for r in panel if r['aiVisits']])

# AI companies' own sites (how much of their traffic comes from other chatbots) — reported, not ranked
ai_sites = [site(r) for r in rows if is_ai_own(r)]

res = {
    'source': {'tranco': {'listId': cand['listId'], 'listDate': cand.get('listDate')}, 'actor': raw['actor'],
               'runs': raw['runs'], 'apifyUsageUsd': round(raw['usageTotalUsd'], 3),
               'trafficMonth': Counter(by_domain[r['domain']].get('trafficMonth') for r in rows).most_common(1)[0][0],
               'data': 'Similarweb estimates'},
    'sample': {'candidates': len(cand['domains']), 'excludedInfra': len(cand['excluded']), 'returned': len(raw['items']),
               'notReturned': missing, 'noEstimate': len(no_estimate), 'noEstimateWithinSampleRange': len(no_estimate_in_range),
               'domains': len(rows), 'lastTrancoRank': last_rank, 'withChannelSplit': len(with_split),
               'withAiShare': len(with_ai), 'withChatbotRanking': len(with_bots), 'withAiPanel': len(panel),
               'bigSites': len(big), 'minVisitsForShareRanking': MIN_VISITS_FOR_SHARE_RANKING},
    'totals': {'visits': total_visits, 'visitsWithSplit': split_visits, 'aiVisits': round(ai_total),
               'aiPooledPct': pct(ai_total, split_visits),
               'aiMedianPct': round(100 * st.median([r['aiShare'] for r in with_split]), 3),
               'aiMeanPct': round(100 * st.mean([r['aiShare'] for r in with_split]), 3),
               'aiPanelVisits': panel_total, 'panelToSplitMedianRatio': round(panel_vs_split, 3)},
    'channels': channels,
    'aiVsOther': {'organicPerAi': round(channel_visits['organicSearch'] / ai_total, 1),
                  'socialPerAi': round(channel_visits['socialOrganic'] / ai_total, 2),
                  'medianSiteOrganicPerAi': round(st.median(ratio_org), 1),
                  'aiBeatsSocialSites': ai_beats_social, 'aiBeatsOrganicSites': len(ai_beats_organic),
                  'aiBeatsOrganic': [site(r) for r in sorted(ai_beats_organic, key=lambda r: -r['visits'])],
                  'aiBeatsEmailSites': ai_beats_email, 'aiBeatsPaidSearchSites': ai_beats_paid},
    'distribution': {k: dist[k] for k in ['0', '<0.5', '0.5-1', '1-2', '2-5', '5+']},
    'topByShare': [site(r) for r in top_by_share],
    'topByShareInclAiOwn': [site(r) for r in top_by_share_all],
    'uncategorised': len(cats.get('Unknown', [])),
    'googleShareOfAiPct': pct(next(r['aiVisits'] for r in rows if r['domain'] == 'google.com'), ai_total),
    'topByVisits': [site(r) for r in top_by_visits],
    'top20VisitsShareOfAi': pct(sum(r['aiVisits'] for r in top_by_visits), ai_total),
    'bots': bots, 'pairs': pairs,
    'notChatgptFirst': [{**site(r), 'bots': [BOTS.get(b, b) for b in r['bots'][:3]]} for r in not_chatgpt],
    'categories': categories,
    'aiCompanySites': ai_sites,
}
(OUT / 'results.json').write_text(json.dumps(res, indent=1, ensure_ascii=False))
(OUT / 'sample.json').write_text(json.dumps({'trafficMonth': res['source']['trafficMonth'], 'tranco': cand['listId'], 'rows': [
    {k: r[k] for k in ('rank', 'domain', 'visits', 'category', 'aiShare', 'aiVisits', 'bots', 'sources')} for r in rows]}, separators=(',', ':')))

# ---------------------------------------------------------------- tables.md (for checking by eye)
f = lambda n: f'{n:,.0f}'  # noqa: E731
md = [f"# ai-traffic-2026 tables (Tranco {cand['listId']}, Similarweb month {res['source']['trafficMonth']})", '',
      json.dumps(res['sample']), '', json.dumps(res['totals']), '', json.dumps(res['aiVsOther'] | {'aiBeatsOrganic': [s['domain'] for s in res['aiVsOther']['aiBeatsOrganic']]}), '',
      '## Channels', '| Channel | Visits | Pooled % | Median % |', '|---|---:|---:|---:|']
md += [f"| {c['label']} | {f(c['visits'])} | {c['pooledPct']} | {c['medianPct']} |" for c in channels]
md += ['', '## Distribution of AI share', json.dumps(res['distribution']), '', "## Top 20 by AI share (>= 1M visits, AI companies' own sites excluded)", '| # | Domain | Visits | AI % | AI visits | Organic % | Top bot | Category |', '|---:|---|---:|---:|---:|---:|---|---|']
md += [f"| {s['rank']} | {s['domain']} | {f(s['visits'])} | {s['aiSharePct']} | {f(s['aiVisits'])} | {s['organicPct']} | {s['topBot']} | {s['category']} |" for s in res['topByShare']]
md += ['', f"## Top 20 by AI visits ({res['top20VisitsShareOfAi']}% of all AI visits)", '| # | Domain | Visits | AI % | AI visits | Top bot | Category |', '|---:|---|---:|---:|---:|---|---|']
md += [f"| {s['rank']} | {s['domain']} | {f(s['visits'])} | {s['aiSharePct']} | {f(s['aiVisits'])} | {s['topBot']} | {s['category']} |" for s in res['topByVisits']]
md += ['', '## Chatbots', '| Bot | First on | First % | Listed on | Listed % | AI visits at sites where first % |', '|---|---:|---:|---:|---:|---:|']
md += [f"| {b['bot']} | {b['firstSites']} | {b['firstPct']} | {b['listedSites']} | {b['listedPct']} | {b['aiVisitsAtSitesWhereFirstPct']} |" for b in bots]
md += ['', json.dumps(pairs), '', '## Not ChatGPT first', *[f"- {s['domain']} ({f(s['visits'])}, {s['category']}): {', '.join(s['bots'])}" for s in res['notChatgptFirst'][:40]]]
md += ['', '## Categories', '| Category | Sites | AI pooled % | AI median % | Index | Organic % | Social % | AI per organic | ChatGPT first % | Top bots |', '|---|---:|---:|---:|---:|---:|---:|---:|---:|---|']
md += [f"| {c['category']} | {c['sites']} | {c['aiPooledPct']} | {c['aiMedianPct']} | {c['index']} | {c['organicPooledPct']} | {c['socialPooledPct']} | {c['aiPerOrganic']} | {c['chatgptFirstPct']} | {c['topBot']} |" for c in categories]
md += ['', '## AI company sites', *[f"- {s}" for s in ai_sites]]
(OUT / 'tables.md').write_text('\n'.join(md) + '\n')
print('\n'.join(md))