# Copyright 2020 Ross Wightman
# Various utility functions

import torch
import torch.nn as nn
from functools import partial
import math
import warnings
import torch.nn.functional as F

#from timesformer.models.helpers import load_pretrained
#from .build import MODEL_REGISTRY
from itertools import repeat
import collections.abc as container_abcs

DEFAULT_CROP_PCT = 0.875
IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5)
IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5)
IMAGENET_DPN_MEAN = (124 / 255, 117 / 255, 104 / 255)
IMAGENET_DPN_STD = tuple([1 / (.0167 * 255)] * 3)

def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor

def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.
    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

#From PyTorch internals
def _ntuple(n):
    def parse(x):
        if isinstance(x, container_abcs.Iterable):
            return x
        return tuple(repeat(x, n))
    return parse
to_2tuple = _ntuple(2)

# Calculate symmetric padding for a convolution
def get_padding(kernel_size: int, stride: int = 1, dilation: int = 1, **_) -> int:
    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
    return padding

def get_padding_value(padding, kernel_size, **kwargs):
    dynamic = False
    if isinstance(padding, str):
        # for any string padding, the padding will be calculated for you, one of three ways
        padding = padding.lower()
        if padding == 'same':
            # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact
            if is_static_pad(kernel_size, **kwargs):
                # static case, no extra overhead
                padding = get_padding(kernel_size, **kwargs)
            else:
                # dynamic 'SAME' padding, has runtime/GPU memory overhead
                padding = 0
                dynamic = True
        elif padding == 'valid':
            # 'VALID' padding, same as padding=0
            padding = 0
        else:
            # Default to PyTorch style 'same'-ish symmetric padding
            padding = get_padding(kernel_size, **kwargs)
    return padding, dynamic

# Calculate asymmetric TensorFlow-like 'SAME' padding for a convolution
def get_same_padding(x: int, k: int, s: int, d: int):
    return max((int(math.ceil(x // s)) - 1) * s + (k - 1) * d + 1 - x, 0)


# Can SAME padding for given args be done statically?
def is_static_pad(kernel_size: int, stride: int = 1, dilation: int = 1, **_):
    return stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0


# Dynamically pad input x with 'SAME' padding for conv with specified args
#def pad_same(x, k: List[int], s: List[int], d: List[int] = (1, 1), value: float = 0):
def pad_same(x, k, s, d=(1, 1), value= 0):
    ih, iw = x.size()[-2:]
    pad_h, pad_w = get_same_padding(ih, k[0], s[0], d[0]), get_same_padding(iw, k[1], s[1], d[1])
    if pad_h > 0 or pad_w > 0:
        x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2], value=value)
    return x

def adaptive_pool_feat_mult(pool_type='avg'):
    if pool_type == 'catavgmax':
        return 2
    else:
        return 1

def drop_path(x, drop_prob: float = 0., training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.
    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output

class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)

def hf_weight_translation(state_dict):
    new_state_dict = {}
    for key, value in state_dict.items():
        new_key = None
        
        # Handle embeddings
        if key == 'vit.embeddings.cls_token':
            new_key = 'cls_token'
        elif key == 'vit.embeddings.position_embeddings':
            new_key = 'pos_embed'
        elif key.startswith('vit.embeddings.patch_embeddings.projection'):
            new_key = key.replace('vit.embeddings.patch_embeddings.projection', 'patch_embed.proj')
        
        # Handle encoder layers
        elif key.startswith('vit.encoder.layer.'):
            # Extract layer number
            parts = key.split('.')
            layer_num = parts[3]  # vit.encoder.layer.{N}
            
            # Handle attention weights - need to combine q,k,v into qkv
            if 'attention.attention.query' in key:
                # Store for later combination
                q_key = f'blocks.{layer_num}.attn.qkv.weight' if 'weight' in key else f'blocks.{layer_num}.attn.qkv.bias'
                if q_key not in new_state_dict:
                    new_state_dict[q_key] = {'q': None, 'k': None, 'v': None}
                new_state_dict[q_key]['q'] = value
                continue
            elif 'attention.attention.key' in key:
                q_key = f'blocks.{layer_num}.attn.qkv.weight' if 'weight' in key else f'blocks.{layer_num}.attn.qkv.bias'
                if q_key not in new_state_dict:
                    new_state_dict[q_key] = {'q': None, 'k': None, 'v': None}
                new_state_dict[q_key]['k'] = value
                continue
            elif 'attention.attention.value' in key:
                q_key = f'blocks.{layer_num}.attn.qkv.weight' if 'weight' in key else f'blocks.{layer_num}.attn.qkv.bias'
                if q_key not in new_state_dict:
                    new_state_dict[q_key] = {'q': None, 'k': None, 'v': None}
                new_state_dict[q_key]['v'] = value
                continue
            elif 'attention.output.dense' in key:
                new_key = key.replace('vit.encoder.layer.', 'blocks.').replace('.attention.output.dense', '.attn.proj')
            
            # Handle layer norms
            elif 'layernorm_before' in key:
                new_key = key.replace('vit.encoder.layer.', 'blocks.').replace('.layernorm_before', '.norm1')
            elif 'layernorm_after' in key:
                new_key = key.replace('vit.encoder.layer.', 'blocks.').replace('.layernorm_after', '.norm2')
            
            # Handle MLP layers
            elif 'intermediate.dense' in key:
                new_key = key.replace('vit.encoder.layer.', 'blocks.').replace('.intermediate.dense', '.mlp.fc1')
            elif 'output.dense' in key and 'attention' not in key:
                new_key = key.replace('vit.encoder.layer.', 'blocks.').replace('.output.dense', '.mlp.fc2')
        
        # Handle final layer norm
        elif key == 'vit.layernorm.weight':
            new_key = 'norm.weight'
        elif key == 'vit.layernorm.bias':
            new_key = 'norm.bias'
        
        # Skip classifier layers as they're typically not needed for video models
        elif key.startswith('classifier'):
            continue
            
        # Add the mapped key-value pair
        if new_key is not None:
            new_state_dict[new_key] = value
    
    final_state_dict = {}
    
    for key, value in new_state_dict.items():
        if isinstance(value, dict) and 'q' in value:
            # This is a qkv tensor that needs to be combined
            if value['q'] is not None and value['k'] is not None and value['v'] is not None:
                # Concatenate along the first dimension (output dimension)
                combined = torch.cat([value['q'], value['k'], value['v']], dim=0)
                final_state_dict[key] = combined
            else:
                print(f"Warning: Incomplete qkv tensor for {key}")
        else:
            final_state_dict[key] = value
    
    return final_state_dict
    

def load_pretrained(model, state_dict, num_classes=1000, in_chans=3, filter_fn=None, img_size=224, num_frames=8, num_patches=196, attention_type='divided_space_time', strict=True):

    if filter_fn is not None:
        state_dict = filter_fn(state_dict)

    if in_chans == 1:
        conv1_name = 'patch_embed.proj'
        _logger.info('Converting first conv (%s) pretrained weights from 3 to 1 channel' % conv1_name)
        conv1_weight = state_dict[conv1_name + '.weight']
        conv1_type = conv1_weight.dtype
        conv1_weight = conv1_weight.float()
        O, I, J, K = conv1_weight.shape
        if I > 3:
            assert conv1_weight.shape[1] % 3 == 0
            # For models with space2depth stems
            conv1_weight = conv1_weight.reshape(O, I // 3, 3, J, K)
            conv1_weight = conv1_weight.sum(dim=2, keepdim=False)
        else:
            conv1_weight = conv1_weight.sum(dim=1, keepdim=True)
        conv1_weight = conv1_weight.to(conv1_type)
        state_dict[conv1_name + '.weight'] = conv1_weight
    elif in_chans != 3:
        conv1_name = 'patch_embed.proj'
        conv1_weight = state_dict[conv1_name + '.weight']
        conv1_type = conv1_weight.dtype
        conv1_weight = conv1_weight.float()
        O, I, J, K = conv1_weight.shape
        if I != 3:
            _logger.warning('Deleting first conv (%s) from pretrained weights.' % conv1_name)
            del state_dict[conv1_name + '.weight']
            strict = False
        else:
            _logger.info('Repeating first conv (%s) weights in channel dim.' % conv1_name)
            repeat = int(math.ceil(in_chans / 3))
            conv1_weight = conv1_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
            conv1_weight *= (3 / float(in_chans))
            conv1_weight = conv1_weight.to(conv1_type)
            state_dict[conv1_name + '.weight'] = conv1_weight


    classifier_name = 'head'

    ## Resizing the positional embeddings in case they don't match
    if num_patches + 1 != state_dict['pos_embed'].size(1):
        pos_embed = state_dict['pos_embed']
        cls_pos_embed = pos_embed[0,0,:].unsqueeze(0).unsqueeze(1)
        other_pos_embed = pos_embed[0,1:,:].unsqueeze(0).transpose(1, 2)
        new_pos_embed = F.interpolate(other_pos_embed, size=(num_patches), mode='nearest')
        new_pos_embed = new_pos_embed.transpose(1, 2)
        new_pos_embed = torch.cat((cls_pos_embed, new_pos_embed), 1)
        state_dict['pos_embed'] = new_pos_embed

    ## Resizing time embeddings in case they don't match
    if 'time_embed' in state_dict and num_frames != state_dict['time_embed'].size(1):
        time_embed = state_dict['time_embed'].transpose(1, 2)
        new_time_embed = F.interpolate(time_embed, size=(num_frames), mode='nearest')
        state_dict['time_embed'] = new_time_embed.transpose(1, 2)

    ## Initializing temporal attention
    if attention_type == 'divided_space_time':
        new_state_dict = state_dict.copy()
        for key in state_dict:
            if 'blocks' in key and 'attn' in key:
                new_key = key.replace('attn','temporal_attn')
                if not new_key in state_dict:
                   new_state_dict[new_key] = state_dict[key]
                else:
                   new_state_dict[new_key] = state_dict[new_key]
            if 'blocks' in key and 'norm1' in key:
                new_key = key.replace('norm1','temporal_norm1')
                if not new_key in state_dict:
                   new_state_dict[new_key] = state_dict[key]
                else:
                   new_state_dict[new_key] = state_dict[new_key]
        state_dict = new_state_dict

    new_state_dict = {}
    for key, val in state_dict.items():
        new_key = 'model.' + key
        new_state_dict[new_key] = val
    
    state_dict = new_state_dict
    ## Loading the weights
    return model.load_state_dict(state_dict, strict=False)