# -*- coding: utf-8 -*-
"""ラーメンスープ評価 後段スコアリング・パイプライン（自己完結版）

打ち手1（判断不能画像の分岐）を恒久化した本番想定モジュール。
元ファイル（検証/scripts/run_downstream_model_comparison.py 等）は変更せず、
必要な特徴量生成・グレード定義ロジックをここに複製している。

パイプライン:
    LLM画像評価JSON(+濃度) ─┬─ assessable=0 → 「要確認/撮り直し」（スコアしない）
                            └─ assessable=1 → 線形回帰でスコア回帰 → グレード付与

採用構成（CV検証済み, 5-fold）:
    - total(human_total_score): 線形回帰 / feature_set=visual_plus_brix
    - soup (human_soup_score) : 線形回帰 / feature_set=visual_only
    - 学習対象は assessable=1 の行のみ

公開API:
    prepare(raw_df, require_targets=False)  -> 特徴量DataFrame
    train(input_csv, model_dir)             -> 学習して .pkl 保存
    load_models(model_dir)                  -> 学習済みモデル読込
    score_record(record, models)            -> 単件推論（dict in -> dict out）
    score_batch(raw_df, models)             -> バッチ推論（DataFrame in -> DataFrame out）
    grade(score)                            -> A/B/C/D グレード
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# ====================================================================
# 定数・スコアマッピング（元 run_downstream_model_comparison.py から複製）
# ====================================================================
UNASSESSABLE = "判定不可"

WATER_SCORE = {"水位良好": 1.0, "水位不足": 0.0, UNASSESSABLE: 0.0}
OIL_SCORE = {
    "油感と乳化の状態が良い": 1.0, "乳化が進んでいる": 0.85,
    "油感が強い": 0.45, "乳化しすぎ": 0.35, UNASSESSABLE: 0.0,
}
BOILING_SCORE = {"良い": 1.0, "弱い": 0.45, "強すぎる": 0.35, UNASSESSABLE: 0.0}
PHOTO_SCORE = {
    "良い": 1.0, "湯気が多い": 0.45, "白飛び": 0.35,
    "暗い": 0.35, "角度不良": 0.35, "不良": 0.0, UNASSESSABLE: 0.0,
}

# 採用モデル構成
# 本番入力 = 濃度(brix) ＋ 動画(→画像→LLM評価) ＋ 時間(slot)。
# 濃度はBrix/濃度推定AIが担う設計のため、LLMの濃度予測(visual_density)系は
# 後段特徴量から除外する（§4.11。精度ほぼ不変・設計整合のため）。
TARGETS = {
    "human_total_score": "brix_slot_no_llm_density",
    "human_soup_score": "brix_slot_no_llm_density",
}

# raw入力で期待する列（training_dataset.csv 互換）
RAW_COLUMNS = [
    "video_path", "frame_path", "store", "slot", "brix_raw",
    "visual_density", "water_level", "oil_emulsification",
    "boiling_heat_state", "photo_quality", "image_condition_score",
]


@dataclass(frozen=True)
class FeatureSet:
    name: str
    numeric: tuple[str, ...]
    categorical: tuple[str, ...]


# ====================================================================
# グレード定義（複製）
# ====================================================================
def grade(score) -> np.ndarray:
    values = np.asarray(score, dtype=float)
    out = np.full(values.shape, "D", dtype="<U1")
    out[values >= 21] = "C"
    out[values >= 60] = "B"
    out[values >= 80] = "A"
    return out


def grade_one(score: float) -> str:
    return str(grade(np.array([score]))[0])


# ====================================================================
# 補助
# ====================================================================
def video_id_from_path(path: object) -> str:
    name = Path(str(path)).stem
    match = re.search(r"\d+", name)
    return match.group(0) if match else name


def _to_density(value: object) -> float:
    try:
        return float(value)
    except Exception:
        return np.nan


def _map_score(series: pd.Series, mapping: dict[str, float]) -> pd.Series:
    return series.fillna(UNASSESSABLE).astype(str).str.strip().map(mapping).fillna(0.0)


# ====================================================================
# 特徴量生成（build_features を複製。inference時は人間点数不要）
# ====================================================================
def prepare(raw: pd.DataFrame, require_targets: bool = False) -> pd.DataFrame:
    raw = raw.copy()
    df = pd.DataFrame()
    df["sample_id"] = raw.get("video_path", pd.Series(range(len(raw)))).map(video_id_from_path) \
        if "video_path" in raw else pd.Series(range(len(raw)))
    df["video_path"] = raw.get("video_path")
    df["frame_path"] = raw.get("frame_path")
    df["store"] = raw.get("store", pd.Series(["unknown"] * len(raw))).fillna("unknown").astype(str)
    df["slot"] = raw.get("slot", pd.Series(["unknown"] * len(raw))).fillna("unknown").astype(str)
    df["brix"] = pd.to_numeric(raw.get("brix_raw"), errors="coerce")

    if require_targets:
        df["human_soup_score"] = pd.to_numeric(raw["human_soup_score"], errors="coerce")
        df["human_total_score"] = pd.to_numeric(raw["human_total_score"], errors="coerce")

    for col in ["visual_density", "water_level", "oil_emulsification", "boiling_heat_state", "photo_quality"]:
        df[col] = raw.get(col, pd.Series([UNASSESSABLE] * len(raw))).fillna(UNASSESSABLE).astype(str).str.strip()

    df["image_condition_score"] = pd.to_numeric(raw.get("image_condition_score"), errors="coerce").fillna(0)
    df["visual_density_num"] = df["visual_density"].map(_to_density)
    df["visual_density_assessable"] = df["visual_density_num"].notna().astype(int)
    # assessable は 4軸が全て判定可能か（学習時と同じ導出。LLM出力のassessableに依存しない）
    df["assessable"] = (
        (df["visual_density"] != UNASSESSABLE)
        & (df["water_level"] != UNASSESSABLE)
        & (df["oil_emulsification"] != UNASSESSABLE)
        & (df["boiling_heat_state"] != UNASSESSABLE)
    ).astype(int)
    df["retake_recommended"] = ((df["photo_quality"] != "良い") | (df["assessable"] == 0)).astype(int)
    df["density_brix_abs_diff"] = (df["visual_density_num"] - df["brix"]).abs()
    df["density_brix_signed_diff"] = df["visual_density_num"] - df["brix"]
    df["density_below_brix"] = (df["visual_density_num"] < df["brix"]).fillna(False).astype(int)

    df["water_score"] = _map_score(df["water_level"], WATER_SCORE)
    df["oil_score"] = _map_score(df["oil_emulsification"], OIL_SCORE)
    df["boiling_score"] = _map_score(df["boiling_heat_state"], BOILING_SCORE)
    df["photo_score"] = _map_score(df["photo_quality"], PHOTO_SCORE)

    density_match = (1.0 - (df["density_brix_abs_diff"].fillna(4).clip(0, 4) / 4.0)).clip(0, 1)
    df["density_match_score"] = density_match

    df["slot_is_opening"] = df["slot"].str.contains("開店", na=False).astype(int)
    df["slot_is_evening"] = df["slot"].str.contains("17|18|19|夜|夕", regex=True, na=False).astype(int)
    df["low_quality_flag"] = (df["image_condition_score"] <= 3).astype(int)
    df["unassessable_axis_count"] = (
        (df[["visual_density", "water_level", "oil_emulsification", "boiling_heat_state", "photo_quality"]] == UNASSESSABLE)
        .sum(axis=1).astype(int)
    )

    if require_targets:
        df = df.dropna(subset=["human_soup_score", "human_total_score"]).reset_index(drop=True)
        # 100点スケール外の異常値を除外（元ロジック踏襲）
        in_range = df["human_soup_score"].between(0, 100) & df["human_total_score"].between(0, 100)
        df = df[in_range].reset_index(drop=True)
    return df


def feature_sets() -> dict[str, FeatureSet]:
    visual_numeric = (
        "visual_density_num", "image_condition_score", "water_score", "oil_score",
        "boiling_score", "photo_score", "visual_density_assessable", "assessable",
        "retake_recommended", "low_quality_flag", "unassessable_axis_count",
    )
    visual_cat = ("visual_density", "water_level", "oil_emulsification", "boiling_heat_state", "photo_quality")
    brix_num = ("brix", "density_brix_abs_diff", "density_brix_signed_diff", "density_below_brix", "density_match_score")
    slot_num = ("slot_is_opening", "slot_is_evening")

    # LLMの濃度予測(visual_density)に由来する特徴量を除いた構成。
    #   除外: visual_density_num / visual_density_assessable / visual_density(cat)
    #         density_brix_abs_diff / density_brix_signed_diff / density_below_brix / density_match_score
    #   濃度信号は brix（実測 or 濃度推定AI）のみが担う。
    nod_numeric = (
        "image_condition_score", "water_score", "oil_score", "boiling_score", "photo_score",
        "assessable", "retake_recommended", "low_quality_flag", "unassessable_axis_count",
    )
    nod_cat = ("water_level", "oil_emulsification", "boiling_heat_state", "photo_quality")
    return {
        "visual_only": FeatureSet("visual_only", visual_numeric, visual_cat),
        "visual_plus_brix": FeatureSet("visual_plus_brix", visual_numeric + brix_num, visual_cat),
        "visual_plus_brix_slot": FeatureSet(
            "visual_plus_brix_slot", visual_numeric + brix_num + slot_num, visual_cat + ("slot",)),
        # 本番採用: LLM濃度を除外し、濃度はbrixのみ。動画(状態評価)＋時間(slot)を使う。
        "brix_slot_no_llm_density": FeatureSet(
            "brix_slot_no_llm_density", nod_numeric + ("brix",) + slot_num, nod_cat + ("slot",)),
    }


def _preprocessor(numeric: list[str], categorical: list[str]) -> ColumnTransformer:
    try:
        encoder = OneHotEncoder(handle_unknown="ignore", sparse_output=False)
    except TypeError:
        encoder = OneHotEncoder(handle_unknown="ignore", sparse=False)
    return ColumnTransformer(
        transformers=[
            ("num", Pipeline([("imputer", SimpleImputer(strategy="median")), ("scaler", StandardScaler())]), numeric),
            ("cat", Pipeline([("imputer", SimpleImputer(strategy="most_frequent")), ("onehot", encoder)]), categorical),
        ],
        remainder="drop", verbose_feature_names_out=False,
    )


# ====================================================================
# 学習
# ====================================================================
def train(input_csv: str | Path, model_dir: str | Path) -> dict:
    import joblib
    model_dir = Path(model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)
    raw = pd.read_csv(input_csv)
    df = prepare(raw, require_targets=True)
    # 打ち手1: assessable=1 の行のみで学習（判断不能行はスコア対象外）
    scoreable = df[df["assessable"] == 1].reset_index(drop=True)

    fsets = feature_sets()
    meta = {"n_total": int(len(df)), "n_scoreable": int(len(scoreable)), "models": {}}
    for target, fsname in TARGETS.items():
        fs = fsets[fsname]
        numeric, categorical = list(fs.numeric), list(fs.categorical)
        X = scoreable[numeric + categorical]
        y = scoreable[target].to_numpy()
        pipe = Pipeline([("prep", _preprocessor(numeric, categorical)), ("model", LinearRegression())])
        pipe.fit(X, y)
        path = model_dir / f"{target}.pkl"
        joblib.dump({"pipeline": pipe, "feature_set": fsname,
                     "numeric": numeric, "categorical": categorical, "target": target}, path)
        meta["models"][target] = {"feature_set": fsname, "path": path.name, "model": "LinearRegression"}
    (model_dir / "metadata.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    return meta


def load_models(model_dir: str | Path) -> dict:
    import joblib
    model_dir = Path(model_dir)
    models = {}
    for target in TARGETS:
        models[target] = joblib.load(model_dir / f"{target}.pkl")
    return models


# ====================================================================
# 推論
# ====================================================================
def _predict_scores(feat_df: pd.DataFrame, models: dict) -> dict[str, np.ndarray]:
    preds = {}
    for target, bundle in models.items():
        X = feat_df[bundle["numeric"] + bundle["categorical"]]
        preds[target] = np.clip(bundle["pipeline"].predict(X), 0, 100)
    return preds


def score_batch(raw: pd.DataFrame, models: dict) -> pd.DataFrame:
    """raw（training_dataset互換の列）を受けて、スコア／要確認判定を返す。

    判定可否ゲート（重要）は2段で構成する:
      ① 上流: RF-DETRスープ可読性 `readable`（任意入力。Falseなら評価不能）
      ② LLM側: 評価4軸が判定不可でない `assessable`
    どちらかが不可なら「要確認」（スコアしない）。
    `readable` 列が無い場合は ② のみで判定（後方互換）。
    """
    feat = prepare(raw, require_targets=False)
    preds = _predict_scores(feat, models)
    # 上流可読性（任意）。指定が無ければ全Trueとみなす。
    if "readable" in raw.columns:
        readable = raw["readable"].fillna(True).astype(bool).to_numpy()
    else:
        readable = np.ones(len(feat), dtype=bool)
    llm_ok = (feat["assessable"] == 1).to_numpy()
    scoreable = llm_ok & readable
    out = pd.DataFrame({
        "sample_id": feat["sample_id"],
        "video_path": feat["video_path"],
        "frame_path": feat["frame_path"],
        "assessable": feat["assessable"],
        "readable": readable,
        "retake_recommended": feat["retake_recommended"],
    })
    reason = np.where(~readable, "readable=False（スープ表面が抽出/確認できない）→撮り直し/職人確認へ",
                      np.where(~llm_ok, "assessable=0（LLMが画像評価不能）→撮り直し/職人確認へ", ""))
    out["status"] = np.where(scoreable, "scored", "要確認")
    out["reason"] = reason
    for target in TARGETS:
        col = "total" if target == "human_total_score" else "soup"
        s = preds[target].astype(float)
        out[f"{col}_score"] = np.where(scoreable, np.round(s, 1), np.nan)
        out[f"{col}_grade"] = np.where(scoreable, grade(s), "")
    return out


def score_record(record: dict, models: dict) -> dict:
    """単件推論。record は LLM画像評価JSON ＋ brix_raw/slot/store 等を含む dict。

    例:
        {
          "video_path": "123.MOV", "brix_raw": 12.0, "slot": "開店前", "store": "豚山X",
          "visual_density": "12", "water_level": "水位良好",
          "oil_emulsification": "乳化が進んでいる", "boiling_heat_state": "良い",
          "photo_quality": "良い", "image_condition_score": 8
        }
    返り値:
        assessable=0 → {"status":"要確認", ...}
        assessable=1 → {"status":"scored", "total_score":.., "total_grade":.., "soup_score":.., "soup_grade":..}
    """
    raw = pd.DataFrame([record])
    res = score_batch(raw, models).iloc[0].to_dict()
    # NaN を None に正規化
    for k, v in list(res.items()):
        if isinstance(v, float) and np.isnan(v):
            res[k] = None
    return res