# openbookqa.py

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

# ─── 1) Tokenizer ───────────────────────────────────────────────────────────────
ENC = tiktoken.get_encoding("gpt2")

# ─── 2) Data paths & URLs ───────────────────────────────────────────────────────
DATA_DIR = os.path.join(os.path.dirname(__file__), "openbookqa_data")
ZIP_URL  = "https://s3-us-west-2.amazonaws.com/ai2-website/data/OpenBookQA-V1-Sep2018.zip"
SPLIT_PATHS = {
    "train": "OpenBookQA-V1-Sep2018/Data/Main/train.jsonl",
    "val":   "OpenBookQA-V1-Sep2018/Data/Main/dev.jsonl",
    "test":  "OpenBookQA-V1-Sep2018/Data/Main/test.jsonl",
}

# ─── 3) Download & extract ─────────────────────────────────────────────────────
def _ensure_data():
    os.makedirs(DATA_DIR, exist_ok=True)
    zip_fp = os.path.join(DATA_DIR, "openbookqa.zip")
    if not os.path.exists(zip_fp):
        print("▶ Downloading OpenBookQA…")
        resp = requests.get(ZIP_URL, stream=True)
        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))
        print("▶ Extracting OpenBookQA…")
        with zipfile.ZipFile(zip_fp) as zf:
            zf.extractall(DATA_DIR)

# call on import
_ensure_data()

# ─── 4) Public API ─────────────────────────────────────────────────────────────
def iterate_examples(split="val"):
    """
    Yields one JSON object per line from train/dev/test.
    """
    assert split in SPLIT_PATHS, f"split must be one of {list(SPLIT_PATHS)}"
    path = os.path.join(DATA_DIR, SPLIT_PATHS[split])
    with open(path, encoding="utf-8") as f:
        for line in f:
            yield json.loads(line)

def render_example(example):
    """
    Converts an OpenBookQA example into (tokens, mask, label):
      - tokens: LongTensor (4, L)
      - mask:   FloatTensor (4, L) where 1 marks the answer span
      - label:  int in [0..3]
    """
    # question stem
    stem = (
        example.get("question_stem")
        or example.get("question", {}).get("stem", "")
    )

    # choices list
    raw = example.get("choices") or example.get("question", {}).get("choices", [])
    choices = sorted(raw, key=lambda c: c["label"])

    # answer key
    key   = example.get("answerKey") or example.get("answer_key")
    label = {"A":0,"B":1,"C":2,"D":3}.get(key, 0)

    # tokenize & mask
    stem_tokens = ENC.encode(stem)
    tok_rows, mask_rows = [], []
    for c in choices:
        full = stem + " " + c["text"]
        toks = ENC.encode(full)
        mask = [0]*len(stem_tokens) + [1]*(len(toks)-len(stem_tokens))
        tok_rows.append(toks)
        mask_rows.append(mask)

    # pad to same length
    L = max(len(r) for r in tok_rows)
    tokens = torch.zeros((4, L), dtype=torch.long)
    mask   = torch.zeros((4, L), dtype=torch.float)
    for i, (r, m) in enumerate(zip(tok_rows, mask_rows)):
        tokens[i, :len(r)] = torch.tensor(r, dtype=torch.long)
        mask[i,   :len(m)] = torch.tensor(m, dtype=torch.float)

    return tokens, mask, label

# ─── 5) CLI entrypoint ─────────────────────────────────────────────────────────
if __name__ == "__main__":
    # simply trigger the download & report
    print(f"OpenBookQA data directory is ready at: {DATA_DIR}")
    sizes = {split: os.path.getsize(os.path.join(DATA_DIR, path))
             for split, path in SPLIT_PATHS.items()}
    for split, sz in sizes.items():
        print(f"  {split}.jsonl -> {sz/1024/1024:.2f} MB")