"""Gate primitives — boolean masks over bars (doc 04 Rule 6).

A gate is a callable that takes the bars DataFrame and returns a boolean
``numpy.ndarray`` (``True`` = entry allowed on that bar). Gates never open or
close trades; they only mask the signal array the caller hands to the engine.
"""
from __future__ import annotations

from typing import Protocol, runtime_checkable

import numpy as np
import pandas as pd


@runtime_checkable
class Gate(Protocol):
    """A callable producing a boolean mask over bars (``True`` = allow entry)."""

    def __call__(self, bars: pd.DataFrame) -> np.ndarray: ...


def time_of_day_gate(
    bars: pd.DataFrame,
    *,
    start_hour: int,
    end_hour: int,
    timezone: str | None = None,
) -> np.ndarray:
    """Allow entries only within ``[start_hour, end_hour)`` (hours, 0–24).

    Useful for sessions that only trade London/NY open. ``timezone`` is
    applied to ``bars['timestamp']`` if given; otherwise the timestamp's
    existing tz is used (or naive local time).
    """
    ts = pd.to_datetime(bars["timestamp"])
    if timezone is not None:
        ts = ts.dt.tz_localize(None).dt.tz_localize(timezone) if ts.dt.tz is None else ts.dt.tz_convert(timezone)
    hours = ts.dt.hour
    if start_hour <= end_hour:
        mask = (hours >= start_hour) & (hours < end_hour)
    else:
        # Wrap past midnight, e.g. 22 → 6.
        mask = (hours >= start_hour) | (hours < end_hour)
    return mask.to_numpy()


def regime_gate(
    bars: pd.DataFrame,
    *,
    trend_filter: np.ndarray,
    direction: int,
) -> np.ndarray:
    """Allow entries only when ``trend_filter`` agrees with ``direction``.

    ``trend_filter`` is a +1/-1 array (e.g. from an EMA slope or ADX sign).
    ``direction=+1`` keeps bars where the trend is up; ``-1`` keeps downtrend.
    """
    tf = np.asarray(trend_filter)
    mask = tf == direction
    return mask


def exhaustion_gate(
    rsi_arr: np.ndarray,
    *,
    overbought: float = 70.0,
    oversold: float = 30.0,
) -> np.ndarray:
    """Block entries when RSI is in the exhaustion zone for the direction.

    Returns ``True`` where entry is *allowed* (i.e. not exhausted). Block longs
    when ``rsi >= overbought`` and shorts when ``rsi <= oversold`` — combine
    with the directional signal in the caller.
    """
    rsi_arr = np.asarray(rsi_arr, dtype=float)
    allow = (rsi_arr < overbought) & (rsi_arr > oversold)
    allow = np.where(np.isnan(rsi_arr), False, allow)
    return allow