from .blocks import (
                    VisionEmbeddingModel,
                    KeylessAttention,
                    Encoder,
                    EncoderArgs,
                    InputEmbeddings
)
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
import torch
from ..LLM_pieces import SwiGLU, DynamicTanh, Dynamic_erf
from dataclasses import dataclass

@dataclass
class CMAArgs(EncoderArgs):
    in_channels:int = 3
    patch_size:int = 14
    H:int = 16
    W:int = 16
    n_tasks:int = 2
    n_classes:int = [2, 3]
    cls_text_index:int = 0
    cls_image_index:int = 10

class CMA(
        nn.Module,
        PyTorchModelHubMixin,
        pipeline_tag="text-generation",
        library_name="pytorch",
        license="mit"
    ):
    def __init__(self, args:CMAArgs=None):
        super().__init__()
        self.args = args
        self._prepare_model()

    def _prepare_model(self):
        
        self.vision_emb = VisionEmbeddingModel(self.args.in_channels,
                                                self.args.dim,
                                                self.args.patch_size,
                                                self.args.H,
                                                self.args.W
                                                )
        self.text_emb = InputEmbeddings(self.args)
        if self.args.layer_norm == "RMSNorm":
            self.layer_norm = nn.RMSNorm(self.args.dim)
        elif self.args.layer_norm == "Derf":
            self.layer_norm = Dynamic_erf(self.args.dim)
        elif self.args.layer_norm == "DyT":
            self.layer_norm = DynamicTanh(self.args.dim)
        else:
            raise ValueError(f"allowed layer norms: 'RMSNorm', 'Derf', 'DyT' ; got: {self.args.layer_norm}")
        self.encoder = Encoder(self.args)

        assert len(self.args.n_classes) == self.args.n_tasks
        self.outs = nn.ModuleList([nn.Sequential(
                                        KeylessAttention(self.args.dim),
                                        nn.Sequential(
                                            SwiGLU(self.args),
                                            nn.SiLU(),
                                            nn.Linear(self.args.dim, self.args.n_classes[i])
                                                )
                                            )
                                for i in range(self.args.n_tasks)])

        self.n_params = sum(p.numel() for p in self.parameters() if p.requires_grad)

        self.to(self.args.device, dtype=self.args.dtype)

    def pred(self, texts, images, cls_image_token_index=None) -> list[torch.Tensor]:
        texts = self.text_emb(texts)
        images = self.vision_emb(images)
        if cls_image_token_index is not None:
            cls_i_em = self.text_emb(torch.tensor([cls_image_token_index]).to(self.args.device)).unsqueeze(0).expand(images.shape[0], -1, -1)
            x = torch.cat([texts, cls_i_em, images], dim=1)
        else:
            x = torch.cat([texts, images], dim=1)
        x = self.encoder(x)
        x = self.layer_norm(x)
        cls_text, cls_image = x[:, self.args.cls_text_index], x[:, self.args.cls_image_index]
        tasks = [out(cls_text, cls_image) for out in self.outs]
        return tasks