import numpy as np

# Normalization parameters
NORM_PARAMS = {
    'u_min': -5.785097340233893,
    'u_max': 12.700397285986645,
    'v_min': -2.22016497137004,
    'v_max': 22.115114215511067,
    'u_15_min': -5.351205997176121,
    'u_15_max': 12.59459668486086,
    'v_15_min': -3.653644730479408,
    'v_15_max': 21.82925627521016,
    'u_30_min': -5.643262967391627,
    'u_30_max': 12.55918134433948,
    'v_30_min': -4.606536463767895,
    'v_30_max': 21.94683289060832,
    'u_45_min': -5.80703228733597,
    'u_45_max': 12.59988012306978,
    'v_45_min': -1.346510193821202,
    'v_45_max': 21.94214151514658,
    'x_min': -97.40959,
    'x_max': -96.55169890366584,
    'y_min': 32.587689999999995,
    'y_max': 33.067024421728206
}

def normalize_uv(u, v, timestep=""):
    """Normalize u and v components for a specific timestep."""
    u_min_key = f'u{timestep}_min'
    u_max_key = f'u{timestep}_max'
    v_min_key = f'v{timestep}_min'
    v_max_key = f'v{timestep}_max'
    
    u_norm = (u - NORM_PARAMS[u_min_key]) / (NORM_PARAMS[u_max_key] - NORM_PARAMS[u_min_key])
    v_norm = (v - NORM_PARAMS[v_min_key]) / (NORM_PARAMS[v_max_key] - NORM_PARAMS[v_min_key])
    
    return u_norm, v_norm

def denormalize_uv(u_norm, v_norm, timestep=""):
    """Denormalize normalized u and v components for a specific timestep."""
    u_min_key = f'u{timestep}_min'
    u_max_key = f'u{timestep}_max'
    v_min_key = f'v{timestep}_min'
    v_max_key = f'v{timestep}_max'
    
    u = u_norm * (NORM_PARAMS[u_max_key] - NORM_PARAMS[u_min_key]) + NORM_PARAMS[u_min_key]
    v = v_norm * (NORM_PARAMS[v_max_key] - NORM_PARAMS[v_min_key]) + NORM_PARAMS[v_min_key]
    
    return u, v

def uv_to_velocity_direction(u, v):
    """Convert u, v components to velocity magnitude and meteorological direction."""
    velocity = np.sqrt(u**2 + v**2)
    # Calculate meteorological wind direction (direction wind is coming from)
    # Convert from mathematical angle to meteorological wind direction
    direction = (270 - np.degrees(np.arctan2(v, u))) % 360
    return velocity, direction

def velocity_direction_to_uv(velocity, direction):
    """Convert velocity magnitude and meteorological direction to u, v components."""
    # Convert meteorological direction to mathematical angle
    theta = np.radians((270 - direction) % 360)
    u = velocity * np.cos(theta)
    v = velocity * np.sin(theta)
    return u, v

def denormalize_predictions(pred):
    """Denormalize model predictions to get velocity and direction."""
    velocity = pred[:, 0] * (NORM_PARAMS['v_max'] - NORM_PARAMS['v_min']) + NORM_PARAMS['v_min']
    direction = np.rad2deg(np.arctan2(pred[:, 1], pred[:, 2])) % 360
    return velocity, direction