import ast
import json
import math
import os
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd

from config import (
    CHUNK_SIZE,
    DATA_PATH,
    MAX_LAT,
    MAX_LNG,
    MIN_LAT,
    MIN_LNG,
    MIN_POINTS,
    PREFIX_RATIOS,
    REF_LAT,
    REF_LNG,
    RANDOM_STATE,
)


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


def ensure_dir(path):
    Path(path).mkdir(parents=True, exist_ok=True)


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 []

    return points if isinstance(points, list) else []


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 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 calc_distance_errors(true_lng, true_lat, pred_lng, pred_lat):
    return np.array([
        haversine(a, b, c, d)
        for a, b, c, d in zip(true_lng, true_lat, pred_lng, pred_lat)
    ], 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 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 angle_diff(angle1, angle2):
    diff = abs(angle1 - angle2)
    return 360 - diff if diff > 180 else diff


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


def get_sample_point(points, ratio):
    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 segment_distance_features(points):
    if len(points) < 2:
        return 0.0, 0.0, 0.0, 0.0

    p0 = get_sample_point(points, 0.00)
    p25 = get_sample_point(points, 0.25)
    p50 = get_sample_point(points, 0.50)
    p75 = get_sample_point(points, 0.75)
    p100 = get_sample_point(points, 1.00)

    d1 = haversine(p0[0], p0[1], p25[0], p25[1])
    d2 = haversine(p25[0], p25[1], p50[0], p50[1])
    d3 = haversine(p50[0], p50[1], p75[0], p75[1])
    d4 = haversine(p75[0], p75[1], p100[0], p100[1])

    return d1, d2, d3, d4


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 sanitize_numeric_df(df):
    df = df.copy()
    df = df.replace([np.inf, -np.inf], np.nan)
    for col in df.columns:
        if pd.api.types.is_numeric_dtype(df[col]):
            med = df[col].median()
            if pd.isna(med):
                med = 0.0
            df[col] = df[col].fillna(med)
    return df


def build_prefix_features(points, row_info, prefix_ratio):
    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
    prefix_avg_speed = prefix_distance_km / (prefix_duration_sec / 3600) if prefix_duration_sec > 0 else 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

    tortuosity = prefix_distance_km / direct_distance_km if direct_distance_km > 1e-8 else 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_lng, mid_lat = get_sample_point(prefix_points, 0.50)
    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 = angle_diff(first_half_angle, second_half_angle)

    seg_d1, seg_d2, seg_d3, seg_d4 = segment_distance_features(prefix_points)

    return {
        "TRIP_ID": row_info["TRIP_ID"],
        "prefix_ratio": prefix_ratio,
        "hour": row_info["hour"],
        "weekday": row_info["weekday"],
        "is_weekend": row_info["is_weekend"],
        "num_points": n,

        "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,

        "seg_d1": seg_d1,
        "seg_d2": seg_d2,
        "seg_d3": seg_d3,
        "seg_d4": seg_d4,

        "target_lng": end_lng,
        "target_lat": end_lat,
    }


def prefix_file_path(output_dir, prefix_ratio, full=False):
    prefix_dir = Path(output_dir) / "prefix_data"
    ensure_dir(prefix_dir)
    tag = "full_dataset" if full else "dataset"
    return prefix_dir / f"{tag}_prefix_{int(prefix_ratio * 100)}.csv"


def all_prefix_files_exist(output_dir, full=False):
    return all(prefix_file_path(output_dir, r, full=full).exists() for r in PREFIX_RATIOS)


def build_prefix_files(output_dir, n_rows=50000, full=False, force=False, save_clean=True):
    """
    从 data/train.csv 生成 20/40/60 前缀数据。
    full=False 时文件名为 dataset_prefix_20.csv。
    full=True 时文件名为 full_dataset_prefix_20.csv。
    """
    output_dir = Path(output_dir)
    ensure_dir(output_dir)
    ensure_dir(output_dir / "prefix_data")

    if all_prefix_files_exist(output_dir, full=full) and not force:
        print("检测到前缀数据已存在，跳过预处理：", output_dir)
        return

    for ratio in PREFIX_RATIOS:
        path = prefix_file_path(output_dir, ratio, full=full)
        if path.exists():
            path.unlink()

    clean_path = output_dir / "clean_trips.csv"
    if clean_path.exists() and force:
        clean_path.unlink()

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

    total_read = 0
    total_missing_removed = 0
    total_short_removed = 0
    total_valid = 0

    header_written = {ratio: False for ratio in PREFIX_RATIOS}
    clean_header_written = False

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

    print("开始预处理：", DATA_PATH)

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

        before = len(chunk)
        chunk = chunk[chunk["MISSING_DATA"].astype(str).str.upper() == "FALSE"].copy()
        total_missing_removed += before - len(chunk)

        buffers = {ratio: [] for ratio in PREFIX_RATIOS}
        clean_rows = []

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

            if len(points) < MIN_POINTS:
                total_short_removed += 1
                continue

            hour, weekday, is_weekend = timestamp_features(row.TIMESTAMP)

            row_info = {
                "TRIP_ID": row.TRIP_ID,
                "hour": hour,
                "weekday": weekday,
                "is_weekend": is_weekend,
            }

            if save_clean:
                clean_rows.append({
                    "TRIP_ID": row.TRIP_ID,
                    "TIMESTAMP": row.TIMESTAMP,
                    "hour": hour,
                    "weekday": weekday,
                    "is_weekend": is_weekend,
                    "num_points": len(points),
                    "duration_sec": (len(points) - 1) * 15,
                    "total_distance_km": trajectory_distance(points),
                    "start_lng": points[0][0],
                    "start_lat": points[0][1],
                    "end_lng": points[-1][0],
                    "end_lat": points[-1][1],
                    "points": points,
                })

            for ratio in PREFIX_RATIOS:
                buffers[ratio].append(build_prefix_features(points, row_info, ratio))

            total_valid += 1

        for ratio in PREFIX_RATIOS:
            if not buffers[ratio]:
                continue

            out_df = sanitize_numeric_df(pd.DataFrame(buffers[ratio]))
            out_path = prefix_file_path(output_dir, ratio, full=full)

            out_df.to_csv(
                out_path,
                mode="a",
                header=not header_written[ratio],
                index=False,
                encoding="utf-8-sig",
            )
            header_written[ratio] = True

        if save_clean and clean_rows:
            clean_df = pd.DataFrame(clean_rows)
            clean_df.to_csv(
                clean_path,
                mode="a",
                header=not clean_header_written,
                index=False,
                encoding="utf-8-sig",
            )
            clean_header_written = True

        print(
            f"chunk={chunk_id}, 已读={total_read}, 有效={total_valid}, "
            f"删缺失={total_missing_removed}, 删短轨迹={total_short_removed}"
        )

    print("预处理完成。")
    print("有效轨迹数：", total_valid)
    for ratio in PREFIX_RATIOS:
        print("已保存：", prefix_file_path(output_dir, ratio, full=full))
