#!/usr/bin/env python3
"""Inspect Julia-1 model repo + Open-Jev dataset (no assumptions - read actual files)."""
import json
import os
import sys
from collections import Counter
from pathlib import Path

import pyarrow.parquet as pq

BASE = Path(os.environ.get("QYVOS_HOME", "/home/z/my-project/download/qyvos"))
DATA = BASE / "data" / "release-v2-redistributable"


def inspect_model() -> None:
    jdir = BASE / "Julia-1"
    print("=" * 70)
    print("JULIA-1 REPO FILES")
    print("=" * 70)
    for p in sorted(jdir.rglob("*")):
        if p.is_file() and ".cache" not in p.parts:
            print(f"  {p.relative_to(jdir)}  ({p.stat().st_size:,}B)")

    print("\n" + "=" * 70)
    print("config.json / julia_config.json / encoder config")
    print("=" * 70)
    for name in ("config.json", "julia_config.json", "encoder/config.json"):
        f = jdir / name
        if f.exists():
            obj = json.loads(f.read_text())
            print(f"--- {name} ---")
            print(json.dumps(obj, indent=2)[:2000])

    print("\n" + "=" * 70)
    print("provenance.json")
    print("=" * 70)
    f = jdir / "provenance.json"
    if f.exists():
        print(f.read_text()[:1500])

    print("\n" + "=" * 70)
    print("safetensors header (tensor names/shapes, first 30)")
    print("=" * 70)
    from safetensors import safe_open
    with safe_open(jdir / "model.safetensors", framework="pt") as sf:
        keys = list(sf.keys())
        print("total tensors:", len(keys))
        head_keys = [k for k in keys if not k.startswith("encoder.")]
        enc_keys = [k for k in keys if k.startswith("encoder.")]
        print("encoder.* tensors:", len(enc_keys), "| non-encoder (head) tensors:", len(head_keys))
        for k in head_keys[:30]:
            t = sf.get_slice(k)
            print(f"  {k}  {t.get_shape()}")
        # a few encoder keys
        for k in enc_keys[:5]:
            t = sf.get_slice(k)
            print(f"  {k}  {t.get_shape()}")


def inspect_dataset() -> None:
    print("\n" + "=" * 70)
    print("OPEN-JEV: release-v2-redistributable")
    print("=" * 70)
    shards = ["train", "validation", "test", "calibration", "ood"]
    for shard in shards:
        f = DATA / f"{shard}-00000-of-00001.parquet"
        if not f.exists():
            print(f"[missing] {shard}")
            continue
        pf = pq.ParquetFile(f)  # mmap, no full read
        md = pf.metadata
        print(f"\n--- {shard}: {md.num_rows} rows, {md.num_row_groups} row-groups, {f.stat().st_size:,}B ---")
        print("  columns:", pf.schema_arrow.names)
        if shard == "train":
            tbl = pf.read_row_group(0)  # first row-group only
            kinds = Counter(tbl.column("kind").to_pylist())
            print("  row-group-0 kind counts:", dict(kinds))
            row = {c: tbl.column(c)[0].as_py() for c in tbl.column_names}
            for k, v in row.items():
                s = str(v)
                print(f"  sample[{k}] = {s[:300]}")
            # target / state sizes across a few rows
            tg = tbl.column("target").to_pylist()[:5]
            print("  targets[:5]:", [t if not isinstance(t, (list, str)) else str(t)[:120] for t in tg])
            st = tbl.column("state_json").to_pylist()[:3]
            for s in st:
                print("  state_json len:", len(str(s)))


if __name__ == "__main__":
    inspect_model()
    inspect_dataset()