from ..base.control_base import ControlNodeBase
from enum import Enum
import numpy as np
import torch
import cv2
import time

class ROIAction(Enum):
    # Behavioral actions
    TOGGLE = "toggle"      # Toggles between min/max values
    MOMENTARY = "momentary"  # Outputs max while motion detected
    TRIGGER = "trigger"    # Triggers once per motion event
    COUNTER = "counter"    # Counts motion events
    # Mathematical actions
    ADD = "add"          # Add value when motion detected
    SUBTRACT = "subtract" # Subtract value when motion detected
    MULTIPLY = "multiply" # Multiply by value when motion detected
    DIVIDE = "divide"     # Divide by value when motion detected
    SET = "set"          # Set to specific value when motion detected

class ROINode(ControlNodeBase):
    """Defines a single region of interest and its action"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mask": ("MASK",),  # ROI defined by mask
                "action": (list(action.value for action in ROIAction),),
                "value": ("FLOAT", {
                    "default": 0.1,
                    "step": 0.01,
                    "tooltip": "Value to use for mathematical operations"
                }),
            },
            "optional": {
                "next_roi": ("ROI",)  # Chain to next ROI
            }
        }

    RETURN_TYPES = ("ROI",)
    FUNCTION = "define_roi"
    CATEGORY = "real-time/control/motion"

    def define_roi(self, mask, action, value, next_roi=None):
        # Convert mask to numpy for bounding box calculation
        mask_np = mask[0].cpu().numpy()
        
        # Find non-zero coordinates in mask
        coords = np.nonzero(mask_np)
        if len(coords[0]) == 0:  # Empty mask
            y_min = x_min = y_max = x_max = 0
        else:
            y_min, y_max = coords[0].min(), coords[0].max()
            x_min, x_max = coords[1].min(), coords[1].max()
        
        roi_data = {
            "mask": mask_np,
            "bounds": (y_min, x_min, y_max, x_max),
            "action": action,
            "value": value,  # Store the operation value
            "next": next_roi
        }
        return (roi_data,)

    def update(self, *args, **kwargs):
        return self.define_roi(*args, **kwargs)

class MotionController(ControlNodeBase):
    """Processes motion detection for a chain of ROIs and manages their states"""
    
    @classmethod
    def INPUT_TYPES(cls):
        inputs = super().INPUT_TYPES()
        inputs["required"].update({
            "image": ("IMAGE",),
            "roi_chain": ("ROI",),
            "threshold": ("FLOAT", {
                "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01,
                "tooltip": "Motion detection threshold"
            }),
            "blur_size": ("INT", {
                "default": 5, "min": 1, "max": 21, "step": 2,
                "tooltip": "Size of Gaussian blur kernel"
            }),
            "minimum_value": ("FLOAT", {
                "default": 0.0,
                "tooltip": "Minimum output value"
            }),
            "maximum_value": ("FLOAT", {
                "default": 1.0,
                "tooltip": "Maximum output value"
            }),
            "starting_value": ("FLOAT", {
                "default": 0.0,
                "tooltip": "Initial output value"
            })
        })
        return inputs

    RETURN_TYPES = ("FLOAT", "MASK")
    FUNCTION = "process_motion"
    CATEGORY = "real-time/control/motion"

    def process_motion(self, image, roi_chain, threshold, blur_size, 
                      minimum_value, maximum_value, starting_value, always_execute=True):
        # Get or initialize state with new caching fields
        state = self.get_state({
            "prev_frame": None,
            "prev_frame_blurred": None,  # Cache blurred previous frame
            "roi_states": {},
            "current_value": starting_value,
            "last_cleanup": time.time()  # Track last state cleanup
        })
        
        # Convert image tensor to numpy array and prepare current frame
        current_frame = (image[0] * 255).cpu().numpy().astype(np.uint8)
        current_gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY)
        current_blurred = cv2.GaussianBlur(current_gray, (blur_size, blur_size), 0)
        
        # Initialize motion mask
        motion_mask = np.zeros_like(current_gray, dtype=np.float32)
        
        if state["prev_frame"] is not None:
            # Collect active ROIs for state cleanup
            active_rois = set()
            
            # Process each ROI in the chain
            current_roi = roi_chain
            while current_roi is not None:
                mask = current_roi["mask"]
                bounds = current_roi["bounds"]
                action = current_roi["action"]
                roi_id = str(bounds)
                active_rois.add(roi_id)
                
                # Extract ROI regions once
                y_min, x_min, y_max, x_max = bounds
                roi_region = slice(y_min, y_max+1), slice(x_min, x_max+1)
                
                # Process ROI using cached blurred frames
                roi_current = current_blurred[roi_region]
                roi_prev = state["prev_frame_blurred"][roi_region]
                roi_mask = mask[roi_region]
                
                # Detect motion in ROI
                diff = cv2.absdiff(roi_current, roi_prev)
                _, thresh = cv2.threshold(diff, threshold * 255, 255, cv2.THRESH_BINARY)
                thresh = thresh * (roi_mask > 0.5)
                
                # Update motion mask
                motion_mask[roi_region] = np.maximum(
                    motion_mask[roi_region],
                    thresh / 255.0
                )
                
                # Calculate motion magnitude
                roi_area = np.sum(roi_mask > 0.5)
                motion_detected = np.sum(thresh) / (255 * roi_area) > threshold if roi_area > 0 else False
                
                # Get or initialize ROI state
                roi_state = state["roi_states"].get(roi_id, {"active": False, "count": 0})
                value = current_roi["value"]
                
                # Process actions with optimized value updates
                if motion_detected and not roi_state["active"]:
                    roi_state["active"] = True
                    
                    # Update current value based on action type
                    if action == ROIAction.ADD.value:
                        state["current_value"] = min(maximum_value, state["current_value"] + value)
                    elif action == ROIAction.SUBTRACT.value:
                        state["current_value"] = max(minimum_value, state["current_value"] - value)
                    elif action == ROIAction.MULTIPLY.value:
                        state["current_value"] = min(maximum_value, state["current_value"] * value)
                    elif action == ROIAction.DIVIDE.value:
                        if value != 0:
                            state["current_value"] = max(minimum_value, state["current_value"] / value)
                    elif action == ROIAction.SET.value:
                        state["current_value"] = max(minimum_value, min(maximum_value, value))
                    elif action == ROIAction.TOGGLE.value:
                        state["current_value"] = maximum_value if state["current_value"] == minimum_value else minimum_value
                    elif action == ROIAction.TRIGGER.value:
                        state["current_value"] = maximum_value
                    elif action == ROIAction.COUNTER.value:
                        roi_state["count"] += 1
                        state["current_value"] = minimum_value + (
                            roi_state["count"] % (int((maximum_value - minimum_value) + 1))
                        )
                
                elif action == ROIAction.MOMENTARY.value:
                    state["current_value"] = maximum_value if motion_detected else minimum_value
                
                elif not motion_detected:
                    roi_state["active"] = False
                    if action == ROIAction.TRIGGER.value:
                        state["current_value"] = minimum_value
                
                state["roi_states"][roi_id] = roi_state
                current_roi = current_roi["next"]
            
            # Cleanup inactive ROI states periodically (every 60 seconds)
            if time.time() - state["last_cleanup"] > 60:
                state["roi_states"] = {k: v for k, v in state["roi_states"].items() 
                                     if k in active_rois}
                state["last_cleanup"] = time.time()
        
        # Update state with current frame data
        state["prev_frame"] = current_gray
        state["prev_frame_blurred"] = current_blurred
        self.set_state(state)
        
        # Convert mask to tensor
        mask_tensor = torch.from_numpy(motion_mask).unsqueeze(0)
        
        return (state["current_value"], mask_tensor) 

    def update(self, *args, **kwargs):
        return self.process_motion(*args, **kwargs) 

class RoundingMode(Enum):
    ROUND = "round"    # Standard rounding
    FLOOR = "floor"    # Round down
    CEIL = "ceil"      # Round up
    TRUNCATE = "truncate"  # Remove decimal part

class IntegerMotionController(MotionController):
    """Processes motion detection for ROIs and outputs integer values"""
    
    @classmethod
    def INPUT_TYPES(cls):
        inputs = super().INPUT_TYPES()
        # Convert float parameters to int
        inputs["required"]["minimum_value"] = ("INT", {
            "default": 0,
            "tooltip": "Minimum output value"
        })
        inputs["required"]["maximum_value"] = ("INT", {
            "default": 100,
            "tooltip": "Maximum output value"
        })
        inputs["required"]["starting_value"] = ("INT", {
            "default": 0,
            "tooltip": "Initial output value"
        })
        # Add rounding mode parameter
        inputs["required"]["rounding_mode"] = (
            [mode.value for mode in RoundingMode],
            {"default": RoundingMode.ROUND.value}
        )
        return inputs

    RETURN_TYPES = ("INT", "MASK")
    FUNCTION = "process_motion"
    CATEGORY = "real-time/control/motion"

    def _apply_rounding(self, value, mode):
        """Apply the specified rounding mode to convert float to int"""
        if mode == RoundingMode.ROUND.value:
            return int(round(value))
        elif mode == RoundingMode.FLOOR.value:
            return int(np.floor(value))
        elif mode == RoundingMode.CEIL.value:
            return int(np.ceil(value))
        else:  # TRUNCATE
            return int(value)

    def process_motion(self, image, roi_chain, threshold, blur_size,
                      minimum_value, maximum_value, starting_value, 
                      rounding_mode, always_execute=True):
        # Call parent class method with float conversion of bounds
        float_result, mask = super().process_motion(
            image, roi_chain, threshold, blur_size,
            float(minimum_value), float(maximum_value), float(starting_value),
            always_execute
        )
        
        # Convert result to integer using specified rounding mode
        int_result = self._apply_rounding(float_result, rounding_mode)
        
        # Ensure result stays within bounds
        int_result = max(minimum_value, min(maximum_value, int_result))
        
        return (int_result, mask) 