"""
數據分析輔助函數

提供常用的技術分析和數據處理函數，方便在 Jupyter Notebook 中使用。
"""

import polars as pl
from typing import List, Optional, Tuple


def calculate_returns(
    df: pl.DataFrame,
    price_col: str = "close"
) -> pl.DataFrame:
    """計算收益率

    Args:
        df: 輸入 DataFrame
        price_col: 價格列名

    Returns:
        添加了 returns 和 log_returns 列的 DataFrame
    """
    return df.with_columns([
        ((pl.col(price_col) / pl.col(price_col).shift(1)) - 1).alias("returns"),
        (pl.col(price_col) / pl.col(price_col).shift(1)).log().alias("log_returns")
    ])


def calculate_volatility(
    df: pl.DataFrame,
    returns_col: str = "returns",
    window: int = 20
) -> pl.DataFrame:
    """計算滾動波動率

    Args:
        df: 輸入 DataFrame
        returns_col: 收益率列名
        window: 滾動窗口大小

    Returns:
        添加了 volatility 列的 DataFrame
    """
    return df.with_columns([
        pl.col(returns_col).rolling_std(window).alias(f"volatility_{window}d")
    ])


def resample_to_daily(
    df: pl.DataFrame,
    timestamp_col: str = "ts"
) -> pl.DataFrame:
    """重採樣為日線數據

    Args:
        df: 輸入 DataFrame（分鐘線或其他高頻數據）
        timestamp_col: 時間戳列名

    Returns:
        日線 DataFrame
    """
    required_cols = ["open", "high", "low", "close"]
    if not all(col in df.columns for col in required_cols):
        raise ValueError(f"DataFrame 必須包含列: {required_cols}")

    daily = df.group_by_dynamic(
        timestamp_col,
        every="1d"
    ).agg([
        pl.col("open").first().alias("open"),
        pl.col("high").max().alias("high"),
        pl.col("low").min().alias("low"),
        pl.col("close").last().alias("close"),
        pl.col("volume").sum().alias("volume") if "volume" in df.columns else pl.lit(0).alias("volume"),
        pl.count().alias("num_bars")
    ]).sort(timestamp_col)

    return daily


def find_indicator_columns(
    df: pl.DataFrame,
    patterns: Optional[List[str]] = None
) -> List[str]:
    """查找技術指標列

    Args:
        df: 輸入 DataFrame
        patterns: 要搜索的模式列表（默認為常見技術指標）

    Returns:
        匹配的列名列表
    """
    if patterns is None:
        patterns = ["SMA", "EMA", "ATR", "RSI", "MACD", "BB", "BOLL"]

    indicator_cols = [
        col for col in df.columns
        if any(pattern in col for pattern in patterns)
    ]

    return indicator_cols


def filter_by_date_range(
    df: pl.DataFrame,
    start_date: str,
    end_date: str,
    timestamp_col: str = "ts"
) -> pl.DataFrame:
    """按日期範圍過濾數據

    Args:
        df: 輸入 DataFrame
        start_date: 起始日期（ISO 8601 格式）
        end_date: 結束日期（ISO 8601 格式）
        timestamp_col: 時間戳列名

    Returns:
        過濾後的 DataFrame
    """
    return df.filter(
        (pl.col(timestamp_col) >= start_date) &
        (pl.col(timestamp_col) <= end_date)
    )


def get_ohlcv_stats(df: pl.DataFrame) -> pl.DataFrame:
    """獲取 OHLCV 統計信息

    Args:
        df: 輸入 DataFrame

    Returns:
        統計信息 DataFrame
    """
    ohlcv_cols = ["open", "high", "low", "close", "volume"]
    available_cols = [col for col in ohlcv_cols if col in df.columns]

    if not available_cols:
        raise ValueError("DataFrame 不包含任何 OHLCV 列")

    return df.select(available_cols).describe()


def detect_outliers(
    df: pl.DataFrame,
    column: str,
    n_std: float = 3.0
) -> Tuple[pl.DataFrame, int]:
    """使用標準差方法檢測異常值

    Args:
        df: 輸入 DataFrame
        column: 要檢測的列名
        n_std: 標準差倍數（默認 3.0）

    Returns:
        (標記了異常值的 DataFrame, 異常值數量)
    """
    mean = df[column].mean()
    std = df[column].std()

    lower_bound = mean - n_std * std
    upper_bound = mean + n_std * std

    df_with_outliers = df.with_columns([
        (
            (pl.col(column) < lower_bound) |
            (pl.col(column) > upper_bound)
        ).alias(f"{column}_is_outlier")
    ])

    outlier_count = df_with_outliers.filter(pl.col(f"{column}_is_outlier")).height

    return df_with_outliers, outlier_count


def calculate_drawdown(
    df: pl.DataFrame,
    price_col: str = "close"
) -> pl.DataFrame:
    """計算回撤

    Args:
        df: 輸入 DataFrame
        price_col: 價格列名

    Returns:
        添加了 cummax 和 drawdown 列的 DataFrame
    """
    return df.with_columns([
        pl.col(price_col).cum_max().alias("cummax"),
        ((pl.col(price_col) / pl.col(price_col).cum_max()) - 1).alias("drawdown")
    ])


def get_max_drawdown(
    df: pl.DataFrame,
    price_col: str = "close"
) -> float:
    """獲取最大回撤

    Args:
        df: 輸入 DataFrame
        price_col: 價格列名

    Returns:
        最大回撤值（百分比）
    """
    df_with_dd = calculate_drawdown(df, price_col)
    return df_with_dd["drawdown"].min()


def split_train_test(
    df: pl.DataFrame,
    train_ratio: float = 0.8,
    timestamp_col: str = "ts"
) -> Tuple[pl.DataFrame, pl.DataFrame]:
    """按時間順序分割訓練集和測試集

    Args:
        df: 輸入 DataFrame
        train_ratio: 訓練集比例
        timestamp_col: 時間戳列名

    Returns:
        (訓練集, 測試集)
    """
    df_sorted = df.sort(timestamp_col)
    split_idx = int(len(df_sorted) * train_ratio)

    train = df_sorted[:split_idx]
    test = df_sorted[split_idx:]

    return train, test


def summary_stats(df: pl.DataFrame, price_col: str = "close") -> dict:
    """計算常用的總結統計量

    Args:
        df: 輸入 DataFrame
        price_col: 價格列名

    Returns:
        統計量字典
    """
    df_with_returns = calculate_returns(df, price_col)

    stats = {
        "總行數": len(df),
        "起始價格": df[price_col].first(),
        "結束價格": df[price_col].last(),
        "最高價": df[price_col].max(),
        "最低價": df[price_col].min(),
        "平均價格": df[price_col].mean(),
        "價格標準差": df[price_col].std(),
        "總收益率": (df[price_col].last() / df[price_col].first() - 1) * 100,
    }

    if "returns" in df_with_returns.columns:
        stats.update({
            "平均日收益率": df_with_returns["returns"].mean() * 100,
            "收益率標準差": df_with_returns["returns"].std() * 100,
            "夏普比率（假設 Rf=0）": (
                df_with_returns["returns"].mean() / df_with_returns["returns"].std()
                if df_with_returns["returns"].std() != 0 else 0
            ),
        })

    # 計算最大回撤
    max_dd = get_max_drawdown(df, price_col)
    stats["最大回撤"] = max_dd * 100

    return stats


def print_summary(df: pl.DataFrame, price_col: str = "close"):
    """打印數據摘要

    Args:
        df: 輸入 DataFrame
        price_col: 價格列名
    """
    stats = summary_stats(df, price_col)

    print("=" * 50)
    print("數據摘要統計")
    print("=" * 50)

    for key, value in stats.items():
        if isinstance(value, float):
            if "率" in key or "撤" in key:
                print(f"{key:<20}: {value:>10.2f}%")
            else:
                print(f"{key:<20}: {value:>10.2f}")
        else:
            print(f"{key:<20}: {value:>10}")

    print("=" * 50)
