import cv2
import numpy as np
import time


class Tracker:
    def __init__(self,
                 adaptive,
                 cuda,
                 config_tracker=None,
                 lower_threshold=12.0,
                 upper_threshold=18.0,
                 loop_closure_interval=1,
                #  dis_preset='medium',
                 verbose=False):
        # Check if OpenCV is built with CUDA support
        # if cuda and not cv2.cuda.getCudaEnabledDeviceCount():
        #     # print("Warning: CUDA requested but OpenCV not built with CUDA support or no CUDA device found. Falling back to CPU.")
        #     cuda = False
        
        self.cuda = cuda
        self.verbose = verbose
        self.__loop_closure_interval = loop_closure_interval
        self.__frame_count = 0
        dis_preset = config_tracker["dis_preset"]
        loop_closure_interval = config_tracker["loop_closure_interval"]
        self.__algorithm = config_tracker["algorithm"]
        
        # Conditionally import torch only if RAFT algorithm is used
        self.torch = None
        if self.__algorithm == "raft":
            import torch
            self.torch = torch


        if self.__algorithm == "raft" and (not self.cuda or not self.torch.cuda.is_available()):
            self.__algorithm = "dis"
            if self.verbose:
                print("No CUDA available, falling back to OPENCV DIS optical flow")


        if self.__algorithm == "raft":
            ### RAFT
            self.__raft_device = "cuda"
            print(f"RAFT device: {self.__raft_device}")
            self.__raft_model = self.setup_simple_raft()
            self.__raft_ref_frame = None
            self.__raft_ref_tensor = None

        elif self.__algorithm == "dis":
            # DIS method - optimize for speed
            if dis_preset == 'ultrafast':
                preset = cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST
            elif dis_preset == 'fast':
                preset = cv2.DISOPTICAL_FLOW_PRESET_FAST
            else:
                preset = cv2.DISOPTICAL_FLOW_PRESET_MEDIUM
            
            self.__tracking_inst = cv2.DISOpticalFlow_create(preset)            
            
            # Aggressive optimization for CPU
            self.__tracking_inst.setFinestScale(config_tracker["dis_mode"])      # Skip finest pyramid level (2x faster)
            # self.__tracking_inst.setUseSpatialPropagation(False)
            # self.__tracking_inst.setUseMeanNormalization(False)
            # self.__tracking_inst.setGradientDescentIterations(16) #16
            # self.__tracking_inst.setPatchSize(12) # 8
            # self.__tracking_inst.setPatchStride(5) # 3
            # self.__tracking_inst.setVariationalRefinementIterations(5)#5

            # else:
            #     # CUDA flow method with stream support
            #     self.__tracking_inst = cv2.cuda.FarnebackOpticalFlow_create()
            #     # Create CUDA stream for async operations
            #     self.__cuda_stream = cv2.cuda.Stream()
            
            self.__fine_base_frame_queue = []
            self.__fine_flow_queue = []
            self.__prev_fine_flow = None
            self.__adaptive = adaptive
            self.__lower_threshold = config_tracker["lower_threshold"]  # error lower than threshold will close the loop
            self.__upper_threshold = config_tracker["upper_threshold"]  # error greater than threshold will add new reference frame



    def track(self, frame):
        """Main tracking function that dispatches to CUDA or CPU implementation"""
        if self.__algorithm == "raft":
            return self.track_raft(frame)
        elif self.__algorithm == "dis":
            if frame.ndim == 3:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            return self.track_cpu(frame)
    
    def track_cpu(self, frame):
        """CPU-based optical flow tracking"""
        frame_height, frame_width = frame.shape[:2]
        
        if self.__adaptive:
            if not self.__fine_base_frame_queue:
                self.__fine_base_frame_queue.append(frame)
                # Pre-compute meshgrid for remap operations
                x = np.linspace(0, frame_width - 1, frame_width, dtype='float32')
                y = np.linspace(0, frame_height - 1, frame_height, dtype='float32')
                self.map = np.array(np.meshgrid(x, y)).transpose((1, 2, 0))
                
                # Pre-allocate with correct dtype
                self.__fine_flow_queue.append(
                    np.zeros((frame_height, frame_width, 2), dtype=np.float32))
                self.__prev_fine_flow = np.zeros(
                    (frame_height, frame_width, 2), dtype=np.float32)
                return self.__fine_flow_queue[-1]
            else:
                self.__frame_count += 1
                
                # Only check loop closure periodically (expensive operation)
                if self.__frame_count % self.__loop_closure_interval == 0:
                    self.__checkLoopClosure(frame)
                
                # CPU optical flow computation
                self.__prev_fine_flow = self.__tracking_inst.calc(
                    self.__fine_base_frame_queue[-1], frame,
                    self.__prev_fine_flow)
                # Direct numpy addition is faster than cv2.add
                final_flow = self.__fine_flow_queue[-1] + self.__prev_fine_flow
                
                # Use fine flow for error estimation
                error = self.__getForwardBackwardError(
                    self.__fine_base_frame_queue[-1], frame,
                    self.__prev_fine_flow)
                
                if error > self.__upper_threshold:
                    self.__fine_base_frame_queue.append(frame)
                    self.__fine_flow_queue.append(final_flow)
                    self.__prev_fine_flow = np.zeros(
                        (frame_height, frame_width, 2), dtype=np.float32)
                return final_flow
        else:
            if not self.__fine_base_frame_queue:
                self.__fine_base_frame_queue.append(frame)
            
            self.__prev_fine_flow = self.__tracking_inst.calc(
                self.__fine_base_frame_queue[-1], frame,
                self.__prev_fine_flow)
            return self.__prev_fine_flow
        
    # def track_2dto3d(self, flow):
    #     decompose_obj = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18))
    #     flag,output = decompose_obj.get_force(flow, verbose=0)
    #     return output
        
    def __getForwardBackwardError(self, I0, I1, flow):
        # if self.cuda:
        #     flow_map = cv2.cuda.add(self.map, flow)
        #     map_x, map_y = cv2.cuda.split(flow_map)
        #     z = cv2.cuda.remap(I1,
        #                        map_x,
        #                        map_y,
        #                        interpolation=cv2.INTER_CUBIC,
        #                        borderMode=cv2.BORDER_REPLICATE)
        #     t = I0.clone()
        #     t = t.convertTo(rtype=cv2.CV_32FC1, dst=t)
        #     z = z.convertTo(rtype=cv2.CV_32FC1, dst=z)
        #     rev = cv2.cuda.absSum(cv2.cuda.subtract(
        #         t, z))[0] / t.size()[0] / t.size()[1]
        # else:
        # Optimized CPU path - avoid redundant type conversions
        flow_map = self.map + flow
        z = cv2.remap(I1,
                        flow_map[:, :, 0],
                        flow_map[:, :, 1],
                        interpolation=cv2.INTER_CUBIC,  # INTER_LINEAR faster than INTER_CUBIC
                        borderMode=cv2.BORDER_REPLICATE)
        # Use in-place operations when possible
        if I0.dtype != np.float32:
            diff = z.astype('float32') - I0.astype('float32')
        else:
            diff = z - I0
        rev = np.mean(np.abs(diff))
        return rev

    def __checkLoopClosure(self, frame):  # update (prune) the frame queue with loop closure detection
        for i in range(len(self.__fine_base_frame_queue) - 2, -1, -1):
            f = self.__tracking_inst.calc(
                self.__fine_base_frame_queue[i], frame, None)
            error = self.__getForwardBackwardError(
                self.__fine_base_frame_queue[i], frame, f)
            if error < self.__lower_threshold:
                self.__fine_base_frame_queue = self.__fine_base_frame_queue[:i+1]
                self.__fine_flow_queue = self.__fine_flow_queue[:i + 1]
                return True
        return False

    def reset(self):
        if self.__algorithm == "raft":
            self.__raft_ref_frame = None
            self.__raft_ref_tensor = None

        elif self.__algorithm == "dis":
            self.__fine_base_frame_queue = []
            self.__fine_flow_queue = []
            self.__prev_fine_flow = None
            self.__frame_count = 0

    def getFrameQueueLength(self):
        return len(self.__fine_base_frame_queue)
    
### RAFT ###

    def setup_simple_raft(self):
        """
        Setup a simpler version using available models with performance optimizations
        """
        import torchvision.transforms as T
        from torchvision.models.optical_flow import raft_large
        # Load model and move to GPU
        model = raft_large(pretrained=True, progress=False).to(self.__raft_device)
        model = model.eval()
        
        if self.verbose:
            print(f"RAFT model loaded on {self.__raft_device}")
        return model
    
    def _preprocess_frame_to_tensor(self, frame):
        """Efficiently convert numpy frame to RAFT-ready tensor"""
        # Ensure contiguous memory for faster tensor conversion
        if not frame.flags['C_CONTIGUOUS']:
            frame = np.ascontiguousarray(frame)
        
        # Upload to device as uint8 first to save bandwidth (4x less data than float32)
        tensor = self.torch.from_numpy(frame).permute(2, 0, 1)
        
        # Pin memory for faster CPU-GPU transfer (similar to DataLoader pin_memory=True)
        if self.__raft_device == 'cuda':
            # Pin memory before transferring to GPU (must be on CPU first)
            if not tensor.is_pinned():
                tensor = tensor.pin_memory()
            tensor = tensor.to(self.__raft_device, non_blocking=True)
        else:
            tensor = tensor.to(self.__raft_device)
        
        # Convert to float and normalize on device
        tensor = tensor.float() / 255.0

        # Normalize to [-1, 1]
        tensor = (tensor - 0.5) * 2.0
        
        # Add batch dimension
        tensor = tensor.unsqueeze(0)
        
        return tensor
    
    def track_raft(self, frame, gpu = False):
        # No need to copy if we just read
        # frame is expected to be (H, W, 3) or (H, W)
        
        if frame.ndim == 2:
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
        elif frame.shape[2] == 3:
             # Assume BGR input from OpenCV
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        if self.__raft_ref_frame is None:
            self.__raft_ref_frame = frame
            self.__raft_ref_tensor = self._preprocess_frame_to_tensor(frame)

        prev_tensor = self.__raft_ref_tensor
        curr_tensor = self._preprocess_frame_to_tensor(frame)
        if self.verbose:
            t0 = time.perf_counter()

        with self.torch.inference_mode():
            if self.__raft_device == 'cuda':
                with self.torch.amp.autocast('cuda'):
                    flow = self.__raft_model(prev_tensor, curr_tensor)
            else:
                flow = self.__raft_model(prev_tensor, curr_tensor)
            predicted_flows = flow[-1]

        if self.verbose:
            t1 = time.perf_counter()        
            print(f"RAFT time: {t1 - t0}")
        # Optimize tensor->numpy conversion: contiguous() ensures optimal memory layout
        flow_np = predicted_flows[0].contiguous().cpu().numpy()
        
        # Transpose to (H, W, 2) format expected by draw_arrows
        flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)
        return flow_np


    # def track_cuda(self, frame):
    #     """CUDA-accelerated optical flow tracking"""
    #     frame_height, frame_width = frame.shape[:2]
    #     frame = cv2.cuda_GpuMat(frame)
        
    #     if self.__adaptive:
    #         if not self.__fine_base_frame_queue:
    #             self.__fine_base_frame_queue.append(frame)
    #             # Pre-compute meshgrid for remap operations
    #             x = np.linspace(0, frame_width - 1, frame_width, dtype='float32')
    #             y = np.linspace(0, frame_height - 1, frame_height, dtype='float32')
    #             self.map = np.array(np.meshgrid(x, y)).transpose((1, 2, 0))
                
    #             cuMat = cv2.cuda_GpuMat()
    #             cuMat.upload(self.map)
    #             self.map = cuMat
    #             self.__fine_flow_queue.append(
    #                 cv2.cuda_GpuMat(frame.size(), cv2.CV_32FC2, (0.0, 0.0)))
    #             self.__prev_fine_flow = cv2.cuda_GpuMat(
    #                 frame.size(), cv2.CV_32FC2, (0.0, 0.0))
    #             return self.__fine_flow_queue[-1]
    #         else:
    #             self.__frame_count += 1
                
    #             # Only check loop closure periodically (expensive operation)
    #             if self.__frame_count % self.__loop_closure_interval == 0:
    #                 self.__checkLoopClosure(frame)
                
    #             # Use CUDA stream for async optical flow computation
    #             self.__prev_fine_flow = self.__tracking_inst.calc(
    #                 self.__fine_base_frame_queue[-1], frame,
    #                 self.__prev_fine_flow, self.__cuda_stream)
                
    #             final_flow = cv2.cuda.add(self.__fine_flow_queue[-1],
    #                                       self.__prev_fine_flow,
    #                                       stream=self.__cuda_stream)
                
    #             # Use fine flow for error estimation
    #             error = self.__getForwardBackwardError(
    #                 self.__fine_base_frame_queue[-1], frame,
    #                 self.__prev_fine_flow)
                
    #             if error > self.__upper_threshold:
    #                 self.__fine_base_frame_queue.append(frame)
    #                 self.__fine_flow_queue.append(final_flow)
    #                 self.__prev_fine_flow = None
    #             return final_flow
    #     else:
    #         if not self.__fine_base_frame_queue:
    #             self.__fine_base_frame_queue.append(frame)
            
    #         self.__prev_fine_flow = self.__tracking_inst.calc(
    #             self.__fine_base_frame_queue[-1], frame,
    #             self.__prev_fine_flow, self.__cuda_stream)
    #         return self.__prev_fine_flow