# -*- coding: utf-8 -*-
"""Base export CSV の読込・正規化・train/val/test 分割（全Track共通）。"""
import pandas as pd, numpy as np, re, os, json

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_CSV = os.path.join(ROOT, "data", "raw_csv", "base_export.csv")

def parse_brix(x):
    """'～10'→9.5, '14～'→14.5, '13'→13.0（順序尺度の代表値）"""
    x = str(x)
    if '10' in x and ('～' in x or '〜' in x): return 9.5
    if '14' in x and ('～' in x or '〜' in x): return 14.5
    m = re.findall(r'\d+', x); return float(m[0]) if m else np.nan

def _ai_pred(s):
    try:
        return json.loads(s)['results'][0].get('pred_score')
    except Exception:
        return np.nan

def load(csv_path=DEFAULT_CSV, human_only=True):
    """正規化済みDataFrameを返す。
    列: dt, store, slot, brix_raw, brix, video, soup, total, comment, evaluator, ai_pred, dinner, target
    """
    df = pd.read_csv(csv_path)
    ren = {'テキスト':'dt','提出者':'store','時間帯':'slot','スープ濃度':'brix_raw',
           '添付ファイル':'video','スープのみの点数':'soup','濃度込みの点数':'total',
           '評価備考':'comment','通知するブロック長':'evaluator','結果':'result'}
    df = df.rename(columns={k:v for k,v in ren.items() if k in df.columns})
    df = df.dropna(subset=['dt']).copy()
    if 'result' in df.columns: df['ai_pred'] = df['result'].apply(_ai_pred)
    for c in ['soup','total']:
        if c in df.columns: df[c] = pd.to_numeric(df[c], errors='coerce')
    df['brix'] = df['brix_raw'].apply(parse_brix)
    df['dinner'] = ~df['slot'].astype(str).str.contains('開店前')
    df['target'] = np.where(df['dinner'], 12.0, 13.0)
    if human_only:
        df = df[df['soup'].notna() & df['total'].notna() & (df['total'] <= 100)].copy()
    return df.reset_index(drop=True)

def split(df, test=0.2, val=0.1, seed=42, label='total'):
    """グレード層別で train/val/test に分割（BytePlusと同じ70/10/20・seed固定）。"""
    from collections import defaultdict
    rng = np.random.default_rng(seed)
    # グレードで層別
    def g(s): return 'A' if s>=80 else 'B' if s>=60 else 'C' if s>=21 else 'D'
    buckets = defaultdict(list)
    for i, s in enumerate(df[label].values): buckets[g(s)].append(i)
    tr, va, te = [], [], []
    for _, idx in buckets.items():
        idx = np.array(idx); rng.shuffle(idx)
        n = len(idx); n_te = int(round(n*test)); n_va = int(round(n*val))
        te += list(idx[:n_te]); va += list(idx[n_te:n_te+n_va]); tr += list(idx[n_te+n_va:])
    return (df.iloc[sorted(tr)].reset_index(drop=True),
            df.iloc[sorted(va)].reset_index(drop=True),
            df.iloc[sorted(te)].reset_index(drop=True))

if __name__ == "__main__":
    df = load()
    print("件数:", len(df), "| 期間:", df['dt'].min(), "〜", df['dt'].max())
    tr, va, te = split(df)
    print(f"分割 train={len(tr)} val={len(va)} test={len(te)}")
    print("動画ファイル名あり:", int(df['video'].notna().sum()) if 'video' in df else 0)