import math
import os

import pandas as pd

from sklearn.ensemble import (
    ExtraTreesRegressor,
    HistGradientBoostingRegressor,
    RandomForestRegressor,
)
from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

try:
    from xgboost import XGBRegressor
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False


# =========================
# 1. 基础配置
# =========================

OUTPUT_DIR = "output"
os.makedirs(OUTPUT_DIR, exist_ok=True)

RANDOM_STATE = 42

PREFIX_FILES = {
    20: os.path.join(OUTPUT_DIR, "dataset_prefix_20.csv"),
    40: os.path.join(OUTPUT_DIR, "dataset_prefix_40.csv"),
    60: os.path.join(OUTPUT_DIR, "dataset_prefix_60.csv"),
}

FEATURE_COLUMNS = [
    # 时间特征
    "hour",
    "weekday",
    "is_weekend",

    # 起点和当前点
    "start_lng",
    "start_lat",
    "current_lng",
    "current_lat",

    # 前缀基础特征
    "prefix_len",
    "prefix_duration_sec",
    "prefix_distance_km",
    "direct_distance_km",
    "prefix_avg_speed",
    "direction_angle",
    "lng_delta",
    "lat_delta",

    # 轨迹形状采样特征
    "p25_lng",
    "p25_lat",
    "p50_lng",
    "p50_lat",
    "p75_lng",
    "p75_lat",

    # 轨迹形态和方向变化特征
    "tortuosity",
    "recent_direction_angle",
    "recent_distance_km",
    "lng_range",
    "lat_range",
    "first_half_angle",
    "second_half_angle",
    "angle_change",

    # 分段距离特征
    "seg_d1",
    "seg_d2",
    "seg_d3",
    "seg_d4",
]

TARGET_COLUMNS = [
    "target_lng",
    "target_lat",
]


# =========================
# 2. 距离误差函数
# =========================

def haversine(lng1, lat1, lng2, lat2):
    """
    计算两个经纬度点之间的球面距离，单位：公里。
    """
    R = 6371.0

    lng1 = math.radians(lng1)
    lat1 = math.radians(lat1)
    lng2 = math.radians(lng2)
    lat2 = math.radians(lat2)

    dlng = lng2 - lng1
    dlat = lat2 - lat1

    a = math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlng / 2) ** 2
    c = 2 * math.asin(math.sqrt(a))

    return R * c


def calc_distance_errors(true_lng, true_lat, pred_lng, pred_lat):
    """
    批量计算预测终点和真实终点之间的距离误差。
    """
    errors = []

    for a, b, c, d in zip(true_lng, true_lat, pred_lng, pred_lat):
        errors.append(haversine(a, b, c, d))

    return errors


def evaluate_errors(errors):
    """
    汇总误差指标。
    """
    s = pd.Series(errors)

    return {
        "mean_error_km": s.mean(),
        "median_error_km": s.median(),
        "p90_error_km": s.quantile(0.90),
        "max_error_km": s.max(),
    }


# =========================
# 3. 模型列表
# =========================

def get_models():
    models = {
        "RandomForest": RandomForestRegressor(
            n_estimators=300,
            max_depth=30,
            min_samples_leaf=2,
            random_state=RANDOM_STATE,
            n_jobs=-1,
        ),

        "ExtraTrees": ExtraTreesRegressor(
            n_estimators=500,
            max_depth=None,
            min_samples_leaf=2,
            random_state=RANDOM_STATE,
            n_jobs=-1,
        ),

        "HistGradientBoosting": MultiOutputRegressor(
            HistGradientBoostingRegressor(
                max_iter=500,
                learning_rate=0.04,
                max_leaf_nodes=31,
                l2_regularization=0.01,
                random_state=RANDOM_STATE,
            )
        ),

        "KNN": Pipeline([
            ("scaler", StandardScaler()),
            ("knn", KNeighborsRegressor(
                n_neighbors=15,
                weights="distance",
            ))
        ]),
    }

    if HAS_XGBOOST:
        models["XGBoost"] = MultiOutputRegressor(
            XGBRegressor(
                n_estimators=800,
                max_depth=6,
                learning_rate=0.03,
                subsample=0.85,
                colsample_bytree=0.85,
                reg_lambda=1.0,
                reg_alpha=0.0,
                objective="reg:squarederror",
                random_state=RANDOM_STATE,
                n_jobs=-1,
                tree_method="hist",
            )
        )

        # 算力充足时可以保留这个加强版。
        # 如果跑得太慢，可以把 XGBoost_Deep 这一段注释掉。
        models["XGBoost_Deep"] = MultiOutputRegressor(
            XGBRegressor(
                n_estimators=1200,
                max_depth=8,
                learning_rate=0.02,
                subsample=0.90,
                colsample_bytree=0.90,
                reg_lambda=1.5,
                reg_alpha=0.05,
                objective="reg:squarederror",
                random_state=RANDOM_STATE,
                n_jobs=-1,
                tree_method="hist",
            )
        )
    else:
        print("未检测到 xgboost，已跳过 XGBoost 模型。")
        print("如需使用，请先运行：pip install xgboost")

    return models


# =========================
# 4. 单个前缀比例实验
# =========================

def check_columns(df, file_path):
    """
    检查增强特征是否存在。
    如果报错，通常说明你还没有重新运行增强版 01_preprocess_enhanced.py。
    """
    missing = [col for col in FEATURE_COLUMNS + TARGET_COLUMNS if col not in df.columns]
    if missing:
        raise ValueError(
            f"{file_path} 缺少字段：{missing}\n"
            f"请先重新运行 01_preprocess_enhanced.py，生成带增强特征的新 dataset_prefix 文件。"
        )


def run_one_prefix(prefix_ratio, file_path):
    print("=" * 70)
    print(f"开始比较 {prefix_ratio}% 前缀下的模型")
    print("=" * 70)

    df = pd.read_csv(file_path)
    check_columns(df, file_path)

    X = df[FEATURE_COLUMNS]
    y = df[TARGET_COLUMNS]

    X_train, X_test, y_train, y_test, train_idx, test_idx = train_test_split(
        X,
        y,
        df.index,
        test_size=0.2,
        random_state=RANDOM_STATE,
    )

    test_df = df.loc[test_idx].copy()

    true_lng = y_test["target_lng"].values
    true_lat = y_test["target_lat"].values

    results = []

    # =========================
    # 4.1 基线方法
    # =========================

    baseline_errors = calc_distance_errors(
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=test_df["current_lng"].values,
        pred_lat=test_df["current_lat"].values,
    )

    baseline_metrics = evaluate_errors(baseline_errors)

    results.append({
        "prefix_ratio": prefix_ratio,
        "model_name": "Baseline_CurrentPoint",
        "mean_error_km": baseline_metrics["mean_error_km"],
        "median_error_km": baseline_metrics["median_error_km"],
        "p90_error_km": baseline_metrics["p90_error_km"],
        "max_error_km": baseline_metrics["max_error_km"],
        "improvement_vs_baseline_percent": 0.0,
    })

    print(f"Baseline 平均误差/km：{baseline_metrics['mean_error_km']:.4f}")

    # =========================
    # 4.2 正式模型
    # =========================

    models = get_models()

    best_model_name = None
    best_model_error = 999999.0
    best_pred_df = None
    best_feature_importance_df = None

    for model_name, model in models.items():
        print(f"正在训练模型：{model_name}")

        model.fit(X_train, y_train)

        y_pred = model.predict(X_test)

        pred_lng = y_pred[:, 0]
        pred_lat = y_pred[:, 1]

        model_errors = calc_distance_errors(
            true_lng=true_lng,
            true_lat=true_lat,
            pred_lng=pred_lng,
            pred_lat=pred_lat,
        )

        metrics = evaluate_errors(model_errors)

        improvement = (
            (baseline_metrics["mean_error_km"] - metrics["mean_error_km"])
            / baseline_metrics["mean_error_km"]
            * 100
        )

        print(
            f"{model_name} 平均误差/km：{metrics['mean_error_km']:.4f}，"
            f"相比基线提升：{improvement:.2f}%"
        )

        results.append({
            "prefix_ratio": prefix_ratio,
            "model_name": model_name,
            "mean_error_km": metrics["mean_error_km"],
            "median_error_km": metrics["median_error_km"],
            "p90_error_km": metrics["p90_error_km"],
            "max_error_km": metrics["max_error_km"],
            "improvement_vs_baseline_percent": improvement,
        })

        if metrics["mean_error_km"] < best_model_error:
            best_model_error = metrics["mean_error_km"]
            best_model_name = model_name

            temp_df = test_df.copy()
            temp_df["pred_lng"] = pred_lng
            temp_df["pred_lat"] = pred_lat
            temp_df["model_error_km"] = model_errors
            temp_df["baseline_error_km"] = baseline_errors
            temp_df["best_model_name"] = model_name
            best_pred_df = temp_df

            # 保存部分模型的特征重要性
            best_feature_importance_df = extract_feature_importance(model, model_name)

    print(f"{prefix_ratio}% 前缀下最优模型：{best_model_name}")
    print(f"{prefix_ratio}% 前缀下最优平均误差/km：{best_model_error:.4f}")

    best_pred_path = os.path.join(
        OUTPUT_DIR,
        f"best_predictions_prefix_{prefix_ratio}.csv"
    )
    best_pred_df.to_csv(best_pred_path, index=False, encoding="utf-8-sig")

    if best_feature_importance_df is not None:
        importance_path = os.path.join(
            OUTPUT_DIR,
            f"best_feature_importance_prefix_{prefix_ratio}.csv"
        )
        best_feature_importance_df.to_csv(importance_path, index=False, encoding="utf-8-sig")

    return results


def extract_feature_importance(model, model_name):
    """
    提取树模型特征重要性。
    对 MultiOutputRegressor 包装的模型，取两个输出模型的重要性平均值。
    KNN 没有特征重要性，返回 None。
    """
    try:
        if hasattr(model, "feature_importances_"):
            importances = model.feature_importances_
        elif hasattr(model, "estimators_"):
            # MultiOutputRegressor
            importances_list = []
            for estimator in model.estimators_:
                if hasattr(estimator, "feature_importances_"):
                    importances_list.append(estimator.feature_importances_)

            if not importances_list:
                return None

            importances = sum(importances_list) / len(importances_list)
        else:
            return None

        importance_df = pd.DataFrame({
            "feature": FEATURE_COLUMNS,
            "importance": importances,
        }).sort_values("importance", ascending=False)

        importance_df["model_name"] = model_name

        return importance_df

    except Exception:
        return None


# =========================
# 5. 主程序
# =========================

def main():
    all_results = []

    for prefix_ratio, file_path in PREFIX_FILES.items():
        prefix_results = run_one_prefix(prefix_ratio, file_path)
        all_results.extend(prefix_results)

    result_df = pd.DataFrame(all_results)

    result_path = os.path.join(OUTPUT_DIR, "model_compare_results.csv")
    result_df.to_csv(result_path, index=False, encoding="utf-8-sig")

    print("\n====== 所有模型对比结果 ======")
    print(result_df)

    print("\n模型对比结果已保存：", result_path)

    best_df = (
        result_df[result_df["model_name"] != "Baseline_CurrentPoint"]
        .sort_values(["prefix_ratio", "mean_error_km"])
        .groupby("prefix_ratio")
        .head(1)
    )

    best_path = os.path.join(OUTPUT_DIR, "best_model_results.csv")
    best_df.to_csv(best_path, index=False, encoding="utf-8-sig")

    print("\n====== 每个前缀比例下的最优模型 ======")
    print(best_df)

    print("\n最优模型结果已保存：", best_path)


if __name__ == "__main__":
    main()
