"""Shared utilities for indicator computation.

Centralises array extraction, constants, and common operations
so individual indicator modules stay focused on math.
"""

from __future__ import annotations

from typing import Optional, Tuple, Union

import numpy as np

# ── Constants ────────────────────────────────────────────────────────

EPSILON: float = 1e-15
TRADING_DAYS_PER_YEAR: int = 252
ANNUALIZATION_FACTOR: float = np.sqrt(TRADING_DAYS_PER_YEAR)

# Standard OHLCV column names (overridable per-call)
DEFAULT_COL_OPEN: str = "open"
DEFAULT_COL_HIGH: str = "high"
DEFAULT_COL_LOW: str = "low"
DEFAULT_COL_CLOSE: str = "close"
DEFAULT_COL_VOLUME: str = "volume"

# Type alias for anything that can carry price data
PriceInput = Union["pd.DataFrame", "pd.Series", np.ndarray, list]


# ── Array extraction ─────────────────────────────────────────────────

def _is_polars_frame(obj) -> bool:
    """Check for Polars DataFrame/Series without importing polars at module level."""
    typ = type(obj).__module__
    return typ.startswith("polars")


def ensure_array(data: PriceInput, col: Optional[str] = None) -> np.ndarray:
    """Convert any supported input into a float64 numpy array.

    Accepts pandas DataFrames/Series, **Polars DataFrames/Series**,
    numpy arrays, and plain lists.

    Args:
        data: DataFrame, Series, ndarray, or list of floats.
        col: Column name to extract when *data* is a DataFrame.
             Defaults to ``"close"`` when *col* is ``None`` and
             *data* is a DataFrame.
    """
    import pandas as pd

    if isinstance(data, pd.DataFrame):
        col = col or DEFAULT_COL_CLOSE
        if col not in data.columns:
            raise KeyError(f"Column '{col}' not found. Available: {list(data.columns)}")
        return data[col].values.astype(np.float64)
    if isinstance(data, pd.Series):
        return data.values.astype(np.float64)

    if _is_polars_frame(data):
        cls_name = type(data).__name__
        if cls_name == "DataFrame":
            col = col or DEFAULT_COL_CLOSE
            return data[col].to_numpy().astype(np.float64)
        if cls_name in ("Series", "Expr"):
            return data.to_numpy().astype(np.float64)

    return np.asarray(data, dtype=np.float64)


def _col_to_numpy(df, col: str) -> np.ndarray:
    """Extract a column as float64 numpy array from pandas or Polars DataFrame."""
    if _is_polars_frame(df):
        return df[col].to_numpy().astype(np.float64)
    return df[col].values.astype(np.float64)


def ensure_ohlcv(
    bars,
    *,
    col_open: str = DEFAULT_COL_OPEN,
    col_high: str = DEFAULT_COL_HIGH,
    col_low: str = DEFAULT_COL_LOW,
    col_close: str = DEFAULT_COL_CLOSE,
    col_volume: str = DEFAULT_COL_VOLUME,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Extract OHLCV arrays from a pandas or Polars DataFrame.

    Column names are fully configurable so users with non-standard
    schemas don't have to rename anything.
    """
    return (
        _col_to_numpy(bars, col_open),
        _col_to_numpy(bars, col_high),
        _col_to_numpy(bars, col_low),
        _col_to_numpy(bars, col_close),
        _col_to_numpy(bars, col_volume),
    )


def ensure_ohlc(
    bars,
    *,
    col_open: str = DEFAULT_COL_OPEN,
    col_high: str = DEFAULT_COL_HIGH,
    col_low: str = DEFAULT_COL_LOW,
    col_close: str = DEFAULT_COL_CLOSE,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Extract OHLC arrays (no volume). Accepts pandas or Polars."""
    return (
        _col_to_numpy(bars, col_open),
        _col_to_numpy(bars, col_high),
        _col_to_numpy(bars, col_low),
        _col_to_numpy(bars, col_close),
    )


def ensure_hlc(
    bars,
    *,
    col_high: str = DEFAULT_COL_HIGH,
    col_low: str = DEFAULT_COL_LOW,
    col_close: str = DEFAULT_COL_CLOSE,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Extract High / Low / Close arrays. Accepts pandas or Polars."""
    return (
        _col_to_numpy(bars, col_high),
        _col_to_numpy(bars, col_low),
        _col_to_numpy(bars, col_close),
    )


# ── Smoothing primitives (reused across indicators) ─────────────────

def wilder_smooth(values: np.ndarray, period: int) -> np.ndarray:
    """Wilder's smoothing (equivalent to EMA with alpha = 1/period).

    Used by RSI, ATR, ADX, and other Wilder-family indicators.
    """
    n = len(values)
    out = np.full(n, np.nan)
    if n < period:
        return out
    out[period - 1] = np.mean(values[:period])
    alpha = 1.0 / period
    for i in range(period, n):
        out[i] = alpha * values[i] + (1.0 - alpha) * out[i - 1]
    return out


# ── Source price helpers ─────────────────────────────────────────────

def source_price(
    bars,
    source: str = "close",
    *,
    col_open: str = DEFAULT_COL_OPEN,
    col_high: str = DEFAULT_COL_HIGH,
    col_low: str = DEFAULT_COL_LOW,
    col_close: str = DEFAULT_COL_CLOSE,
) -> np.ndarray:
    """Compute a derived price source from OHLC data.

    Accepts pandas or Polars DataFrames.

    Supported *source* values:
        ``"close"``  — close price (default)
        ``"open"``   — open price
        ``"hl2"``    — (high + low) / 2
        ``"hlc3"``   — (high + low + close) / 3  (typical price)
        ``"ohlc4"``  — (open + high + low + close) / 4
    """
    import pandas as pd

    is_df = isinstance(bars, pd.DataFrame) or _is_polars_frame(bars)
    if not is_df:
        return ensure_array(bars)

    if source == "close":
        return _col_to_numpy(bars, col_close)
    if source == "open":
        return _col_to_numpy(bars, col_open)
    if source == "hl2":
        h = _col_to_numpy(bars, col_high)
        l = _col_to_numpy(bars, col_low)
        return (h + l) / 2.0
    if source == "hlc3":
        h = _col_to_numpy(bars, col_high)
        l = _col_to_numpy(bars, col_low)
        c = _col_to_numpy(bars, col_close)
        return (h + l + c) / 3.0
    if source == "ohlc4":
        o = _col_to_numpy(bars, col_open)
        h = _col_to_numpy(bars, col_high)
        l = _col_to_numpy(bars, col_low)
        c = _col_to_numpy(bars, col_close)
        return (o + h + l + c) / 4.0
    raise ValueError(f"Unknown source '{source}'. Use: close, open, hl2, hlc3, ohlc4")
