import torch
import torch.nn.functional as F
from torchmetrics.regression import PearsonCorrCoef
import numpy as np


def masked_cos_sim(y_true, y_pred):
    """Masked, cosine similarity between true and pred vectors

    """

    # To avoid numerical instability during training on GPUs,
    # we add a fuzzing constant epsilon of 1×10−7 to all vectors
    epsilon = 1e-7
    # Masking: we multiply values by (true + 1) because then the peaks that cannot
    # be there (and have value of -1 as explained above) won't be considered
    pred_masked = ((y_true + 1) * y_pred) / (y_true + 1 + epsilon)
    true_masked = ((y_true + 1) * y_true) / (y_true + 1 + epsilon)
    pred_masked = F.normalize(pred_masked, p=2, dim=1)
    true_masked = F.normalize(true_masked, p=2, dim=1)
    return -(pred_masked * true_masked).sum(dim=1).mean()


def masked_spectral_angle(y_true, y_pred):
    """Masked, cosine similarity between true and pred vectors

    """

    # To avoid numerical instability during training on GPUs,
    # we add a fuzzing constant epsilon of 1×10−7 to all vectors
    epsilon = 1e-7
    # Masking: we multiply values by (true + 1) because then the peaks that cannot
    # be there (and have value of -1 as explained above) won't be considered
    pred_masked = ((y_true + 1) * y_pred) / (y_true + 1 + epsilon)
    true_masked = ((y_true + 1) * y_true) / (y_true + 1 + epsilon)
    pred_masked = F.normalize(pred_masked, p=2, dim=1)
    true_masked = F.normalize(true_masked, p=2, dim=1)
    # print(pred_masked.sum(dim=1))
    # print((pred_masked * true_masked).sum(dim=1).shape)
    return 1 -2 * torch.acos((pred_masked * true_masked).sum(dim=1)).mean() / np.pi


def masked_pearson_correlation_distance(y_true, y_pred, reduce='mean'):
    """
    Calculates the masked Pearson correlation distance between true and predicted intensity vectors.

    The masked Pearson correlation distance is a metric for comparing the similarity between two intensity vectors,
    taking into account only the non-negative values in the true values tensor (which represent valid peaks).

    Parameters:
    -----------
    y_true : Tensor
        A tensor containing the true values, with shape `(batch_size, num_values)`.
    y_pred : Tensor
        A tensor containing the predicted values, with the same shape as `y_true`.

    Returns:
    --------
    Tensor
        A tensor containing the masked Pearson correlation distance between `y_true` and `y_pred`.

    Raises:
    -------
    ValueError
        If `y_true` and `y_pred` have different shapes.
    """
    epsilon = 1e-7

    # Masking: we multiply values by (true + 1) because then the peaks that cannot
    # be there (and have value of -1 as explained above) won't be considered
    pred_masked = ((y_true + 1) * y_pred) / (y_true + 1 + epsilon)
    true_masked = ((y_true + 1) * y_true) / (y_true + 1 + epsilon)
    loss = PearsonCorrCoef()
    if reduce == 'mean':
        return torch.mean(loss(pred_masked, true_masked))
    if reduce == 'sum':
        return torch.sum(loss(pred_masked, true_masked))
    if reduce is None:
        return loss(pred_masked, true_masked)

def distance(x, y):
    return torch.mean(torch.abs(x - y))


def cos_sim_to_sa(cos):
    return 1 - (2 * np.arccos(cos) / np.pi)