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 = []

    # 필요한 영업일수가 채워질 때까지 하루씩 진행
    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()