
import torch
import torch.nn as nn
from transformers import AutoTokenizer, CLIPProcessor

from medpalm.transformer import (
    AutoregressiveWrapper,
    Decoder,
    Encoder,
    Transformer,
    ViTransformerWrapper,
)


class MedPalmTokenizer:
    def __init__(self):
        try:

            self.processor = CLIPProcessor.from_pretrained("laion/CLIP-ViT-L-14-laion2B-s32B-b82K")
            self.tokenizer = AutoTokenizer.from_pretrained(
                "EleutherAI/gpt-neox-20b",
                additional_special_tokens=["<image>", "</image>"],
                eos_token ="<eos>",
                pad_token="<pad>",
                extra_ids=0,
                model_max_length=8192
            )

            self.im_idx, self.im_end_idx = self.tokenizer.convert_tokens_to_ids(["<image>", "</image>"])
        except Exception as e:
            print(f"Error init tokenizer: {e}")


    def tokenize_texts(self, texts):
        try:

            texts =  self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True).input_ids
            image_tokens = torch.tensor([[self.im_idx, self.im_end_idx]] * texts.shape[0])
            return torch.cat([texts[:, 0:1], image_tokens, texts[:, 1:]], dim=1), texts
        except Exception as e:
            print(f"Error tokenizing texts: {e}")

        

    def tokenize_images(self, images):
        try:
                
            tokenized_images = self.processor(images=images, return_tensors="pt").pixel_values
            print(f"Tokenized image: {tokenized_images.shape}")
            return tokenized_images
        
        except Exception as e:
            print(f"Error tokenizing texts: {e}")

    def tokenize(self, sample):
        try:
            
            text_tokens, only_text_tokens = self.tokenize_texts(sample["target_text"])
            attention_mask = text_tokens != self.tokenizer.pad_token_id
            dummy_image_features = torch.ones((text_tokens.shape[0], 64))
            attention_mask = torch.cat([dummy_image_features, attention_mask], dim=1)
            return {
                "text_tokens": text_tokens,
                "images": self.tokenize_images(sample["image"]),
                "labels": only_text_tokens,
                "attention_mask": attention_mask,
            }
        
        except Exception as e:
            print(f"Error during tokenization {e}")
        

# class MedPalm(nn.Module):
#     """
#     MedPalm is a transformer-based model architecture. It initializes with 
#     a Transformer and AutoregressiveWrapper with default or user-specified parameters.

#     Initialize the model with specified or default parameters.
#         Args:
#         - num_tokens: Number of tokens in the vocabulary
#         - max_seq_len: Maximum sequence length
#         - dim: Dimension of the model
#         - depth: Depth of the model
#         - dim_head: Dimension of the model head
#         - heads: Number of heads
#         - use_abs_pos_emb: Whether to use absolute position embedding
#         - alibi_pos_bias: Alibi position bias
#         - alibi_num_heads: Number of alibi heads
#         - rotary_xpos: Rotary position
#         - attn_flash: Attention flash
#         - deepnorm: Deep normalization
#         - shift_tokens: Number of tokens to shift
#         - attn_one_kv_head: Attention one key/value head
#         - qk_norm: Query-key normalization
#         - attn_qk_norm: Attention query-key normalization
#         - attn_qk_norm_dim_scale: Attention query-key normalization dimension scale
#         - embedding_provider: Embedding provider module
#     """
#     def __init__(self, 
#                  num_tokens=20000, 
#                  max_seq_len=4096, 
#                  dim=2560, 
#                  depth=32, 
#                  dim_head=128, 
#                  heads=24,
#                  use_abs_pos_emb=False, 
#                  alibi_pos_bias=True, 
#                  alibi_num_heads=12, 
#                  rotary_xpos=True,
#                  attn_flash=True, 
#                  image_size=256,
#                  patch_size=32,
#                  attn_one_kv_head=False,  # multiquery attention
#                  qk_norm=True, 
#                  attn_qk_norm=False, 
#                  attn_qk_norm_dim_scale=False, 
#                  ):
#         super(MedPalm, self).__init__()

#         self.encoder = ViTransformerWrapper(
#             image_size=image_size,
#             patch_size=patch_size,
#             attn_layers=Encoder(
#                 dim=dim,
#                 depth=depth,
#                 dim_head=dim_head,
#                 heads=heads
#             )
#         )

#         self.decoder = Transformer(
#             num_tokens=num_tokens,
#             max_seq_len=max_seq_len,
#             use_abs_pos_emb=use_abs_pos_emb,
#             attn_layers=Decoder(
#                 dim=dim,
#                 depth=depth,
#                 dim_head=dim_head,
#                 heads=heads,
#                 alibi_pos_bias=alibi_pos_bias,
#                 alibi_num_heads=alibi_num_heads,
#                 rotary_xpos=rotary_xpos,
#                 attn_flash=attn_flash,
#                 attn_one_kv_head=False,
#                 qk_norm=qk_norm,
#                 attn_qk_norm=False,
#                 attn_qk_norm_dim_scale=False,
#                 cross_attend=True
#             )
#         )

#         # self.decoder = AutoregressiveWrapper(self.decoder)

#     def forward(self, text_tokens, img, **kwargs):
#         """
#         Forward pass through the model. It expects the input text_tokens.
#         Args:
#         - text_tokens: Input tokens
#         - kwargs: Other arguments
#         Returns:
#         - output from the decoder
#         """
#         try:
#             print(f"Text tokens shape: {text_tokens.shape}")
#             encoded = self.encoder(img, return_embeddings=True)
#             print(encoded.shape)
#             return self.decoder(text_tokens, context=encoded)
#         except Exception as error:
#             print(f"Failed in forward method: {error}")
#             raise


class MedPalm(torch.nn.Module):
    def __init__(self, 
                 image_size=256, 
                 patch_size=32, 
                 encoder_dim=512, 
                 encoder_depth=6, 
                 encoder_heads=8,
                 num_tokens=20000, 
                 max_seq_len=1024, 
                 decoder_dim=512, 
                 decoder_depth=6, 
                 decoder_heads=8, 
                 alibi_num_heads=4,
                 use_abs_pos_emb=False,
                 cross_attend=True,
                 alibi_pos_bias=True,
                 rotary_xpos=True,
                 attn_flash=True,
                 qk_norm=True):
        
        super(MedPalm, self).__init__()
        
        self.encoder = ViTransformerWrapper(
            image_size=image_size,
            patch_size=patch_size,
            attn_layers=Encoder(
                dim=encoder_dim,
                depth=encoder_depth,
                heads=encoder_heads
            )
        )

        self.decoder = Transformer(
            num_tokens=num_tokens,
            max_seq_len=max_seq_len,
            use_abs_pos_emb=use_abs_pos_emb,
            attn_layers=Decoder(
                dim=decoder_dim,
                depth=decoder_depth,
                heads=decoder_heads,
                cross_attend=cross_attend,
                alibi_pos_bias=alibi_pos_bias,
                alibi_num_heads=alibi_num_heads,
                rotary_xpos=rotary_xpos,
                attn_flash=attn_flash,
                qk_norm=qk_norm,
            )
        )

    def forward(self, img, text):
        try:    
            encoded = self.encoder(img, return_embeddings=True)
            return self.decoder(text, context=encoded)
        except Exception as error:
            print(f"Failed in forward method: {error}")
            raise


