# --------------------------------------------------------
# Copyright (c) 2025 NVIDIA
# Licensed under customized NSCLv1 [see LICENSE.md for details]
# --------------------------------------------------------

# Based on https://github.com/OpenGVLab/InternVL/blob/main/streamlit_demo/model_worker.py
# https://github.com/OpenGVLab/InternVL/?tab=MIT-1-ov-file#readme

# Importing torch before transformers can cause `segmentation fault`

from transformers import  AutoTokenizer,  AutoConfig
from transformers.modeling_outputs import SequenceClassifierOutputWithPast
import base64
import os
from io import BytesIO
from typing import Tuple
import math
import requests
import torch
from torch import Tensor
import torchvision.transforms as T
from PIL import Image
from torchvision.transforms.functional import InterpolationMode
from typing import Optional, Any, Union, Dict, List

from tqdm import tqdm
import torch.nn.functional as F
from datasets import Dataset
from torch.utils.data import DataLoader

from .modeling_eagle_chat import Eagle2ChatModel
from .configuration_eagle_chat import Eagle2ChatConfig
from .conversation import get_conv_template 

from .configuration_siglip import SiglipVisionConfig
from .modeling_siglip import SiglipVisionModel
from .flash_attention import *

from .llama_bidirectional_model import LlamaBidirectionalModel
from transformers import PreTrainedModel

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

SIGLIP_MEAN = (0.5, 0.5, 0.5)
SIGLIP_STD = (0.5, 0.5, 0.5)

def load_image(image):
    if isinstance(image, Image.Image):
        return image
    elif isinstance(image, str) and os.path.exists(image):
        return Image.open(image)
    elif isinstance(image, dict):
        if 'disk_path' in image:
            return Image.open(image['disk_path'])
        elif 'base64' in image:
            return Image.open(BytesIO(base64.b64decode(image['base64'])))
        elif 'url' in image:
            response = requests.get(image['url'])
            return Image.open(BytesIO(response.content))
        elif 'bytes' in image:
            return Image.open(BytesIO(image['bytes']))
        else:
            raise ValueError(f'Invalid image: {image}')
    else:
        raise ValueError(f'Invalid image: {image}')

def build_transform(input_size, norm_type='imagenet'):
    if norm_type == 'imagenet':
        MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    elif norm_type == 'siglip':
        MEAN, STD = SIGLIP_MEAN, SIGLIP_STD
        
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform

def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    """
    previous version mainly foucs on ratio.
    We also consider area ratio here.
    """
    best_factor = float('-inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        area_ratio = (ratio[0]*ratio[1]*image_size*image_size)/ area
        """
        new area > 60% of original image area is enough.
        """
        factor_based_on_area_n_ratio = min((ratio[0]*ratio[1]*image_size*image_size)/ area, 0.6)* \
                                     min(target_aspect_ratio/aspect_ratio, aspect_ratio/target_aspect_ratio)
        
        if factor_based_on_area_n_ratio > best_factor:
            best_factor = factor_based_on_area_n_ratio
            best_ratio = ratio
        
    return best_ratio

def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    # calculate the existing image aspect ratio
    target_ratios = set(
        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
        i * j <= max_num and i * j >= min_num)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    # calculate the target width and height
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    # resize the image
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size
        )
        # split the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images

def split_model(model_path, device):

    device_map = {}
    world_size = torch.cuda.device_count()
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    num_layers = config.llm_config.num_hidden_layers

    print('world_size', world_size)
    num_layers_per_gpu_ = math.floor(num_layers / (world_size - 1))
    num_layers_per_gpu = [num_layers_per_gpu_] * world_size
    num_layers_per_gpu[device] = num_layers - num_layers_per_gpu_ * (world_size-1)
    print(num_layers_per_gpu)
    layer_cnt = 0
    for i, num_layer in enumerate(num_layers_per_gpu):
        for j in range(num_layer):
            device_map[f'language_model.model.layers.{layer_cnt}'] = i
            layer_cnt += 1
    device_map['vision_model'] = device
    device_map['mlp1'] = device
    device_map['language_model.model.tok_embeddings'] = device
    device_map['language_model.model.embed_tokens'] = device
    device_map['language_model.output'] = device
    device_map['language_model.model.norm'] = device
    device_map['language_model.lm_head'] = device
    device_map['language_model.model.rotary_emb'] = device
    device_map[f'language_model.model.layers.{num_layers - 1}'] = device
    return device_map

class llama_NemoRetrieverColEmbedConfig(Eagle2ChatConfig):
    model_type = "llama_nemoretrievercolembed"

    q_max_length: Optional[int] 
    p_max_length: Optional[int]     
    query_prefix: str 
    passage_prefix: str
    pooling: str
    bidirectional_attention: bool

    def __init__(
        self, 
        q_max_length: Optional[int] = 512,
        p_max_length: Optional[int] = 10240,        
        query_prefix: str = "query:",
        passage_prefix: str = "passage:",
        pooling: str = "last", 
        bidirectional_attention: bool = False,
        max_input_tiles: int = 2,
        img_context_token_id: int = 128258, #tokenizer.convert_tokens_to_ids("<IMG_CONTEXT>")
        out_dimension: int = -1,
        **kwargs,
    ):
        self.q_max_length = q_max_length
        self.p_max_length = p_max_length        
        self.query_prefix = query_prefix
        self.passage_prefix = passage_prefix
        self.pooling = pooling
        self.bidirectional_attention = bidirectional_attention
        self.img_context_token_id = img_context_token_id
        self.max_input_tiles = max_input_tiles
        self.out_dimension = out_dimension
        super().__init__(**kwargs)

class llama_NemoRetrieverColEmbed(Eagle2ChatModel):

    config_class = llama_NemoRetrieverColEmbedConfig
    _supports_flash_attn_2 = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.padding = True
        self.q_max_length = 512
        self.p_max_length = 10240
        self.pad_to_multiple_of = None
        self.query_prefix = 'query:'
        self.passage_prefix = 'passage:'

        if isinstance(args[0], llama_NemoRetrieverColEmbedConfig):
            tokenizer = AutoTokenizer.from_pretrained(args[0]._name_or_path, trust_remote_code=True)
            tokens_to_keep = ['<box>', '</box>', '<ref>', '</ref>']
            tokenizer.additional_special_tokens = [item for item in tokenizer.additional_special_tokens if item not in tokens_to_keep]
            tokenizer.padding_side = 'left'
            self.tokenizer = tokenizer

        self.norm_type = 'siglip'
        self.image_size = self.config.force_image_size
        self.max_input_tiles = 6
        self.system_message = ""
        self.use_visual_embedding = True
    
    def process_documents(self, documents: Union[Dict,List[Dict]], **kwargs):
        if isinstance(documents, dict):
            images = documents["images"]
            texts = documents["texts"]
            assert len(texts) == len(images)
        elif isinstance(documents, list):
            images = [pair['image'] for pair in documents ]
            texts = [pair['text'] for pair in documents ]
        else:
            raise ValueError("The documents need to be a dict or list of dicts")
        
        if self.passage_prefix:
            texts = [self.passage_prefix + ' ' + t for t in texts]

        contents, pil_images, max_input_tile_list, llm_onlys = [], [], [], []
        for image, text in zip(images, texts):
            prefix = ''
            llm_only = True
            if image != '':
                pil_images.append(load_image(image))
                prefix = '<image>'
                max_input_tile_list.append(self.max_input_tiles)
                llm_only = False
            else:
                pil_images.append(None)
                max_input_tile_list.append(self.max_input_tiles)
            
            llm_onlys.append(llm_only)

            content = text
            if prefix!='':
                content = prefix + ' ' + content
            if self.passage_prefix:
                content = self.passage_prefix + ' ' + content
            contents.append(content)
        
        transform = build_transform(input_size=self.image_size, norm_type=self.norm_type)

        template = get_conv_template(self.config.template)
        template.system_message = self.system_message
        
        content_prompts = []
        pixel_values_list = []
        for content, pil_image, max_input_tiles, llm_only in zip(contents, pil_images, max_input_tile_list, llm_onlys):
            if pil_image is not None:
                if self.config.dynamic_image_size:
                    image_tiles = dynamic_preprocess(
                        pil_image, image_size=self.image_size, max_num=max_input_tiles,
                        use_thumbnail=self.config.use_thumbnail)
                else:
                    image_tiles = [pil_image]
                        
                pixel_values = [transform(item) for item in image_tiles]
                pixel_values = torch.stack(pixel_values).to(dtype=torch.bfloat16)
                pixel_values_list.append(pixel_values)
            else:
                pixel_values = None

            IMG_START_TOKEN='<img>'
            IMG_END_TOKEN='</img>'
            IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'        

            if pixel_values is not None and '<image>' not in content and not llm_only:
                content = '<image> ' + content

            # Reseting conversation messages
            template.messages.clear()

            # TODO: do we need this template?
            template.append_message(template.roles[0], content) # user
            template.append_message(template.roles[1], None)     # assistant
            content_prompt = template.get_prompt()

            if '<image>' not in content:
                content_prompt = content_prompt
            else:
                num_patches = pixel_values.shape[0]
                image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
                content_prompt = content_prompt.replace('<image>', image_tokens, 1)

            content_prompts.append(content_prompt)

        model_inputs = self.tokenizer(content_prompts, 
                                      truncation=True,
                                      max_length=self.p_max_length,
                                      padding=self.padding,
                                      pad_to_multiple_of=self.pad_to_multiple_of,
                                      return_tensors='pt')

        if len(pixel_values_list)>1:
            pixel_values_squeezed = torch.concat(pixel_values_list, axis=0)
        elif len(pixel_values_list)==1:
            pixel_values_squeezed = pixel_values_list[0]
        else:
            pixel_values_squeezed = None

        batch_docs = {
            "input_ids": model_inputs['input_ids'],
            "attention_mask": model_inputs['attention_mask'],
            "pixel_values": None
        }
        if pixel_values_squeezed is not None:
            batch_docs["pixel_values"] = pixel_values_squeezed
        
        return batch_docs

    def process_queries(self, queries: List[str], **kwargs):

        template = get_conv_template(self.config.template)
        template.system_message = self.system_message

        query_prompts = []
        for query in queries:
            if self.query_prefix:
                query = f"{self.query_prefix} {query}"

            # Reseting conversation messages
            template.messages.clear()

            template.append_message(template.roles[0], query)    # user
            template.append_message(template.roles[1], None)     # assistant
            query_prompt = template.get_prompt()

            query_prompts.append(query_prompt)

        
        batch_query = self.tokenizer(
            query_prompts, 
            truncation=True,
            max_length=self.q_max_length,
            padding=self.padding,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors='pt'
        )

        return batch_query

    def get_scores(
        self, 
        query_embeddings: Union[torch.Tensor, List[torch.Tensor]],
        passage_embeddings: Union[torch.Tensor, List[torch.Tensor]],
        batch_size: Optional[int] = 8,
    ) -> torch.Tensor:
        """Dot-product similarity between queries and passages."""        
        if isinstance(query_embeddings, list):
            if len(query_embeddings[0].shape)==2:
                # Expend Batch Dimension as ViDoRe Framework remove it
                query_embeddings = [q.unsqueeze(0) for q in query_embeddings]
            query_embeddings = self.padding_various_shape_tensor(query_embeddings)
        if isinstance(passage_embeddings, list):
            if len(passage_embeddings[0].shape)==2:
                # Expend Batch Dimension as ViDoRe Framework remove it
                passage_embeddings = [p.unsqueeze(0) for p in passage_embeddings]
            passage_embeddings = self.padding_various_shape_tensor(passage_embeddings)
        
        return self.colbert_score(query_embeddings, passage_embeddings, batch_size)

    def colbert_score(
        self,
        qs: Union[torch.Tensor, List[torch.Tensor]],
        ps: Union[torch.Tensor, List[torch.Tensor]],
        batch_size: int = 128,
        device: Optional[Union[str, torch.device]] = None,
    ) -> torch.Tensor:
        if batch_size is None:
            batch_size = 128
        if len(qs) == 0:
            raise ValueError("No queries provided")
        if len(ps) == 0:
            raise ValueError("No passages provided")
    
        scores_list: List[torch.Tensor] = []
        for i in range(0, len(qs), batch_size):
            scores_batch = []
            qs_batch = torch.nn.utils.rnn.pad_sequence(qs[i : i + batch_size].cuda(), batch_first=True, padding_value=0)
            for j in range(0, len(ps), batch_size):
                ps_batch = torch.nn.utils.rnn.pad_sequence(
                    ps[j : j + batch_size].cuda(), batch_first=True, padding_value=0
                )
                scores_batch.append(torch.einsum("bnd,csd->bcns", qs_batch, ps_batch).max(dim=3)[0].sum(dim=2))
            # Keep scores_batch on the GPU
            scores_batch = torch.cat(scores_batch, dim=1)
            scores_list.append(scores_batch)
    
        scores = torch.cat(scores_list, dim=0)
        return(scores)
    
    def _extract_embeddings(self, dataloader: DataLoader, is_query: bool) -> List[torch.Tensor]:
        qs = []
        message = "query" if is_query else "document"
        for batch in tqdm(dataloader, desc=f"Extracting {message} embeddings..."):
            with torch.inference_mode():
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    if 'pixel_values' in batch and batch['pixel_values'] is None:
                        batch.pop('pixel_values')
                    batch = {k: v.to(self.device) for k, v in batch.items()}
                    embeddings = self(**batch, output_hidden_states=True).hidden_states[-1]
                    embeddings = embeddings*batch['attention_mask'].unsqueeze(-1)
                    embeddings = F.normalize(embeddings, dim=-1)
                    
            # Detecting abnormal outputs
            assert torch.sum(embeddings).float().item() not in [float(0.), float("inf")]
            qs.append(embeddings.contiguous())
        
        qs_tensor = self.padding_various_shape_tensor(qs)
        all_embeddings_tensor = qs_tensor.detach().cpu()
        return all_embeddings_tensor

    def forward_passages(self, passages, batch_size=8, **kwargs) -> Union[torch.Tensor, List[torch.Tensor]]:
        """Forward passages as image-only documents."""
        corpus = []
        for image in passages:
            corpus.append({
                "image": image,
                "text": ''
            })
        return self.forward_documents(corpus, batch_size)
    
    def forward_queries(self, queries: List, batch_size=8) -> List[torch.Tensor]:
        dataset = ListDataset[str](queries)
        dataloader = DataLoader(
            dataset=dataset,
            batch_size=batch_size,
            collate_fn=self.process_queries,
            shuffle=False,
            num_workers=8,
            pin_memory=True,
            drop_last=False,
        )
        return self._extract_embeddings(dataloader=dataloader, is_query=True)

    def forward_documents(self, corpus: List, batch_size=8) -> List[torch.Tensor]:    
        images = []
        texts = []
        for doc in corpus:
            text = doc["text"]
            image = doc.get("image", "")
            if image.mode != "RGB":
                image = image.convert("RGB")
            images.append(image)
            texts.append(text)

        dataset = Dataset.from_dict({"image": images, "text": texts})
        dataloader = DataLoader(
            dataset=dataset,
            batch_size=batch_size,
            collate_fn=self.process_documents,
            shuffle=False,
            num_workers=8,
            pin_memory=True,
            drop_last=False,
        )
        return self._extract_embeddings(dataloader=dataloader, is_query=False)

    def padding_various_shape_tensor(self, tensors: List[torch.Tensor]) -> torch.Tensor:
        """Pad tensors of various shapes for colbert-like scoring"""
        max_seq_len = max(t.shape[1] for t in tensors)
        padded_tensors = [F.pad(t, (0, 0, 0, max_seq_len - t.shape[1]), mode="constant", value=0) for t in tensors]
        return torch.cat(padded_tensors, dim=0)


from typing import TypeVar
from torch.utils.data import Dataset as TorchDataset
TV = TypeVar("T")
class ListDataset(TorchDataset[TV]):
    def __init__(self, elements: List[TV]):
        self.elements = elements

    def __len__(self) -> int:
        return len(self.elements)

    def __getitem__(self, idx: int) -> TV:
        return self.elements[idx]