# -*- coding: utf-8 -*-
"""
OpenAlgo Technical Indicators - Statistical Indicators
"""

import numpy as np
import pandas as pd
from numba import jit
from typing import Union, Tuple, Optional
from .base import BaseIndicator


class LINREG(BaseIndicator):
    """
    Linear Regression
    
    Linear Regression calculates the linear regression line for the given period.
    
    Formula: y = mx + b (least squares method)
    """
    
    def __init__(self):
        super().__init__("Linear Regression")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_linearreg(data: np.ndarray, period: int) -> np.ndarray:
        """Numba optimized Linear Regression calculation"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            # Extract window
            y = data[i - period + 1:i + 1]
            x = np.arange(period)
            
            # Calculate linear regression
            sum_x = np.sum(x)
            sum_y = np.sum(y)
            sum_xy = np.sum(x * y)
            sum_x2 = np.sum(x * x)
            
            # Calculate slope and intercept
            denominator = period * sum_x2 - sum_x * sum_x
            if denominator != 0:
                slope = (period * sum_xy - sum_x * sum_y) / denominator
                intercept = (sum_y - slope * sum_x) / period
                
                # Calculate value at the end of the period
                result[i] = slope * (period - 1) + intercept
            else:
                result[i] = y[-1]  # Fallback to last value
        
        return result
    
    def calculate(self, data: Union[np.ndarray, pd.Series, list], period: int = 14) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Linear Regression
        
        Parameters:
        -----------
        data : Union[np.ndarray, pd.Series, list]
            Price data (typically closing prices)
        period : int, default=14
            Period for linear regression calculation
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Linear Regression values in the same format as input
        """
        validated_data, input_type, index = self.validate_input(data)
        self.validate_period(period, len(validated_data))
        
        result = self._calculate_linearreg(validated_data, period)
        return self.format_output(result, input_type, index)


class LRSLOPE(BaseIndicator):
    """
    Linear Regression Slope - matches TradingView exactly
    
    TradingView calculates slope as the difference between consecutive 
    linear regression values divided by interval:
    linear_reg = linreg(close_price, len, 0)
    linear_reg_prev = linreg(close[1], len, 0)  
    slope = ((linear_reg - linear_reg_prev) / interval)
    """
    
    def __init__(self):
        super().__init__("Linear Regression Slope")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_linreg_value(data: np.ndarray, period: int, offset: int) -> float:
        """Calculate linear regression value at given offset"""
        y = data
        x = np.arange(len(y))
        
        sum_x = np.sum(x)
        sum_y = np.sum(y)
        sum_xy = np.sum(x * y)
        sum_x2 = np.sum(x * x)
        
        denominator = period * sum_x2 - sum_x * sum_x
        if denominator != 0:
            slope = (period * sum_xy - sum_x * sum_y) / denominator
            intercept = (sum_y - slope * sum_x) / period
            
            # Calculate value at offset position (TradingView uses offset 0 for current)
            return slope * (period - 1 - offset) + intercept
        else:
            return y[-1 - offset] if offset < len(y) else y[-1]
    
    def calculate(self, data: Union[np.ndarray, pd.Series, list], 
                 period: int = 100, interval: int = 1) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Linear Regression Slope - matches TradingView exactly
        
        Parameters:
        -----------
        data : Union[np.ndarray, pd.Series, list]
            Price data (typically closing prices)
        period : int, default=100
            Period for linear regression calculation (TradingView default)
        interval : int, default=1
            Interval divisor (TradingView uses timeframe interval)
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Slope values in the same format as input
        """
        validated_data, input_type, index = self.validate_input(data)
        self.validate_period(period + 1, len(validated_data))  # Need period+1 for TradingView method
        
        if interval <= 0:
            raise ValueError(f"Interval must be positive, got {interval}")
        
        result = _calculate_slope_tv(validated_data, period, interval)
        return self.format_output(result, input_type, index)


@jit(nopython=True)
def _calculate_linreg_value_standalone(data: np.ndarray, period: int, offset: int) -> float:
    """Calculate linear regression value at given offset (standalone function)"""
    y = data
    x = np.arange(len(y))
    
    sum_x = np.sum(x)
    sum_y = np.sum(y)
    sum_xy = np.sum(x * y)
    sum_x2 = np.sum(x * x)
    
    denominator = period * sum_x2 - sum_x * sum_x
    if denominator != 0:
        slope = (period * sum_xy - sum_x * sum_y) / denominator
        intercept = (sum_y - slope * sum_x) / period
        
        # Calculate value at offset position (TradingView uses offset 0 for current)
        return slope * (period - 1 - offset) + intercept
    else:
        return y[-1 - offset] if offset < len(y) else y[-1]


@jit(nopython=True) 
def _calculate_slope_tv(data: np.ndarray, period: int, interval: int = 1) -> np.ndarray:
    """
    Calculate slope using TradingView method:
    slope = ((linreg(close, len, 0) - linreg(close[1], len, 0)) / interval)
    """
    n = len(data)
    result = np.full(n, np.nan)
    
    for i in range(period, n):  # Start from period (not period-1) to have previous value
        # Current window
        current_window = data[i - period + 1:i + 1]
        # Previous window (shifted by 1)
        prev_window = data[i - period:i] 
        
        # Calculate linear regression values
        linear_reg = _calculate_linreg_value_standalone(current_window, period, 0)
        linear_reg_prev = _calculate_linreg_value_standalone(prev_window, period, 0)
        
        # Calculate slope as TradingView does
        result[i] = (linear_reg - linear_reg_prev) / interval
    
    return result


class CORREL(BaseIndicator):
    """
    Pearson Correlation Coefficient
    
    Measures the correlation between two data series.
    
    Formula: r = Σ((x - x̄)(y - ȳ)) / √(Σ(x - x̄)² × Σ(y - ȳ)²)
    """
    
    def __init__(self):
        super().__init__("Correlation")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_correl(data1: np.ndarray, data2: np.ndarray, period: int) -> np.ndarray:
        """Numba optimized correlation calculation"""
        n = len(data1)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            x = data1[i - period + 1:i + 1]
            y = data2[i - period + 1:i + 1]
            
            # Calculate means
            mean_x = np.mean(x)
            mean_y = np.mean(y)
            
            # Calculate correlation coefficient
            numerator = np.sum((x - mean_x) * (y - mean_y))
            sum_sq_x = np.sum((x - mean_x) ** 2)
            sum_sq_y = np.sum((y - mean_y) ** 2)
            
            denominator = np.sqrt(sum_sq_x * sum_sq_y)
            
            if denominator > 0:
                result[i] = numerator / denominator
            else:
                result[i] = 0
        
        return result
    
    def calculate(self, data1: Union[np.ndarray, pd.Series, list],
                 data2: Union[np.ndarray, pd.Series, list],
                 period: int = 20) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Pearson Correlation Coefficient
        
        Parameters:
        -----------
        data1 : Union[np.ndarray, pd.Series, list]
            First data series
        data2 : Union[np.ndarray, pd.Series, list]
            Second data series
        period : int, default=20
            Period for correlation calculation
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Correlation values in the same format as input
        """
        data1_validated, input_type, index = self.validate_input(data1)
        data2_validated, _, _ = self.validate_input(data2)
        
        data1_validated, data2_validated = self.align_arrays(data1_validated, data2_validated)
        self.validate_period(period, len(data1_validated))
        
        result = self._calculate_correl(data1_validated, data2_validated, period)
        return self.format_output(result, input_type, index)


class BETA(BaseIndicator):
    """
    Beta Coefficient
    
    Measures the volatility of a security relative to the market.
    
    Formula: β = Cov(asset, market) / Var(market)
    """
    
    def __init__(self):
        super().__init__("Beta")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_beta(asset: np.ndarray, market: np.ndarray, period: int) -> np.ndarray:
        """Numba optimized Beta calculation"""
        n = len(asset)
        result = np.full(n, np.nan)
        
        for i in range(period, n):
            # Calculate returns
            asset_returns = np.diff(asset[i - period:i + 1])
            market_returns = np.diff(market[i - period:i + 1])
            
            # Calculate means
            mean_asset = np.mean(asset_returns)
            mean_market = np.mean(market_returns)
            
            # Calculate covariance and variance
            covariance = np.mean((asset_returns - mean_asset) * (market_returns - mean_market))
            market_variance = np.mean((market_returns - mean_market) ** 2)
            
            if market_variance > 0:
                result[i] = covariance / market_variance
            else:
                result[i] = 0
        
        return result
    
    def calculate(self, asset: Union[np.ndarray, pd.Series, list],
                 market: Union[np.ndarray, pd.Series, list],
                 period: int = 252) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Beta Coefficient
        
        Parameters:
        -----------
        asset : Union[np.ndarray, pd.Series, list]
            Asset price data
        market : Union[np.ndarray, pd.Series, list]
            Market price data
        period : int, default=252
            Period for beta calculation (typically 1 year = 252 trading days)
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Beta values in the same format as input
        """
        asset_data, input_type, index = self.validate_input(asset)
        market_data, _, _ = self.validate_input(market)
        
        asset_data, market_data = self.align_arrays(asset_data, market_data)
        self.validate_period(period + 1, len(asset_data))  # +1 for diff
        
        result = self._calculate_beta(asset_data, market_data, period)
        return self.format_output(result, input_type, index)


class VAR(BaseIndicator):
    """
    Variance - TradingView Pine Script v4 Implementation
    
    Calculates variance with support for both logarithmic returns and price modes,
    including EMA smoothing and z-score analysis with signal generation.
    
    TradingView Pine Script v4 Implementation by moot-al-cabal:
    
    Variance modes:
    - "LR" (Logarithmic Returns): variance of log(close/close[1])*100
    - "PR" (Price): variance of price values
    
    Formula:
    1. source = if mode == "LR" then log(close/close[1])*100 else close
    2. mean = sma(source, lookback)
    3. For each i in lookback: array.push(s, pow(source[i]-mean, 2))
    4. variance = array.sum(s) / (lookback-1)
    5. stdev = sqrt(variance)
    6. ema_variance = ema(variance, ema_period)
    7. zscore = (variance - sma(variance, filter_lookback)) / stdev(variance, filter_lookback)
    8. ema_zscore = ema(zscore, ema_length)
    """
    
    def __init__(self):
        super().__init__("Variance")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_sma_safe(data: np.ndarray, period: int) -> np.ndarray:
        """Safe SMA calculation without Numba issues"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            window = data[i - period + 1:i + 1]
            valid_count = 0
            sum_val = 0.0
            for j in range(len(window)):
                if not np.isnan(window[j]):
                    sum_val += window[j]
                    valid_count += 1
            
            if valid_count == period:
                result[i] = sum_val / period
        
        return result
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_ema_safe(data: np.ndarray, period: int) -> np.ndarray:
        """Safe EMA calculation without Numba issues"""
        n = len(data)
        result = np.full(n, np.nan)
        alpha = 2.0 / (period + 1)
        
        # Find first valid value
        first_valid_idx = -1
        for i in range(n):
            if not np.isnan(data[i]):
                first_valid_idx = i
                break
        
        if first_valid_idx == -1:
            return result
            
        # Initialize with first valid value
        result[first_valid_idx] = data[first_valid_idx]
        
        # Continue EMA calculation
        for i in range(first_valid_idx + 1, n):
            if not np.isnan(data[i]):
                result[i] = alpha * data[i] + (1 - alpha) * result[i - 1]
            else:
                result[i] = result[i - 1]
        
        return result
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_stdev_safe(data: np.ndarray, period: int) -> np.ndarray:
        """Safe standard deviation calculation"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            window = data[i - period + 1:i + 1]
            valid_count = 0
            sum_val = 0.0
            for j in range(len(window)):
                if not np.isnan(window[j]):
                    sum_val += window[j]
                    valid_count += 1
            
            if valid_count == period:
                mean_val = sum_val / period
                sum_sq_diff = 0.0
                for j in range(len(window)):
                    if not np.isnan(window[j]):
                        diff = window[j] - mean_val
                        sum_sq_diff += diff * diff
                
                result[i] = np.sqrt(sum_sq_diff / period)
        
        return result
    
    @staticmethod
    def _calculate_variance_tv(data: np.ndarray, lookback: int, use_log_returns: bool) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """TradingView Variance calculation with all components"""
        n = len(data)
        
        # Step 1: Calculate source based on mode
        source = np.full(n, np.nan)
        
        if use_log_returns:  # Logarithmic Returns
            for i in range(1, n):
                if data[i] > 0 and data[i-1] > 0:
                    source[i] = np.log(data[i] / data[i-1]) * 100
        else:  # Price mode
            for i in range(n):
                source[i] = data[i]
        
        # Step 2: Calculate variance using TradingView method
        variance = np.full(n, np.nan)
        stdev = np.full(n, np.nan)
        
        for i in range(lookback - 1, n):
            # Calculate mean for the window
            window = source[i - lookback + 1:i + 1]
            valid_count = 0
            sum_val = 0.0
            for j in range(len(window)):
                if not np.isnan(window[j]):
                    sum_val += window[j]
                    valid_count += 1
            
            if valid_count == lookback:
                mean_val = sum_val / lookback
                
                # TradingView method: sum of squared differences divided by (lookback-1)
                sum_sq_diff = 0.0
                for j in range(len(window)):
                    if not np.isnan(window[j]):
                        diff = window[j] - mean_val
                        sum_sq_diff += diff * diff
                
                variance[i] = sum_sq_diff / (lookback - 1)  # TradingView uses (n-1)
                stdev[i] = np.sqrt(variance[i])
        
        return source, variance, stdev
    
    def calculate(self, data: Union[np.ndarray, pd.Series, list], 
                 lookback: int = 20, mode: str = "PR", ema_period: int = 20,
                 filter_lookback: int = 20, ema_length: int = 14,
                 return_components: bool = False) -> Union[np.ndarray, pd.Series, Tuple]:
        """
        Calculate Variance - TradingView Pine Script v4 Implementation
        
        Parameters:
        -----------
        data : Union[np.ndarray, pd.Series, list]
            Price data (close prices)
        lookback : int, default=20
            Variance lookback period
        mode : str, default="PR"
            Variance mode: "LR" for Logarithmic Returns, "PR" for Price
        ema_period : int, default=20
            EMA period for variance smoothing
        filter_lookback : int, default=20
            Lookback period for variance filter (z-score calculation)
        ema_length : int, default=14
            EMA length for z-score smoothing
        return_components : bool, default=False
            If True, returns (variance, ema_variance, zscore, ema_zscore, stdev)
            
        Returns:
        --------
        Union[np.ndarray, pd.Series, Tuple]
            Variance values or tuple of all components
        """
        validated_data, input_type, index = self.validate_input(data)
        self.validate_period(lookback, len(validated_data))
        
        # Validate mode
        if mode not in ["LR", "PR"]:
            raise ValueError("Mode must be 'LR' (Logarithmic Returns) or 'PR' (Price)")
        
        # Calculate base variance components
        use_log_returns = (mode == "LR")
        source, variance, stdev = self._calculate_variance_tv(validated_data, lookback, use_log_returns)
        
        # Calculate EMA of variance
        ema_variance = self._calculate_ema_safe(variance, ema_period)
        
        # Calculate z-score components
        variance_sma = self._calculate_sma_safe(variance, filter_lookback)
        variance_stdev = self._calculate_stdev_safe(variance, filter_lookback)
        
        # Calculate z-score: (variance - sma(variance)) / stdev(variance)
        zscore = np.full(len(variance), np.nan)
        for i in range(len(variance)):
            if (not np.isnan(variance[i]) and not np.isnan(variance_sma[i]) and 
                not np.isnan(variance_stdev[i]) and variance_stdev[i] > 0):
                zscore[i] = (variance[i] - variance_sma[i]) / variance_stdev[i]
        
        # Calculate EMA of z-score
        ema_zscore = self._calculate_ema_safe(zscore, ema_length)
        
        if return_components:
            results = (variance, ema_variance, zscore, ema_zscore, stdev)
            return self.format_multiple_outputs(results, input_type, index)
        else:
            return self.format_output(variance, input_type, index)


class TSF(BaseIndicator):
    """
    Time Series Forecast
    
    Forecasts the next value using linear regression.
    """
    
    def __init__(self):
        super().__init__("Time Series Forecast")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_tsf(data: np.ndarray, period: int) -> np.ndarray:
        """Numba optimized TSF calculation"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            y = data[i - period + 1:i + 1]
            x = np.arange(period)
            
            sum_x = np.sum(x)
            sum_y = np.sum(y)
            sum_xy = np.sum(x * y)
            sum_x2 = np.sum(x * x)
            
            denominator = period * sum_x2 - sum_x * sum_x
            if denominator != 0:
                slope = (period * sum_xy - sum_x * sum_y) / denominator
                intercept = (sum_y - slope * sum_x) / period
                
                # Forecast next value
                result[i] = slope * period + intercept
            else:
                result[i] = y[-1]
        
        return result
    
    def calculate(self, data: Union[np.ndarray, pd.Series, list], period: int = 14) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Time Series Forecast
        
        Parameters:
        -----------
        data : Union[np.ndarray, pd.Series, list]
            Price data
        period : int, default=14
            Period for forecast calculation
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Time Series Forecast values in the same format as input
        """
        validated_data, input_type, index = self.validate_input(data)
        self.validate_period(period, len(validated_data))
        
        result = self._calculate_tsf(validated_data, period)
        return self.format_output(result, input_type, index)


class MEDIAN(BaseIndicator):
    """
    Rolling Median (Pine Script v6)
    
    Calculates the median value over a rolling window using percentile_nearest_rank.
    Pine Script v6 Formula:
    median = ta.percentile_nearest_rank(source, length, 50)
    
    This is equivalent to the 50th percentile of the data.
    """
    
    def __init__(self):
        super().__init__("Median")
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_median(data: np.ndarray, period: int) -> np.ndarray:
        """Numba optimized median calculation (percentile_nearest_rank with 50th percentile)"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            window = data[i - period + 1:i + 1].copy()
            
            # Sort the window
            for j in range(period):
                for k in range(period - 1 - j):
                    if window[k] > window[k + 1]:
                        window[k], window[k + 1] = window[k + 1], window[k]
            
            # Get median (50th percentile) - nearest rank method
            # For percentile_nearest_rank, we take the actual value at the position
            if period % 2 == 1:
                result[i] = window[period // 2]
            else:
                # For even length, take the average of the two middle values
                result[i] = (window[period // 2 - 1] + window[period // 2]) / 2.0
        
        return result
    
    def calculate(self, data: Union[np.ndarray, pd.Series, list], period: int = 3) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Rolling Median (Pine Script v6)
        
        Pine Script v6 Formula:
        median = ta.percentile_nearest_rank(source, length, 50)
        
        Parameters:
        -----------
        data : Union[np.ndarray, pd.Series, list]
            Price data (default hl2 in Pine Script)
        period : int, default=3
            Period for median calculation (Pine Script default)
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Median values in the same format as input
        """
        validated_data, input_type, index = self.validate_input(data)
        self.validate_period(period, len(validated_data))
        
        result = self._calculate_median(validated_data, period)
        return self.format_output(result, input_type, index)


class MedianBands(BaseIndicator):
    """
    Median Bands (Pine Script v6 Complete Implementation)
    
    Implements the complete Pine Script v6 Median indicator with:
    - Median line using percentile_nearest_rank
    - ATR-based upper and lower bands
    - EMA of the median
    
    Pine Script v6 Formula:
    median = ta.percentile_nearest_rank(source, length, 50)
    atr_ = atr_mult * ta.atr(atr_length)
    upper = median + atr_
    lower = median - atr_
    median_ema = ta.ema(median, length)
    """
    
    def __init__(self):
        super().__init__("Median Bands")
    
    def calculate(self, *args, **kwargs):
        """Wrapper for calculate_with_bands to satisfy BaseIndicator abstract method"""
        return self.calculate_with_bands(*args, **kwargs)
        
    @staticmethod
    @jit(nopython=True)
    def _calculate_median_percentile(data: np.ndarray, period: int) -> np.ndarray:
        """Calculate median using percentile_nearest_rank method"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            window = data[i - period + 1:i + 1].copy()
            
            # Sort the window
            for j in range(period):
                for k in range(period - 1 - j):
                    if window[k] > window[k + 1]:
                        window[k], window[k + 1] = window[k + 1], window[k]
            
            # percentile_nearest_rank for 50th percentile
            if period % 2 == 1:
                result[i] = window[period // 2]
            else:
                # For even length, take the lower middle value
                result[i] = window[period // 2 - 1]
        
        return result
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_ema(data: np.ndarray, period: int) -> np.ndarray:
        """Calculate EMA with NaN handling"""
        n = len(data)
        result = np.full(n, np.nan)
        alpha = 2.0 / (period + 1)
        
        # Find first valid value
        first_valid_idx = -1
        for i in range(n):
            if not np.isnan(data[i]):
                first_valid_idx = i
                break
        
        if first_valid_idx == -1:
            return result
            
        # Initialize with first valid value
        result[first_valid_idx] = data[first_valid_idx]
        
        # Continue EMA calculation
        for i in range(first_valid_idx + 1, n):
            if not np.isnan(data[i]):
                result[i] = alpha * data[i] + (1 - alpha) * result[i - 1]
            else:
                result[i] = result[i - 1]
        
        return result
    
    @staticmethod
    @jit(nopython=True)
    def _calculate_atr(high: np.ndarray, low: np.ndarray, close: np.ndarray, period: int) -> np.ndarray:
        """Calculate ATR"""
        n = len(close)
        tr = np.full(n, np.nan)
        atr = np.full(n, np.nan)
        
        # Calculate True Range
        tr[0] = high[0] - low[0]
        for i in range(1, n):
            hl = high[i] - low[i]
            hc = abs(high[i] - close[i-1])
            lc = abs(low[i] - close[i-1])
            tr[i] = max(hl, hc, lc)
        
        # Calculate ATR (Wilder's smoothing)
        # Initial ATR
        if n >= period:
            sum_tr = 0.0
            for i in range(period):
                sum_tr += tr[i]
            atr[period-1] = sum_tr / period
            
            # Subsequent ATR values
            for i in range(period, n):
                atr[i] = (atr[i-1] * (period - 1) + tr[i]) / period
        
        return atr
    
    def calculate_with_bands(self, high: Union[np.ndarray, pd.Series, list],
                            low: Union[np.ndarray, pd.Series, list],
                            close: Union[np.ndarray, pd.Series, list],
                            source: Optional[Union[np.ndarray, pd.Series, list]] = None,
                            median_length: int = 3,
                            atr_length: int = 14,
                            atr_mult: float = 2.0) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], 
                                                           Tuple[pd.Series, pd.Series, pd.Series, pd.Series]]:
        """
        Calculate Median with Bands and EMA (Pine Script v6 Complete)
        
        Parameters:
        -----------
        high : Union[np.ndarray, pd.Series, list]
            High prices
        low : Union[np.ndarray, pd.Series, list]
            Low prices
        close : Union[np.ndarray, pd.Series, list]
            Close prices
        source : Optional[Union[np.ndarray, pd.Series, list]]
            Source data for median (default: hl2 = (high + low) / 2)
        median_length : int, default=3
            Period for median calculation
        atr_length : int, default=14
            Period for ATR calculation
        atr_mult : float, default=2.0
            ATR multiplier for bands
            
        Returns:
        --------
        Tuple of (median, upper_band, lower_band, median_ema)
        """
        high_data, input_type, index = self.validate_input(high)
        low_data, _, _ = self.validate_input(low)
        close_data, _, _ = self.validate_input(close)
        
        # Align arrays
        high_data, low_data, close_data = self.align_arrays(high_data, low_data, close_data)
        
        # Calculate source (hl2 if not provided)
        if source is None:
            source_data = (high_data + low_data) / 2.0
        else:
            source_data, _, _ = self.validate_input(source)
        
        # Calculate median
        median = self._calculate_median_percentile(source_data, median_length)
        
        # Calculate ATR
        atr = self._calculate_atr(high_data, low_data, close_data, atr_length)
        atr_scaled = atr * atr_mult
        
        # Calculate bands
        upper_band = median + atr_scaled
        lower_band = median - atr_scaled
        
        # Calculate EMA of median
        median_ema = self._calculate_ema(median, median_length)
        
        results = (median, upper_band, lower_band, median_ema)
        return self.format_multiple_outputs(results, input_type, index)


class MODE(BaseIndicator):
    """
    Rolling Mode
    
    Calculates the most frequent value over a rolling window.
    """
    
    def __init__(self):
        super().__init__("Mode")
    
    def calculate(self, data: Union[np.ndarray, pd.Series, list], 
                 period: int = 20, bins: int = 10) -> Union[np.ndarray, pd.Series]:
        """
        Calculate Rolling Mode using optimized algorithm
        
        Parameters:
        -----------
        data : Union[np.ndarray, pd.Series, list]
            Price data
        period : int, default=20
            Period for mode calculation
        bins : int, default=10
            Number of bins for discretization
            
        Returns:
        --------
        Union[np.ndarray, pd.Series]
            Mode values in the same format as input
        """
        validated_data, input_type, index = self.validate_input(data)
        self.validate_period(period, len(validated_data))
        
        result = self._calculate_mode_optimized(validated_data, period, bins)
        return self.format_output(result, input_type, index)
    
    @staticmethod
    @jit(nopython=True, cache=True)
    def _calculate_mode_optimized(data: np.ndarray, period: int, bins: int) -> np.ndarray:
        """Optimized rolling mode calculation using vectorized binning"""
        n = len(data)
        result = np.full(n, np.nan)
        
        for i in range(period - 1, n):
            window = data[i - period + 1:i + 1]
            
            # Fast min/max calculation
            min_val = np.min(window)
            max_val = np.max(window)
            
            if max_val > min_val:
                bin_width = (max_val - min_val) / bins
                
                # Vectorized binning
                bin_indices = ((window - min_val) / bin_width).astype(np.int32)
                bin_indices = np.clip(bin_indices, 0, bins - 1)
                
                # Fast histogram using numpy bincount
                counts = np.bincount(bin_indices, minlength=bins)
                
                # Find mode bin
                mode_bin = np.argmax(counts)
                result[i] = min_val + (mode_bin + 0.5) * bin_width
            else:
                result[i] = window[0]
        
        return result