import os
from abc import ABCMeta, abstractmethod
from typing import Optional, Union, Dict, List
from termcolor import colored
import random


import numpy as np
import torch
from transformers import (
    AutoProcessor,
    AutoTokenizer,
    LlavaConfig, 
    LlamaForCausalLM,
)
from torchvision.transforms.v2 import (
    ToPILImage,
)
import decord
from decord import VideoReader

decord.bridge.set_bridge("torch")

# TODO: need to use these directly
from tarsier.modeling_tarsier import TarsierForConditionalGeneration
from tarsier.processor import Processor
# from utils.model import transform_pixel_values


EOL_PROMPTS = {
    'text': '<sent>\nSummary above sentence in one word:',
    'image': '<image>\nSummary above image in one word:',
    'video': '<video>\nSummary above video in one word:',
}


def transform_pixel_values(pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
    # NOTE: this function doesn't accept unbatched inputs
    # pixel_values should be uint8 of (B, T, C, H, W)
    if isinstance(pixel_values, list):
        pixel_values = torch.stack(pixel_values)

    if pixel_values.ndim == 4:
        # pixel_values is (B, C, H, W)
        # (B, C, H, W) -> (B, 1, C, H, W)
        pixel_values = pixel_values.unsqueeze(1)
    elif pixel_values.ndim == 5:
        # pixel_values is (B, T, C, H, W)
        pass
    else:
        raise ValueError(f"pixel_values should be 4D or 5D, got {pixel_values.ndim}D")
    return pixel_values


base_registry = {}
class BaseModel(metaclass=ABCMeta):
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # register model architecture
        if hasattr(cls, 'ARCHITECTURE'):
            base_registry[cls.ARCHITECTURE] = cls
    
    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        load_llm: bool = False,
        device_map: Optional[Union[str, Dict[str, int]]] = None,
        **kwargs):
        print(colored(f'[ MODEL ] Loading {cls.__name__} from {model_name_or_path} [..............]', 'yellow'))

        return cls(model_name_or_path, load_llm=load_llm, device_map=device_map, **kwargs)


class BaseModelForTARA(BaseModel):
    
    ARCHITECTURE = "TarsierForConditionalGeneration"
    LLM_CLASS = LlamaForCausalLM
    MLLM_CLASS = TarsierForConditionalGeneration

    @property
    def describe_prompt(self):
        return "Describe the video in detail."

    @property
    def text_eol_prompt(self):
        prompt = f'USER: {EOL_PROMPTS["text"]} ASSISTANT: '
        return prompt
    
    @property
    def image_eol_prompt(self):
        prompt = f'USER: {EOL_PROMPTS["image"]} ASSISTANT: '
        return prompt
    
    @property
    def video_eol_prompt(self):
        prompt = f'USER: {EOL_PROMPTS["video"]} ASSISTANT: '
        return prompt

    def __init__(
            self, 
            model_name_or_path: str,
            load_llm: Optional[bool] = None,
            device_map: Optional[Union[str, Dict[str, int]]] = None,
            **kwargs,
        ):

        MODEL_CLASS = self.LLM_CLASS if load_llm else self.MLLM_CLASS

        if load_llm:
            self.split_weights(model_name_or_path, model_name_or_path + '-llm')
            model_name_or_path += '-llm'
            model_config = None
            self.processor = AutoProcessor.from_pretrained(model_name_or_path, use_fast=False)
        else:
            model_config = LlavaConfig.from_pretrained(
                model_name_or_path,
                # trust_remote_code=True,
            )
            self.processor = Processor(
                model_name_or_path,
                max_n_frames=32,
            )
        
        self.tokenizer = self.processor.tokenizer

        self.model = MODEL_CLASS.from_pretrained(
            model_name_or_path,
            config=model_config,
            torch_dtype=kwargs.get("torch_dtype", torch.bfloat16),
            device_map=device_map,
            # trust_remote_code=True
        )
        
        self.model.eval()

    def split_weights(self, mllm_path, llm_path):
        if os.path.exists(llm_path):
            print(f'{llm_path} already exists. Skip splitting weights.')
            return
        print('Splitting LLM weights from MLLM.')
        model = self.MLLM_CLASS.from_pretrained(mllm_path)
        llm = model.language_model
        processor = AutoProcessor.from_pretrained(mllm_path)
        tokenizer = AutoTokenizer.from_pretrained(mllm_path)
        llm.save_pretrained(llm_path)
        processor.save_pretrained(llm_path)
        tokenizer.save_pretrained(llm_path)


encoder_registry = {}
class EncodeMixin(metaclass=ABCMeta):
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # register model architecture
        if hasattr(cls, 'ARCHITECTURE'):
            encoder_registry[cls.ARCHITECTURE] = cls

    @abstractmethod
    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
        """
        Encodes vision data (images or videos) into a tensor representation.

        Args:
            pixel_values (torch.Tensor | List[torch.Tensor]): The input pixel values. 
                - If a tensor, it should be of shape (B, C, H, W) for images or (B, T, C, H, W) for videos.
                - If a list, it will be stacked into a tensor.

        Returns:
            torch.Tensor: The encoded tensor representation of the input vision data.

        Raises:
            ValueError: If `pixel_values` is not 4D or 5D.

        ## Notes:
            - This function does not accept unbatched inputs.
            - `pixel_values` should be of type uint8.
        """
        raise NotImplementedError

    @abstractmethod
    def encode_text(self, text: str | List[str]) -> torch.Tensor:
        """
        Encodes the given text(s) into a tensor representation using the model.

        Args:
            text (str | List[str]): A single string or a list of strings to be encoded.

        Returns:
            torch.Tensor: The tensor representation of the encoded text(s).

        ## Notes:
            - The method uses a prompt to encode the text.
            - If a single string is provided, it is converted into a list containing that string.
            - The method processes the prompts and generates the tensor representation using the model.
            - The output tensor contains the hidden states of the last token for each input text.
        """
        raise NotImplementedError


class TARA(BaseModelForTARA, EncodeMixin):

    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:

        pixel_values = transform_pixel_values(pixel_values) # [B, T, C, H, W]
        nframes = pixel_values.shape[1]
        prompt = self.image_eol_prompt if nframes == 1 else self.video_eol_prompt
        
        to_image = ToPILImage()
        batched_frames = []
        for batch in pixel_values:
            frames = [to_image(v) for v in batch]
            batched_frames.append(frames)

        generate_kwargs = {
            "max_new_tokens": 1,
            "output_hidden_states": True,
            "return_dict_in_generate": True,
        }

        vision_embs = []

        for frames in batched_frames:
            input_prompt = prompt.replace("<video>", "<image>"*len(frames))
            input_ids = self.processor.get_text_inputs(input_prompt)
            frames = self.processor.get_pixel_values(frames)
            inputs = {
                "input_ids": input_ids,
                "pixel_values": frames
            }
            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
            outputs = self.model.generate(
                **inputs,
                **generate_kwargs,
            )
            vision_embs.append(outputs.hidden_states[0][-1][:, -1, :])
        
        vision_embs = torch.cat(vision_embs)
        return vision_embs
    
    def encode_text(self, text: str | List[str]) -> torch.Tensor:

        prompt = self.text_eol_prompt

        if isinstance(text, str):
            text = [text]
        
        prompts = [prompt.replace('<sent>', t) for t in text]

        generate_kwargs = {
            "max_new_tokens": 1,
            "output_hidden_states": True,
            "return_dict_in_generate": True,
        }

        text_embs = []

        for p in prompts:
            text_inputs = self.processor.get_text_inputs(p)
            inputs = {
                "input_ids": text_inputs,
            }
            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
            outputs = self.model.generate(
                **inputs,
                **generate_kwargs,
            )
            text_embs.append(outputs.hidden_states[0][-1][:, -1, :])
        
        text_embs = torch.cat(text_embs)
        return text_embs





def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1):
    if sample in ["rand", "middle"]: # uniform sampling
        acc_samples = min(num_frames, vlen)
        # split the video into `acc_samples` intervals, and sample from each interval.
        intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
        ranges = []
        for idx, interv in enumerate(intervals[:-1]):
            ranges.append((interv, intervals[idx + 1] - 1))
        if sample == 'rand':
            try:
                frame_indices = [random.choice(range(x[0], x[1])) for x in ranges]
            except (ValueError, IndexError):
                frame_indices = np.random.permutation(vlen)[:acc_samples]
                frame_indices.sort()
                frame_indices = list(frame_indices)
        elif fix_start is not None:
            frame_indices = [x[0] + fix_start for x in ranges]
        elif sample == 'middle':
            frame_indices = [(x[0] + x[1]) // 2 for x in ranges]
        else:
            raise NotImplementedError

        if len(frame_indices) < num_frames:  # padded with last frame
            padded_frame_indices = [frame_indices[-1]] * num_frames
            padded_frame_indices[:len(frame_indices)] = frame_indices
            frame_indices = padded_frame_indices
    elif "fps" in sample:  # fps0.5, sequentially sample frames at 0.5 fps
        output_fps = float(sample[3:])
        duration = float(vlen) / input_fps
        delta = 1 / output_fps  # gap between frames, this is also the clip length each frame represents
        frame_seconds = np.arange(0 + delta / 2, duration + delta / 2, delta)
        frame_indices = np.around(frame_seconds * input_fps).astype(int)
        frame_indices = [e for e in frame_indices if e < vlen]
        if max_num_frames > 0 and len(frame_indices) > max_num_frames:
            frame_indices = frame_indices[:max_num_frames]
            # frame_indices = np.linspace(0 + delta / 2, duration + delta / 2, endpoint=False, num=max_num_frames)
    else:
        raise ValueError
    return frame_indices


def read_frames_decord(
        video_path, num_frames, sample='middle', fix_start=None, 
        max_num_frames=-1, trimmed30=False, height=-1, width=-1
    ):
    decord.bridge.set_bridge('torch')

    # num_threads = 1 if video_path.endswith('.webm') else 0 # make ssv2 happy
    num_threads = 1
    video_reader = VideoReader(video_path, num_threads=num_threads, height=height, width=width)
    try:
        vlen = len(video_reader)

        fps = video_reader.get_avg_fps()
        duration = vlen / float(fps)

        # only use top 30 seconds
        if trimmed30 and duration > 30:
            duration = 30
            vlen = int(30 * float(fps))

        frame_indices = get_frame_indices(
            num_frames, vlen, sample=sample, fix_start=fix_start,
            input_fps=fps, max_num_frames=max_num_frames
        )

        frames = video_reader.get_batch(frame_indices)  # (T, H, W, C), torch.uint8
        if not isinstance(frames, torch.Tensor):
            frames = torch.from_numpy(frames.asnumpy())
        frames = frames.permute(0, 3, 1, 2)  # (T, C, H, W), torch.uint8
        return frames
    finally:
        # Explicitly release underlying resources to avoid file descriptor leaks
        del video_reader


if __name__ == "__main__":

    # Load model
    model = TARA.from_pretrained(
        "/work/piyush/experiments/CaRe/Tarsier-7b/final-10112025/nli_9000+ego_1000+subj_replaced-seed_42/merged_checkpoint",
        device_map='auto',
        dtype=torch.bfloat16,
    )
    n_params = sum(p.numel() for p in model.model.parameters())
    print(f"Number of parameters: {round(n_params/1e9, 3)}B")
    
    # Let's encode a sample video
    print(colored("Testing video encoding...", 'cyan'))
    video_path = "./assets/folding_paper.mp4"
    video_tensor = read_frames_decord(video_path, num_frames=16)
    video_tensor = video_tensor.unsqueeze(0)
    video_tensor = video_tensor.to(model.model.device)
    with torch.no_grad():
        video_emb = model.encode_vision(video_tensor).cpu().squeeze(0).float()
    print("Video shape:", video_tensor.shape) # torch.Size([1, 16, 3, 240, 426])
    print("Video embedding shape:", video_emb.shape) # torch.Size([4096])
    
    # Let's encode a sample text
    print(colored("Testing text encoding...", 'cyan'))
    text = ['someone is folding a paper', 'cutting a paper', 'someone is unfolding a paper']
    # NOTE: It can also take a single string
    with torch.no_grad():
        text_emb = model.encode_text(text).cpu().float()
    print("Text:", text)
    print("Text embedding shape:", text_emb.shape) # torch.Size([3, 4096])