"""
Stratified sampling: chọn 1M users đại diện từ filtered_reviews.jsonl
Chiều stratify: interaction_count_bucket × year_of_last_review

Output:
  sampled_1m_reviews.jsonl   — reviews của 1M users đã chọn
  sampled_1m_users.json      — danh sách user_id + stats
  sampling_report.json       — phân bố cells trước/sau sampling
"""
import json, time, random
from collections import defaultdict
from pathlib import Path

FILTERED   = Path("/workspace/amazon/filtered_reviews.jsonl")
OUT_REVIEWS = Path("/workspace/amazon/sampled_1m_reviews.jsonl")
OUT_USERS   = Path("/workspace/amazon/sampled_1m_users.json")
OUT_REPORT  = Path("/workspace/amazon/sampling_report.json")

TARGET_USERS = 1_000_000
SEED = 42
random.seed(SEED)

# ── Buckets ────────────────────────────────────────────────────────────────
def interaction_bucket(n):
    if n < 10:   return "05-09"
    if n < 20:   return "10-19"
    if n < 50:   return "20-49"
    if n < 100:  return "50-99"
    return "100+"

def year_bucket(ts):
    if ts is None: return "unknown"
    import datetime
    y = datetime.datetime.fromtimestamp(ts).year
    if y <= 2012: return "≤2012"
    if y <= 2014: return "2013-14"
    if y <= 2016: return "2015-16"
    return "2017-18"

# ── Pass 1: build per-user stats ───────────────────────────────────────────
print("Pass 1: building user stats...", flush=True)
t0 = time.time()

user_stats = {}   # uid -> {count, last_ts, asins}

with open(FILTERED, "r", encoding="utf-8", errors="replace") as f:
    for i, line in enumerate(f):
        if i % 10_000_000 == 0 and i > 0:
            print(f"  {i/1e6:.0f}M rows  ({time.time()-t0:.0f}s)", flush=True)
        line = line.strip()
        if not line:
            continue
        try:
            rec = json.loads(line)
        except:
            continue
        uid  = rec.get("reviewerID", "")
        asin = rec.get("asin", "")
        ts   = rec.get("unixReviewTime")
        if not uid:
            continue
        if uid not in user_stats:
            user_stats[uid] = {"count": 0, "last_ts": 0, "n_asins": set()}
        s = user_stats[uid]
        s["count"] += 1
        if ts and ts > s["last_ts"]:
            s["last_ts"] = ts
        if asin:
            s["n_asins"].add(asin)

print(f"Pass 1 done in {time.time()-t0:.1f}s — {len(user_stats):,} users", flush=True)

# convert sets to counts
for uid, s in user_stats.items():
    s["n_asins"] = len(s["n_asins"])

# ── Assign each user to a cell ─────────────────────────────────────────────
cells = defaultdict(list)   # cell_key -> [uid, ...]

for uid, s in user_stats.items():
    ib = interaction_bucket(s["count"])
    yb = year_bucket(s["last_ts"] if s["last_ts"] > 0 else None)
    cells[f"{ib}|{yb}"].append(uid)

print(f"\nCell distribution (before sampling):")
total_users = len(user_stats)
for cell in sorted(cells):
    n = len(cells[cell])
    print(f"  {cell:<20} {n:>8,}  ({n/total_users*100:.1f}%)")

# ── Proportional sampling per cell ────────────────────────────────────────
print(f"\nProportional sampling → {TARGET_USERS:,} users...", flush=True)

sampled_uids = set()
report_cells = {}

for cell, uids in cells.items():
    n_cell    = len(uids)
    n_sample  = round(n_cell / total_users * TARGET_USERS)
    n_sample  = min(n_sample, n_cell)
    chosen    = random.sample(uids, n_sample) if n_sample < n_cell else uids
    sampled_uids.update(chosen)
    report_cells[cell] = {"total": n_cell, "sampled": n_sample}

# adjust to hit exactly TARGET_USERS (rounding may cause off-by-few)
all_uids = list(user_stats.keys())
if len(sampled_uids) < TARGET_USERS:
    remaining = [u for u in all_uids if u not in sampled_uids]
    random.shuffle(remaining)
    sampled_uids.update(remaining[:TARGET_USERS - len(sampled_uids)])
elif len(sampled_uids) > TARGET_USERS:
    sampled_list = list(sampled_uids)
    random.shuffle(sampled_list)
    sampled_uids = set(sampled_list[:TARGET_USERS])

print(f"Sampled {len(sampled_uids):,} users", flush=True)

# ── Save user list ─────────────────────────────────────────────────────────
user_records = [
    {"uid": uid, **{k: v for k, v in user_stats[uid].items()}}
    for uid in sampled_uids
]
with open(OUT_USERS, "w") as f:
    json.dump(user_records, f)
print(f"Saved {OUT_USERS}", flush=True)

# ── Pass 2: write reviews for sampled users ────────────────────────────────
print(f"\nPass 2: writing reviews for sampled users...", flush=True)
t1 = time.time()
written = 0

with open(FILTERED, "r", encoding="utf-8", errors="replace") as fin, \
     open(OUT_REVIEWS, "w", encoding="utf-8") as fout:
    for i, line in enumerate(fin):
        if i % 10_000_000 == 0 and i > 0:
            print(f"  {i/1e6:.0f}M scanned, {written:,} written  ({time.time()-t1:.0f}s)", flush=True)
        line = line.strip()
        if not line:
            continue
        try:
            rec = json.loads(line)
        except:
            continue
        if rec.get("reviewerID", "") in sampled_uids:
            fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
            written += 1

print(f"Pass 2 done in {time.time()-t1:.1f}s — {written:,} reviews written", flush=True)

# ── Report ─────────────────────────────────────────────────────────────────
sampled_stats = [user_stats[u] for u in sampled_uids]
counts = sorted(s["count"] for s in sampled_stats)
n = len(counts)

report = {
    "target_users": TARGET_USERS,
    "sampled_users": len(sampled_uids),
    "reviews_written": written,
    "seed": SEED,
    "cells": report_cells,
    "interaction_dist": {
        "min": counts[0], "p25": counts[n//4], "p50": counts[n//2],
        "p75": counts[3*n//4], "p90": counts[int(n*.9)],
        "p99": counts[int(n*.99)], "max": counts[-1],
    }
}
with open(OUT_REPORT, "w") as f:
    json.dump(report, f, indent=2)

print(f"\n{'='*50}")
print(f"Done in {time.time()-t0:.1f}s total")
print(f"  Sampled users : {len(sampled_uids):,}")
print(f"  Reviews       : {written:,}")
print(f"  interaction p50={counts[n//2]}  p90={counts[int(n*.9)]}  max={counts[-1]}")
print(f"  Output        : {OUT_REVIEWS}")