import math,random
import torch
import config as C

def aug_dihedral(xc, yc):
    """One dihedral-8 transform per crop-batch, image and label together."""
    if torch.rand(()) > 0.5:
        xc, yc = xc.flip(-2), yc.flip(-2)
    if torch.rand(()) > 0.5:
        xc, yc = xc.flip(-1), yc.flip(-1)
    k = int(torch.randint(0, 4, ()))
    if k:
        xc, yc = torch.rot90(xc, k, (-2, -1)), torch.rot90(yc, k, (-2, -1))
    return xc.contiguous(), yc.contiguous()


def aug_temporal_subsample(xc, pos, p=0.5, drop_max=0.2, keep_min=24):
    """Drop up to `drop_max` of the frames, order preserved. Positions keep
    their original index values -> emulates missing acquisitions of an
    otherwise identical season. No-op with probability 1-p."""
    if torch.rand(()) > p:
        return xc, pos
    T = xc.shape[1]
    keep = int(round(T * (1.0 - random.uniform(0.0, drop_max))))
    keep = max(keep_min, min(T, keep))
    if keep >= T:
        return xc, pos
    idx = torch.randperm(T)[:keep].sort().values.to(xc.device)
    return (xc.index_select(1, idx).contiguous(),
            pos.index_select(1, idx).contiguous())


def aug_frame_dropout(xc, p=0.25, max_frames=2):
    """Zero 1-2 whole frames (0 == dataset mean after normalisation) —
    emulates cloudy/lost acquisitions."""
    if torch.rand(()) > p:
        return xc
    T = xc.shape[1]
    n = random.randint(1, max_frames)
    idx = torch.randperm(T)[:n].to(xc.device)
    xc = xc.clone()
    xc[:, idx] = 0.0
    return xc


def aug_spectral(xc, gain_std=0.05, bias_std=0.03, noise_std=0.01):
    """Per-sample, per-band gain/bias jitter + small gaussian noise (data is
    z-normalised, so these are fractions of one std)."""
    B, T, Ch, H, W = xc.shape
    dev = xc.device
    gain = 1.0 + gain_std * torch.randn(B, 1, Ch, 1, 1, device=dev)
    bias = bias_std * torch.randn(B, 1, Ch, 1, 1, device=dev)
    xc = xc * gain + bias
    if noise_std > 0:
        xc = xc + noise_std * torch.randn_like(xc)
    return xc.contiguous()


def lr_factor(step, total_steps, warmup_steps, min_ratio):
    """Linear warmup then cosine decay to min_ratio."""
    if step < warmup_steps:
        return (step + 1) / max(1, warmup_steps)
    p = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    p = min(1.0, p)
    cos = 0.5 * (1.0 + math.cos(math.pi * p))
    return min_ratio + (1.0 - min_ratio) * cos