"""
K线数据聚合器
从1H K线聚合出4H和日线数据
"""

import pandas as pd
import numpy as np
from typing import Dict


def aggregate_to_4h(df_1h: pd.DataFrame) -> pd.DataFrame:
    """
    1H K线聚合为4H K线
    4H聚合方法: floor到4小时整点 (0,4,8,12,16,20 UTC)
    """
    df = df_1h.copy()
    df['group'] = df['timestamp'].dt.floor('4h')

    agg = df.groupby('group').agg({
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last',
        'volume': 'sum',
    }).reset_index()
    agg.rename(columns={'group': 'timestamp'}, inplace=True)
    # 只保留完整的4根1H组成的4H bar
    counts = df.groupby('group').size()
    complete = counts[counts == 4].index
    agg = agg[agg['timestamp'].isin(complete)].reset_index(drop=True)
    return agg


def aggregate_to_daily(df_1h: pd.DataFrame) -> pd.DataFrame:
    """
    1H K线聚合为日线（UTC日期）
    """
    df = df_1h.copy()
    df['group'] = df['timestamp'].dt.floor('1D')

    agg = df.groupby('group').agg({
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last',
        'volume': 'sum',
    }).reset_index()
    agg.rename(columns={'group': 'timestamp'}, inplace=True)
    # 只保留完整24根1H的日线
    counts = df.groupby('group').size()
    complete = counts[counts == 24].index
    agg = agg[agg['timestamp'].isin(complete)].reset_index(drop=True)
    return agg


def prepare_all_timeframes(df_1h: pd.DataFrame) -> Dict[str, pd.DataFrame]:
    """
    从1H数据准备所有时间周期
    返回 {'1H': df, '4H': df, '1D': df}
    """
    df_4h = aggregate_to_4h(df_1h)
    df_1d = aggregate_to_daily(df_1h)
    return {
        '1H': df_1h,
        '4H': df_4h,
        '1D': df_1d,
    }