import torch
import torch.nn as nn
import torch.nn.functional as F
from aiRNN.preprocessors import denorm_coords
import math


class HaversineMSEAltitudeLoss(nn.Module):
    def __init__(self, alt_const=1.0, earth_radius_km=6371.0):
        super().__init__()
        self.alt_const = alt_const
        self.R = earth_radius_km

    def forward(self, pred, target):
        """
        pred:   (..., 3) -> [lat, lon, alt]
        target: (..., 3) -> [lat, lon, alt]
        
        lat/lon expected in degrees.
        alt in meters (or same unit in pred/target).
        """
        lat1, lon1, alt1 = denorm_coords(pred[...,0], pred[...,1], pred[...,2])
        lat2, lon2, alt2 = denorm_coords(target[...,0], target[...,1], target[...,2])
        lat1 = torch.deg2rad(lat1)
        lon1 = torch.deg2rad(lon1)
        lat2 = torch.deg2rad(lat2)
        lon2 = torch.deg2rad(lon2)

        dlat = lat2 - lat1
        dlon = lon2 - lon1

        # Haversine formula
        a = (
            torch.sin(dlat / 2) ** 2
            + torch.cos(lat1) * torch.cos(lat2) * torch.sin(dlon / 2) ** 2
        )
        c = 2 * torch.atan2(torch.sqrt(a), torch.sqrt(1 - a))
        dist_km = self.R * c  # great-circle distance

        # MSE on distance
        haversine_mse = torch.mean(dist_km ** 2)

        # Altitude penalty
        delta_alt = alt2 - alt1
        alt_penalty = self.alt_const * torch.mean(delta_alt ** 2)

        return haversine_mse + alt_penalty
    
class HaversineAltitudeLoss(nn.Module):
    def __init__(self, alt_const=1.0, earth_radius_km=6371.0):
        super().__init__()
        self.alt_const = alt_const
        self.R = earth_radius_km

    def forward(self, pred, target):
        """
        pred:   (..., 3) -> [lat, lon, alt]
        target: (..., 3) -> [lat, lon, alt]
        
        lat/lon expected in degrees.
        alt in meters (or same unit in pred/target).
        """
        lat1, lon1, alt1 = denorm_coords(pred[...,0], pred[...,1], pred[...,2])
        lat2, lon2, alt2 = denorm_coords(target[...,0], target[...,1], target[...,2])
        lat1 = torch.deg2rad(lat1)
        lon1 = torch.deg2rad(lon1)
        lat2 = torch.deg2rad(lat2)
        lon2 = torch.deg2rad(lon2)

        dlat = lat2 - lat1
        dlon = lon2 - lon1

        # Haversine formula
        a = (
            torch.sin(dlat / 2) ** 2
            + torch.cos(lat1) * torch.cos(lat2) * torch.sin(dlon / 2) ** 2
        )
        c = 2 * torch.atan2(torch.sqrt(a), torch.sqrt(1 - a))
        dist_km = self.R * c  # great-circle distance

        # Mean distance
        haversine_loss = torch.mean(dist_km)

        # Altitude penalty
        delta_alt = alt2 - alt1
        alt_penalty = self.alt_const * torch.mean(delta_alt ** 2)

        return haversine_loss + alt_penalty
    
class WeightedMSELoss(nn.Module):
    def __init__(self, weights):
        super().__init__()
        self.weights = torch.tensor(weights).float()

    def forward(self, pred, target):
        """
        pred:   (..., D)
        target: (..., D)
        """
        diff = pred - target
        weighted_diff = diff * self.weights.to(diff.device)
        return torch.mean(weighted_diff ** 2)
