"""
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 sixtysix.ta import pivot_structure
"""

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

if TYPE_CHECKING:
    from .helpers import HasBarIndex


class PivotStructureDict(TypedDict):
    """Return type of pivot_structure(). All values are np.ndarray of length n."""
    higher_low: np.ndarray       # 1.0 if last pivot low > previous pivot low
    higher_high: np.ndarray      # 1.0 if last pivot high > previous pivot high
    bars_since_pivot: np.ndarray # Bars since last pivot low (capped at 200)
    last_pivot_low: np.ndarray   # Price of most recent pivot low (NaN if none yet)
    prev_pivot_low: np.ndarray   # Price of previous pivot low (NaN if none yet)
    last_pivot_high: np.ndarray  # Price of most recent pivot high (NaN if none yet)
    prev_pivot_high: np.ndarray  # Price of previous pivot high (NaN if none yet)
    consecutive_hl: np.ndarray   # Count of consecutive higher lows (0–4)
    consecutive_hh: np.ndarray   # Count of consecutive higher highs (0–4)


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)


def pivot_structure(
    high: np.ndarray, low: np.ndarray, lookback: int = 30,
) -> PivotStructureDict:
    """
    Full pivot structure analysis — higher lows, higher highs, consecutive
    counts, distances, and slopes.

    Detects pivot highs/lows using a window of `lookback` bars on each side,
    then walks forward tracking structure evolution at each bar.

    Args:
        high: High prices array
        low: Low prices array
        lookback: Bars on each side for pivot confirmation

    Returns:
        Dict with keys:
            'higher_low': 1.0 if last pivot low > previous pivot low
            'higher_high': 1.0 if last pivot high > previous pivot high
            'bars_since_pivot': Bars since last pivot low (capped at 200)
            'last_pivot_low': Price of most recent pivot low
            'prev_pivot_low': Price of previous pivot low
            'last_pivot_high': Price of most recent pivot high
            'prev_pivot_high': Price of previous pivot high
            'consecutive_hl': Count of consecutive higher lows (0-4)
            'consecutive_hh': Count of consecutive higher highs (0-4)

    Example:
        ps = pivot_structure(df['high'].values, df['low'].values, 30)
        # Higher lows forming = bullish structure
        bullish = ps['higher_low'] > 0.5
        # Strong trend = 3+ consecutive higher lows
        strong = ps['consecutive_hl'] >= 3
    """
    n = len(high)
    higher_low = np.zeros(n)
    higher_high = np.zeros(n)
    bars_since = np.zeros(n)
    last_pl_arr = np.full(n, np.nan)
    prev_pl_arr = np.full(n, np.nan)
    last_ph_arr = np.full(n, np.nan)
    prev_ph_arr = np.full(n, np.nan)
    consec_hl = np.zeros(n)
    consec_hh = np.zeros(n)

    # Detect raw pivots
    pivot_lows = []
    pivot_highs = []

    for i in range(lookback, n - lookback):
        if low[i] == np.min(low[max(0, i - lookback):i + lookback + 1]):
            pivot_lows.append((i, low[i]))
            if len(pivot_lows) > 4:
                pivot_lows = pivot_lows[-4:]
        if high[i] == np.max(high[max(0, i - lookback):i + lookback + 1]):
            pivot_highs.append((i, high[i]))
            if len(pivot_highs) > 4:
                pivot_highs = pivot_highs[-4:]

    # Walk forward, tracking structure
    pl_idx = 0
    ph_idx = 0
    cur_last_pl = np.nan
    cur_prev_pl = np.nan
    cur_last_ph = np.nan
    cur_prev_ph = np.nan

    for i in range(n):
        # Advance pivot low pointer
        while pl_idx < len(pivot_lows) - 1 and pivot_lows[pl_idx + 1][0] <= i:
            pl_idx += 1
        if pl_idx > 0 and pivot_lows[pl_idx][0] <= i:
            cur_last_pl = pivot_lows[pl_idx][1]
            cur_prev_pl = pivot_lows[pl_idx - 1][1]
            higher_low[i] = 1.0 if cur_last_pl > cur_prev_pl else 0.0

            consec = 0
            for k in range(pl_idx, 0, -1):
                if pivot_lows[k][0] <= i and pivot_lows[k][1] > pivot_lows[k - 1][1]:
                    consec += 1
                else:
                    break
            consec_hl[i] = min(consec, 4)
        elif pl_idx == 0 and pivot_lows and pivot_lows[0][0] <= i:
            cur_last_pl = pivot_lows[0][1]

        # Advance pivot high pointer
        while ph_idx < len(pivot_highs) - 1 and pivot_highs[ph_idx + 1][0] <= i:
            ph_idx += 1
        if ph_idx > 0 and pivot_highs[ph_idx][0] <= i:
            cur_last_ph = pivot_highs[ph_idx][1]
            cur_prev_ph = pivot_highs[ph_idx - 1][1]
            higher_high[i] = 1.0 if cur_last_ph > cur_prev_ph else 0.0

            consec = 0
            for k in range(ph_idx, 0, -1):
                if pivot_highs[k][0] <= i and pivot_highs[k][1] > pivot_highs[k - 1][1]:
                    consec += 1
                else:
                    break
            consec_hh[i] = min(consec, 4)
        elif ph_idx == 0 and pivot_highs and pivot_highs[0][0] <= i:
            cur_last_ph = pivot_highs[0][1]

        last_pl_arr[i] = cur_last_pl
        prev_pl_arr[i] = cur_prev_pl
        last_ph_arr[i] = cur_last_ph
        prev_ph_arr[i] = cur_prev_ph

        if pivot_lows and pl_idx < len(pivot_lows) and pivot_lows[min(pl_idx, len(pivot_lows) - 1)][0] <= i:
            bars_since[i] = i - pivot_lows[min(pl_idx, len(pivot_lows) - 1)][0]

    return {
        'higher_low': higher_low,
        'higher_high': higher_high,
        'bars_since_pivot': np.clip(bars_since, 0, 200),
        'last_pivot_low': last_pl_arr,
        'prev_pivot_low': prev_pl_arr,
        'last_pivot_high': last_ph_arr,
        'prev_pivot_high': prev_ph_arr,
        'consecutive_hl': consec_hl,
        'consecutive_hh': consec_hh,
    }
