# -*- coding: utf-8 -*-
"""Track A: 各動画(フレーム)を4軸でVLM採点 → 特徴量CSVを作る。
出力: trackA_byteplus/features.csv（軸スコア + brix + slot + 派生特徴 + 正解total）
実行（データ・キー受領後）: python trackA_byteplus/score_dimensions.py
"""
import os, sys, glob, json
import pandas as pd, numpy as np
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "common"))
import data_loader, vlm_client
from prompts import DIMENSIONS, build_prompt, SCHEMA_HINT

FRAMES = os.path.join(ROOT, "data", "frames")
ANCHORS = os.path.join(ROOT, "data", "anchors")   # high.jpg / mid.jpg / low.jpg

def video_key(name):  # 添付ファイル名→フレームのベース名
    return os.path.splitext(str(name))[0]

def frames_for(name, n=3):
    base = video_key(name)
    fs = sorted(glob.glob(os.path.join(FRAMES, f"{base}_f*.jpg")))[:n]
    if not fs:
        single = os.path.join(FRAMES, f"{base}.jpg")
        fs = [single] if os.path.exists(single) else []
    return fs

def anchor_imgs():
    return [os.path.join(ANCHORS, f) for f in ("high.jpg","mid.jpg","low.jpg")
            if os.path.exists(os.path.join(ANCHORS, f))]

def run():
    df = data_loader.load()
    cfg = vlm_client.load_config()
    anchors = anchor_imgs()
    rows = []
    for _, r in df.iterrows():
        imgs = frames_for(r.get("video"))
        if not imgs:
            continue
        slot = "開店前" if not r["dinner"] else "ディナー"
        rec = {"video": r["video"], "brix": r["brix"], "slot": int(r["dinner"]),
               "total": r["total"], "soup": r["soup"]}
        ok = True
        for d in DIMENSIONS:
            try:
                out = vlm_client.score(anchors + imgs, build_prompt(d, r["brix"], slot),
                                       SCHEMA_HINT, cfg=cfg)
                rec[d["key"]] = float(out.get("score", np.nan))
            except Exception as e:
                print("ERR", r["video"], d["key"], e); ok = False; break
        if ok:
            rows.append(rec); print("OK", video_key(r["video"]))
    out = pd.DataFrame(rows)
    # 派生特徴（BytePlusの重要特徴に倣う）
    dim_cols = [d["key"] for d in DIMENSIONS]
    out["raw_total"] = out[dim_cols].sum(axis=1)
    out["max_dim"] = out[dim_cols].max(axis=1)
    out["n_low"] = (out[dim_cols] < out[dim_cols].quantile(0.25)).sum(axis=1)
    out.to_csv(os.path.join(ROOT, "trackA_byteplus", "features.csv"),
               index=False, encoding="utf-8-sig")
    print(f"\n{len(out)}件 → trackA_byteplus/features.csv")

if __name__ == "__main__":
    run()