from __future__ import annotations
import torch
import torch.nn as nn

class SquaredReLU(nn.Module):

    def forward(self, x):
        return torch.pow(torch.nn.functional.relu(x), 2)

class RMSNorm(nn.Module):

    def __init__(self, hidden_size, eps=1e-05):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)

class RadioHProjector(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.downsample_ratio = 0.5
        self.ps_version = 'v2'
        vit_hidden_size = 1280
        vision_projection_hidden_size = getattr(config, 'projector_hidden_size', config.hidden_size)
        llm_hidden_size = config.hidden_size
        in_dim = int(vit_hidden_size * (1 / self.downsample_ratio) ** 2)
        self.mlp = nn.Sequential(RMSNorm(in_dim, eps=1e-05), nn.Linear(in_dim, vision_projection_hidden_size, bias=False), SquaredReLU(), nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False))
        self.radio_multiply = True

    def pixel_shuffle(self, x, scale_factor=0.5):
        (n, w, h, c) = x.size()
        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
        x = x.permute(0, 2, 1, 3).contiguous()
        x = x.view(n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor)))
        if self.ps_version != 'v1':
            x = x.permute(0, 2, 1, 3).contiguous()
        return x

    def forward(self, image_features: torch.Tensor, image_grid_thw: torch.Tensor):
        device = image_features.device
        indices = []
        offset = 0
        for (_, H, W) in image_grid_thw.tolist():
            (H, W) = (int(H), int(W))
            (H2, W2) = (H // 2, W // 2)
            k = torch.arange(H2 * W2, device=device)
            if self.ps_version == 'v1':
                (out_r, out_c) = (k % H2, k // H2)
            else:
                (out_r, out_c) = (k // W2, k % W2)
            tl = offset + 2 * out_r * W + 2 * out_c
            indices.append(torch.stack([tl, tl + 1, tl + W, tl + W + 1], dim=1).reshape(-1))
            offset += H * W
        gather_idx = torch.cat(indices)
        merged = image_features[gather_idx].view(-1, 4 * image_features.shape[-1])
        if self.radio_multiply:
            return 0.01 * self.mlp(merged)
        return self.mlp(merged)