import polars as pl
from typing import TypedDict


class IchimokuValues(TypedDict):
    conversion_line: pl.Series
    base_line: pl.Series
    leading_span1: pl.Series
    leading_span2: pl.Series
    lagging_span: pl.Series


def get_ichimoku_values(df: pl.DataFrame) -> IchimokuValues:
    # Decimal型はrolling操作がサポートされていないため、先にFloat64に変換
    high = df["high.amount"].cast(pl.Float64)
    low = df["low.amount"].cast(pl.Float64)
    close = df["close.amount"].cast(pl.Float64)

    # 転換線: 過去9日間の (Max + Min) / 2
    conversion_line = (high.rolling_max(9) + low.rolling_min(9)) / 2

    # 基準線: 過去26日間の (Max + Min) / 2
    base_line = (high.rolling_max(26) + low.rolling_min(26)) / 2

    # 先行スパン1: (転換線 + 基準線) / 2 を26日未来にずらす
    # Polarsのshiftはデフォルトで空いた部分をnullで埋めます
    leading_span1 = ((conversion_line + base_line) / 2).shift(26)

    # 先行スパン2: 過去52日間の (Max + Min) / 2 を26日未来にずらす
    leading_span2 = ((high.rolling_max(52) + low.rolling_min(52)) / 2).shift(26)

    # 遅行スパン: 今日の終値を26日過去にずらす
    lagging_span = close.shift(-26)

    return {
        "conversion_line": conversion_line,
        "base_line": base_line,
        "leading_span1": leading_span1,
        "leading_span2": leading_span2,
        "lagging_span": lagging_span,
    }

from datetime import datetime, timedelta

def generate_future_business_days(last_date: str, n_days: int = 26) -> list[str]:
    """
    最後の日付から未来の営業日（平日）を生成する

    Args:
        last_date: 最後のデータの日付（文字列形式: "YYYY-MM-DD"）
        n_days: 生成する営業日数（デフォルト26日）

    Returns:
        未来の営業日のリスト（文字列形式: "YYYY-MM-DD"）
    """
    # 最後の日付を日付オブジェクトに変換
    current = datetime.strptime(last_date, "%Y-%m-%d")
    future_dates = []

    # 必要な営業日数が揃うまで1日ずつ進める
    while len(future_dates) < n_days:
        current += timedelta(days=1)
        # 営業日（月〜金）のみ追加（weekday(): 月=0, 金=4）
        if current.weekday() < 5:
            future_dates.append(current.strftime("%Y-%m-%d"))

    return future_dates

import numpy as np
import plotly.graph_objects as go

def create_cloud_segments(dates, span1, span2):
    """
    一目均衡表の雲を陰転・陽転に応じて色分けしたセグメントに分割する
    """
    # Polars SeriesをNumPy配列に変換
    if hasattr(span1, "to_numpy"):
        span1 = span1.to_numpy()
    if hasattr(span2, "to_numpy"):
        span2 = span2.to_numpy()

    # 陽転判定（先行スパン1 > 先行スパン2）
    is_bullish = np.nan_to_num(span1 > span2, nan=False)

    # クロスオーバーポイントを検出（陰転⇔陽転が切り替わる箇所）
    crossovers = np.where(np.diff(is_bullish.astype(int)) != 0)[0] + 1

    # セグメントの境界を作成
    boundaries = np.concatenate([[0], crossovers, [len(dates)]])

    traces = []

    # 各セグメントごとに処理
    for i in range(len(boundaries) - 1):
        start_idx = boundaries[i]
        end_idx = boundaries[i + 1]

        segment_dates = dates[start_idx:end_idx]
        segment_span1 = span1[start_idx:end_idx]
        segment_span2 = span2[start_idx:end_idx]

        # NaNのみのセグメントはスキップ
        if np.all(np.isnan(segment_span1)) or np.all(np.isnan(segment_span2)):
            continue

        # セグメントの陽転・陰転を判定して色を設定
        segment_is_bullish = is_bullish[start_idx]
        fillcolor = (
            "rgba(135, 206, 250, 0.3)"  # 陽転: ライトブルー
            if segment_is_bullish
            else "rgba(255, 165, 0, 0.3)"  # 陰転: オレンジ
        )

        # 先行スパン1のライン（塗りつぶしなし）
        traces.append(
            go.Scatter(
                x=segment_dates,
                y=segment_span1,
                mode="lines",
                line=dict(width=0.5, color="green"),
                showlegend=False,
                hoverinfo="skip",
            )
        )

        # 先行スパン2のライン（先行スパン1との間を塗りつぶし）
        traces.append(
            go.Scatter(
                x=segment_dates,
                y=segment_span2,
                mode="lines",
                line=dict(width=0.5, color="orange"),
                fill="tonexty",  # 前のトレース（先行スパン1）との間を塗りつぶし
                fillcolor=fillcolor,
                showlegend=False,
                hoverinfo="skip",
            )
        )

    return traces

def create_cloud_segments_with_future(dates, span1, span2, num_future_days=26):
    """
    一目均衡表の雲（未来の雲を含む）を陰転・陽転に応じて色分けしたセグメントに分割する
    """
    # 未来の日付を生成
    last_date = dates[-1]
    future_dates = generate_future_business_days(last_date, num_future_days)

    # 日付を結合（過去 + 未来）
    all_dates = dates + future_dates

    # Polars SeriesをNumPy配列に変換
    if hasattr(span1, "to_numpy"):
        span1_values = span1.to_numpy()
    else:
        span1_values = span1
    if hasattr(span2, "to_numpy"):
        span2_values = span2.to_numpy()
    else:
        span2_values = span2

    # 先行スパンの配列を拡張（先頭26個をNaNで埋める）
    span1_extended = np.concatenate([[np.nan] * num_future_days, span1_values])
    span2_extended = np.concatenate([[np.nan] * num_future_days, span2_values])

    # 過去部分と未来部分を分離
    past_dates = dates
    future_span1 = span1_extended[-num_future_days:]
    future_span2 = span2_extended[-num_future_days:]
    past_span1 = span1_extended[:-num_future_days]
    past_span2 = span2_extended[:-num_future_days]

    # 過去部分の陽転判定
    is_bullish = np.nan_to_num(past_span1 > past_span2, nan=False)

    # クロスオーバーポイントを検出
    crossovers = np.where(np.diff(is_bullish.astype(int)) != 0)[0] + 1

    # セグメントの境界を作成
    boundaries = np.concatenate([[0], crossovers, [len(past_dates)]])

    traces = []
    last_segment_color = None

    # 過去部分の各セグメントを処理
    for i in range(len(boundaries) - 1):
        start_idx = boundaries[i]
        end_idx = boundaries[i + 1]

        segment_dates = past_dates[start_idx:end_idx]
        segment_span1 = past_span1[start_idx:end_idx]
        segment_span2 = past_span2[start_idx:end_idx]

        if np.all(np.isnan(segment_span1)) or np.all(np.isnan(segment_span2)):
            continue

        # セグメントの陽転・陰転を判定して色を設定
        segment_is_bullish = is_bullish[start_idx]
        fillcolor = (
            "rgba(135, 206, 250, 0.3)"  # 陽転: ライトブルー
            if segment_is_bullish
            else "rgba(255, 165, 0, 0.3)"  # 陰転: オレンジ
        )
        last_segment_color = fillcolor

        # 先行スパン1のライン
        traces.append(
            go.Scatter(
                x=segment_dates,
                y=segment_span1,
                mode="lines",
                line=dict(width=0.5, color="green"),
                showlegend=False,
                hoverinfo="skip",
            )
        )

        # 先行スパン2のライン
        traces.append(
            go.Scatter(
                x=segment_dates,
                y=segment_span2,
                mode="lines",
                line=dict(width=0.5, color="orange"),
                fill="tonexty",
                fillcolor=fillcolor,
                showlegend=False,
                hoverinfo="skip",
            )
        )

    # 未来の雲を追加（最後のセグメントの色を継続）
    if last_segment_color and not np.all(np.isnan(future_span1)):
        traces.append(
            go.Scatter(
                x=future_dates,
                y=future_span1,
                mode="lines",
                line=dict(width=0.5, color="green", dash="dot"),  # 点線で未来を示す
                showlegend=False,
                hoverinfo="skip",
            )
        )

        traces.append(
            go.Scatter(
                x=future_dates,
                y=future_span2,
                mode="lines",
                line=dict(width=0.5, color="orange", dash="dot"),  # 点線で未来を示す
                fill="tonexty",
                fillcolor=last_segment_color,  # 最後のセグメントの色を継続
                showlegend=False,
                hoverinfo="skip",
            )
        )

    return traces

import polars as pl
import yfinance_pl as yf

# データ取得
ticker = yf.Ticker("8381.T")
hist = ticker.history(period="6mo")

# 一目均衡表の計算
ichimoku = get_ichimoku_values(hist)

# DataFrameに追加
df_plot = hist.with_columns(
    conversion_line=ichimoku["conversion_line"],
    base_line=ichimoku["base_line"],
    leading_span1=ichimoku["leading_span1"],
    leading_span2=ichimoku["leading_span2"],
    lagging_span=ichimoku["lagging_span"],
).with_columns(
    [
        pl.col("open.amount").cast(pl.Float64),
        pl.col("high.amount").cast(pl.Float64),
        pl.col("low.amount").cast(pl.Float64),
        pl.col("close.amount").cast(pl.Float64),
    ]
)

dates = df_plot["date"].to_list()

# チャート作成
data = []

# ローソク足
data.append(
    go.Candlestick(
        x=dates,
        open=df_plot["open.amount"],
        high=df_plot["high.amount"],
        low=df_plot["low.amount"],
        close=df_plot["close.amount"],
        increasing_line_color="red",
        decreasing_line_color="green",
        name="株価",
    )
)

# 転換線
data.append(
    go.Scatter(
        x=dates,
        y=df_plot["conversion_line"],
        mode="lines",
        name="転換線",
        line=dict(color="blue", width=1),
    )
)

# 基準線
data.append(
    go.Scatter(
        x=dates,
        y=df_plot["base_line"],
        mode="lines",
        name="基準線",
        line=dict(color="red", width=1),
    )
)

# 雲（陰転・陽転の色分け + 未来の雲）
cloud_traces = create_cloud_segments_with_future(
    dates,
    df_plot["leading_span1"],
    df_plot["leading_span2"],
    num_future_days=26,
)
data.extend(cloud_traces)

# 遅行スパン
data.append(
    go.Scatter(
        x=dates,
        y=df_plot["lagging_span"],
        mode="lines",
        name="遅行スパン",
        line=dict(color="purple", width=1, dash="dot"),
    )
)

# レイアウト設定
layout = go.Layout(
    title="株価と一目均衡表",
    xaxis_title="日付",
    yaxis_title="価格",
    xaxis_rangeslider_visible=False,
)

fig = go.Figure(data=data, layout=layout)
fig.show()