import math

import numpy as np
import pandas as pd
from sklearn.neighbors import NearestNeighbors

from config import K_NEIGHBORS, REF_LAT, REF_LNG, WEIGHT_POWER
from dm_taxi.common import haversine


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


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


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)
            x, y = lonlat_to_xy_one(lng, lat)
            sw = math.sqrt(float(weights[i]))
            values.extend([x * sw, y * 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 make_similarity_features_from_neighbors(distances, indices, train_targets):
    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 = NearestNeighbors(
        n_neighbors=K_NEIGHBORS + 1,
        metric="euclidean",
        algorithm="auto",
        n_jobs=-1,
    )
    nn.fit(train_vectors)

    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)

    return make_similarity_features_from_neighbors(clean_distances, clean_indices, train_targets)


def make_similarity_features_query(train_vectors, query_vectors, train_targets):
    nn = NearestNeighbors(
        n_neighbors=K_NEIGHBORS,
        metric="euclidean",
        algorithm="auto",
        n_jobs=-1,
    )
    nn.fit(train_vectors)

    distances, indices = nn.kneighbors(query_vectors, return_distance=True)

    return make_similarity_features_from_neighbors(distances, indices, train_targets)


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
