import polars as pl
from limen.data.bars._bar_aggregation import aggregate_bars



MIN_DATA_LENGTH = 2


def _standard_bars(data: pl.DataFrame,
                   threshold: int | float,
                   column_name: str) -> pl.DataFrame:

    '''
    Compute standard bars from klines data using fixed threshold on
    column value accumulation.

    Args:
        data (pl.DataFrame): Klines data
        threshold (Union[int, float]): Threshold value for bar creation
        column_name (str): Column name to accumulate for threshold detection

    Returns:
        pl.DataFrame:  Data bars with columns 'datetime', 'open', 'high',
                      'low', 'close', 'volume', 'no_of_trades', 'liquidity_sum',
                      'maker_ratio', 'maker_volume', 'maker_liquidity', 'mean'
                      'bar_count', 'base_interval'

    '''

    base_interval = (
          (data['datetime'][1] - data['datetime'][0]).total_seconds()
          if len(data) >= MIN_DATA_LENGTH
          else 0
      )

    df_with_cumsum = data.with_row_index().with_columns([
        pl.col(column_name).alias('threshold_col')
    ])

    def create_bar_groups(df: pl.DataFrame) -> pl.DataFrame:
        cumsum = 0.0
        bar_group = 0
        groups: list[int] = []

        for row in df.iter_rows(named=True):
            cumsum += row['threshold_col']
            groups.append(bar_group)

            if cumsum >= threshold:
                bar_group += 1
                cumsum = 0.0

        return df.with_columns(pl.Series('bar_group', groups))

    df_with_groups = create_bar_groups(df_with_cumsum)

    max_bar_group = df_with_groups['bar_group'].max()
    last_bar_sum = (
        df_with_groups
        .filter(pl.col('bar_group') == max_bar_group)[column_name]
        .sum())

    if last_bar_sum < threshold:
        df_with_groups = (
            df_with_groups
            .filter(pl.col('bar_group') < max_bar_group)
        )

    return aggregate_bars(df_with_groups.drop(['index', 'threshold_col']), base_interval)




def volume_bars(data: pl.DataFrame, volume_threshold: float) -> pl.DataFrame:

    '''
    Compute volume bars with fixed volume size sampling.

    Args:
        data (pl.DataFrame): Klines data
        volume_threshold (float): Volume threshold per bar

    Returns:
        pl.DataFrame: Standard volume data bars
    '''

    return _standard_bars(data, volume_threshold, 'volume')


def trade_bars(data: pl.DataFrame, trade_threshold: int) -> pl.DataFrame:

    '''
    Compute trade bars with fixed trade count sampling.

    Args:
        data (pl.DataFrame): Klines data
        trade_threshold (int): Number of trades per bar

    Returns:
        pl.DataFrame: Standard trade count data bars
    '''

    return _standard_bars(data, trade_threshold, 'no_of_trades')


def liquidity_bars(data: pl.DataFrame,
                   liquidity_threshold: float) -> pl.DataFrame:

    '''
    Compute liquidity bars with fixed liquidity sampling.

    Args:
        data (pl.DataFrame): Klines data
        liquidity_threshold (float): Liquidity threshold per bar

    Returns:
        pl.DataFrame: Standard liquidity data bars
    '''

    return _standard_bars(data, liquidity_threshold, 'liquidity_sum')