# -*- coding: utf-8 -*-
"""固定配点・線形回帰・Ridge・Random Forest・XGBoost・LightGBMを比較する。"""
import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "common"))
import compare_models  # noqa: E402

DEFAULT_FEATURES = ROOT / "検証" / "data" / "features.csv"
DEFAULT_SPLITS = ROOT / "検証" / "data" / "splits.csv"
DEFAULT_OUTPUT = ROOT / "検証" / "results" / "comparison_latest.csv"


def _load_split(features_path, splits_path):
    df = pd.read_csv(features_path)
    splits = pd.read_csv(splits_path)
    if "sample_id" not in df.columns or "sample_id" not in splits.columns:
        raise ValueError("features と splits には sample_id が必要です。")
    return df.merge(splits, on="sample_id", how="inner")


def _numeric_feature_cols(df, target_col, dim_cols, exclude):
    reserved = {"sample_id", "split", target_col} | set(exclude)
    cols = []
    for col in df.columns:
        if col in reserved:
            continue
        if pd.api.types.is_numeric_dtype(df[col]):
            cols.append(col)
    for col in dim_cols:
        if col not in cols and col in df.columns and pd.api.types.is_numeric_dtype(df[col]):
            cols.append(col)
    return cols


def run(features_path, splits_path, output_path, dim_cols, target_col, feature_cols, exclude, seed, data_version, feature_set):
    df = _load_split(features_path, splits_path)
    if target_col not in df.columns:
        raise ValueError(f"目的変数列がありません: {target_col}")
    missing_dim = [c for c in dim_cols if c not in df.columns]
    if missing_dim:
        raise ValueError(f"固定配点用の列がありません: {missing_dim}")

    if not feature_cols:
        feature_cols = _numeric_feature_cols(df, target_col, dim_cols, exclude)
    if not feature_cols:
        raise ValueError("特徴量列がありません。--feature-cols で指定してください。")

    df[feature_cols] = df[feature_cols].replace([np.inf, -np.inf], np.nan)
    fill_values = df[df["split"] == "train"][feature_cols].median(numeric_only=True)
    df[feature_cols] = df[feature_cols].fillna(fill_values).fillna(0)

    train = df[df["split"] == "train"].copy()
    val = df[df["split"] == "val"].copy()
    test = df[df["split"] == "test"].copy()
    if len(train) == 0 or len(test) == 0:
        raise ValueError("train と test の両方が必要です。")

    eval_sets = [("test", test)]
    if len(val) > 0:
        eval_sets.insert(0, ("val", val))

    rows = []
    by_model = {}
    for split_name, part in eval_sets:
        results = compare_models.run(
            train[feature_cols],
            train[target_col].values,
            part[feature_cols],
            part[target_col].values,
            dim_cols,
            feature_cols,
            seed=seed,
        )
        for model, metrics in results:
            by_model.setdefault(model, {"model": model})
            for key in ["rho", "mae", "grade_match", "anomaly_recall"]:
                by_model[model][f"{split_name}_{key}"] = metrics.get(key)

    for model, values in by_model.items():
        rows.append({
            "data_version": data_version,
            "feature_set": feature_set,
            "target": target_col,
            "split_seed": seed,
            "model": model,
            "val_rho": values.get("val_rho"),
            "val_mae": values.get("val_mae"),
            "val_grade_match": values.get("val_grade_match"),
            "val_anomaly_recall": values.get("val_anomaly_recall"),
            "test_rho": values.get("test_rho"),
            "test_mae": values.get("test_mae"),
            "test_grade_match": values.get("test_grade_match"),
            "test_anomaly_recall": values.get("test_anomaly_recall"),
            "notes": "",
        })

    output_path.parent.mkdir(parents=True, exist_ok=True)
    out = pd.DataFrame(rows)
    out.to_csv(output_path, index=False, encoding="utf-8-sig")
    return out


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--features", type=Path, default=DEFAULT_FEATURES)
    parser.add_argument("--splits", type=Path, default=DEFAULT_SPLITS)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--dim-cols", nargs="+", required=True, help="固定配点に使う列。例: visual_density water_level oil_emulsification boiling_heat_state")
    parser.add_argument("--target", default="target")
    parser.add_argument("--feature-cols", nargs="*", default=[])
    parser.add_argument("--exclude", nargs="*", default=[])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--data-version", default="learning_data")
    parser.add_argument("--feature-set", default="default")
    args = parser.parse_args()

    out = run(args.features, args.splits, args.output, args.dim_cols, args.target, args.feature_cols, args.exclude, args.seed, args.data_version, args.feature_set)
    print(out.to_string(index=False))
    print(f"saved: {args.output}")


if __name__ == "__main__":
    main()