import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score


def mse(real_data: pd.Series, predictions: pd.Series) -> float:
    """
    Calculates the mean squared error between the real data and the predictions.
    """
    return round(
        mean_squared_error(real_data.astype(np.float64), predictions.astype(np.float64)), 2
    )


def mae(real_data: pd.Series, predictions: pd.Series) -> float:
    """
    Calculates the mean absolute error between the real data and the predictions.
    """
    return round(
        mean_absolute_error(real_data.astype(np.float64), predictions.astype(np.float64)), 2
    )


def r2(real_data: pd.Series, predictions: pd.Series) -> float:
    """
    Calculates the R2 score between the real data and the predictions.
    """
    return round(r2_score(real_data.astype(np.float64), predictions.astype(np.float64)), 2)


def silverman_radius(data: np.ndarray) -> float:
    """
    Calculate the Silverman bandwidth (radius) for a given dataset.

    Args:
        data (np.ndarray): Input data (1D array)

    Returns:
        float: Silverman bandwidth (radius)
    """
    n = len(data)
    sigma = np.std(data)
    iqr = np.percentile(data, 75) - np.percentile(data, 25)
    radius = 0.9 * min(sigma, iqr / 1.34) * n ** (-1 / 5)
    return radius


def rce_train(training_set: pd.DataFrame, radius: float | None = None) -> pd.DataFrame:
    """
    Get the Reduced Coulomb Energy (RCE) prototypes.

    Args:
        training_set (pd.DataFrame): The training set
        radius (float | None): The radius of the RCE prototypes. If None, computed using Silverman's rule.

    Returns:
        pd.DataFrame: The RCE prototypes
    """
    train_vectors = training_set.values

    # Vectorized distance computation for the radius calculation
    diff_vectors = train_vectors[:, np.newaxis] - train_vectors[np.newaxis, :]
    distances = np.linalg.norm(diff_vectors, axis=-1)

    # Non-parametric radius: Silverman Radius (compute if not provided)
    effective_radius = radius if radius is not None else silverman_radius(distances.flatten())

    # Initialize prototypes with the first vector
    prototypes = [train_vectors[0]]

    for vector in train_vectors[1:]:
        # Vectorized distance check between current vector and all prototypes
        distances_to_prototypes = np.linalg.norm(prototypes - vector, axis=1)

        # If no prototype is close, add the current vector as a new prototype
        if np.all(distances_to_prototypes > effective_radius):
            prototypes.append(vector)

    return pd.DataFrame(prototypes)


def rce_test(test_set: pd.DataFrame, prototypes: pd.DataFrame) -> pd.Series:
    """
    Get the signed Reduced Coulomb Energy (RCE) predictions.

    Args:
        test_set (pd.DataFrame): The test set
        prototypes (pd.DataFrame): The RCE prototypes

    Returns:
        pd.Series: The signed distances to the closest prototype for each test vector
    """
    test_vectors = test_set.values
    prototype_vectors = prototypes.values

    # Vectorized computation of distances between test vectors and all prototypes
    diff_vectors = test_vectors[:, np.newaxis] - prototype_vectors[np.newaxis, :]
    distances = np.linalg.norm(diff_vectors, axis=-1)

    # Find the closest prototype for each test vector
    min_distances = np.min(distances, axis=1)
    closest_prototypes = prototype_vectors[np.argmin(distances, axis=1)]

    # Compute the signed distance for each test vector
    signed_distances = np.sqrt(min_distances**2) * np.sign(
        np.mean(test_vectors - closest_prototypes, axis=1)
    )

    return pd.Series(signed_distances)


def rce_drift(reference_data: pd.DataFrame, real_data: pd.DataFrame, column: str) -> pd.Series:
    """
    Detect drift using the Reduced Coulomb Energy (RCE) method.

    Args:
        reference_data (pd.DataFrame): The reference data
        real_data (pd.DataFrame): The real data
        column (str): The target column to be analyzed. 'target' or 'prediction'

    Returns:
        pd.Series: Normalized drift distances
    """
    common_columns = list(set(reference_data.columns).intersection(real_data.columns))
    reference_data = reference_data[common_columns]
    real_data = real_data[common_columns]

    # Get prototypes
    if column == 'target':
        prototypes = rce_train(reference_data.drop(columns=['prediction']), 0.1)
    else:
        prototypes = rce_train(reference_data.drop(columns=['target']), 0.1)

    # Distances to prototypes
    if column == 'target':
        distances_train = rce_test(reference_data.drop(columns=['prediction']), prototypes)
        distances_test = rce_test(real_data.drop(columns=['prediction']), prototypes)
    else:
        distances_train = rce_test(reference_data.drop(columns=['target']), prototypes)
        distances_test = rce_test(real_data.drop(columns=['target']), prototypes)

    # Find the maximum absolute distance in the training set
    max_abs_distance = max(abs(distances_train.max()), abs(distances_train.min()))

    # Normalize while preserving sign
    distances = distances_test / max_abs_distance

    return distances