import numpy as np
import polars as pl
import limen._feature_constraints as _check

CDLTHRUSTING_BODY_LONG_AVG_PERIOD = 10
CDLTHRUSTING_BODY_LONG_PERIOD_TOTAL = 0.0
CDLTHRUSTING_EQUAL_AVG_PERIOD = 5
CDLTHRUSTING_EQUAL_FACTOR = 0.05
CDLTHRUSTING_EQUAL_PERIOD_TOTAL = 0.0


def _cdlthrusting_impl(
    data: pl.DataFrame,
    open_col: str = 'open',
    high_col: str = 'high',
    low_col: str = 'low',
    close_col: str = 'close',
) -> pl.DataFrame:

    '''
    Compute Thrusting candlestick pattern.

    Args:
        data (pl.DataFrame): Dataset with OHLC columns
        open_col (str): Column name for open prices
        high_col (str): Column name for high prices
        low_col (str): Column name for low prices
        close_col (str): Column name for close prices

    Returns:
        pl.DataFrame: The input data with a new column 'cdlthrusting'
    '''

    open_values = data[open_col].to_numpy().astype(float, copy=False).tolist()
    high_values = data[high_col].to_numpy().astype(float, copy=False).tolist()
    low_values = data[low_col].to_numpy().astype(float, copy=False).tolist()
    close_values = data[close_col].to_numpy().astype(float, copy=False).tolist()
    n = len(data)


    equal_avg_period = CDLTHRUSTING_EQUAL_AVG_PERIOD
    equal_factor = CDLTHRUSTING_EQUAL_FACTOR
    body_long_avg_period = CDLTHRUSTING_BODY_LONG_AVG_PERIOD
    lookback_total = max(equal_avg_period, body_long_avg_period) + 1

    out = np.zeros(n, dtype=np.int32)
    if n <= lookback_total:
        return data.with_columns(pl.Series(name='cdlthrusting', values=out))

    equal_period_total = CDLTHRUSTING_EQUAL_PERIOD_TOTAL
    body_long_period_total = CDLTHRUSTING_BODY_LONG_PERIOD_TOTAL
    equal_trailing_idx = lookback_total - equal_avg_period
    body_long_trailing_idx = lookback_total - body_long_avg_period

    i = equal_trailing_idx
    while i < lookback_total:
        equal_period_total += high_values[i - 1] - low_values[i - 1]
        i += 1

    i = body_long_trailing_idx
    while i < lookback_total:
        body_long_period_total += abs(close_values[i - 1] - open_values[i - 1])
        i += 1

    i = lookback_total
    while i < n:
        color_i1 = 1 if close_values[i - 1] >= open_values[i - 1] else -1
        color_i0 = 1 if close_values[i] >= open_values[i] else -1
        real_body_i1 = abs(close_values[i - 1] - open_values[i - 1])
        equal_avg_i1 = equal_factor * (equal_period_total / equal_avg_period)
        body_long_avg_i1 = body_long_period_total / body_long_avg_period

        if (
            color_i1 == -1
            and real_body_i1 > body_long_avg_i1
            and color_i0 == 1
            and open_values[i] < low_values[i - 1]
            and close_values[i] > close_values[i - 1] + equal_avg_i1
            and close_values[i] <= close_values[i - 1] + (real_body_i1 * 0.5)
        ):
            out[i] = -100

        equal_period_total += (
            (high_values[i - 1] - low_values[i - 1])
            - (high_values[equal_trailing_idx - 1] - low_values[equal_trailing_idx - 1])
        )
        body_long_period_total += (
            abs(close_values[i - 1] - open_values[i - 1])
            - abs(close_values[body_long_trailing_idx - 1] - open_values[body_long_trailing_idx - 1])
        )

        i += 1
        equal_trailing_idx += 1
        body_long_trailing_idx += 1

    return data.with_columns(pl.Series(name='cdlthrusting', values=out))


def cdlthrusting(
    data: pl.DataFrame,
    open_col: str = 'open',
    high_col: str = 'high',
    low_col: str = 'low',
    close_col: str = 'close',
) -> pl.DataFrame:

    _check.names(open_col=open_col, high_col=high_col, low_col=low_col, close_col=close_col)

    out_col = 'cdlthrusting'
    input_cols = [open_col, high_col, low_col, close_col]
    return data.with_columns(
        pl.struct(input_cols).map_batches(
            lambda s: _cdlthrusting_impl(
                pl.DataFrame({col: s.struct.field(col) for col in input_cols}),
                open_col=open_col,
                high_col=high_col,
                low_col=low_col,
                close_col=close_col,
            ).get_column(out_col),
            return_dtype=pl.Int32,
        ).alias(out_col)
    )