"""
Is something wrong in production?  Steps 3 to 5.
================================================
Data: WildChat-1M (allenai/WildChat-1M, ODC-BY), 837,989 real conversations with
a free ChatGPT-based chatbot run by the Allen Institute for AI, April 2023 to
April 2024. Zhao et al., 2024, "WildChat: 1M ChatGPT Interaction Logs in the Wild."
This script reads only a few small columns (time, model, language, country,
number of turns, and a hashed sender address), never the conversations.

Setup (once):  pip install pandas pyarrow
Run:           python monitor_lesson.py                      (month by month, with and without heavy senders)
               python monitor_lesson.py alerts               (four daily alert rules over the whole year)
               python monitor_lesson.py weekly 2023-05-15 2023-08-13
                                                             (week by week between two dates)
               python monitor_lesson.py senders 2023-10-29   (the busiest senders on one day)
The first run reads about 35 MB from the 14 data files (a few minutes) and keeps it in ./data.
"""
import io
import os
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

BASE = 'https://huggingface.co/datasets/allenai/WildChat-1M/resolve/main/data/train-{:05d}-of-00014.parquet'
COLS = ['timestamp', 'model', 'language', 'country', 'turn', 'hashed_ip']
LOCAL = 'data/wildchat_meta.parquet'
HEAVY = 1000        # a sender with this many conversations in a month counts as heavy


class Remote(io.RawIOBase):
    """Lets pyarrow read parts of a file over the web, so we fetch only the columns we need."""
    def __init__(self, url, size):
        self.url, self.size, self.pos = url, size, 0

    def seekable(self):
        return True

    def readable(self):
        return True

    def tell(self):
        return self.pos

    def seek(self, off, whence=0):
        self.pos = off if whence == 0 else self.pos + off if whence == 1 else self.size + off
        return self.pos

    def readinto(self, b):
        if self.pos >= self.size or not len(b):
            return 0
        end = min(self.size, self.pos + len(b)) - 1
        for attempt in range(4):
            try:
                req = urllib.request.Request(self.url, headers={'Range': f'bytes={self.pos}-{end}'})
                data = urllib.request.urlopen(req, timeout=60).read()
                break
            except OSError:
                if attempt == 3:
                    raise
                time.sleep(2)
        b[:len(data)] = data
        self.pos += len(data)
        return len(data)


def fetch_file(i):
    head = urllib.request.urlopen(urllib.request.Request(BASE.format(i), method='HEAD'))
    url, size = head.geturl(), int(head.headers['Content-Length'])
    groups = pq.ParquetFile(Remote(url, size)).metadata.num_row_groups
    with ThreadPoolExecutor(12) as ex:
        parts = ex.map(lambda g: pq.ParquetFile(Remote(url, size)).read_row_group(g, columns=COLS), range(groups))
        return pa.concat_tables(list(parts))


def load():
    if not os.path.exists(LOCAL):
        os.makedirs('data', exist_ok=True)
        tables = []
        for i in range(14):
            print(f'reading file {i + 1} of 14 ...', flush=True)
            tables.append(fetch_file(i))
        pq.write_table(pa.concat_tables(tables), LOCAL)
    d = pd.read_parquet(LOCAL)
    d['timestamp'] = d.timestamp.dt.tz_localize(None)
    d['month'] = d.timestamp.dt.strftime('%Y-%m')
    d['day'] = d.timestamp.dt.floor('D')
    d['sender_month'] = d.groupby(['month', 'hashed_ip']).hashed_ip.transform('size')
    return d.sort_values('timestamp')


def monthly(d):
    print(f"Month by month. Heavy sender = {HEAVY:,}+ conversations that month.\n")
    print(f"  {'month':<8} {'conversations':>13} {'senders':>8} {'from heavy':>11}"
          f" {'Chinese, all':>13} {'without heavy':>14} {'turns, all':>11} {'without heavy':>14}")
    for m, g in d.groupby('month'):
        rest = g[g.sender_month < HEAVY]
        zh = lambda x: (x.language == 'Chinese').mean()
        print(f"  {m:<8} {len(g):>13,} {g.hashed_ip.nunique():>8,} {1 - len(rest) / len(g):>11.0%}"
              f" {zh(g):>13.1%} {zh(rest):>14.1%} {g.turn.mean():>11.2f} {rest.turn.mean():>14.2f}")


def daily(d):
    days = d.groupby('day')
    t = pd.DataFrame({'n': days.size(),
                      'top': days.hashed_ip.agg(lambda s: s.value_counts(normalize=True).iloc[0]),
                      'zh': days.language.agg(lambda s: (s == 'Chinese').mean()),
                      'turns': days.turn.mean()}).iloc[1:-1]          # first and last day are partial
    t['usual_n'] = t.n.shift(1).rolling(7).median()                  # median of the previous 7 days
    t['usual_turns'] = t.turns.shift(1).rolling(7).median()
    return t


def alerts(d):
    t = daily(d)
    rules = {
        'volume below half of usual': t.n < 0.5 * t.usual_n,
        'volume above twice usual': t.n > 2 * t.usual_n,
        'Chinese share moves 5+ points in a day': (t.zh - t.zh.shift(1)).abs() > 0.05,
        'average turns 20%+ below usual': t.turns < 0.8 * t.usual_turns,
    }
    print(f"Four alert rules run on each of {len(t)} days. 'Usual' = median of the previous 7 days.\n")
    print(f"  {'rule':<40} {'days it fires':>13} {'of those, one sender sent 5%+':>31}")
    for name, fired in rules.items():
        print(f"  {name:<40} {fired.sum():>13} {(t.top[fired] > 0.05).sum():>31}")
    print("\nDays with volume below half of usual:")
    for day, r in t[rules['volume below half of usual']].iterrows():
        print(f"  {day:%Y-%m-%d}  {r.n:>5,.0f} conversations (usual {r.usual_n:,.0f}), busiest sender {r.top:.0%}")


def weekly(d, start, end):
    w = d[(d.timestamp >= start) & (d.timestamp < pd.Timestamp(end) + pd.Timedelta(days=1))].copy()
    w['week'] = w.timestamp.dt.to_period('W').dt.start_time
    w['sender_week'] = w.groupby(['week', 'hashed_ip']).hashed_ip.transform('size')
    print(f"  {'week of':<10} {'conversations':>13} {'senders':>8} {'turns, all':>11} {'senders <50 a week':>19}  models")
    for wk, g in w.groupby('week'):
        models = ', '.join(f"{m.replace('gpt-', '').replace('-turbo', '')} {n / len(g):.0%}"
                           for m, n in g.model.value_counts().items() if n / len(g) >= 0.01)
        print(f"  {wk:%Y-%m-%d} {len(g):>13,} {g.hashed_ip.nunique():>8,} {g.turn.mean():>11.2f}"
              f" {g[g.sender_week < 50].turn.mean():>19.2f}  {models}")


def senders(d, day):
    g = d[d.day == pd.Timestamp(day)]
    print(f"{day}: {len(g):,} conversations from {g.hashed_ip.nunique():,} senders. Busiest five:\n")
    for ip, n in g.hashed_ip.value_counts().head(5).items():
        s = g[g.hashed_ip == ip]
        print(f"  sender {ip[:8]}  {n:>5,} conversations ({n / len(g):.0%})  {s.language.mode()[0]:<10}"
              f" {s.country.mode()[0]:<15} turns {s.turn.mean():.2f}  {s.timestamp.min():%H:%M} to {s.timestamp.max():%H:%M}")


if __name__ == '__main__':
    data = load()
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    if cmd == 'alerts':
        alerts(data)
    elif cmd == 'weekly' and len(sys.argv) == 4:
        weekly(data, sys.argv[2], sys.argv[3])
    elif cmd == 'senders' and len(sys.argv) == 3:
        senders(data, sys.argv[2])
    else:
        monthly(data)