#!/usr/bin/env python3
"""Step 02 — Draw the stratified extraction sample from the frame.

Stratification: cma × age_band × price_tier, proportional allocation with a floor per non-empty
stratum, fixed seed. A nested random subsample receives the third (Opus) extractor tier.
Images: the first `max_images_per_listing` URLs in the seller's presentation order (the order is
itself an object of study: image-count sensitivity is simulated by truncation at K = 1, 2, 4, 8, 12).

Outputs:
  data/processed/sample.parquet                 sampled listings with `image_urls` (list) and tier flags
  data/processed/sample_images.parquet          one row per (uid, image_index, url)
  results/raw/sample_summary.json
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import yaml

ROOT = Path(__file__).resolve().parents[2]
CFG = yaml.safe_load(open(ROOT / "configs" / "data.yaml"))["sample"]


def main() -> None:
    rng = np.random.default_rng(CFG["seed"])
    df = pd.read_parquet(ROOT / "data" / "processed" / "frame.parquet")
    df["stratum"] = df[CFG["strata"]].astype(str).agg("|".join, axis=1)
    sizes = df["stratum"].value_counts()
    n_target = CFG["n_listings"]
    floor = 20
    # proportional allocation, then enforce floor where the stratum is large enough, then rescale
    alloc = (sizes / sizes.sum() * n_target).round().astype(int)
    alloc = alloc.clip(lower=np.minimum(floor, sizes))
    alloc = np.minimum(alloc, sizes)
    # trim largest strata to hit the target exactly
    while alloc.sum() > n_target:
        k = alloc.idxmax()
        alloc[k] -= 1
    while alloc.sum() < n_target:
        room = (sizes - alloc)
        k = room.idxmax()
        alloc[k] += 1
    parts = []
    for strat, n in alloc.items():
        sub = df[df["stratum"] == strat]
        idx = rng.choice(sub.index.values, size=int(n), replace=False)
        parts.append(df.loc[idx])
    smp = pd.concat(parts).sample(frac=1.0, random_state=int(rng.integers(1e9))).reset_index(drop=True)
    # design weight = stratum size in frame / stratum size in sample (for frame-representative statistics)
    nsamp = smp["stratum"].value_counts()
    smp["weight"] = smp["stratum"].map(lambda k: sizes[k] / nsamp[k])
    smp["tier_haiku"] = True
    smp["tier_sonnet"] = True
    opus_idx = rng.choice(smp.index.values, size=CFG["subsample_opus"], replace=False)
    smp["tier_opus"] = False
    smp.loc[opus_idx, "tier_opus"] = True
    K = CFG["max_images_per_listing"]
    smp["image_urls_used"] = smp["image_urls"].map(lambda l: list(l)[:K])
    smp["n_images_used"] = smp["image_urls_used"].map(len)
    smp.to_parquet(ROOT / "data" / "processed" / "sample.parquet", index=False)
    rows = [{"uid": r.uid, "image_index": i, "url": u} for r in smp.itertuples() for i, u in enumerate(r.image_urls_used)]
    img = pd.DataFrame(rows)
    img.to_parquet(ROOT / "data" / "processed" / "sample_images.parquet", index=False)
    summary = {
        "seed": CFG["seed"], "n_listings": int(len(smp)), "n_images": int(len(img)),
        "n_strata": int(len(alloc)), "opus_subsample": int(smp["tier_opus"].sum()),
        "by_cma": smp["cma"].value_counts().to_dict(),
        "by_age_band": smp["age_band"].astype(str).value_counts().sort_index().to_dict(),
        "by_price_tier": smp["price_tier"].astype(str).value_counts().sort_index().to_dict(),
        "by_source": smp["source"].value_counts().to_dict(),
        "images_per_listing_used": smp["n_images_used"].describe().to_dict(),
        "frame_vs_sample_price_median": [float(df["price"].median()), float(smp["price"].median())],
    }
    json.dump(summary, open(ROOT / "results" / "raw" / "sample_summary.json", "w"), indent=2, ensure_ascii=False, default=str)
    print(json.dumps(summary, indent=1, ensure_ascii=False, default=str))


if __name__ == "__main__":
    sys.exit(main())