# social_iqa.py

import os, json, torch, tiktoken, requests, zipfile
from tqdm import tqdm

# 1) GPT-2 tokenizer
ENC = tiktoken.get_encoding("gpt2")

# 2) Where we cache & extract
DATA_DIR = os.path.join(os.path.dirname(__file__), "social_iqa_data")
ZIP_URL  = "https://storage.googleapis.com/ai2-mosaic/public/socialiqa/socialiqa-train-dev.zip"
SPLIT_FILES = {
    "train": ("train.jsonl",    "train-labels.lst"),
    "val":   ("dev.jsonl",      "dev-labels.lst"),
}

def _ensure_data():
    """Download & unzip the Social IQA train/dev split once."""
    os.makedirs(DATA_DIR, exist_ok=True)
    zip_fp = os.path.join(DATA_DIR, "socialiqa-train-dev.zip")
    if not os.path.exists(zip_fp):
        # Download
        resp = requests.get(ZIP_URL, stream=True)
        resp.raise_for_status()
        total = int(resp.headers.get("content-length", 0))
        with open(zip_fp, "wb") as f, tqdm(total=total, unit="B", unit_scale=True) as p:
            for chunk in resp.iter_content(8192):
                f.write(chunk); p.update(len(chunk))
        # Extract
        with zipfile.ZipFile(zip_fp) as zf:
            zf.extractall(DATA_DIR)

def iterate_examples(split="val"):
    """
    Yields each example dict from the given split.
    split ∈ {"train","val"}.
    """
    assert split in SPLIT_FILES, "split must be 'train' or 'val'"
    _ensure_data()
    # zip contains a folder named 'socialiqa-train-dev'
    base = os.path.join(DATA_DIR, "socialiqa-train-dev")
    jsonl_name, labels_name = SPLIT_FILES[split]
    jsonl_fp  = os.path.join(base, jsonl_name)
    labels_fp = os.path.join(base, labels_name)
    with open(jsonl_fp, "r", encoding="utf-8") as jf, \
         open(labels_fp, "r", encoding="utf-8") as lf:
        for jline in jf:
            lline = lf.readline()
            if not lline:
                break
            ex = json.loads(jline)
            ex["label"] = int(lline.strip())
            yield ex

def render_example(example):
    """
    Converts a Social IQA example into:
      - tokens: LongTensor of shape (3, L)
      - mask:   FloatTensor of shape (3, L) where 1 marks choice tokens
      - label:  int in {0,1,2}
    """
    # 1) Build the prompt: context + question
    context  = example.get("context", "")
    question = example.get("question", "")
    prompt = context.strip() + " " + question.strip()

    # 2) Gather the three answers
    choices = [
        example.get("answerA", ""),
        example.get("answerB", ""),
        example.get("answerC", ""),
    ]
    label = example["label"]  # 0-based index loaded from the labels.lst

    # 3) Tokenize & mask
    p_tokens = ENC.encode(prompt)
    p_len    = len(p_tokens)
    toks_list, mask_list = [], []
    for ch in choices:
        full = prompt + " " + ch
        toks = ENC.encode(full)
        mask = torch.zeros(len(toks), dtype=torch.float)
        if len(toks) > p_len:
            mask[p_len:] = 1.0
        toks_list.append(torch.tensor(toks, dtype=torch.long))
        mask_list.append(mask)

    # 4) Pad to common length
    L = max(t.size(0) for t in toks_list)
    padded_toks, padded_masks = [], []
    for t, m in zip(toks_list, mask_list):
        pad = L - t.size(0)
        if pad > 0:
            t = torch.cat([t, torch.zeros(pad, dtype=torch.long)])
            m = torch.cat([m, torch.zeros(pad, dtype=torch.float)])
        padded_toks.append(t)
        padded_masks.append(m)

    tokens = torch.stack(padded_toks, dim=0)
    mask   = torch.stack(padded_masks, dim=0)
    return tokens, mask, label