import numpy as np
import pandas as pd
from sklearn.cluster import MiniBatchKMeans

from config import N_CLUSTERS, RANDOM_STATE, TOP_K_CLUSTERS
from dm_taxi.common import BASE_FEATURE_COLUMNS, lonlat_to_xy_km_array, xy_km_to_lonlat_array
from dm_taxi.models import make_xgboost_classifier, make_xgboost_regressor


def fit_train_kmeans(train_df, n_clusters=N_CLUSTERS):
    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=10,
    )

    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 train_offset_regressor(train_df, train_labels, centers_df, use_gpu=False):
    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_xgboost_regressor(use_gpu=use_gpu, deep=False)
    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(test_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_pred_lng = np.zeros(len(test_df), dtype=np.float64)
    final_pred_lat = np.zeros(len(test_df), dtype=np.float64)

    base_part = test_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_pred_lng += candidate_lng * candidate_probs
        final_pred_lat += candidate_lat * candidate_probs

    return final_pred_lng, final_pred_lat


def run_cluster_topk_offset(train_df, test_df, n_clusters=N_CLUSTERS, top_k=TOP_K_CLUSTERS, use_gpu=False):
    km, train_labels, centers_df = fit_train_kmeans(train_df, n_clusters=n_clusters)
    test_true_labels = assign_cluster_labels(km, test_df)

    clf = make_xgboost_classifier(n_clusters, use_gpu=use_gpu)
    clf.fit(train_df[BASE_FEATURE_COLUMNS].values, train_labels)

    prob_matrix = clf.predict_proba(test_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,
    )

    top1_acc = float(np.mean(top_clusters[:, 0] == test_true_labels))
    top3_hit = float(np.mean([
        test_true_labels[i] in set(top_clusters[i, :min(3, top_k)].tolist())
        for i in range(len(test_true_labels))
    ]))
    top10_hit = float(np.mean([
        test_true_labels[i] in set(top_clusters[i, :min(10, top_k)].tolist())
        for i in range(len(test_true_labels))
    ]))
    topk_hit = float(np.mean([
        test_true_labels[i] in set(top_clusters[i, :top_k].tolist())
        for i in range(len(test_true_labels))
    ]))

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

    pred_lng, pred_lat = predict_topk_offset_weighted(
        test_df=test_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,
    )

    info = {
        "cluster_top1_acc": top1_acc,
        "cluster_top3_hit_rate": top3_hit,
        "cluster_top10_hit_rate": top10_hit,
        "cluster_topk_hit_rate": topk_hit,
    }

    return pred_lng, pred_lat, info, centers_df
