#!/usr/bin/env python3
from pathlib import Path
import json
import torch

root = Path(__file__).resolve().parents[1]

print("package root:", root)

ckpt = root / "checkpoint"
assert ckpt.exists(), f"missing checkpoint dir: {ckpt}"

lang_root = root / "lang_embed"
assert lang_root.exists(), f"missing lang_embed dir: {lang_root}"

task_dirs = sorted([p for p in lang_root.iterdir() if p.is_dir()])
print("num task dirs:", len(task_dirs))
assert len(task_dirs) > 0, "no task dirs under lang_embed"

num_embeds = 0
bad = []

for td in task_dirs:
    js = td / "expanded_instruction_gpt-4-turbo.json"
    if not js.exists():
        bad.append(f"missing json: {js}")
        continue

    data = json.load(open(js, "r", encoding="utf-8"))
    for k, v in data.items():
        if not isinstance(v, str):
            bad.append(f"{js}: value for {k} is not str: {type(v)}")

    embeds = sorted(td.glob("lang_embed_*.pt"))
    num_embeds += len(embeds)
    if not embeds:
        bad.append(f"missing lang_embed_*.pt in {td}")

    if embeds:
        x = torch.load(embeds[0], map_location="cpu")
        print("sample embed:", embeds[0], type(x))

print("num lang embeds:", num_embeds)

if bad:
    print("BAD:")
    for x in bad:
        print(" ", x)
    raise SystemExit(1)

print("Package looks OK.")