import ast
import math
import os
import pandas as pd
from datetime import datetime

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

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

# 第一次建议只跑 50000 条，确认没问题后再改大
N_ROWS = 5000

# 至少保留 10 个轨迹点
# 太短的轨迹不适合做终点预测
MIN_POINTS = 10

# Porto 大致经纬度范围，设置得宽松一点
# 主要用于过滤明显错误的 GPS 点
MIN_LNG, MAX_LNG = -8.80, -8.45
MIN_LAT, MAX_LAT = 41.00, 41.35


# =========================
# 2. 工具函数
# =========================

def parse_polyline(polyline_str):
    """
    将 POLYLINE 字符串解析成 Python 列表。
    原始格式类似：
    [[-8.628399,41.157288],[-8.628327,41.157351]]
    """
    try:
        points = ast.literal_eval(polyline_str)
        if not isinstance(points, list):
            return []
        return points
    except Exception:
        return []


def is_valid_point(point):
    """
    判断单个 GPS 点是否合法。
    数据中每个点格式是 [经度, 纬度]。
    """
    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

    if not (MIN_LNG <= lng <= MAX_LNG):
        return False

    if not (MIN_LAT <= lat <= MAX_LAT):
        return False

    return True


def clean_points(points):
    """
    清洗轨迹点。
    这里不删除重复点，因为重复点可能表示车辆等待或红灯停车。
    只删除格式错误或明显超出范围的点。
    """
    return [p for p in points if is_valid_point(p)]


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 trajectory_distance(points):
    """
    计算整条轨迹的实际行驶距离，单位：公里。
    """
    if len(points) < 2:
        return 0.0

    total = 0.0
    for i in range(1, len(points)):
        lng1, lat1 = points[i - 1]
        lng2, lat2 = points[i]
        total += haversine(lng1, lat1, lng2, lat2)

    return total


def bearing(lng1, lat1, lng2, lat2):
    """
    计算从起点到当前点的大致方向角。
    方向角可以反映车辆行驶方向。
    """
    lng1, lat1, lng2, lat2 = map(math.radians, [lng1, lat1, lng2, lat2])

    dlng = lng2 - lng1
    x = math.sin(dlng) * math.cos(lat2)
    y = math.cos(lat1) * math.sin(lat2) - math.sin(lat1) * math.cos(lat2) * math.cos(dlng)

    angle = math.degrees(math.atan2(x, y))
    return (angle + 360) % 360


def timestamp_features(ts):
    """
    从时间戳中提取小时、星期、是否周末。
    """
    dt = datetime.fromtimestamp(int(ts))

    hour = dt.hour
    weekday = dt.weekday()
    is_weekend = 1 if weekday >= 5 else 0

    return hour, weekday, is_weekend


# =========================
# 3. 读取数据
# =========================

print("正在读取数据...")

df = pd.read_csv(DATA_PATH, nrows=N_ROWS)

print("原始读取数据量：", len(df))
print("字段列表：", df.columns.tolist())


# =========================
# 4. 删除缺失轨迹
# =========================

before = len(df)

# MISSING_DATA 在原始文件中可能是布尔值，也可能是字符串
df = df[df["MISSING_DATA"].astype(str).str.upper() == "FALSE"].copy()

after = len(df)
print("删除 MISSING_DATA=True 的数量：", before - after)


# =========================
# 5. 解析并清洗 POLYLINE
# =========================

print("正在解析 POLYLINE...")

df["points"] = df["POLYLINE"].apply(parse_polyline)
df["points"] = df["points"].apply(clean_points)
df["num_points"] = df["points"].apply(len)

before = len(df)

# 删除轨迹点太少的样本
df = df[df["num_points"] >= MIN_POINTS].copy()

after = len(df)
print("删除短轨迹数量：", before - after)
print("清洗后剩余数据量：", len(df))


# =========================
# 6. 提取完整轨迹基础信息
# =========================

print("正在提取完整轨迹特征...")

df["start_lng"] = df["points"].apply(lambda x: x[0][0])
df["start_lat"] = df["points"].apply(lambda x: x[0][1])
df["end_lng"] = df["points"].apply(lambda x: x[-1][0])
df["end_lat"] = df["points"].apply(lambda x: x[-1][1])

# Porto 数据每 15 秒一个点
# 如果有 n 个点，行程持续时间约为 (n - 1) * 15 秒
df["duration_sec"] = (df["num_points"] - 1) * 15

df["total_distance_km"] = df["points"].apply(trajectory_distance)

df[["hour", "weekday", "is_weekend"]] = df["TIMESTAMP"].apply(
    lambda ts: pd.Series(timestamp_features(ts))
)

# 保存清洗后的完整轨迹信息
clean_cols = [
    "TRIP_ID",
    "TAXI_ID",
    "CALL_TYPE",
    "TIMESTAMP",
    "hour",
    "weekday",
    "is_weekend",
    "num_points",
    "duration_sec",
    "total_distance_km",
    "start_lng",
    "start_lat",
    "end_lng",
    "end_lat",
    "points"
]

df[clean_cols].to_csv(
    os.path.join(OUTPUT_DIR, "clean_trips.csv"),
    index=False,
    encoding="utf-8-sig"
)

print("已保存 clean_trips.csv")


# =========================
# 7. 构造前缀数据集
# =========================

def get_sample_point(points, ratio):
    """
    从轨迹中按照比例取一个采样点。
    ratio=0.25 表示取前缀轨迹 25% 位置附近的点。
    """
    if len(points) == 0:
        return [0.0, 0.0]

    idx = int((len(points) - 1) * ratio)
    idx = max(0, min(idx, len(points) - 1))
    return points[idx]


def build_prefix_features(row, prefix_ratio):
    """
    根据前缀比例构造一条样本的模型特征。
    增强版：在原有位置、时间、距离、方向特征基础上，
    加入轨迹形状采样点、弯曲度、最近一段方向等特征。
    """
    points = row["points"]
    n = len(points)

    prefix_len = max(2, int(n * prefix_ratio))
    prefix_points = points[:prefix_len]

    start_lng, start_lat = prefix_points[0]
    current_lng, current_lat = prefix_points[-1]

    end_lng, end_lat = points[-1]

    prefix_distance_km = trajectory_distance(prefix_points)
    direct_distance_km = haversine(start_lng, start_lat, current_lng, current_lat)

    prefix_duration_sec = (prefix_len - 1) * 15

    if prefix_duration_sec > 0:
        prefix_avg_speed = prefix_distance_km / (prefix_duration_sec / 3600)
    else:
        prefix_avg_speed = 0.0

    direction_angle = bearing(start_lng, start_lat, current_lng, current_lat)

    lng_delta = current_lng - start_lng
    lat_delta = current_lat - start_lat

    # 轨迹弯曲度：实际行驶距离 / 起点到当前点直线距离
    # 如果接近 1，说明轨迹比较直；如果明显大于 1，说明存在绕行或转弯。
    if direct_distance_km > 0:
        tortuosity = prefix_distance_km / direct_distance_km
    else:
        tortuosity = 1.0

    # 对前缀轨迹做形状采样
    # 让模型不仅看到起点和当前点，也能看到中间路径形状。
    p25_lng, p25_lat = get_sample_point(prefix_points, 0.25)
    p50_lng, p50_lat = get_sample_point(prefix_points, 0.50)
    p75_lng, p75_lat = get_sample_point(prefix_points, 0.75)

    # 最近一小段轨迹方向
    if len(prefix_points) >= 3:
        prev_lng, prev_lat = prefix_points[-3]
        recent_direction_angle = bearing(prev_lng, prev_lat, current_lng, current_lat)
        recent_distance_km = haversine(prev_lng, prev_lat, current_lng, current_lat)
    else:
        recent_direction_angle = direction_angle
        recent_distance_km = direct_distance_km

    # 前缀轨迹经纬度范围，用于描述行驶区域大小
    lng_values = [p[0] for p in prefix_points]
    lat_values = [p[1] for p in prefix_points]

    lng_range = max(lng_values) - min(lng_values)
    lat_range = max(lat_values) - min(lat_values)

    # 轨迹前半段和后半段的方向变化
    mid_point = get_sample_point(prefix_points, 0.50)
    mid_lng, mid_lat = mid_point

    first_half_angle = bearing(start_lng, start_lat, mid_lng, mid_lat)
    second_half_angle = bearing(mid_lng, mid_lat, current_lng, current_lat)

    angle_change = abs(second_half_angle - first_half_angle)
    if angle_change > 180:
        angle_change = 360 - angle_change

    return {
        "TRIP_ID": row["TRIP_ID"],
        "prefix_ratio": prefix_ratio,

        # 时间特征
        "hour": row["hour"],
        "weekday": row["weekday"],
        "is_weekend": row["is_weekend"],

        # 起点与当前点
        "start_lng": start_lng,
        "start_lat": start_lat,
        "current_lng": current_lng,
        "current_lat": current_lat,

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

        # 增强轨迹形状特征
        "p25_lng": p25_lng,
        "p25_lat": p25_lat,
        "p50_lng": p50_lng,
        "p50_lat": p50_lat,
        "p75_lng": p75_lng,
        "p75_lat": p75_lat,
        "tortuosity": tortuosity,
        "recent_direction_angle": recent_direction_angle,
        "recent_distance_km": recent_distance_km,
        "lng_range": lng_range,
        "lat_range": lat_range,
        "first_half_angle": first_half_angle,
        "second_half_angle": second_half_angle,
        "angle_change": angle_change,

        # 预测目标
        "target_lng": end_lng,
        "target_lat": end_lat,

        # 后续画图用
        "true_num_points": n,
    }


for ratio in [0.2, 0.4, 0.6]:
    print(f"正在构造 {int(ratio * 100)}% 前缀数据集...")

    rows = []
    for _, row in df.iterrows():
        rows.append(build_prefix_features(row, ratio))

    prefix_df = pd.DataFrame(rows)

    output_path = os.path.join(
        OUTPUT_DIR,
        f"dataset_prefix_{int(ratio * 100)}.csv"
    )

    prefix_df.to_csv(output_path, index=False, encoding="utf-8-sig")

    print(f"已保存：{output_path}，样本数：{len(prefix_df)}")

print("====== 数据统计结果 ======")
print("清洗后样本数：", len(df))
print("最短轨迹点数：", df["num_points"].min())
print("最长轨迹点数：", df["num_points"].max())
print("平均轨迹点数：", df["num_points"].mean())
print("平均行程距离/km：", df["total_distance_km"].mean())
print("平均行程时长/秒：", df["duration_sec"].mean())

print("预处理完成。")