import ast
import json
import math
import os

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.cluster import MiniBatchKMeans
from sklearn.metrics import silhouette_score


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

DATA_PATH = "data/train.csv"
OUTPUT_DIR = "output"
FIGURE_DIR = os.path.join(OUTPUT_DIR, "figures")
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(FIGURE_DIR, exist_ok=True)

# None 表示读取全部数据。
# 如果你只是先调试，可以改成 100000。
MAX_ROWS = None

CHUNK_SIZE = 50000

MIN_POINTS = 2

MIN_LNG, MAX_LNG = -8.80, -8.45
MIN_LAT, MAX_LAT = 41.00, 41.35

# 聚类数量。可以先用 50，后面根据肘部图调整。
N_CLUSTERS = 50

# 肘部图候选 K
K_CANDIDATES = [10, 20, 30, 50, 80, 120]

# 肘部图和轮廓系数采样数量，避免全量计算太慢
ELBOW_SAMPLE_SIZE = 50000
SILHOUETTE_SAMPLE_SIZE = 10000

# 可视化时采样点数，避免图太密
VIS_SAMPLE_SIZE = 80000

RANDOM_STATE = 42

REF_LNG = -8.61
REF_LAT = 41.16

plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "Arial Unicode MS"]
plt.rcParams["axes.unicode_minus"] = False


# ============================================================
# 2. 基础函数
# ============================================================

def parse_polyline(polyline_str):
    if not isinstance(polyline_str, str) or len(polyline_str) < 2:
        return []

    try:
        points = json.loads(polyline_str)
    except Exception:
        try:
            points = ast.literal_eval(polyline_str)
        except Exception:
            return []

    if not isinstance(points, list):
        return []

    return points


def is_valid_point(point):
    if not isinstance(point, list) or len(point) != 2:
        return False

    lng, lat = point

    if not isinstance(lng, (int, float)) or not isinstance(lat, (int, float)):
        return False

    return MIN_LNG <= lng <= MAX_LNG and MIN_LAT <= lat <= MAX_LAT


def lonlat_to_xy_km_array(lng, lat):
    x = (lng - REF_LNG) * 111.320 * math.cos(math.radians(REF_LAT))
    y = (lat - REF_LAT) * 110.574
    return x, y


def xy_km_to_lonlat_array(x, y):
    lng = x / (111.320 * math.cos(math.radians(REF_LAT))) + REF_LNG
    lat = y / 110.574 + REF_LAT
    return lng, lat


def load_destination_points():
    """
    从 train.csv 中提取所有有效轨迹的终点。
    结果会缓存到 output/destination_points_cache.csv，后面重复运行会快很多。
    """
    cache_path = os.path.join(OUTPUT_DIR, "destination_points_cache.csv")

    if os.path.exists(cache_path):
        print("检测到缓存文件，直接读取：", cache_path)
        cached = pd.read_csv(cache_path)
        return cached

    print("开始从原始数据中提取终点。全量 100w 数据可能需要一些时间。")

    usecols = ["TRIP_ID", "MISSING_DATA", "POLYLINE"]

    all_rows = []
    total_read = 0
    total_valid = 0

    reader = pd.read_csv(
        DATA_PATH,
        usecols=lambda c: c in usecols,
        chunksize=CHUNK_SIZE,
        nrows=MAX_ROWS,
    )

    for chunk_id, chunk in enumerate(reader, start=1):
        total_read += len(chunk)

        chunk = chunk[chunk["MISSING_DATA"].astype(str).str.upper() == "FALSE"].copy()

        for row in chunk.itertuples(index=False):
            points = parse_polyline(row.POLYLINE)

            if len(points) < MIN_POINTS:
                continue

            end_point = points[-1]

            if not is_valid_point(end_point):
                continue

            end_lng, end_lat = end_point

            all_rows.append({
                "TRIP_ID": row.TRIP_ID,
                "end_lng": end_lng,
                "end_lat": end_lat,
            })
            total_valid += 1

        print(f"已读取原始行数：{total_read}，有效终点数：{total_valid}")

    df = pd.DataFrame(all_rows)

    df.to_csv(cache_path, index=False, encoding="utf-8-sig")
    print("终点缓存已保存：", cache_path)

    return df


def make_xy_matrix(df):
    x_list = []
    y_list = []

    for lng, lat in zip(df["end_lng"].values, df["end_lat"].values):
        x, y = lonlat_to_xy_km_array(lng, lat)
        x_list.append(x)
        y_list.append(y)

    X = np.column_stack([x_list, y_list]).astype(np.float32)
    return X


# ============================================================
# 3. 聚类与密度统计
# ============================================================

def run_elbow_analysis(X):
    """
    通过不同 K 的 inertia 变化观察是否存在明显聚类结构。
    """
    n = len(X)
    sample_n = min(ELBOW_SAMPLE_SIZE, n)

    rng = np.random.default_rng(RANDOM_STATE)
    sample_idx = rng.choice(n, size=sample_n, replace=False)
    X_sample = X[sample_idx]

    rows = []

    print("正在计算肘部图数据...")

    for k in K_CANDIDATES:
        print(f"正在测试 K={k}")

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

        km.fit(X_sample)

        rows.append({
            "k": k,
            "inertia": km.inertia_,
        })

    elbow_df = pd.DataFrame(rows)

    elbow_path = os.path.join(OUTPUT_DIR, "destination_kmeans_elbow.csv")
    elbow_df.to_csv(elbow_path, index=False, encoding="utf-8-sig")
    print("肘部图数据已保存：", elbow_path)

    plt.figure(figsize=(8, 5))
    plt.plot(elbow_df["k"], elbow_df["inertia"], marker="o")
    plt.xlabel("聚类数量 K")
    plt.ylabel("Inertia")
    plt.title("终点聚类 KMeans 肘部图")
    plt.grid(True)
    plt.tight_layout()

    fig_path = os.path.join(FIGURE_DIR, "destination_kmeans_elbow.png")
    plt.savefig(fig_path, dpi=300, bbox_inches="tight")
    plt.show()
    print("肘部图已保存：", fig_path)

    return elbow_df


def cluster_destinations(X, df):
    print(f"正在使用 MiniBatchKMeans 聚类，K={N_CLUSTERS}")

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

    labels = km.fit_predict(X)

    df = df.copy()
    df["cluster_id"] = labels

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

    # 每个点到自己聚类中心的距离，单位 km
    center_for_point = centers_xy[labels]
    dist_to_center = np.sqrt(((X - center_for_point) ** 2).sum(axis=1))
    df["dist_to_cluster_center_km"] = dist_to_center

    total_n = len(df)
    rows = []

    for cluster_id in range(N_CLUSTERS):
        sub = df[df["cluster_id"] == cluster_id]
        count = len(sub)

        if count == 0:
            continue

        distances = sub["dist_to_cluster_center_km"].values

        mean_radius = float(np.mean(distances))
        p50_radius = float(np.quantile(distances, 0.50))
        p90_radius = float(np.quantile(distances, 0.90))

        # 使用 p90 半径估算覆盖面积，得到一个粗略密度。
        if p90_radius > 1e-6:
            area_km2 = math.pi * (p90_radius ** 2)
            density_p90 = count / area_km2
        else:
            area_km2 = 0.0
            density_p90 = 0.0

        rows.append({
            "cluster_id": cluster_id,
            "count": count,
            "percent": count / total_n * 100,
            "center_lng": center_lng[cluster_id],
            "center_lat": center_lat[cluster_id],
            "mean_radius_km": mean_radius,
            "p50_radius_km": p50_radius,
            "p90_radius_km": p90_radius,
            "area_by_p90_km2": area_km2,
            "density_by_p90_count_per_km2": density_p90,
        })

    cluster_df = pd.DataFrame(rows)

    cluster_df = cluster_df.sort_values("count", ascending=False)

    cluster_path = os.path.join(OUTPUT_DIR, "destination_cluster_density.csv")
    cluster_df.to_csv(cluster_path, index=False, encoding="utf-8-sig")

    labeled_path = os.path.join(OUTPUT_DIR, "destination_points_with_clusters.csv")
    df.to_csv(labeled_path, index=False, encoding="utf-8-sig")

    print("聚类密度表已保存：", cluster_path)
    print("带聚类标签的终点数据已保存：", labeled_path)

    return df, cluster_df, km


def calc_silhouette(X, labels):
    """
    抽样计算轮廓系数。
    轮廓系数越接近 1，聚类结构越明显；接近 0 表示边界不明显。
    """
    n = len(X)

    if n < 3:
        return None

    sample_n = min(SILHOUETTE_SAMPLE_SIZE, n)

    rng = np.random.default_rng(RANDOM_STATE)
    idx = rng.choice(n, size=sample_n, replace=False)

    print(f"正在抽样计算轮廓系数，样本数：{sample_n}")

    score = silhouette_score(X[idx], labels[idx])

    score_path = os.path.join(OUTPUT_DIR, "destination_cluster_silhouette.txt")
    with open(score_path, "w", encoding="utf-8") as f:
        f.write(f"silhouette_score={score}\n")
        f.write("说明：轮廓系数越接近 1，聚类越明显；接近 0 表示类间边界不清楚。\n")

    print("轮廓系数：", score)
    print("轮廓系数已保存：", score_path)

    return score


# ============================================================
# 4. 可视化
# ============================================================

def plot_cluster_scatter(df, cluster_df):
    """
    终点聚类散点图。
    """
    n = len(df)
    sample_n = min(VIS_SAMPLE_SIZE, n)

    sample_df = df.sample(n=sample_n, random_state=RANDOM_STATE)

    plt.figure(figsize=(8, 7))
    plt.scatter(
        sample_df["end_lng"],
        sample_df["end_lat"],
        c=sample_df["cluster_id"],
        s=2,
        alpha=0.45,
        cmap="tab20",
    )

    plt.scatter(
        cluster_df["center_lng"],
        cluster_df["center_lat"],
        s=np.sqrt(cluster_df["count"]) * 3,
        marker="x",
        linewidths=1.5,
    )

    plt.xlabel("经度")
    plt.ylabel("纬度")
    plt.title(f"出租车终点聚类可视化（K={N_CLUSTERS}，抽样 {sample_n} 点）")
    plt.grid(True)
    plt.tight_layout()

    fig_path = os.path.join(FIGURE_DIR, "destination_clusters_scatter.png")
    plt.savefig(fig_path, dpi=300, bbox_inches="tight")
    plt.show()

    print("终点聚类散点图已保存：", fig_path)


def plot_cluster_count_bar(cluster_df):
    """
    按样本数量显示前 20 个聚类。
    """
    top_df = cluster_df.sort_values("count", ascending=False).head(20).sort_values("count")

    plt.figure(figsize=(9, 6))
    plt.barh(top_df["cluster_id"].astype(str), top_df["count"])
    plt.xlabel("终点数量")
    plt.ylabel("聚类编号")
    plt.title("终点数量最多的前 20 个聚类")
    plt.tight_layout()

    fig_path = os.path.join(FIGURE_DIR, "destination_cluster_count_top20.png")
    plt.savefig(fig_path, dpi=300, bbox_inches="tight")
    plt.show()

    print("聚类数量 Top20 图已保存：", fig_path)


def plot_cluster_density_bar(cluster_df):
    """
    按 p90 半径估算密度，显示密度最高的前 20 个聚类。
    """
    top_df = (
        cluster_df.sort_values("density_by_p90_count_per_km2", ascending=False)
        .head(20)
        .sort_values("density_by_p90_count_per_km2")
    )

    plt.figure(figsize=(9, 6))
    plt.barh(
        top_df["cluster_id"].astype(str),
        top_df["density_by_p90_count_per_km2"],
    )
    plt.xlabel("密度（终点数 / km²，基于 p90 半径估算）")
    plt.ylabel("聚类编号")
    plt.title("终点密度最高的前 20 个聚类")
    plt.tight_layout()

    fig_path = os.path.join(FIGURE_DIR, "destination_cluster_density_top20.png")
    plt.savefig(fig_path, dpi=300, bbox_inches="tight")
    plt.show()

    print("聚类密度 Top20 图已保存：", fig_path)


def plot_radius_distribution(cluster_df):
    """
    聚类 p90 半径分布。
    半径越小、数量越多，说明终点越集中。
    """
    plt.figure(figsize=(8, 5))
    plt.hist(cluster_df["p90_radius_km"], bins=20)
    plt.xlabel("聚类 p90 半径 / km")
    plt.ylabel("聚类数量")
    plt.title("终点聚类 p90 半径分布")
    plt.tight_layout()

    fig_path = os.path.join(FIGURE_DIR, "destination_cluster_radius_distribution.png")
    plt.savefig(fig_path, dpi=300, bbox_inches="tight")
    plt.show()

    print("聚类半径分布图已保存：", fig_path)


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

def main():
    dest_df = load_destination_points()

    print("有效终点数量：", len(dest_df))
    print(dest_df.head())

    X = make_xy_matrix(dest_df)

    elbow_df = run_elbow_analysis(X)

    labeled_df, cluster_df, km = cluster_destinations(X, dest_df)

    silhouette = calc_silhouette(X, labeled_df["cluster_id"].values)

    print("\n====== 终点数量最多的前 10 个聚类 ======")
    print(cluster_df.sort_values("count", ascending=False).head(10))

    print("\n====== 终点密度最高的前 10 个聚类 ======")
    print(cluster_df.sort_values("density_by_p90_count_per_km2", ascending=False).head(10))

    if silhouette is not None:
        print("\n====== 聚类结构初步判断 ======")
        if silhouette >= 0.35:
            print("轮廓系数较高，说明终点存在较明显的聚类结构。")
        elif silhouette >= 0.15:
            print("轮廓系数中等，说明终点有一定聚集趋势，但边界不是特别清楚。")
        else:
            print("轮廓系数偏低，说明终点总体更像连续空间分布，局部有高密度区域，但类间边界不明显。")

    plot_cluster_scatter(labeled_df, cluster_df)
    plot_cluster_count_bar(cluster_df)
    plot_cluster_density_bar(cluster_df)
    plot_radius_distribution(cluster_df)

    print("\n全部完成。")
    print("请重点查看：")
    print("1. output/destination_cluster_density.csv")
    print("2. output/destination_kmeans_elbow.csv")
    print("3. output/destination_cluster_silhouette.txt")
    print("4. output/figures/destination_clusters_scatter.png")
    print("5. output/figures/destination_cluster_density_top20.png")


if __name__ == "__main__":
    main()
