import math
from types import SimpleNamespace
import numpy as np
import torch
from torch import nn
from timm.models.vision_transformer import Block
from .radio_patch import ViTPatchGenerator
from .radio_cls import ClsToken as _bundle_radio_cls

class RadioBackbone(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch_generator = ViTPatchGenerator(16,1280,224,cls_token=True,max_input_dims=2048,num_cls_tokens=4,register_multiple=10)
        self.blocks = nn.Sequential(*[Block(1280,16,mlp_ratio=4,qkv_bias=True) for _ in range(32)])
    def forward(self, pixels):
        return self.blocks(self.patch_generator(pixels))[:,10:]

class RadioModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = RadioBackbone()
        self.register_buffer('summary_idxs',torch.tensor([0,1,2],dtype=torch.int64))
    def forward(self, pixels):
        return self.model(pixels)

class RadioWrapper(nn.Module):
    def __init__(self):
        super().__init__()
        self.radio_model = RadioModel()
    def forward(self, pixels):
        return self.radio_model(pixels)

class RadioVisionTower(nn.Module):
    def __init__(self):
        super().__init__()
        self.vision_tower = RadioWrapper()
        self.image_processor = RadioImageProcessor()
        self.is_loaded = True
        self.config = SimpleNamespace(hidden_size=1280,patch_size=16)
    def forward(self, images):
        reference = next(self.parameters())
        if isinstance(images, (tuple,list)):
            return torch.cat([self.vision_tower(x.to(reference)) for x in images], dim=1)
        return self.vision_tower(images.to(reference))

class RadioImageProcessor:
    radio_downsample_ratio = 0.5
    def __call__(self, images, return_tensors='pt'):
        if len(images)!=1:
            raise ValueError('Process one image at a time')
        image = images[0].convert('RGB')
        ph,pw = round(image.height/16+0.5), round(image.width/16+0.5)
        factor=min(math.sqrt(12000/(ph*pw)),1.)
        h,w=math.floor(factor*ph),math.floor(factor*pw)
        if h*w<1024:
            up=math.sqrt(1024/(h*w)); h,w=math.ceil(up*h),math.ceil(up*w)
        if h%2:h=h+1 if (h+1)*w<=12000 else max(2,h-1)
        if w%2:w=w+1 if h*(w+1)<=12000 else max(2,w-1)
        pixels=torch.from_numpy(np.asarray(image,dtype=np.uint8).copy()).permute(2,0,1).unsqueeze(0).float()
        if pixels.shape[-2:]!=(h*16,w*16):
            pixels=torch.nn.functional.interpolate(pixels,size=(h*16,w*16),mode='bicubic',align_corners=False,antialias=True)
        mean=torch.tensor([.48145466,.4578275,.40821073]).view(1,3,1,1)
        std=torch.tensor([.26862954,.26130258,.27577711]).view(1,3,1,1)
        return SimpleNamespace(pixel_values=(pixels/255.-mean)/std)