import torch
import torch.nn as nn
import torch.nn.functional as F

from functools import wraps

from einops import rearrange, reduce
import math


def once(fn):
    called = False

    @wraps(fn)
    def inner(x):
        nonlocal called
        if called:
            return
        called = True
        return fn(x)

    return inner


print_once = once(print)

# functions

def exists(val):
    return val is not None




def default(val, d):
    return val if exists(val) else d






def Sequential(*modules):
    return nn.Sequential(*filter(exists, modules))


# tensor functions

def log(t, eps=1e-20):
    return torch.log(t.clamp(min=eps))


def l2norm(t):
    return F.normalize(t, p=2, dim=-1)


def matrix_diag(t):
    device = t.device
    i, j = t.shape[-2:]
    num_diag_el = min(i, j)
    i_range = torch.arange(i, device=device)
    j_range = torch.arange(j, device=device)
    diag_mask = rearrange(i_range, 'i -> i 1') == rearrange(j_range, 'j -> 1 j')
    diag_el = t.masked_select(diag_mask)
    return rearrange(diag_el, '(b d) -> b d', d=num_diag_el)


# 2d sinusoidal positional embedding
# simple vit paper shows it is good enough compared to learned


class LayerNorm(nn.Module):
    def __init__(self, dim, scale=True):
        super().__init__()
        self.learned_gamma = nn.Parameter(torch.ones(dim)) if scale else None

        self.register_buffer('gamma', torch.ones(dim), persistent=False)
        self.register_buffer('beta', torch.zeros(dim), persistent=False)

    def forward(self, x):
        return F.layer_norm(x, x.shape[-1:], default(self.learned_gamma, self.gamma), self.beta)


def freeze(model):
    for n, p in model.named_parameters():
        p.requires_grad = False









def print_trainable_parameters(model):
    trainable_params = 0
    all_param = 0
    for k, param in model.named_parameters():

        num_params = param.numel()
        # if using DS Zero 3 and the weights are initialized empty
        if num_params == 0 and hasattr(param, "ds_numel"):
            num_params = param.ds_numel

        all_param += num_params
        if param.requires_grad:
            print(k)
            trainable_params += num_params

    print(
        f"trainable params: {trainable_params:,d} || all params: {all_param:,d} || trainable%: {100 * trainable_params / all_param}"
    )