import marimo

__generated_with = "0.19.6"
app = marimo.App(width="full", auto_download=["ipynb"], sql_output="pandas")


@app.cell
def _():
    import marimo as mo
    return (mo,)


@app.cell
def _():
    import urllib
    import sqlalchemy

    host = "81.71.3.24"
    port = 6785
    username = "leopard"
    password = urllib.parse.quote_plus("9NEzFzovnddf@PyEP?e*AYAWnCyd7UhYwQK$pJf>7?ccFiN^x4$eKEZ5~E<7<+~X")
    database = "leopard_dev"
    engine = sqlalchemy.create_engine(f"postgresql://{username}:{password}@{host}:{port}/{database}")
    return (engine,)


@app.cell
def _(engine, mo):
    dailies_df = mo.sql(
        f"""
        select trade_date, open * factor as "Open", close * factor as "Close", high * factor as "High", low * factor as "Low", volume as "Volume", coalesce(factor, 1.0) as factor
        from leopard_daily daily
                 left join leopard_stock stock on stock.id = daily.stock_id
        where stock.code = '000001.SZ'
          and daily.trade_date between '2024-01-01 00:00:00' and '2025-12-31 23:59:59'
        order by daily.trade_date
        """,
        engine=engine
    )
    return (dailies_df,)


@app.cell
def _(dailies_df):
    import pandas as pd

    dailies_df["trade_date"] = pd.to_datetime(dailies_df["trade_date"], format='%Y-%m-%d')
    dailies_df.set_index("trade_date", inplace=True)
    dailies_df
    return


@app.cell
def _(dailies_df):
    import talib

    dailies_df['sma30'] = talib.SMA(dailies_df['Close'], timeperiod=30)
    dailies_df['sma60'] = talib.SMA(dailies_df['Close'], timeperiod=60)
    dailies_df['sma120'] = talib.SMA(dailies_df['Close'], timeperiod=120)

    macd, signal, hist = talib.MACD(dailies_df["Close"], fastperiod=10, slowperiod=20, signalperiod=9)
    dailies_df['macd'] = macd
    dailies_df['signal'] = signal
    dailies_df['hist'] = hist

    target_dailies_df = dailies_df.loc['2025-01-01':'2025-12-31']
    target_dailies_df
    return (target_dailies_df,)


@app.cell
def _(target_dailies_df):
    target_dailies_df.reset_index(inplace=True)
    target_dailies_df_inc = target_dailies_df[target_dailies_df["Close"] > target_dailies_df["Open"]]
    target_dailies_df_dec = target_dailies_df[target_dailies_df["Close"] < target_dailies_df["Open"]]
    return target_dailies_df_dec, target_dailies_df_inc


@app.cell
def daily_chart(
        target_dailies_df,
        target_dailies_df_dec,
        target_dailies_df_inc,
):
    from bokeh.io import show, output_notebook
    from bokeh.layouts import column
    from bokeh.plotting import figure

    output_notebook()

    width = 1200
    up_color = "black"
    down_color = "grey"

    price_figure = figure(width=width, height=400, tools='pan,wheel_zoom,box_zoom,reset')

    price_figure.min_border = 0

    x_padding = 1
    y_padding = 10
    price_figure.x_range.bounds = (min(target_dailies_df.index) - x_padding, max(target_dailies_df.index) + x_padding)
    price_figure.y_range.bounds = (min(target_dailies_df['Low']) - y_padding, max(target_dailies_df['High']) + y_padding)
    price_figure.xaxis.major_label_overrides = {i: date.strftime('%Y-%m-%d') for i, date in zip(target_dailies_df.index, target_dailies_df['trade_date'])}
    price_figure.segment(target_dailies_df.index, target_dailies_df['High'], target_dailies_df.index, target_dailies_df['Low'], color='black')
    price_figure.vbar(target_dailies_df_inc.index, 0.6, target_dailies_df_inc['Open'], target_dailies_df_inc['Close'], color=up_color)
    price_figure.vbar(target_dailies_df_dec.index, 0.6, target_dailies_df_dec['Open'], target_dailies_df_dec['Close'], color=down_color)
    price_figure.line(target_dailies_df.index, target_dailies_df['sma30'], color='orange')
    price_figure.line(target_dailies_df.index, target_dailies_df['sma60'], color='red')
    # 控制图表只能放大不能缩小
    macd_figure = figure(width=width, height=200, tools='pan,wheel_zoom,box_zoom,reset')
    macd_figure.x_range.bounds = (min(target_dailies_df.index) - x_padding, max(target_dailies_df.index) + x_padding)
    macd_figure.y_range.bounds = (min(target_dailies_df['macd']) - y_padding, max(target_dailies_df['macd']) + y_padding)
    macd_figure.line(target_dailies_df.index, target_dailies_df['macd'], color='orange')
    macd_figure.line(target_dailies_df.index, target_dailies_df['signal'], color='red')
    # Add MACD histogram bars for positive and negative values
    macd_positive = target_dailies_df[target_dailies_df['macd'] > 0]
    macd_negative = target_dailies_df[target_dailies_df['macd'] < 0]

    macd_figure.vbar(macd_positive.index, 0.6, 0, macd_positive['macd'], color=up_color)
    macd_figure.vbar(macd_negative.index, 0.6, 0, macd_negative['macd'], color=down_color)

    # Add zero line
    macd_figure.line(target_dailies_df.index, [0] * len(target_dailies_df), color='black', line_dash='dashed')
    # show(price_figure)
    # show(macd_figure)
    show(column(price_figure, macd_figure))
    return


if __name__ == "__main__":
    app.run()