import torch
import sys
import torch.nn.functional as F
import numpy as np

import utils.loss
import utils.samp
import utils.data
import utils.improc
import utils.misc
import utils.saveload
import cv2
from nets.blocks import InputPadder

import torch
from PIL import Image, ImageDraw
import numpy as np

torch.set_float32_matmul_precision('medium')

def run_example(processor, model, task_prompt, image, text_input=None):
    if text_input is None:
        prompt = task_prompt
    else:
        prompt = task_prompt + text_input
    inputs = processor(text=prompt, images=image, return_tensors="pt").to('cuda', torch.float32)
    generated_ids = model.generate(
      input_ids=inputs["input_ids"],
      pixel_values=inputs["pixel_values"],
      max_new_tokens=1024,
      num_beams=3
    )
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]

    parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))

    return parsed_answer


def polygons_to_mask(image, prediction, fill_value=255):
    """
    Converts polygons into a mask.
    
    Parameters:
    - image: A PIL Image instance whose size will be used for the mask.
    - prediction: Dictionary containing 'polygons' and 'labels'. 
                  'polygons' is a list where each element is a list of sub-polygons.
    - fill_value: The pixel value used to fill the polygon areas (default 255 for a binary mask).
    
    Returns:
    - A NumPy array representing the mask (same width and height as the input image).
    """
    # Create a blank grayscale mask image with the same size as the original image.
    mask = Image.new('L', image.size, 0)
    draw = ImageDraw.Draw(mask)
    
    # Iterate over each set of polygons
    for polygons in prediction['polygons']:
        # Each element in "polygons" can be a sub-polygon
        for poly in polygons:
            # Ensure the polygon is in the right shape and has at least 3 points.
            poly_arr = np.array(poly).reshape(-1, 2)
            if poly_arr.shape[0] < 3:
                print('Skipping invalid polygon:', poly_arr)
                continue
            # Convert the polygon vertices into a list for drawing.
            poly_list = poly_arr.reshape(-1).tolist()
            # Draw the polygon on the mask with the fill_value.
            draw.polygon(poly_list, fill=fill_value)
    
    # Convert the PIL mask image to a NumPy array and return it.
    return np.array(mask)


class Tracker:
    def __init__(self, model, mean, std, S, stride, inference_iters, target_res, device='cuda'):
        """
        Initializes the Tracker.

        Args:
            model: The model used to compute feature maps and forward window flow.
            mean: Tensor or value used for normalizing the input.
            std: Tensor or value used for normalizing the input.
            S: Window size for the tracker.
            stride: The stride used when updating the window.
            inference_iters: Number of inference iterations.
            device: Torch device, defaults to 'cuda'.
        """
        self.model = model
        self.mean = mean
        self.std = std
        self.S = S
        self.stride = stride
        self.inference_iters = inference_iters
        self.device = device
        self.target_res = target_res

        self.padder = None
        self.cnt = 0
        self.fmap_anchor = None
        self.fmaps2 = None
        self.flows8 = None
        self.visconfs8 = None
        self.flows = []  # List to store computed flows
        self.visibs = []  # List to store visibility confidences
        self.rgbs = []  # List to store RGB frames

    def reset(self):
        """Reset the tracker state."""
        self.padder = None
        self.cnt = 0
        self.fmap_anchor = None
        self.fmaps2 = None
        self.flows8 = None
        self.visconfs8 = None
        self.flows = []
        self.visibs = []
        self.rgbs = []
        
    def preprocess(self, rgb_frame):
        # Resize frame (scale to keep maximum dimension ~1024)
        scale = min(self.target_res / rgb_frame.shape[0], self.target_res / rgb_frame.shape[1])
        rgb_resized = cv2.resize(rgb_frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
        
        # Convert to tensor, normalize and move to device.
        rgb_tensor = torch.from_numpy(rgb_resized).permute(2, 0, 1).float().unsqueeze(0).to(self.device)
        rgb_tensor = rgb_tensor / 255.0
        
        self.rgbs.append(rgb_tensor.cpu())
            
        # import pdb; pdb.set_trace()
        rgb_tensor = (rgb_tensor - self.mean) / self.std
        return rgb_tensor

    @torch.no_grad()
    def track(self, rgb_frame):
        """
        Process a single RGB frame and return the computed flow when available.

        Args:
            rgb_frame: A NumPy array containing the RGB frame.
                       (Assumed to be in RGB; if coming from OpenCV, convert it before passing.)
        
        Returns:
            flow_predictions: The predicted flow for the current frame (or None if not enough frames have been processed).
        """
        torch.cuda.empty_cache()

        rgb_tensor = self.preprocess(rgb_frame)
        
        # Initialize padder on the first frame.
        if self.cnt == 0:
            self.padder = InputPadder(rgb_tensor.shape)
        rgb_padded = self.padder.pad(rgb_tensor)[0]
        _, _, H_pad, W_pad = rgb_padded.shape
        C = 256  # Feature map channel dimension (could be parameterized if needed)
        H8, W8 = H_pad // 8, W_pad // 8

        # Accumulate feature maps until the window is full.
        if self.cnt == 0:
            self.fmap_anchor = self.model.get_fmaps(rgb_padded, 1, 1, None, False, False).reshape(1, C, H8, W8)
            self.fmaps2 = self.fmap_anchor[:, None]
            self.cnt += 1
            return None
            
        new_fmap = self.model.get_fmaps(rgb_padded, 1, 1, None, False, False).reshape(1, 1, C, H8, W8)
        self.fmaps2 = torch.cat([self.fmaps2[:, (1 if self.fmaps2.shape[1] >= self.S else 0):].detach().clone(), new_fmap], dim=1)
        
        # need to track
        if self.cnt - self.S + 1 >= 0 and (self.cnt - self.S + 1) % self.stride == 0:
            # Initialize or update temporary flow buffers.
            iter_num = self.inference_iters
            if self.flows8 is None:
                self.flows8 = torch.zeros((self.S, 2, H_pad // 8, W_pad // 8), device=self.device)
                self.visconfs8 = torch.zeros((self.S, 2, H_pad // 8, W_pad // 8), device=self.device)
                # iter_num = self.inference_iters
            else:
                self.flows8 = torch.cat([
                    self.flows8[self.stride:self.stride + self.S // 2].detach().clone(),
                    self.flows8[self.stride + self.S // 2 - 1:self.stride + self.S // 2].detach().clone().repeat(self.S // 2, 1, 1, 1)
                ])
                self.visconfs8 = torch.cat([
                    self.visconfs8[self.stride:self.stride + self.S // 2].detach().clone(),
                    self.visconfs8[self.stride + self.S // 2 - 1:self.stride + self.S // 2].detach().clone().repeat(self.S // 2, 1, 1, 1)
                ])

            # import pdb; pdb.set_trace()
            # Compute flow predictions using the model's forward window.
            flow_predictions, visconf_predictions, self.flows8, self.visconfs8, _ = self.model.forward_window(
                self.fmap_anchor,
                self.fmaps2,
                self.visconfs8,
                iters=iter_num,
                flowfeat=None,
                flows8=self.flows8,
                is_training=False
            )
            flow_predictions = self.padder.unpad(flow_predictions[-1][0 if self.cnt == self.S - 1 else -self.stride:])
            visconf_predictions = self.padder.unpad(torch.sigmoid(visconf_predictions[-1][0 if self.cnt == self.S - 1 else -self.stride:]))

            self.cnt += 1
            self.flows.append(flow_predictions.cpu())
            self.visibs.append(visconf_predictions.cpu())

            return flow_predictions, visconf_predictions
        
        self.cnt += 1
        return None