"""
Pivot Detection Functions

Vectorized and bar-by-bar pivot high/low detection for price analysis.

Usage:
    from sixtysix.ta import detect_pivots, detect_pivot_lows, pivothigh, pivotlow
"""

from __future__ import annotations
from typing import Optional, Tuple
import numpy as np
import pandas as pd

from .helpers import HasBarIndex


def detect_pivots(high: np.ndarray, low: np.ndarray, length: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    Detect all pivot highs and lows in a price series (vectorized).

    Returns arrays with pivot values at their confirmation bars (length bars after the actual pivot).
    Non-pivot bars contain NaN.

    Args:
        high: High prices array
        low: Low prices array
        length: Number of bars on each side to confirm pivot

    Returns:
        Tuple of (pivot_highs, pivot_lows) arrays

    Example:
        pivot_highs, pivot_lows = detect_pivots(df['high'].values, df['low'].values, 10)
    """
    n = len(high)
    pivot_highs = np.full(n, np.nan)
    pivot_lows = np.full(n, np.nan)

    for i in range(length, n - length):
        # Pivot high: center is higher than all neighbors
        if all(high[i] > high[i - j] and high[i] > high[i + j] for j in range(1, length + 1)):
            pivot_highs[i + length] = high[i]

        # Pivot low: center is lower than all neighbors
        if all(low[i] < low[i - j] and low[i] < low[i + j] for j in range(1, length + 1)):
            pivot_lows[i + length] = low[i]

    return pivot_highs, pivot_lows


def detect_pivot_lows(low: np.ndarray, length: int) -> np.ndarray:
    """
    Vectorized pivot low detection using rolling operations.

    Returns boolean mask where True = pivot confirmed at this bar.
    The pivot value is at (index - length), confirmed at index.

    Args:
        low: Low prices array
        length: Number of bars on each side to confirm pivot

    Returns:
        Boolean mask array

    Example:
        pivot_mask = detect_pivot_lows(df['low'].values, 10)
        pivot_indices = np.where(pivot_mask)[0]
        pivot_values = low[pivot_indices - 10]  # Values at actual pivot bars
    """
    n = len(low)
    s = pd.Series(low)

    # Left: min of previous `length` bars (excluding current)
    left_min = s.shift(1).rolling(length, min_periods=length).min().values
    # Right: min of next `length` bars (excluding current)
    right_min = s.shift(-1).iloc[::-1].rolling(length, min_periods=length).min().iloc[::-1].values

    # Pivot where current is strictly less than both sides
    is_pivot = (low < left_min) & (low < right_min)

    # Shift forward by `length` bars (pivot confirmed `length` bars later)
    is_pivot = np.roll(is_pivot, length)
    is_pivot[:length] = False
    is_pivot[n - length:] = False

    return is_pivot


def detect_pivot_highs(high: np.ndarray, length: int) -> np.ndarray:
    """
    Vectorized pivot high detection using rolling operations.

    Returns boolean mask where True = pivot confirmed at this bar.

    Args:
        high: High prices array
        length: Number of bars on each side to confirm pivot

    Returns:
        Boolean mask array
    """
    n = len(high)
    s = pd.Series(high)

    left_max = s.shift(1).rolling(length, min_periods=length).max().values
    right_max = s.shift(-1).iloc[::-1].rolling(length, min_periods=length).max().iloc[::-1].values

    is_pivot = (high > left_max) & (high > right_max)

    is_pivot = np.roll(is_pivot, length)
    is_pivot[:length] = False
    is_pivot[n - length:] = False

    return is_pivot


def pivothigh(ctx: HasBarIndex, source: pd.Series, left: int, right: int) -> Optional[float]:
    """
    Detect pivot high at confirmation bar.

    Similar to PineScript's ta.pivothigh().

    The pivot is confirmed 'right' bars after it occurs. At confirmation,
    returns the pivot value. Otherwise returns None.

    Args:
        ctx: Trading context (provides current bar index)
        source: Price series (typically df['high'])
        left: Number of bars to the left that must be lower
        right: Number of bars to the right that must be lower

    Returns:
        Pivot high value if confirmed at current bar, None otherwise

    Example:
        ph = pivothigh(self.ctx, df['high'], 10, 10)
        if ph is not None:
            self.last_high = ph
    """
    idx = ctx.bar_index
    pivot_idx = idx - right

    if pivot_idx < left or pivot_idx < 0:
        return None

    pivot_val = source.iloc[pivot_idx]
    if pd.isna(pivot_val):
        return None

    # Check left side (must be lower, or equal for double tops)
    for i in range(1, left + 1):
        if source.iloc[pivot_idx - i] > pivot_val:
            return None

    # Check right side (must be lower, or equal for double tops)
    for i in range(1, right + 1):
        if source.iloc[pivot_idx + i] > pivot_val:
            return None

    return float(pivot_val)


def pivotlow(ctx: HasBarIndex, source: pd.Series, left: int, right: int) -> Optional[float]:
    """
    Detect pivot low at confirmation bar.

    Similar to PineScript's ta.pivotlow().

    The pivot is confirmed 'right' bars after it occurs. At confirmation,
    returns the pivot value. Otherwise returns None.

    Args:
        ctx: Trading context (provides current bar index)
        source: Price series (typically df['low'])
        left: Number of bars to the left that must be higher
        right: Number of bars to the right that must be higher

    Returns:
        Pivot low value if confirmed at current bar, None otherwise

    Example:
        pl = pivotlow(self.ctx, df['low'], 10, 10)
        if pl is not None:
            self.prev_low = self.last_low
            self.last_low = pl
    """
    idx = ctx.bar_index
    pivot_idx = idx - right

    if pivot_idx < left or pivot_idx < 0:
        return None

    pivot_val = source.iloc[pivot_idx]
    if pd.isna(pivot_val):
        return None

    # Check left side (must be higher, or equal for double bottoms)
    for i in range(1, left + 1):
        if source.iloc[pivot_idx - i] < pivot_val:
            return None

    # Check right side (must be higher, or equal for double bottoms)
    for i in range(1, right + 1):
        if source.iloc[pivot_idx + i] < pivot_val:
            return None

    return float(pivot_val)
