# 検証環境
# Polars 1.x系（APIは変更が多いためバージョンを必ず固定してください）
# pip install polars==1.8.0

import polars as pl
import numpy as np
from datetime import datetime, timezone

print(pl.__version__)  # 1.8.0 などを確認

# ダミーOHLCVデータの生成（1分足 × 200本）
np.random.seed(42)
n = 200
base_price = 50000.0

timestamps_ms = [
    int(datetime(2024, 1, 15, 9, 0, tzinfo=timezone.utc).timestamp() * 1000)
    + i * 60_000
    for i in range(n)
]
prices = base_price + np.cumsum(np.random.randn(n) * 100)

df_1m = pl.DataFrame({
    "timestamp_ms": timestamps_ms,
    "open":   prices,
    "high":   prices + np.abs(np.random.randn(n) * 50),
    "low":    prices - np.abs(np.random.randn(n) * 50),
    "close":  prices + np.random.randn(n) * 30,
    "volume": np.abs(np.random.randn(n) * 10 + 5),
    "symbol": ["BTC/USDT"] * n,
})

# CCXTが返すエポックミリ秒を Datetime(time_unit="ms", UTC) に変換
df = df_1m.with_columns(
    pl.from_epoch("timestamp_ms", time_unit="ms")
      .dt.replace_time_zone("UTC")
      .alias("datetime")
)

print(df.schema)
# {'timestamp_ms': Int64, 'datetime': Datetime(time_unit='ms', time_zone='UTC'), ...}

# 文字列タイムスタンプのパース例
df_str = pl.DataFrame({"ts": ["2024-01-15 09:00:00", "2024-01-15 09:01:00"]})

df_parsed = df_str.with_columns(
    pl.col("ts")
      .str.to_datetime(format="%Y-%m-%d %H:%M:%S", time_unit="ms")
      .dt.replace_time_zone("UTC")
      .alias("datetime")
)

# CCXTが返すエポックミリ秒を Datetime(time_unit="ms", UTC) に変換
df = df_1m.with_columns(
    pl.from_epoch("timestamp_ms", time_unit="ms")
      .dt.replace_time_zone("UTC")
      .alias("datetime")
)

print(df.schema)
# {'timestamp_ms': Int64, 'datetime': Datetime(time_unit='ms', time_zone='UTC'), ...}

# 文字列タイムスタンプのパース例
df_str = pl.DataFrame({"ts": ["2024-01-15 09:00:00", "2024-01-15 09:01:00"]})

df_parsed = df_str.with_columns(
    pl.col("ts")
      .str.to_datetime(format="%Y-%m-%d %H:%M:%S", time_unit="ms")
      .dt.replace_time_zone("UTC")
      .alias("datetime")
)

# sort は必ず先に行う
df = df.sort("datetime")

df_5m = (
    df
    .sort("datetime")
    .group_by_dynamic(
        index_column="datetime",
        every="5m",       # 5分ごとに窓を切り出す
        closed="left",    # 窓の左端を含む（右端は含まない）
        label="left",     # 窓の開始時刻をラベルとする
        group_by=["symbol"],
    )
    .agg([
        pl.col("open").first(),
        pl.col("high").max(),
        pl.col("low").min(),
        pl.col("close").last(),
        pl.col("volume").sum(),
    ])
    .sort(["symbol", "datetime"])
)

print(df_5m.head(3))

def resample_ohlcv(df: pl.DataFrame, every: str) -> pl.DataFrame:
    """OHLCV 1分足を指定間隔にリサンプリングする"""
    return (
        df
        .sort("datetime")
        .group_by_dynamic(
            index_column="datetime",
            every=every,
            closed="left",
            label="left",
            group_by=["symbol"],
        )
        .agg([
            pl.col("open").first(),
            pl.col("high").max(),
            pl.col("low").min(),
            pl.col("close").last(),
            pl.col("volume").sum(),
            # VWAP = sum(close * volume) / sum(volume)
            (pl.col("close") * pl.col("volume")).sum().alias("price_vol"),
        ])
        .with_columns(
            (pl.col("price_vol") / pl.col("volume")).alias("vwap")
        )
        .drop("price_vol")
        .sort(["symbol", "datetime"])
    )

df_5m  = resample_ohlcv(df, "5m")
df_1h  = resample_ohlcv(df, "1h")

# 1分グリッドに揃える（欠落した分はnullで埋まる）
df_up = (
    df
    .sort("datetime")
    .upsample(time_column="datetime", every="1m", group_by=["symbol"])
)
print(df_up.filter(pl.col("open").is_null()).head(3))
# 欠損バーはopenからvolumeまで全カラムがnullになっている

df_filled = (
    df_up
    .sort(["symbol", "datetime"])
    .with_columns([
        # 価格系：直前の終値を引き継ぐ（forward_fill）
        pl.col("open").forward_fill().over("symbol"),
        pl.col("high").forward_fill().over("symbol"),
        pl.col("low").forward_fill().over("symbol"),
        pl.col("close").forward_fill().over("symbol"),
        # 出来高・約定数：ゼロで埋める（forward_fillすると出来高を捏造する）
        pl.col("volume").fill_null(0),
    ])
)

df_filled = df_filled.with_columns(
    pl.col("close").is_null().alias("is_gap_bar")  # 欠損フラグを残す
)

# over を付け忘れると銘柄をまたいで平均が漏れる（バグの温床）
df_features = df_filled.with_columns([
    # ✅ 正しい：銘柄ごとに移動平均を計算
    pl.col("close").rolling_mean(window_size=20).over("symbol").alias("ma20"),
    pl.col("close").rolling_mean(window_size=60).over("symbol").alias("ma60"),

    # ✅ リターン（1期前比）
    (pl.col("close") / pl.col("close").shift(1).over("symbol") - 1).alias("return_1m"),

    # ✅ 実現ボラティリティ（20期）
    pl.col("return_1m").rolling_std(window_size=20).over("symbol").alias("realized_vol"),

    # ✅ 出来高比（直近20本の平均比）
    (pl.col("volume") / pl.col("volume").rolling_mean(20).over("symbol")).alias("vol_ratio"),
])

約定時刻: 09:01:03

板スナップショット: 09:00:55 / 09:01:05 / 09:01:30

backward（過去の最近値）: 09:00:55  ← ✅ 実務のデフォルト
forward （未来の最近値）: 09:01:05  ← ❌ 未来を参照（look-ahead bias）
nearest （前後で近い方） : 09:01:05  ← ❌ 未来を含みうる

# サンプル：約定データと板スナップショットを結合
df_trades = pl.DataFrame({
    "datetime": [
        datetime(2024, 1, 15, 9, 1, 3, tzinfo=timezone.utc),
        datetime(2024, 1, 15, 9, 2, 45, tzinfo=timezone.utc),
    ],
    "symbol":   ["BTC/USDT", "BTC/USDT"],
    "trade_price": [50100.0, 50200.0],
}).with_columns(pl.col("datetime").dt.cast_time_unit("ms"))

df_orderbook = pl.DataFrame({
    "datetime": [
        datetime(2024, 1, 15, 9, 0, 55, tzinfo=timezone.utc),
        datetime(2024, 1, 15, 9, 1, 5,  tzinfo=timezone.utc),
        datetime(2024, 1, 15, 9, 2, 30, tzinfo=timezone.utc),
    ],
    "symbol":   ["BTC/USDT"] * 3,
    "mid_price": [50090.0, 50110.0, 50190.0],
}).with_columns(pl.col("datetime").dt.cast_time_unit("ms"))

# ✅ 正しい実装：直前の板を取得、5秒以上古ければnull
result = df_trades.join_asof(
    df_orderbook.sort("datetime"),
    on="datetime",
    by="symbol",
    strategy="backward",   # 過去の最近値のみ
    tolerance="5s",        # 5秒以上古い板はnullにする
)
print(result)

# ❌ 誤った実装：未来の板を参照してしまう
result_bad = df_trades.join_asof(
    df_orderbook.sort("datetime"),
    on="datetime",
    by="symbol",
    strategy="forward",   # ← 未来の値を掴む
)

# ✅ 正しい実装
result_good = df_trades.join_asof(
    df_orderbook.sort("datetime"),
    on="datetime",
    by="symbol",
    strategy="backward",
    tolerance="5s",
)

# 比較：約定時刻09:01:03に対して
# bad  → mid_price=50110.0（09:01:05の板、未来の情報）
# good → mid_price=50090.0（09:00:55の板、正しい過去の情報）
print(result_bad.select(["datetime", "mid_price"]))
print(result_good.select(["datetime", "mid_price"]))

# ❌ label="left"のままリサンプル結果を特徴量に使う（最も気づきにくいleak）
df_5m_bad = (
    df.sort("datetime")
    .group_by_dynamic("datetime", every="5m", closed="left", label="left", group_by=["symbol"])
    .agg(pl.col("close").last().alias("close_5m"))
)

# datetime=09:00 の close_5m は 09:00〜09:04:59 の終値（5分後の情報！）
# これを09:00時点の特徴量として使うと完全にリーク

# ✅ 正しい対処：label="right"にして窓の終端時刻をラベルにする
df_5m_good = (
    df.sort("datetime")
    .group_by_dynamic("datetime", every="5m", closed="left", label="right", group_by=["symbol"])
    .agg(pl.col("close").last().alias("close_5m"))
)
# datetime=09:05 の close_5m は 09:00〜09:04:59 の終値（正しい）

# 簡易リークテスト：全特徴量を1期シフトして成績差を確認
# シフト前後で大幅に成績が変わる = リークの疑い大
df_no_leak = df_features.with_columns([
    pl.col("ma20").shift(1).over("symbol").alias("ma20_safe"),
    pl.col("close_5m").shift(1).over("symbol").alias("close_5m_safe"),
])

def build_feature_table(df_raw: pl.DataFrame) -> pl.DataFrame:
    """
    CCXT取得済みの1分足OHLCVから特徴量テーブルを構築する
    記事232〜237の予測モデルへそのまま渡せる形で出力
    """
    # 1. 時刻型の統一
    df = (
        df_raw
        .with_columns(
            pl.from_epoch("timestamp_ms", time_unit="ms")
              .dt.replace_time_zone("UTC")
              .alias("datetime")
        )
        .sort(["symbol", "datetime"])
    )

    # 2. 欠損バーの補完
    df = (
        df
        .upsample(time_column="datetime", every="1m", group_by=["symbol"])
        .sort(["symbol", "datetime"])
        .with_columns([
            pl.col("open").forward_fill().over("symbol"),
            pl.col("high").forward_fill().over("symbol"),
            pl.col("low").forward_fill().over("symbol"),
            pl.col("close").forward_fill().over("symbol"),
            pl.col("volume").fill_null(0),
        ])
    )

    # 3. 5分足・1時間足を生成してjoin_asof（label="right"でleak防止）
    df_5m = (
        df.sort("datetime")
        .group_by_dynamic("datetime", every="5m", closed="left", label="right", group_by=["symbol"])
        .agg(pl.col("close").last().alias("close_5m"), pl.col("volume").sum().alias("vol_5m"))
    )

    df = df.join_asof(
        df_5m.sort("datetime"),
        on="datetime",
        by="symbol",
        strategy="backward",  # 過去の5分足のみ参照
        tolerance="5m",
    )

    # 4. 特徴量を追加
    df = df.with_columns([
        pl.col("close").rolling_mean(20).over("symbol").alias("ma20"),
        pl.col("close").rolling_mean(60).over("symbol").alias("ma60"),
        (pl.col("close") / pl.col("close").shift(1).over("symbol") - 1).alias("return_1m"),
        pl.col("close").rolling_std(20).over("symbol").alias("vol_20m"),
    ])

    return df

# 実行
df_features = build_feature_table(df_1m)
print(df_features.head())