# -*- coding: utf-8 -*-
"""後段モデル比較ハーネス（Track A/B共通）。
入力: 特徴量DataFrame（VLM軸スコア + Brix + 時間帯 + 派生特徴）と正解スコア。
出力: 固定配点/線形/Ridge/RandomForest/XGBoost/LightGBM を同一train学習・同一test評価し指標を並置。
検証doc §4 の判断ロジック（どのモデルが優位か／横並びか／全滅か）に使う。
"""
import numpy as np, pandas as pd, warnings
warnings.filterwarnings("ignore")
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.ensemble import RandomForestRegressor
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from metrics import evaluate

def _try_import(name):
    try:
        return __import__(name)
    except Exception:
        return None
xgb = _try_import("xgboost"); lgb = _try_import("lightgbm")

def fixed_score(X, dim_cols):
    """固定配点: 軸スコアの単純合計（各軸の満点設計に依存）。"""
    return X[dim_cols].sum(axis=1).values

def run(train_X, train_y, test_X, test_y, dim_cols, feature_cols, seed=42):
    """各モデルを学習・評価して結果dictのリストを返す。"""
    results = []
    # 0) 固定配点（学習不要）
    results.append(("固定配点", evaluate(test_y, fixed_score(test_X, dim_cols))))
    Xtr, Xte = train_X[feature_cols].values, test_X[feature_cols].values
    models = [
        ("線形回帰", LinearRegression()),
        ("Ridge", Ridge(alpha=1.0, random_state=seed) if 'random_state' in Ridge().get_params() else Ridge(alpha=1.0)),
        ("RandomForest", RandomForestRegressor(n_estimators=300, random_state=seed, n_jobs=-1)),
    ]
    if xgb:
        models.append(("XGBoost", xgb.XGBRegressor(n_estimators=300, max_depth=3,
                        learning_rate=0.05, subsample=0.8, random_state=seed)))
    if lgb:
        models.append(("LightGBM", lgb.LGBMRegressor(n_estimators=300, max_depth=3,
                        learning_rate=0.05, subsample=0.8, random_state=seed, verbose=-1)))
    for name, m in models:
        m.fit(Xtr, train_y)
        pred = np.clip(m.predict(Xte), 0, 100)
        results.append((name, evaluate(test_y, pred)))
    return results

def print_table(results, baseline=None):
    cols = ["rho", "mae", "grade_match", "anomaly_recall"]
    print(f"{'モデル':<14}" + "".join(f"{c:>14}" for c in cols))
    print("-" * 70)
    for name, r in results:
        print(f"{name:<14}" + "".join(f"{r.get(c,'-'):>14}" for c in cols))
    if baseline:
        print("-" * 70)
        print(f"{'BytePlus公表':<14}" + "".join(f"{baseline.get(c,'-'):>14}" for c in cols))

BYTEPLUS_BASELINE = {"rho": 0.53, "mae": 11.7, "grade_match": 0.74, "anomaly_recall": 0.44}

if __name__ == "__main__":
    # ダミーデータで動作確認（4軸 + brix + slot + 派生）
    rng = np.random.default_rng(0); n = 400
    dim = rng.integers(5, 30, (n, 4)).astype(float)
    brix = rng.integers(10, 14, n).astype(float)
    slot = rng.integers(0, 2, n).astype(float)
    y = np.clip(dim.sum(1) + (brix-12)*4 + rng.normal(0, 8, n), 0, 100)
    cols = ["d0","d1","d2","d3"]
    df = pd.DataFrame(dim, columns=cols); df["brix"]=brix; df["slot"]=slot
    df["max_dim"]=dim.max(1); df["n_low"]=(dim<10).sum(1); df["raw_total"]=dim.sum(1)
    feat = cols+["brix","slot","max_dim","n_low","raw_total"]
    tr, te = df.iloc[:300], df.iloc[300:]; ytr, yte = y[:300], y[300:]
    res = run(tr, ytr, te, yte, cols, feat)
    print_table(res, BYTEPLUS_BASELINE)