"""
Sensor - Core sensor class for optical tactile sensing.

This module provides the Sensor class for capturing, processing, and analyzing
optical tactile sensor data with support for stitching, optical flow tracking,
and Helmholtz-Hodge decomposition.
"""

import os
import numpy as np
import cv2
import time
from enum import Enum, auto
from typing import Tuple
from threading import Thread, Event
from queue import Queue, Empty

from .Stitcher import Stitcher
from .tracker import Tracker
from .nHHD import nHHD
# from .getDensity import getFilteredDensity
from . import configs 
from .get_depth import Contact3D


class Sensor:
    """
    Optical tactile sensor interface with deformation analysis.
    
    Provides image capture, stitching, optical flow tracking, and force decomposition.
    Supports both synchronous and asynchronous (pipelined) processing modes.
    """
    
    # =========================================================================
    # Info Enum (for read_info API)
    # =========================================================================
    
    class info(Enum):
        """Data types available through read_info()."""
        VRAW = auto()
        VNORMAL = auto()
        VSHEAR = auto()
        FPS = auto()
        TIMESTAMP = auto()
        FNORMAL = auto()
        FSHEARX = auto()
        FSHEARY = auto()
        CENTROID = auto()
        CONTOUR = auto()
        DEPTH = auto()

    # =========================================================================
    # Initialization
    # =========================================================================
    
    def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True, 
                 config_name="prototype", cal_path=None, override=False, box=None,mirror=True):
        """
        Initialize Sensor with video source and configuration.
        
        Args:
            vid_src: Video source (int for camera ID, str for file path)
            isstitch: Enable stitching mode for multi-tile sensor (default: False)
            verbose: Print initialization logs (default: True)
            cuda: Enable CUDA acceleration (default: True)
            config_name: Configuration name (default: "prototype")
            cal_path: Path to calibration file (default: None)
            override: Skip auto-detection, use provided box (default: False)
            box: Manual bounding boxes when override=True (default: None)
        """
        # Check CUDA availability
        try:
            import cupy as cp
            self.CUPY_AVAILABLE = True
        except ImportError:
            self.CUPY_AVAILABLE = False

        if vid_src is None:
            raise Exception("User must specify a video source.")

        # Store settings
        self.vid_src = vid_src
        self.isstitch = isstitch
        self.cuda = cuda
        self.verbose = verbose
        self.config_name = config_name
        self.cal_path = cal_path
        self.mirror = mirror

        # State variables
        self.stitch_valid = False
        self.imgupdate = False
        self.vupdate = False
        self.stitch_calibration = []
        
        # Data Variables
        self.img = None
        self.img_size_raw = None
        self.img_filtered = None
        self.flow = None
        self.vnormal = None
        self.vshear = None
        # self.depth = None
        self.n_boxes = None
        self.fnormal = self.fshearx = self.fsheary = self.ftor = 0.0

        # Contact Data Variables
        self.centroid = None
        self.centroid_found = False
        self.contour = None
        self.depth_map = None
        
        # Timing
        self.last_log_time = time.perf_counter()
        self.fps = 0
        
        # Async pipeline state
        self._async_enabled = False
        self._capture_thread = None
        self._flow_thread = None
        self._compute_thread = None
        self._capture_queue = None
        self._flow_queue = None
        self._result_queue = None
        self._stop_event = None

        # Load configuration
        self.config_cam = None
        self.config_stitcher = None
        self.config_contact = None
        self._load_config()
        

        # Initialize components
        self._stitcher = Stitcher(isstitch=isstitch, index=self.config_index, config_stitcher=self.config_stitcher)
        self.tracker = Tracker(adaptive=True, cuda=cuda)
        self.contact3d = Contact3D(config_contact=self.config_contact, verbose=self.verbose)
        self.hhd = None
        
        # Setup video capture
        self._open_video_source(override, box)

        if self.verbose:
            print(f"  Video source: {vid_src}")
            print(f"  Stitch mode: {self.stitch_valid}")
            print(f"  CUDA acceleration: {self.cuda}")
            print(f"  CuPy available: {self.CUPY_AVAILABLE}")
            print(f"  OpenCV CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}")

    # =========================================================================
    # Public API - Core Methods (Most Frequently Used)
    # =========================================================================
    
    def get_img(self, log_timing=False):
        """
        Capture and preprocess a frame.
        
        Returns:
            np.ndarray: Processed grayscale image
        """
        ret, frame = self.cap.read()
        if not ret:
            print("Failed to get_img from video source.")
            return None

        if log_timing:
            timings = {}
            t_start = time.perf_counter()
        
        # Convert to grayscale
        if log_timing: t0 = time.perf_counter()
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        if log_timing: timings['cvt_color'] = (time.perf_counter() - t0) * 1000

        # Extract and stitch tiles
        if log_timing: t0 = time.perf_counter()
        self._stitcher.get_stitch_unit(gray)
        self._extract_tiles(gray)
        if log_timing: timings['get_tiles'] = (time.perf_counter() - t0) * 1000

        if log_timing: t0 = time.perf_counter()
        img = np.ascontiguousarray(self._stitcher.stitch_rt_optimized())
        if log_timing: timings['stitch'] = (time.perf_counter() - t0) * 1000

        self.img_size_raw = img.shape
        # Resize
        img = cv2.resize(img, tuple(self.config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
        
        # Temporal filtering (noise reduction)
        alpha = 0.7
        if self.img_filtered is None:
            self.img_filtered = img.astype(np.float32)
        else:
            self.img_filtered = alpha * img + (1 - alpha) * self.img_filtered
        
        self.img = self.img_filtered.astype(np.uint8)
        self.imgupdate = True

        if log_timing:
            timings['total'] = (time.perf_counter() - t_start) * 1000
            self._print_timings("get_img", timings)
        if self.mirror:
            self.img = cv2.flip(self.img, 1)
            
        return self.img

    def compute_deformation(self, log_timing=False):
        """
        Compute optical flow and decompose into force components.
        
        Returns:
            np.ndarray: Optical flow field
        """
        if log_timing:
            timings = {}
            t_start = time.perf_counter()
        
        # Optical flow
        if log_timing: t0 = time.perf_counter()
        flow = self.tracker.track(self.img)
        if hasattr(flow, 'download'):
            flow = flow.download()
        if log_timing: timings['optical_flow'] = (time.perf_counter() - t0) * 1000
        
        # Initialize HHD solver
        if self.hhd is None:
            self.hhd = nHHD(grid=flow.shape[:2], spacings=(18, 18))

        # Force computation
        if log_timing: t0 = time.perf_counter()
        if self.cuda and self.CUPY_AVAILABLE:
            _, forces = self.hhd.get_force_gpu_v2(flow)
        else:
            _, forces = self.hhd.get_force(flow)
        self.fnormal, self.fshearx, self.fsheary, self.ftor = forces
        if log_timing: timings['hhd_decompose'] = (time.perf_counter() - t0) * 1000
        
        # Get components
        self.vnormal = self.hhd.get_normal()
        self.vshear = self.hhd.get_shear()
        
        # # Depth computation
        # if log_timing: t0 = time.perf_counter()
        # self.depth = getFilteredDensity(self.vnormal, use_cuda=self.cuda)
        # if log_timing: timings['get_density'] = (time.perf_counter() - t0) * 1000
        
        # Update state
        self.flow = flow
        self.imgupdate = False
        self.vupdate = True
        
        # FPS calculation
        self.fps = 1.0 / (time.perf_counter() - self.last_log_time)
        self.last_log_time = time.perf_counter()
        
        if log_timing:
            timings['total'] = (time.perf_counter() - t_start) * 1000
            self._print_timings("compute_deformation", timings)

        return self.flow

    def compute_contact(self, multiple_foe=False, min_contour_area=3000, min_magnitude=1.0, expansion_threshold=0.01, downsample_factor=1, kernel_size=(5,5), denoise_sigma=15):
        
        if multiple_foe:
            # result = compute_multiple_foe(self.vnormal, min_contour_area=min_contour_area, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma)
            result = self.contact3d.compute_multiple_foe(self.vnormal)
        else:
            if self.cuda: 
                # result = compute_foe_gpu(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma)
                result = self.contact3d.compute_foe_gpu(self.vnormal)
            else:
                # result = compute_foe(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma)
                result = self.contact3d.compute_foe(self.vnormal)

        self.centroid = result['centroid']
        self.centroid_found = result['centroid_found']
        if 'contour' in result:
            self.contour = result['contour']
        self.depth_map = result['divergence_map']

        return result
    

    # def get_depth(self, denoise_sigma=15):
    #     return getFilteredDensity(self.vnormal, use_cuda=self.cuda)
    #     self.depth = getFilteredDensity(self.vnormal, use_cuda=self.cuda)   

    def update(self):
        """
        Convenience method: get_img + compute_deformation in one call.
        
        Returns:
            dict: All sensor data
        """
        self.get_img()
        self.compute_deformation()
        return {
            'img': self.img, 'flow': self.flow,
            'normal': self.vnormal, 'shear': self.vshear, 'depth': self.depth,
            'fps': self.fps, 'fnormal': self.fnormal,
            'fshearx': self.fshearx, 'fsheary': self.fsheary, 'ftor': self.ftor,
        }

    def read_info(self, *args) -> Tuple:
        """
        Get sensor data by info type.
        
        Args:
            *args: info enum values (e.g., sensor.info.FPS, sensor.info.FNORMAL)
        
        Returns:
            Tuple: Requested data in order
        """
        if not self.vupdate:
            print("Warning: Call compute_deformation() first.")
        
        mapping = {
            self.info.VRAW: lambda: self.flow if self.flow is not None else np.array([]),
            self.info.VNORMAL: lambda: self.vnormal if self.vnormal is not None else np.array([]),
            self.info.VSHEAR: lambda: self.vshear if self.vshear is not None else np.array([]),
            self.info.FNORMAL: lambda: self.fnormal,
            self.info.FSHEARX: lambda: self.fshearx,
            self.info.FSHEARY: lambda: self.fsheary,
            self.info.FPS: lambda: self.fps,
            self.info.TIMESTAMP: lambda: time.strftime('%H:%M:%S'),
            self.info.DEPTH: lambda: self.depth_map if self.depth_map is not None else np.array([]),
            self.info.CENTROID: lambda: self.centroid if self.centroid is not None else np.array([]),
            self.info.CONTOUR: lambda: self.contour if self.contour is not None else np.array([]),
        }
        
        return tuple(mapping.get(arg, lambda: None)() for arg in args if isinstance(arg, self.info))

    def reset(self):
        """Reset optical flow tracker."""
        self.tracker.reset()

    def disconnect(self):
        """Release resources and stop async threads."""
        if self._async_enabled:
            self.disable_async()
        if self.cap.isOpened():
            self.cap.release()
            if self.verbose:
                print("Sensor disconnected.")
        cv2.destroyAllWindows()

    # =========================================================================
    # Public API - Async Pipeline
    # =========================================================================
    
    def enable_async(self, queue_size=2):
        """
        Enable 3-stage async pipeline: Capture → Flow → HHD.
        
        Each stage runs in parallel for maximum throughput.
        """
        if self._async_enabled:
            return
        
        self._capture_queue = Queue(maxsize=queue_size)
        self._flow_queue = Queue(maxsize=queue_size)
        self._result_queue = Queue(maxsize=queue_size)
        self._contact_queue = Queue(maxsize=queue_size)
        self._stop_event = Event()
        self._async_enabled = True
        
        self._capture_thread = Thread(target=self._capture_worker, daemon=True)
        self._flow_thread = Thread(target=self._flow_worker, daemon=True)
        self._compute_thread = Thread(target=self._hhd_worker, daemon=True)
        self._contact_thread = Thread(target=self._contact_worker, daemon=True)

        self._capture_thread.start()
        self._flow_thread.start()
        self._compute_thread.start()
        
        if self.verbose:
            print("Async pipeline enabled: Capture → Flow → HHD")

    def get_async(self, timeout=0.05):
        """
        Get latest result from async pipeline (non-blocking).
        
        Returns:
            dict or None: Sensor data if available
        """
        if not self._async_enabled:
            return None
        try:
            result = self._result_queue.get(timeout=timeout)
            t = time.perf_counter()
            result['fps'] = 1.0 / (t - self.last_log_time) if self.last_log_time else 0
            self.last_log_time = t
            
            # Sync instance variables
            # for k in ['img', 'flow', 'normal', 'shear', 'depth', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
            for k in ['img', 'flow', 'normal', 'shear', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
                setattr(self, k, result.get(k))
            self.vupdate = True
            return result
        except Empty:
            return None

    def disable_async(self):
        """Stop async pipeline."""
        if not self._async_enabled:
            return
        self._stop_event.set()
        for t in [self._capture_thread, self._flow_thread, self._compute_thread]:
            if t: t.join(timeout=2.0)
        self._async_enabled = False
        self._capture_thread = self._flow_thread = self._compute_thread = None
        self._capture_queue = self._flow_queue = self._result_queue = self._stop_event = None
        if self.verbose:
            print("Async pipeline disabled")

    # =========================================================================
    # Public API - Configuration & Calibration
    # =========================================================================
    
    def export_config(self, path="./config/exported_config.json"):
        """Return current configuration dict."""
        configs.export_config(self.config, path)
        if self.verbose:
            print(f"Config exported to {path}")
        return True

    def stitch_cal(self):
        """
        Load or create stitch calibration.
        
        Returns:
            np.ndarray: Stitched calibration image
        """
        if self.cal_path is None:
            self.cal_path = "./config/cal.npy"
        
        if os.path.exists(self.cal_path):
            print("Loading calibration from", self.cal_path)
            self.stitch_calibration = np.load(self.cal_path, allow_pickle=True)
            self._stitcher.stitch_calibration = self.stitch_calibration
            return self._stitcher.stitch_rt_optimized()
        else:
            print("Creating new calibration...")
            img, self.stitch_calibration = self._stitcher.stitch_cal()
            os.makedirs("./config", exist_ok=True)
            np.save(self.cal_path, self.stitch_calibration)
            print("Calibration saved to", self.cal_path)
            return img

    # =========================================================================
    # Private - Initialization Helpers
    # =========================================================================
    
    def _load_config(self):
        """Load sensor configuration from JSON file."""
        self.config = configs.load_config(self.config_name)
        cam = self.config["camera"]
        self.config_resolution_raw = cam["resolution_raw"]
        self.config_fps = cam["fps"]
        self.config_units = cam["units"]
        self.config_resolution_stitch = cam["resolution_stitch"]
        self.config_index = cam["index"]
        self.config_pos_default = cam["pos_default"]

        self.config_stitcher = self.config["stitcher"]
        self.config_contact = self.config["contact"]

        if self.verbose:
            print(f"Config: {self.config_name}")
            print(f"  Resolution: {self.config_resolution_raw} → {self.config_resolution_stitch}")
            print(f"  FPS: {self.config_fps}, Units: {self.config_units}")

    def _open_video_source(self, override, box):
        """Initialize video capture and calibration."""
        self.cap = cv2.VideoCapture(self.vid_src)
        self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M', 'J', 'P', 'G'))
        self.cap.set(cv2.CAP_PROP_FPS, self.config_fps)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.config_resolution_raw[0])
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.config_resolution_raw[1])
        
        if self.verbose:
            print("Connecting to video source...")
        
        if not self.cap.isOpened():
            raise Exception("Failed to open video source.")
        
        # Wait for first valid frame
        frame = None
        while frame is None:
            ret, frame = self.cap.read()
            if not ret:
                time.sleep(0.1)

        if override:
            if box is not None:
                self._stitcher.isstitch = self.isstitch = False
                self._stitcher.boxes = box
            else:
                print("Override=True but no box provided.")
        else:
            self._find_stitch_boxes(frame)
            self._extract_tiles(frame)
            self.stitch_cal()

    def _find_stitch_boxes(self, img_src, show=False):
        """Detect vision tile boxes with camera stabilization."""
        n_boxes, v_mean0, timeout = 0, 0, 5
        start_time = cv2.getTickCount()
        
        while n_boxes < self.config_units:
            ret, frame = self.cap.read()
            v_mean = np.sum(cv2.mean(frame))
            
            if abs(v_mean - v_mean0) < 5:
                n_boxes = self._stitcher.find_stitch_boxes(frame, show=show)
            
            if (cv2.getTickCount() - start_time) / cv2.getTickFrequency() > timeout:
                print("Timeout finding vision units.")
                self._stitcher.isstitch = self.isstitch = False
                break
            v_mean0 = v_mean
        
        self.n_boxes = n_boxes
        self.stitch_valid = self.isstitch and n_boxes == self.config_units
        
        status = "valid for stitching" if self.stitch_valid else "non-stitch mode"
        print(f"Sensor: {status}")
        return n_boxes

    def _extract_tiles(self, img_src):
        """Extract tiles from source image."""
        self._stitcher.get_stitch_unit(img_src)
        self._stitcher.extract_tiles()
        return True

    # =========================================================================
    # Private - Async Workers
    # =========================================================================
    
    def _queue_put(self, q, item):
        """Put item in queue, dropping old if full."""
        if q.full():
            try: q.get_nowait()
            except Empty: pass
        q.put(item)

    def _capture_worker(self):
        """Stage 1: Capture and preprocess frames."""
        alpha = 0.7
        while not self._stop_event.is_set():
            try:
                ret, frame = self.cap.read()
                if not ret:
                    continue
                
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                self._stitcher.get_stitch_unit(gray)
                self._extract_tiles(gray)
                img = np.ascontiguousarray(self._stitcher.stitch_rt_optimized())
                img = cv2.resize(img, tuple(self.config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
                
                if self.img_filtered is None:
                    self.img_filtered = img.astype(np.float32)
                else:
                    self.img_filtered = alpha * img + (1 - alpha) * self.img_filtered
                
                self._queue_put(self._capture_queue, self.img_filtered.astype(np.uint8))
            except Exception as e:
                if self.verbose: print(f"Capture error: {e}")

    def _flow_worker(self):
        """Stage 2: Compute optical flow."""
        while not self._stop_event.is_set():
            try:
                img = self._capture_queue.get(timeout=0.1)
                flow = self.tracker.track(img)
                if hasattr(flow, 'download'):
                    flow = flow.download()
                self._queue_put(self._flow_queue, (img, flow))
            except Empty:
                continue
            except Exception as e:
                if self.verbose: print(f"Flow error: {e}")

    def _hhd_worker(self):
        """Stage 3: HHD decomposition."""
        while not self._stop_event.is_set():
            try:
                img, flow = self._flow_queue.get(timeout=0.1)
                
                if self.hhd is None:
                    self.hhd = nHHD(grid=flow.shape[:2], spacings=(18, 18))
                
                if self.cuda and self.CUPY_AVAILABLE:
                    _, forces = self.hhd.get_force_gpu_v2(flow)
                else:
                    _, forces = self.hhd.get_force(flow)
                
                self._queue_put(self._result_queue, {
                    'img': img, 'flow': flow,
                    'normal': self.hhd.get_normal(), 'shear': self.hhd.get_shear(),
                    # 'depth': getFilteredDensity(flow, use_cuda=self.cuda),
                    'fnormal': forces[0], 'fshearx': forces[1],
                    'fsheary': forces[2], 'ftor': forces[3],
                })
            except Empty:
                continue
            except Exception as e:
                if self.verbose: print(f"HHD error: {e}")

    def _contact_worker(self):
        """Stage 4: Contact computation."""
        while not self._stop_event.is_set():
            try:
                img, flow = self._flow_queue.get(timeout=0.1)
                contact_result = self.compute_contact(flow)
                self._queue_put(self._result_queue, contact_result)
            except Empty:
                continue
            except Exception as e:
                if self.verbose: print(f"Contact error: {e}")

    # =========================================================================
    # Private - Utilities
    # =========================================================================
    
    def _print_timings(self, name, timings):
        """Print timing breakdown."""
        print(f"\n[{name} timing]")
        for k, v in timings.items():
            if k != 'total':
                print(f"  {k}: {v:7.2f} ms")
        print(f"  ─────────────────────────")
        print(f"  TOTAL: {timings['total']:7.2f} ms ({1000/timings['total']:.1f} FPS)")