import math
import os

import numpy as np
import pandas as pd

from sklearn.cluster import MiniBatchKMeans
from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputRegressor
from sklearn.neighbors import NearestNeighbors

try:
    from xgboost import XGBClassifier, XGBRegressor
except ImportError as e:
    raise ImportError("未检测到 xgboost。请先运行：pip install xgboost") from e


# ============================================================
# 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"),
}

# 聚类数量：沿用你前面的 K=50
N_CLUSTERS = 50

# 聚类候选数：这里使用 Top-30，而不是 Top-3
TOP_K_CLUSTERS = 30

# 轨迹相似度近邻数量
K_NEIGHBORS = 30
WEIGHT_POWER = 1.5

# None 表示使用 dataset_prefix 文件中的全部样本。
# 如果只想快速调试，可以改成 5000。
MAX_ROWS_PER_PREFIX = None

# Porto 附近参考点，用于经纬度和 km 平面坐标互转
REF_LNG = -8.61
REF_LAT = 41.16


BASE_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"]

SIMILARITY_FEATURE_COLUMNS = [
    "sim_pred_lng",
    "sim_pred_lat",
    "sim_top1_lng",
    "sim_top1_lat",
    "sim_top5_pred_lng",
    "sim_top5_pred_lat",
    "sim_min_dist",
    "sim_mean_dist",
    "sim_weighted_dist",
    "sim_endpoint_dispersion_km",
    "sim_current_to_pred_km",
    "sim_start_to_pred_km",
]

CLUSTER_CONFIDENCE_COLUMNS = [
    "cluster_top1_prob",
    "cluster_top3_prob_sum",
    "cluster_top10_prob_sum",
    "cluster_top30_prob_sum",
    "cluster_prob_entropy",
]

META_FEATURE_COLUMNS = [
    # 当前状态
    "current_lng",
    "current_lat",
    "start_lng",
    "start_lat",
    "prefix_len",
    "prefix_distance_km",
    "direct_distance_km",
    "prefix_avg_speed",
    "tortuosity",
    "angle_change",

    # 三类模型预测结果
    "direct_pred_lng",
    "direct_pred_lat",
    "sim_only_pred_lng",
    "sim_only_pred_lat",
    "sim_enhanced_pred_lng",
    "sim_enhanced_pred_lat",
    "cluster_top30_pred_lng",
    "cluster_top30_pred_lat",

    # 预测之间的分歧
    "dist_direct_sim_enhanced",
    "dist_direct_cluster",
    "dist_sim_enhanced_cluster",
    "dist_current_direct",
    "dist_current_sim_enhanced",
    "dist_current_cluster",

    # 轨迹相似度置信特征
    "sim_min_dist",
    "sim_mean_dist",
    "sim_weighted_dist",
    "sim_endpoint_dispersion_km",

    # 聚类分类置信特征
    "cluster_top1_prob",
    "cluster_top3_prob_sum",
    "cluster_top10_prob_sum",
    "cluster_top30_prob_sum",
    "cluster_prob_entropy",
]


# ============================================================
# 2. 距离与坐标函数
# ============================================================

def haversine(lng1, lat1, lng2, lat2):
    R = 6371.0

    lng1, lat1, lng2, lat2 = map(math.radians, [lng1, lat1, lng2, 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 np.array(errors, dtype=np.float32)


def evaluate_errors(errors):
    s = pd.Series(errors)
    return {
        "mean_error_km": float(s.mean()),
        "median_error_km": float(s.median()),
        "p90_error_km": float(s.quantile(0.90)),
        "max_error_km": float(s.max()),
    }


def lonlat_to_xy_km_array(lng_arr, lat_arr):
    lng_arr = np.asarray(lng_arr)
    lat_arr = np.asarray(lat_arr)

    x = (lng_arr - REF_LNG) * 111.320 * math.cos(math.radians(REF_LAT))
    y = (lat_arr - REF_LAT) * 110.574

    return np.column_stack([x, y]).astype(np.float32)


def xy_km_to_lonlat_array(x_arr, y_arr):
    x_arr = np.asarray(x_arr)
    y_arr = np.asarray(y_arr)

    lng = x_arr / (111.320 * math.cos(math.radians(REF_LAT))) + REF_LNG
    lat = y_arr / 110.574 + REF_LAT

    return lng, lat


def pairwise_distance_array(lng1, lat1, lng2, lat2):
    return np.array([
        haversine(a, b, c, d)
        for a, b, c, d in zip(lng1, lat1, lng2, lat2)
    ], dtype=np.float32)


# ============================================================
# 3. XGBoost 模型
# ============================================================

def make_direct_xgboost_regressor():
    base_model = XGBRegressor(
        n_estimators=1000,
        max_depth=8,
        learning_rate=0.025,
        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",
    )
    return MultiOutputRegressor(base_model)


def make_cluster_classifier(num_classes):
    clf = XGBClassifier(
        n_estimators=700,
        max_depth=6,
        learning_rate=0.035,
        subsample=0.90,
        colsample_bytree=0.90,
        reg_lambda=1.5,
        reg_alpha=0.05,
        objective="multi:softprob",
        num_class=num_classes,
        random_state=RANDOM_STATE,
        n_jobs=-1,
        tree_method="hist",
        eval_metric="mlogloss",
    )
    return clf


def make_offset_xgboost_regressor():
    base_model = XGBRegressor(
        n_estimators=900,
        max_depth=7,
        learning_rate=0.025,
        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",
    )
    return MultiOutputRegressor(base_model)


def make_meta_fusion_regressor():
    """
    最终融合模型。
    它不直接从零学习终点，而是学习如何融合：
    单独 XGBoost、相似轨迹、相似度增强 XGBoost、Top-30 聚类偏移模型。
    """
    base_model = XGBRegressor(
        n_estimators=600,
        max_depth=5,
        learning_rate=0.035,
        subsample=0.90,
        colsample_bytree=0.90,
        reg_lambda=2.0,
        reg_alpha=0.05,
        objective="reg:squarederror",
        random_state=RANDOM_STATE,
        n_jobs=-1,
        tree_method="hist",
    )
    return MultiOutputRegressor(base_model)


def train_direct_xgboost(train_df, query_df):
    model = make_direct_xgboost_regressor()
    model.fit(train_df[BASE_FEATURE_COLUMNS].values, train_df[TARGET_COLUMNS].values)
    pred = model.predict(query_df[BASE_FEATURE_COLUMNS].values)
    return pred, model


# ============================================================
# 4. 轨迹相似度特征
# ============================================================

def build_similarity_vectors(df):
    point_columns = [
        ("start_lng", "start_lat"),
        ("p25_lng", "p25_lat"),
        ("p50_lng", "p50_lat"),
        ("p75_lng", "p75_lat"),
        ("current_lng", "current_lat"),
    ]

    weights = np.array([0.7, 0.9, 1.1, 1.3, 1.7], dtype=np.float32)

    vectors = []
    for row in df.itertuples(index=False):
        values = []
        for i, (lng_col, lat_col) in enumerate(point_columns):
            lng = getattr(row, lng_col)
            lat = getattr(row, lat_col)
            xy = lonlat_to_xy_km_array([lng], [lat])[0]
            sw = math.sqrt(float(weights[i]))
            values.extend([xy[0] * sw, xy[1] * sw])
        vectors.append(values)

    return np.array(vectors, dtype=np.float32)


def weighted_endpoint_prediction(distances, indices, train_targets):
    neighbor_targets = train_targets[indices]

    weights = 1.0 / np.power(distances + 0.05, WEIGHT_POWER)
    weights = weights / weights.sum(axis=1, keepdims=True)

    pred = (neighbor_targets * weights[:, :, None]).sum(axis=1)
    weighted_dist = (distances * weights).sum(axis=1)

    return pred, weights, neighbor_targets, weighted_dist


def endpoint_dispersion_km(neighbor_targets, pred_targets):
    dispersions = []

    for neigh, pred in zip(neighbor_targets, pred_targets):
        pred_lng, pred_lat = pred
        ds = [haversine(pred_lng, pred_lat, lng, lat) for lng, lat in neigh]
        dispersions.append(float(np.mean(ds)))

    return np.array(dispersions, dtype=np.float32)


def build_nn_index(train_vectors, n_neighbors):
    nn = NearestNeighbors(
        n_neighbors=n_neighbors,
        metric="euclidean",
        algorithm="auto",
        n_jobs=-1,
    )
    nn.fit(train_vectors)
    return nn


def make_similarity_features_query(nn, query_vectors, train_targets):
    distances, indices = nn.kneighbors(query_vectors, return_distance=True)

    pred, weights, neighbor_targets, weighted_dist = weighted_endpoint_prediction(
        distances=distances,
        indices=indices,
        train_targets=train_targets,
    )

    top1 = neighbor_targets[:, 0, :]

    k5 = min(5, distances.shape[1])
    pred_top5, _, _, _ = weighted_endpoint_prediction(
        distances=distances[:, :k5],
        indices=indices[:, :k5],
        train_targets=train_targets,
    )

    dispersion = endpoint_dispersion_km(neighbor_targets, pred)

    feature_df = pd.DataFrame({
        "sim_pred_lng": pred[:, 0],
        "sim_pred_lat": pred[:, 1],
        "sim_top1_lng": top1[:, 0],
        "sim_top1_lat": top1[:, 1],
        "sim_top5_pred_lng": pred_top5[:, 0],
        "sim_top5_pred_lat": pred_top5[:, 1],
        "sim_min_dist": distances[:, 0],
        "sim_mean_dist": distances.mean(axis=1),
        "sim_weighted_dist": weighted_dist,
        "sim_endpoint_dispersion_km": dispersion,
    })

    return feature_df, pred


def make_similarity_features_train(train_vectors, train_targets):
    """
    训练集相似特征：去掉自身近邻，避免泄漏。
    """
    nn = build_nn_index(train_vectors, K_NEIGHBORS + 1)
    distances, indices = nn.kneighbors(train_vectors, return_distance=True)

    clean_distances = []
    clean_indices = []

    for i in range(len(train_vectors)):
        mask = indices[i] != i
        clean_distances.append(distances[i][mask][:K_NEIGHBORS])
        clean_indices.append(indices[i][mask][:K_NEIGHBORS])

    clean_distances = np.array(clean_distances, dtype=np.float32)
    clean_indices = np.array(clean_indices, dtype=np.int64)

    pred, weights, neighbor_targets, weighted_dist = weighted_endpoint_prediction(
        distances=clean_distances,
        indices=clean_indices,
        train_targets=train_targets,
    )

    top1 = neighbor_targets[:, 0, :]

    k5 = min(5, K_NEIGHBORS)
    pred_top5, _, _, _ = weighted_endpoint_prediction(
        distances=clean_distances[:, :k5],
        indices=clean_indices[:, :k5],
        train_targets=train_targets,
    )

    dispersion = endpoint_dispersion_km(neighbor_targets, pred)

    feature_df = pd.DataFrame({
        "sim_pred_lng": pred[:, 0],
        "sim_pred_lat": pred[:, 1],
        "sim_top1_lng": top1[:, 0],
        "sim_top1_lat": top1[:, 1],
        "sim_top5_pred_lng": pred_top5[:, 0],
        "sim_top5_pred_lat": pred_top5[:, 1],
        "sim_min_dist": clean_distances[:, 0],
        "sim_mean_dist": clean_distances.mean(axis=1),
        "sim_weighted_dist": weighted_dist,
        "sim_endpoint_dispersion_km": dispersion,
    })

    return feature_df, pred


def add_similarity_distance_features(df, sim_feature_df):
    current_to_pred = []
    start_to_pred = []

    for row, sim_row in zip(df.itertuples(index=False), sim_feature_df.itertuples(index=False)):
        current_to_pred.append(
            haversine(
                row.current_lng,
                row.current_lat,
                sim_row.sim_pred_lng,
                sim_row.sim_pred_lat,
            )
        )
        start_to_pred.append(
            haversine(
                row.start_lng,
                row.start_lat,
                sim_row.sim_pred_lng,
                sim_row.sim_pred_lat,
            )
        )

    sim_feature_df = sim_feature_df.copy()
    sim_feature_df["sim_current_to_pred_km"] = current_to_pred
    sim_feature_df["sim_start_to_pred_km"] = start_to_pred

    return sim_feature_df


def train_similarity_enhanced_xgboost(train_df, train_sim_features):
    train_aug = pd.concat(
        [
            train_df[BASE_FEATURE_COLUMNS].reset_index(drop=True),
            train_sim_features[SIMILARITY_FEATURE_COLUMNS].reset_index(drop=True),
        ],
        axis=1,
    )
    train_aug[TARGET_COLUMNS] = train_df[TARGET_COLUMNS].reset_index(drop=True)

    feature_columns = BASE_FEATURE_COLUMNS + SIMILARITY_FEATURE_COLUMNS

    model = make_direct_xgboost_regressor()
    model.fit(train_aug[feature_columns].values, train_aug[TARGET_COLUMNS].values)

    return model, feature_columns


def predict_similarity_enhanced_xgboost(model, query_df, query_sim_features, feature_columns):
    query_aug = pd.concat(
        [
            query_df[BASE_FEATURE_COLUMNS].reset_index(drop=True),
            query_sim_features[SIMILARITY_FEATURE_COLUMNS].reset_index(drop=True),
        ],
        axis=1,
    )

    return model.predict(query_aug[feature_columns].values)


# ============================================================
# 5. 聚类 Top-30 + 类内偏移回归
# ============================================================

def fit_train_kmeans(train_df):
    train_xy = lonlat_to_xy_km_array(
        train_df["target_lng"].values,
        train_df["target_lat"].values,
    )

    km = MiniBatchKMeans(
        n_clusters=N_CLUSTERS,
        batch_size=8192,
        random_state=RANDOM_STATE,
        n_init="auto",
    )

    train_labels = km.fit_predict(train_xy)

    centers_xy = km.cluster_centers_
    center_lng, center_lat = xy_km_to_lonlat_array(centers_xy[:, 0], centers_xy[:, 1])

    centers_df = pd.DataFrame({
        "cluster_id": np.arange(N_CLUSTERS),
        "center_lng": center_lng,
        "center_lat": center_lat,
    })

    return km, train_labels, centers_df


def assign_cluster_labels(km, df):
    xy = lonlat_to_xy_km_array(df["target_lng"].values, df["target_lat"].values)
    return km.predict(xy)


def get_topk_clusters(prob_matrix, class_labels, top_k):
    top_pos = np.argsort(-prob_matrix, axis=1)[:, :top_k]
    top_probs = np.take_along_axis(prob_matrix, top_pos, axis=1)
    top_clusters = class_labels[top_pos]

    prob_sum = top_probs.sum(axis=1, keepdims=True)
    prob_sum[prob_sum == 0] = 1.0
    top_probs_norm = top_probs / prob_sum

    return top_clusters, top_probs, top_probs_norm


def cluster_prob_entropy(prob_matrix):
    eps = 1e-12
    p = np.clip(prob_matrix, eps, 1.0)
    return -np.sum(p * np.log(p), axis=1)


def train_offset_regressor(train_df, train_labels, centers_df):
    train_df = train_df.copy()
    train_df["cluster_id"] = train_labels

    train_with_center = train_df.merge(
        centers_df[["cluster_id", "center_lng", "center_lat"]],
        on="cluster_id",
        how="left",
    )

    train_with_center["offset_lng"] = (
        train_with_center["target_lng"] - train_with_center["center_lng"]
    )
    train_with_center["offset_lat"] = (
        train_with_center["target_lat"] - train_with_center["center_lat"]
    )

    offset_feature_columns = BASE_FEATURE_COLUMNS + ["center_lng", "center_lat"]

    model = make_offset_xgboost_regressor()
    model.fit(
        train_with_center[offset_feature_columns].values,
        train_with_center[["offset_lng", "offset_lat"]].values,
    )

    return model, offset_feature_columns


def predict_topk_offset_weighted(query_df, top_clusters, top_probs_norm, centers_df, offset_model, offset_feature_columns):
    centers_lng = centers_df.set_index("cluster_id")["center_lng"].to_dict()
    centers_lat = centers_df.set_index("cluster_id")["center_lat"].to_dict()

    final_lng = np.zeros(len(query_df), dtype=np.float64)
    final_lat = np.zeros(len(query_df), dtype=np.float64)

    base_part = query_df[BASE_FEATURE_COLUMNS].reset_index(drop=True)

    for rank in range(top_clusters.shape[1]):
        candidate_clusters = top_clusters[:, rank]
        candidate_probs = top_probs_norm[:, rank]

        temp = base_part.copy()
        temp["center_lng"] = [centers_lng[int(cid)] for cid in candidate_clusters]
        temp["center_lat"] = [centers_lat[int(cid)] for cid in candidate_clusters]

        offsets = offset_model.predict(temp[offset_feature_columns].values)

        candidate_lng = temp["center_lng"].values + offsets[:, 0]
        candidate_lat = temp["center_lat"].values + offsets[:, 1]

        final_lng += candidate_lng * candidate_probs
        final_lat += candidate_lat * candidate_probs

    return final_lng, final_lat


def train_cluster_top30_model(train_df):
    km, train_labels, centers_df = fit_train_kmeans(train_df)

    clf = make_cluster_classifier(N_CLUSTERS)
    clf.fit(train_df[BASE_FEATURE_COLUMNS].values, train_labels)

    offset_model, offset_feature_columns = train_offset_regressor(
        train_df=train_df,
        train_labels=train_labels,
        centers_df=centers_df,
    )

    return km, train_labels, centers_df, clf, offset_model, offset_feature_columns


def predict_cluster_top30(query_df, km, clf, centers_df, offset_model, offset_feature_columns):
    prob_matrix = clf.predict_proba(query_df[BASE_FEATURE_COLUMNS].values)
    class_labels = clf.classes_

    top_clusters, top_probs_raw, top_probs_norm = get_topk_clusters(
        prob_matrix=prob_matrix,
        class_labels=class_labels,
        top_k=TOP_K_CLUSTERS,
    )

    pred_lng, pred_lat = predict_topk_offset_weighted(
        query_df=query_df,
        top_clusters=top_clusters,
        top_probs_norm=top_probs_norm,
        centers_df=centers_df,
        offset_model=offset_model,
        offset_feature_columns=offset_feature_columns,
    )

    top3_prob_sum = top_probs_raw[:, :min(3, TOP_K_CLUSTERS)].sum(axis=1)
    top10_prob_sum = top_probs_raw[:, :min(10, TOP_K_CLUSTERS)].sum(axis=1)
    top30_prob_sum = top_probs_raw.sum(axis=1)
    entropy = cluster_prob_entropy(prob_matrix)

    confidence_df = pd.DataFrame({
        "cluster_top1_prob": top_probs_raw[:, 0],
        "cluster_top3_prob_sum": top3_prob_sum,
        "cluster_top10_prob_sum": top10_prob_sum,
        "cluster_top30_prob_sum": top30_prob_sum,
        "cluster_prob_entropy": entropy,
        "cluster_top1": top_clusters[:, 0],
        "cluster_top2": top_clusters[:, 1],
        "cluster_top3": top_clusters[:, 2],
    })

    return pred_lng, pred_lat, confidence_df, prob_matrix, top_clusters


# ============================================================
# 6. 元融合特征
# ============================================================

def build_meta_features(
    df,
    direct_pred,
    sim_only_pred,
    sim_enhanced_pred,
    cluster_pred,
    sim_features,
    cluster_confidence,
):
    meta = df[[
        "current_lng",
        "current_lat",
        "start_lng",
        "start_lat",
        "prefix_len",
        "prefix_distance_km",
        "direct_distance_km",
        "prefix_avg_speed",
        "tortuosity",
        "angle_change",
    ]].reset_index(drop=True).copy()

    meta["direct_pred_lng"] = direct_pred[:, 0]
    meta["direct_pred_lat"] = direct_pred[:, 1]

    meta["sim_only_pred_lng"] = sim_only_pred[:, 0]
    meta["sim_only_pred_lat"] = sim_only_pred[:, 1]

    meta["sim_enhanced_pred_lng"] = sim_enhanced_pred[:, 0]
    meta["sim_enhanced_pred_lat"] = sim_enhanced_pred[:, 1]

    meta["cluster_top30_pred_lng"] = cluster_pred[:, 0]
    meta["cluster_top30_pred_lat"] = cluster_pred[:, 1]

    meta["dist_direct_sim_enhanced"] = pairwise_distance_array(
        meta["direct_pred_lng"].values,
        meta["direct_pred_lat"].values,
        meta["sim_enhanced_pred_lng"].values,
        meta["sim_enhanced_pred_lat"].values,
    )

    meta["dist_direct_cluster"] = pairwise_distance_array(
        meta["direct_pred_lng"].values,
        meta["direct_pred_lat"].values,
        meta["cluster_top30_pred_lng"].values,
        meta["cluster_top30_pred_lat"].values,
    )

    meta["dist_sim_enhanced_cluster"] = pairwise_distance_array(
        meta["sim_enhanced_pred_lng"].values,
        meta["sim_enhanced_pred_lat"].values,
        meta["cluster_top30_pred_lng"].values,
        meta["cluster_top30_pred_lat"].values,
    )

    meta["dist_current_direct"] = pairwise_distance_array(
        meta["current_lng"].values,
        meta["current_lat"].values,
        meta["direct_pred_lng"].values,
        meta["direct_pred_lat"].values,
    )

    meta["dist_current_sim_enhanced"] = pairwise_distance_array(
        meta["current_lng"].values,
        meta["current_lat"].values,
        meta["sim_enhanced_pred_lng"].values,
        meta["sim_enhanced_pred_lat"].values,
    )

    meta["dist_current_cluster"] = pairwise_distance_array(
        meta["current_lng"].values,
        meta["current_lat"].values,
        meta["cluster_top30_pred_lng"].values,
        meta["cluster_top30_pred_lat"].values,
    )

    for col in ["sim_min_dist", "sim_mean_dist", "sim_weighted_dist", "sim_endpoint_dispersion_km"]:
        meta[col] = sim_features[col].values

    for col in CLUSTER_CONFIDENCE_COLUMNS:
        meta[col] = cluster_confidence[col].values

    return meta


# ============================================================
# 7. 实验流程
# ============================================================

def check_columns(df, file_path):
    missing = [c for c in BASE_FEATURE_COLUMNS + TARGET_COLUMNS if c not in df.columns]
    if missing:
        raise ValueError(
            f"{file_path} 缺少字段：{missing}\n"
            f"请先运行 01_preprocess_enhanced.py 生成增强特征。"
        )


def add_result(rows, prefix_ratio, model_name, true_lng, true_lat, pred_lng, pred_lat, baseline_mean, xgb_mean=None, extra=None):
    errors = calc_distance_errors(
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=pred_lng,
        pred_lat=pred_lat,
    )

    metrics = evaluate_errors(errors)

    row = {
        "prefix_ratio": prefix_ratio,
        "model_name": model_name,
        **metrics,
        "improvement_vs_baseline_percent": (baseline_mean - metrics["mean_error_km"]) / baseline_mean * 100,
    }

    if xgb_mean is not None:
        row["improvement_vs_xgboost_percent"] = (xgb_mean - metrics["mean_error_km"]) / xgb_mean * 100

    if extra:
        row.update(extra)

    rows.append(row)

    return errors, metrics


def run_one_prefix(prefix_ratio, file_path):
    print("\n" + "=" * 90)
    print(f"开始混合模型实验：{prefix_ratio}% 前缀")
    print("=" * 90)

    if MAX_ROWS_PER_PREFIX is None:
        df = pd.read_csv(file_path)
    else:
        df = pd.read_csv(file_path, nrows=MAX_ROWS_PER_PREFIX)

    check_columns(df, file_path)
    print("样本数：", len(df))

    # 三段切分：
    # train_base：训练基础模型
    # val_meta：训练融合模型
    # test：最终测试
    train_idx, temp_idx = train_test_split(
        df.index,
        test_size=0.40,
        random_state=RANDOM_STATE,
    )

    val_idx, test_idx = train_test_split(
        temp_idx,
        test_size=0.50,
        random_state=RANDOM_STATE,
    )

    train_df = df.loc[train_idx].reset_index(drop=True)
    val_df = df.loc[val_idx].reset_index(drop=True)
    test_df = df.loc[test_idx].reset_index(drop=True)

    print("基础模型训练集数量：", len(train_df))
    print("融合模型验证集数量：", len(val_df))
    print("最终测试集数量：", len(test_df))

    result_rows = []

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

    # -------------------------
    # 1. Baseline
    # -------------------------
    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)
    baseline_mean = baseline_metrics["mean_error_km"]

    result_rows.append({
        "prefix_ratio": prefix_ratio,
        "model_name": "Baseline_CurrentPoint",
        **baseline_metrics,
        "improvement_vs_baseline_percent": 0.0,
        "improvement_vs_xgboost_percent": np.nan,
    })

    print(f"Baseline 平均误差/km：{baseline_mean:.4f}")

    # -------------------------
    # 2. 单独 XGBoost
    # -------------------------
    print("正在训练单独 XGBoost 直接回归...")
    direct_pred_test, direct_model = train_direct_xgboost(train_df, test_df)
    direct_pred_val = direct_model.predict(val_df[BASE_FEATURE_COLUMNS].values)

    direct_errors, direct_metrics = add_result(
        rows=result_rows,
        prefix_ratio=prefix_ratio,
        model_name="XGBoost_DirectRegression",
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=direct_pred_test[:, 0],
        pred_lat=direct_pred_test[:, 1],
        baseline_mean=baseline_mean,
        xgb_mean=None,
    )
    xgb_mean = direct_metrics["mean_error_km"]

    # 补上自己相对自己的提升
    result_rows[-1]["improvement_vs_xgboost_percent"] = 0.0

    print(
        f"单独 XGBoost 平均误差/km：{xgb_mean:.4f}，"
        f"相比基线提升：{result_rows[-1]['improvement_vs_baseline_percent']:.2f}%"
    )

    # -------------------------
    # 3. 轨迹相似度与相似度增强 XGBoost
    # -------------------------
    print("正在构造轨迹相似度特征...")

    train_vectors = build_similarity_vectors(train_df)
    val_vectors = build_similarity_vectors(val_df)
    test_vectors = build_similarity_vectors(test_df)

    train_targets = train_df[TARGET_COLUMNS].values

    train_sim_features, train_sim_pred = make_similarity_features_train(
        train_vectors=train_vectors,
        train_targets=train_targets,
    )
    train_sim_features = add_similarity_distance_features(train_df, train_sim_features)

    nn = build_nn_index(train_vectors, K_NEIGHBORS)

    val_sim_features, val_sim_pred = make_similarity_features_query(
        nn=nn,
        query_vectors=val_vectors,
        train_targets=train_targets,
    )
    val_sim_features = add_similarity_distance_features(val_df, val_sim_features)

    test_sim_features, test_sim_pred = make_similarity_features_query(
        nn=nn,
        query_vectors=test_vectors,
        train_targets=train_targets,
    )
    test_sim_features = add_similarity_distance_features(test_df, test_sim_features)

    sim_errors, sim_metrics = add_result(
        rows=result_rows,
        prefix_ratio=prefix_ratio,
        model_name="TrajectorySimilarityOnly",
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=test_sim_pred[:, 0],
        pred_lat=test_sim_pred[:, 1],
        baseline_mean=baseline_mean,
        xgb_mean=xgb_mean,
    )

    print(
        f"相似轨迹单独预测 平均误差/km：{sim_metrics['mean_error_km']:.4f}，"
        f"相比 XGBoost：{result_rows[-1]['improvement_vs_xgboost_percent']:.2f}%"
    )

    print("正在训练相似度增强 XGBoost...")
    sim_enhanced_model, sim_enhanced_feature_columns = train_similarity_enhanced_xgboost(
        train_df=train_df,
        train_sim_features=train_sim_features,
    )

    sim_enhanced_pred_val = predict_similarity_enhanced_xgboost(
        model=sim_enhanced_model,
        query_df=val_df,
        query_sim_features=val_sim_features,
        feature_columns=sim_enhanced_feature_columns,
    )

    sim_enhanced_pred_test = predict_similarity_enhanced_xgboost(
        model=sim_enhanced_model,
        query_df=test_df,
        query_sim_features=test_sim_features,
        feature_columns=sim_enhanced_feature_columns,
    )

    sim_enhanced_errors, sim_enhanced_metrics = add_result(
        rows=result_rows,
        prefix_ratio=prefix_ratio,
        model_name="SimilarityEnhancedXGBoost",
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=sim_enhanced_pred_test[:, 0],
        pred_lat=sim_enhanced_pred_test[:, 1],
        baseline_mean=baseline_mean,
        xgb_mean=xgb_mean,
    )

    print(
        f"相似度增强 XGBoost 平均误差/km：{sim_enhanced_metrics['mean_error_km']:.4f}，"
        f"相比 XGBoost：{result_rows[-1]['improvement_vs_xgboost_percent']:.2f}%"
    )

    # -------------------------
    # 4. Top-30 聚类 + 类内偏移 XGBoost
    # -------------------------
    print("正在训练 Top-30 聚类 + 类内偏移模型...")
    km, train_labels, centers_df, clf, offset_model, offset_feature_columns = train_cluster_top30_model(train_df)

    val_true_labels = assign_cluster_labels(km, val_df)
    test_true_labels = assign_cluster_labels(km, test_df)

    cluster_val_lng, cluster_val_lat, cluster_conf_val, prob_val, top_clusters_val = predict_cluster_top30(
        query_df=val_df,
        km=km,
        clf=clf,
        centers_df=centers_df,
        offset_model=offset_model,
        offset_feature_columns=offset_feature_columns,
    )

    cluster_test_lng, cluster_test_lat, cluster_conf_test, prob_test, top_clusters_test = predict_cluster_top30(
        query_df=test_df,
        km=km,
        clf=clf,
        centers_df=centers_df,
        offset_model=offset_model,
        offset_feature_columns=offset_feature_columns,
    )

    top1_acc = float(np.mean(top_clusters_test[:, 0] == test_true_labels))
    top3_hit = float(np.mean([test_true_labels[i] in set(top_clusters_test[i, :3].tolist()) for i in range(len(test_true_labels))]))
    top10_hit = float(np.mean([test_true_labels[i] in set(top_clusters_test[i, :10].tolist()) for i in range(len(test_true_labels))]))
    top30_hit = float(np.mean([test_true_labels[i] in set(top_clusters_test[i, :30].tolist()) for i in range(len(test_true_labels))]))

    cluster_errors, cluster_metrics = add_result(
        rows=result_rows,
        prefix_ratio=prefix_ratio,
        model_name="ClusterTop30OffsetXGBoost",
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=cluster_test_lng,
        pred_lat=cluster_test_lat,
        baseline_mean=baseline_mean,
        xgb_mean=xgb_mean,
        extra={
            "cluster_top1_acc": top1_acc,
            "cluster_top3_hit_rate": top3_hit,
            "cluster_top10_hit_rate": top10_hit,
            "cluster_top30_hit_rate": top30_hit,
        }
    )

    print(
        f"ClusterTop30OffsetXGBoost 平均误差/km：{cluster_metrics['mean_error_km']:.4f}，"
        f"相比 XGBoost：{result_rows[-1]['improvement_vs_xgboost_percent']:.2f}%"
    )
    print(f"聚类 Top-1 Acc={top1_acc:.4f}, Top-3 Hit={top3_hit:.4f}, Top-10 Hit={top10_hit:.4f}, Top-30 Hit={top30_hit:.4f}")

    # -------------------------
    # 5. 元融合模型
    # -------------------------
    print("正在构造融合特征并训练 HybridFusionXGBoost...")

    cluster_pred_val = np.column_stack([cluster_val_lng, cluster_val_lat])
    cluster_pred_test = np.column_stack([cluster_test_lng, cluster_test_lat])

    meta_val = build_meta_features(
        df=val_df,
        direct_pred=direct_pred_val,
        sim_only_pred=val_sim_pred,
        sim_enhanced_pred=sim_enhanced_pred_val,
        cluster_pred=cluster_pred_val,
        sim_features=val_sim_features,
        cluster_confidence=cluster_conf_val,
    )

    meta_test = build_meta_features(
        df=test_df,
        direct_pred=direct_pred_test,
        sim_only_pred=test_sim_pred,
        sim_enhanced_pred=sim_enhanced_pred_test,
        cluster_pred=cluster_pred_test,
        sim_features=test_sim_features,
        cluster_confidence=cluster_conf_test,
    )

    fusion_model = make_meta_fusion_regressor()
    fusion_model.fit(meta_val[META_FEATURE_COLUMNS].values, val_df[TARGET_COLUMNS].values)

    fusion_pred_test = fusion_model.predict(meta_test[META_FEATURE_COLUMNS].values)

    fusion_errors, fusion_metrics = add_result(
        rows=result_rows,
        prefix_ratio=prefix_ratio,
        model_name="HybridFusionXGBoost",
        true_lng=true_lng,
        true_lat=true_lat,
        pred_lng=fusion_pred_test[:, 0],
        pred_lat=fusion_pred_test[:, 1],
        baseline_mean=baseline_mean,
        xgb_mean=xgb_mean,
        extra={
            "cluster_top1_acc": top1_acc,
            "cluster_top3_hit_rate": top3_hit,
            "cluster_top10_hit_rate": top10_hit,
            "cluster_top30_hit_rate": top30_hit,
        }
    )

    print(
        f"HybridFusionXGBoost 平均误差/km：{fusion_metrics['mean_error_km']:.4f}，"
        f"相比 XGBoost：{result_rows[-1]['improvement_vs_xgboost_percent']:.2f}%"
    )

    # -------------------------
    # 6. 保存预测明细
    # -------------------------
    pred_df = test_df[["TRIP_ID", "current_lng", "current_lat", "target_lng", "target_lat"]].copy()

    pred_df["baseline_error_km"] = baseline_errors

    pred_df["direct_pred_lng"] = direct_pred_test[:, 0]
    pred_df["direct_pred_lat"] = direct_pred_test[:, 1]
    pred_df["direct_error_km"] = direct_errors

    pred_df["sim_only_pred_lng"] = test_sim_pred[:, 0]
    pred_df["sim_only_pred_lat"] = test_sim_pred[:, 1]
    pred_df["sim_only_error_km"] = sim_errors

    pred_df["sim_enhanced_pred_lng"] = sim_enhanced_pred_test[:, 0]
    pred_df["sim_enhanced_pred_lat"] = sim_enhanced_pred_test[:, 1]
    pred_df["sim_enhanced_error_km"] = sim_enhanced_errors

    pred_df["cluster_top30_pred_lng"] = cluster_test_lng
    pred_df["cluster_top30_pred_lat"] = cluster_test_lat
    pred_df["cluster_top30_error_km"] = cluster_errors

    pred_df["fusion_pred_lng"] = fusion_pred_test[:, 0]
    pred_df["fusion_pred_lat"] = fusion_pred_test[:, 1]
    pred_df["fusion_error_km"] = fusion_errors

    pred_df["true_cluster"] = test_true_labels
    pred_df["cluster_top1"] = cluster_conf_test["cluster_top1"].values
    pred_df["cluster_top2"] = cluster_conf_test["cluster_top2"].values
    pred_df["cluster_top3"] = cluster_conf_test["cluster_top3"].values
    pred_df["cluster_top1_prob"] = cluster_conf_test["cluster_top1_prob"].values
    pred_df["cluster_top3_prob_sum"] = cluster_conf_test["cluster_top3_prob_sum"].values
    pred_df["cluster_top10_prob_sum"] = cluster_conf_test["cluster_top10_prob_sum"].values
    pred_df["cluster_top30_prob_sum"] = cluster_conf_test["cluster_top30_prob_sum"].values
    pred_df["cluster_prob_entropy"] = cluster_conf_test["cluster_prob_entropy"].values

    pred_path = os.path.join(
        OUTPUT_DIR,
        f"hybrid_cluster_similarity_xgboost_predictions_prefix_{prefix_ratio}.csv"
    )
    pred_df.to_csv(pred_path, index=False, encoding="utf-8-sig")
    print("预测明细已保存：", pred_path)

    meta_path = os.path.join(
        OUTPUT_DIR,
        f"hybrid_meta_features_test_prefix_{prefix_ratio}.csv"
    )
    meta_test.to_csv(meta_path, index=False, encoding="utf-8-sig")
    print("融合特征已保存：", meta_path)

    return result_rows


# ============================================================
# 8. 主程序
# ============================================================

def main():
    all_results = []

    for prefix_ratio, file_path in PREFIX_FILES.items():
        if not os.path.exists(file_path):
            raise FileNotFoundError(
                f"没有找到 {file_path}。\n"
                f"请先运行 01_preprocess_enhanced.py。"
            )

        all_results.extend(run_one_prefix(prefix_ratio, file_path))

    result_df = pd.DataFrame(all_results)

    result_path = os.path.join(OUTPUT_DIR, "hybrid_cluster_similarity_xgboost_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, "hybrid_cluster_similarity_xgboost_best_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()
