import torch


def getVelocityStats(data, quantile=0.95, with_stats=True):
    """
    data: T x D or T x N x D
    """
    assert data.ndim == 2 or data.ndim == 3
    vel = torch.linalg.norm(data[1:] - data[:-1], dim=-1)
    if vel.ndim == 2:
        vel = vel.mean(dim=-1)

    if with_stats:
        errors = {'mean_vel': vel.mean().item(), 'max_vel': vel.max().item(), 'median_vel': vel.median().item(), 'std_vel': vel.std().item(), f'{quantile}q': vel.kthvalue(int(quantile*len(vel))).values.item()}
    else:
        errors = {}
    return vel, errors

def getAccelerationStats(data, quantile=0.95, with_stats=True):
    """
    data: T x D
    """
    assert data.ndim == 2 or data.ndim == 3
    acc = torch.linalg.norm(data[2:] - 2*data[1:-1] + data[:-2], dim=-1)
    if acc.ndim == 2:
        acc = acc.mean(dim=-1)
    if with_stats:
        errors = {'mean_acc': acc.mean().item(), 'max_acc': acc.max().item(), 'median_acc': acc.median().item(), 'std_acc': acc.std().item(), f'{quantile}q': acc.kthvalue(int(quantile*len(acc))).values.item()}
    else:
        errors = {}
    return acc, errors

def getJerkStats(data, quantile=0.95, with_stats=True):
    """
    Measures the 'jerkiness' (3rd derivative). 
    High jerk = robotic/shaky motion.
    """
    assert data.ndim == 2 or data.ndim == 3
    if data.shape[0] < 4: return None, {}
    
    # Third-order finite difference
    jerk = torch.linalg.norm(
        data[3:] - 3*data[2:-1] + 3*data[1:-2] - data[:-3], 
        dim=-1
    )
    if jerk.ndim == 2:
        jerk = jerk.mean(dim=-1)
     
    if not with_stats:
        errors = {}
    else:
        errors = {
            'mean_jerk': jerk.mean().item(),
            'max_jerk': jerk.max().item(),
            'median_jerk': jerk.median().item(),
            'std_jerk': jerk.std().item(),
            f'{quantile}q': jerk.kthvalue(int(quantile*len(jerk))).values.item(),
            # Normalized jerk is often used: mean_jerk / max_vel
            'norm_jerk': (jerk.mean() / (torch.linalg.norm(data[1:] - data[:-1], dim=-1).mean() + 1e-6)).item()
        }
    return jerk, errors

def getHighFreqNoise(data, fps=30, noise_floor_hz=10.0):
    """
    Calculates the ratio of energy present in frequencies above 'noise_floor_hz'.
    
    data: (T, D) tensor
    fps: Frames Per Second of the data
    noise_floor_hz: Frequencies above this are considered 'Noise'. 
                    (Human motion is rarely > 10-12 Hz).
    """
    assert data.ndim == 2

    T = data.shape[0]
    
    # 1. Compute FFT
    # Frequencies range from 0 to fps/2
    fft_coeffs = torch.fft.rfft(data, dim=0)
    power = fft_coeffs.abs()**2
    
    # 2. Calculate the frequency corresponding to each bin
    # frequency step per bin = (fps / 2) / (num_bins - 1)
    num_bins = power.shape[0]
    freqs = torch.linspace(0, fps/2, num_bins)
    
    # 3. Find index for the noise floor
    # We want valid mask where freq > noise_floor_hz
    k = (freqs > noise_floor_hz).nonzero(as_tuple=True)[0]
    
    if len(k) == 0:
        return {'high_freq_noise_ratio': 0.0} # Cutoff is higher than Nyquist
        
    start_idx = k[0].item()
    
    # 4. Sum Energies
    total_energy = power.sum(dim=0)
    high_freq_energy = power[start_idx:].sum(dim=0)
    
    # Avoid div by zero
    noise_ratio = (high_freq_energy / (total_energy + 1e-8)).mean()
    
    return {'high_freq_noise_ratio': noise_ratio.item()}

def getNASC(data, fps=30):
    """
    Number of Acceleration Sign Changes (NASC).
    A robotics metric for "wobble" or "hunting".
    
    data: (T, D) tensor (usually joints or params)
    """
    T = data.shape[0]
    
    # 1. Calculate Velocity (1st derivative)
    vel = data[1:] - data[:-1]
    
    # 2. Calculate Acceleration (2nd derivative)
    acc = vel[1:] - vel[:-1]
    
    # 3. Compute Sign Changes per channel
    # sign(x[t]) != sign(x[t-1])
    # We use (a * b) < 0 to detect sign flip
    sign_changes = (acc[1:] * acc[:-1]) < 0
    
    # Sum changes across time, average across dimensions
    num_changes = sign_changes.float().sum(dim=0).mean()
    
    # Normalize by time (changes per second)
    nasc_score = num_changes / (T / fps)
    
    return {'NASC': nasc_score.item()}
