# -*- coding: utf-8 -*-
"""
---------------------------------------------
Created on 2024/12/9 下午1:29
@author: ZhangYundi
@email: yundi.xxii@outlook.com
---------------------------------------------
"""

import logair
import numpy as np
import pandas as pd
import plotly.figure_factory as ff
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from .options import Options
from .utils import guess_plotly_rangebreaks
from .. import empyrical

logger = logair.get_logger("silars.plot")

_layout = dict(template='gridon',
               hovermode='x unified',
               hoverlabel=dict(bgcolor='rgba(255, 255, 255, 0.5)'),
               legend=dict(
                   orientation="h",  # 垂直排列
                   yanchor="bottom",  # 图例的锚点在上
                   y=1.02,  # 图例的 y 位置
                   # xanchor="center",  #
                   xanchor="right",  # 靠右
                   x=1,
                   bordercolor='lightgrey',  # 边框颜色
                   borderwidth=1,  # 边框宽度
                   bgcolor='rgba(255, 255, 255, 0.8)',  # 背景颜色
               ), )
_axis = dict(xaxis=dict(linecolor="black",
                        linewidth=1,
                        mirror=True,
                        gridcolor='lightgrey',
                        gridwidth=1,
                        griddash='dot',
                        showspikes=True,  # 启用 x 轴的十字光标
                        spikemode='across',  # 十字光标模式
                        spikecolor='black',  # 十字光标颜色
                        spikethickness=1.5,  # 十字光标粗细
                        spikedash='dot',  # 十字光标样式
                        tickformat="%Y-%m-%d"),
             yaxis=dict(linecolor="black",
                        linewidth=1,
                        mirror=True,
                        gridcolor='lightgrey',
                        gridwidth=1.5,
                        griddash='dot',
                        showspikes=True,  # 启用 x 轴的十字光标
                        spikemode='across',  # 十字光标模式
                        spikecolor='black',  # 十字光标颜色
                        spikethickness=1,  # 十字光标粗细
                        spikedash='dot',  # 十字光标样式
                        tickformat="%Y-%m-%d"),
             # height=500,
             )


def table(tb_data: pd.DataFrame, highlight_cols: list[str] = None):
    """绘制表格"""
    cols = tb_data.columns
    cellvalues = [[f'<b>{val}</b>' for val in tb_data[cols[0]]]]
    if highlight_cols is not None:
        col_dict = {k: 1 for k in highlight_cols}
        for col in cols[1:]:
            if col in col_dict:
                cellvalues.append([f'<b>{val}</b>' for val in tb_data[col].tolist()])
            else:
                cellvalues.append(tb_data[col])
    else:
        for col in cols[1:]:
            cellvalues.append(tb_data[col].tolist())
    tb = go.Table(
        header=dict(
            values=[f"<b>{col}</b>" for col in tb_data.columns],
            **Options.Table.header,
        ),
        cells=dict(
            values=cellvalues,  # [d.tolist() for g, d in tb_data.items()],
            **Options.Table.cell,
            fill=dict(color=['gold', 'white']),
        ),
    )
    return go.Figure(data=[tb], )


def distplot(data: pd.Series | pd.DataFrame, bin_size, title: str = None):
    """分布图: 每一列都是一组数据集"""
    if isinstance(data, pd.Series):
        data = data.to_frame()
    hist_data = list()
    group_labels = list()
    for col_name, col_data in data.items():
        hist_data.append(col_data.dropna())
        group_labels.append(col_name)
    fig = ff.create_distplot(hist_data, group_labels, bin_size=bin_size, colors=Options.colors)
    fig.update_layout(**Options.layout, )
    if title is not None:
        fig.update_layout(title=title)
    if data.index.name == "date":
        fig.update_xaxes(hoverformat='%Y-%m-%d',
                         tickformat="%Y-%m-%d", )
    return fig


def violin(data: pd.Series | pd.DataFrame, title: str = None):
    """index是x-axis, columns是datasets"""
    if isinstance(data, pd.Series):
        data = data.to_frame()
    fig = go.Figure()
    x_axis = data.index
    cols = data.columns
    for i, name in enumerate(cols):
        dataset = data[name]
        fig.add_trace(
            go.Violin(
                x=x_axis,
                y=dataset,
                name=name,
                marker_color=Options.colors[i % len(Options.colors)],
                box_visible=True,
                meanline_visible=True,
            )
        )

    fig.update_layout(**Options.layout, )
    if title is not None:
        fig.update_layout(title=title)
    if data.index.name == "date":
        fig.update_xaxes(hoverformat='%Y-%m-%d', tickformat="%Y-%m-%d")
    return fig


def bar(data: pd.Series | pd.DataFrame, title: str = None):
    """index是x-axis, columns是datasets"""
    if isinstance(data, pd.Series):
        data = data.to_frame()
    fig = go.Figure()
    x_axis = data.index
    cols = data.columns
    for i, name in enumerate(cols):
        dataset = data[name]
        fig.add_trace(
            go.Bar(
                x=x_axis,
                y=dataset,
                name=name,
                marker_color=Options.colors[i % len(Options.colors)],
            )
        )

    fig.update_layout(**Options.layout, )
    if title is not None:
        fig.update_layout(title=title)
    if data.index.name == "date":
        data.index = pd.to_datetime(data.index)
        fig.update_xaxes(hoverformat='%Y-%m-%d',
                         tickformat="%Y-%m-%d",
                         rangebreaks=guess_plotly_rangebreaks(data.index))
    return fig


def lines(data: pd.Series | pd.DataFrame, title: str = None):
    if isinstance(data, pd.Series):
        data = data.to_frame()
    fig = go.Figure()
    x_axis = data.index
    cols = data.columns
    for i, name in enumerate(cols):
        dataset = data[name]
        fig.add_trace(
            go.Scatter(
                x=x_axis,
                y=dataset,
                name=name,
                mode="lines",
                marker_color=Options.colors[i % len(Options.colors)],
            )
        )

    fig.update_layout(**Options.layout, )
    if title is not None:
        fig.update_layout(title=title)
    if data.index.name == "date":
        fig.update_xaxes(hoverformat='%Y-%m-%d',
                         tickformat="%Y-%m-%d", )
    return fig


def nv_plot(data: pd.Series, starting_nv: float = 1.0):
    """
    绘制净值曲线
    Parameters
    ----------
    data: pd.Series
        累计净值, index: date
    starting_nv: float
        初始净值，默认 1

    Returns
    -------

    """
    if not isinstance(data, pd.Series):
        logger.error(f"Input data must be pandas.Series, not {type(data)}")
        return

    # 收益数据
    ret = data.pct_change().fillna(data[0] / starting_nv - 1)
    # 计算指标
    ann_ret = empyrical.annual_return(ret)
    cum_ret = data[-1] - starting_nv
    ann_sd = empyrical.annual_volatility(ret)
    running_max = np.maximum.accumulate(data)
    underwater = (data - running_max) / running_max
    underwater = underwater.round(4)
    mdd = underwater.min()
    sharpe = ann_ret / ann_sd
    calmar = np.inf if mdd >= 0 else -ann_ret / mdd

    # if isinstance(data, pd.Series):
    #     data = data.to_frame()
    x_axis = data.index
    fig = go.Figure()
    fig.add_trace(
        go.Scatter(
            x=x_axis,
            y=data,
            name=data.name,
            mode="lines",
            # marker_color="rgba(255, 0, 0, 0.8)",
            line=dict(width=2, color='rgba(200, 0, 0, 1)', ),
            # line=dict(width=3, color='blue', ),
        )
    )
    # 添加最大回撤曲线
    fig.add_trace(
        go.Scatter(x=underwater.index,
                   y=underwater,
                   fill='tozeroy',
                   name='动态回撤',
                   fillcolor='rgba(199, 21, 133, 0.4)',
                   line=dict(width=0.5, color='rgba(199, 21, 133, 0.7)', ),
                   hovertemplate='%{fullData.name}:%{y:.2%}<extra></extra>',
                   yaxis="y2",
                   ),
        # secondary_y=True,
    )
    metric = dict(ann_ret=np.round(ann_ret, 3),
                  cum_ret=np.round(cum_ret, 3),
                  ann_sd=np.round(ann_sd, 3),
                  mdd=np.round(mdd, 5),
                  sharpe=np.round(sharpe, 2),
                  calmar=np.round(calmar, 2), )
    fig.add_annotation(
        xref='paper', yref='paper',
        x=0.02, y=2 / 3,
        xanchor='left', yanchor='middle',
        text=f'年化收益: {ann_ret * 100:.2f}% <br>'
             f'累计收益: {cum_ret * 100:.2f}% <br>'
             f'年化波动: {ann_sd * 100:.2f}% <br>'
             f'最大回撤: {mdd * 100:.2f}% <br>'
             f'sharpe: {sharpe:.2f} <br>'
             f'calmar: {calmar:.2f} <br>',
        font=dict(family='Arial, sans-serif',
                  size=12,
                  color='black'),
        align='left',
        borderwidth=2,
        borderpad=5,
        bordercolor='rgba(220, 220, 220, 0.5)',
        showarrow=False, textangle=0,
    )
    Options.layout.update({"title": dict(text=f"Backtest Report",
                                         x=0.05,
                                         y=0.98,
                                         yanchor="top",
                                         xanchor="left",
                                         font=dict(
                                             family="Courier New, monospace",
                                             size=25,
                                             # color="RebeccaPurple",
                                             variant="small-caps",
                                             weight="bold",
                                         ))})
    Options.layout.update({"legend": dict(orientation="h",
                                          yanchor="bottom",
                                          y=1.02,
                                          xanchor="right",
                                          x=1,
                                          bordercolor='lightgrey',  # 边框颜色
                                          borderwidth=1,  # 边框宽度
                                          bgcolor='rgba(255, 255, 255, 0.8)', )})
    fig.update_layout(**Options.layout, yaxis2=dict(overlaying="y",
                                                    side="right",
                                                    gridcolor='lightgrey',
                                                    gridwidth=1,
                                                    griddash='dot',
                                                    range=[2 * mdd, 0]))
    if data.index.name == "date":
        data.index = pd.to_datetime(data.index)
        fig.update_xaxes(hoverformat='%Y-%m-%d',
                         tickformat="%Y-%m-%d",
                         rangebreaks=guess_plotly_rangebreaks(data.index))
    return fig, metric


def signal_plot(data: pd.DataFrame, title: str = None):
    """
    买卖点绘制
    Parameters
    ----------
    data: pd.DataFrame
        包含 price、signal
    title: 标题

    Returns
    -------

    """
    fig = go.Figure()
    x_axis = data.index
    fig.add_trace(
        go.Scatter(
            x=x_axis,
            y=data["price"],
            name="price",
            mode="lines",
            marker_color="black",
        )
    )
    buy_signal = data[data["signal"] == 1]
    if not buy_signal.empty:
        fig.add_trace(
            go.Scatter(
                x=buy_signal.index,
                y=buy_signal["price"],
                name="B",
                mode="markers",
                marker=dict(
                    symbol="triangle-up",
                    color="red",
                    size=10,
                ),
            )
        )
    sell_signal = data[data["signal"] == -1]
    if not sell_signal.empty:
        fig.add_trace(
            go.Scatter(
                x=sell_signal.index,
                y=sell_signal["price"],
                name="S",
                mode="markers",
                marker=dict(
                    symbol="triangle-down",
                    color="green",
                    size=10,
                ),
            )
        )

    fig.update_layout(**Options.layout, )
    if title is not None:
        fig.update_layout(title=title)
    return fig

def nvs_plot(data: pd.DataFrame,
             color_configs: dict[str, str],
             mdd_name: str,
             starting_nv: float = 1.0):
    """
    绘制多条净值曲线
    Parameters
    ----------
    data: pd.DataFrame
        累计净值, index: date
    color_configs: dict
        颜色配置
    mdd_name: str
        用于绘制最大回撤的序列
    starting_nv: float
        初始净值，默认 1

    Returns
    -------

    """
    if not isinstance(data, pd.DataFrame):
        logger.error(f"Input data must be pandas.DataFrame, not {type(data)}")
        return

    # 收益数据
    s = data[mdd_name]
    ret = s.pct_change().fillna(0.0)
    # 计算指标
    ann_ret = empyrical.annual_return(ret)
    cum_ret = empyrical.cum_returns_final(ret)
    ann_sd = empyrical.annual_volatility(ret)
    running_max = np.maximum.accumulate(s)
    underwater = (s - running_max) / running_max
    underwater = underwater.round(4)
    mdd = underwater.min()
    sharpe = ann_ret / ann_sd
    calmar = np.inf if mdd >= 0 else -ann_ret / mdd

    # if isinstance(data, pd.Series):
    #     data = data.to_frame()
    # fig = go.Figure()
    row_heights = [2, 1, 1]
    rows = len(row_heights)
    fig = make_subplots(
        rows=rows, cols=1,
        vertical_spacing=0.15,
        specs=[
            [{"type": "xy", "secondary_y": True}],
            [{"type": "bar"}],
            [{"type": "table"}, ],
        ],
        subplot_titles=["", "年度收益", "",],
        row_heights=[h / sum(row_heights) for h in row_heights],
    )

    x_axis = data.index
    # visible = {name: "legendonly" if name != mdd_name else True for name in color_configs}
    width = {name: 2 if name != mdd_name else 2.5 for name in color_configs}

    # ======================== 累计收益曲线 ======================== 
    for line_name, line_color in color_configs.items():
        fig.add_trace(
            go.Scatter(
                x=x_axis,
                y=data[line_name],
                name=line_name,
                mode="lines",
                # visible=visible[line_name],
                line=dict(width=width[line_name], color=line_color, ),
                hovertemplate='%{fullData.name}:%{y:.2%}<extra></extra>',
            )
        )
    fig.add_trace(
        go.Scatter(
            x=x_axis,
            y=pd.Series([starting_nv] * len(x_axis)),
            name="base",
            mode="lines",
            line=dict(width=1, color="grey", ),
        ),
        row=1, col=1,
    )
    # 添加最大回撤曲线
    fig.add_trace(
        go.Scatter(x=underwater.index,
                   y=underwater,
                   fill='tozeroy',
                   name='动态回撤',
                   fillcolor='rgba(199, 21, 133, 0.3)', # 'rgba(199, 21, 133, 0.3)',
                   line=dict(width=0.5, color="rgba(199, 21, 133, 0.7)", ), # 'rgba(199, 21, 133, 1)'
                   hovertemplate='%{fullData.name}:%{y:.2%}<extra></extra>',
                   yaxis="y2",
                   ),
        secondary_y=True,
        row=1, col=1,
    )
    metric = dict(ann_ret=np.round(ann_ret, 3),
                  cum_ret=np.round(cum_ret, 3),
                  ann_sd=np.round(ann_sd, 3),
                  mdd=np.round(mdd, 5),
                  sharpe=np.round(sharpe, 2),
                  calmar=np.round(calmar, 2), )
    fig.add_annotation(
        xref='paper', yref='paper',
        x=0.02, y=0.9,
        xanchor='left', yanchor='middle',
        text=f'年化收益: {ann_ret * 100:.2f}% <br>'
             f'累计收益: {cum_ret * 100:.2f}% <br>'
             f'年化波动: {ann_sd * 100:.2f}% <br>'
             f'最大回撤: {mdd * 100:.2f}% <br>'
             f'sharpe: {sharpe:.2f} <br>'
             f'calmar: {calmar:.2f} <br>',
        font=dict(family='Arial, sans-serif',
                  size=12,
                  color='black'),
        align='left',
        borderwidth=2,
        borderpad=5,
        bordercolor='rgba(220, 220, 220, 0.5)',
        showarrow=False, textangle=0,
    )
    # ======================== 年度统计 ========================
    ret_df = data.pct_change().fillna(0.0)
    ret_df.index = pd.to_datetime(ret_df.index)
    # 分年度进行统计
    year_stat: pd.DataFrame = ret_df.groupby(ret_df.index.year).apply(lambda df: (df+1.0).prod()-1.0)
    tb_vals = [year_stat.index.tolist(),]
    for n, col in year_stat.items():
        tb_vals.append((col*100).round(2).tolist())
    # ======================== 年度统计表 ========================
    fig.add_trace(
        go.Table(
            header=dict(
                values=['Year', *year_stat.columns],
                font=dict(size=12, family="Arial, sans-serif", color='#ffffff'),
                fill_color='#3552a2',  # '#545859',
                align='center',
                line_width=0.5,
                line_color='rgba(252, 252, 252, 1)'
            ),
            cells=dict(
                values=tb_vals,
                align='center',
                font_color=[*[['black' for val in col] for col in tb_vals[:-1]],
                            ['green' if val < 0 else 'red' for val in tb_vals[-1]],
                            ],
                font=dict(family="Arial, sans-serif", size=12),
                line_width=0.5,
                line_color='rgba(252, 252, 252, 1)',
                fill_color='#dae1f5',  # 'rgba(0, 0, 0, 0)'
            ),
        ),
        row=3, col=1,
    )
    # ======================== 年度统计直方图 ========================
    for name, s in year_stat.items():
        fig.add_trace(
            go.Bar(
                x=s.index,
                y=s,
                name=name,
                marker_color=color_configs[name],
                hovertemplate='%{fullData.name}:%{y:.2%}<extra></extra>',
            ),
            row=2, col=1,
        )
    Options.layout.update({"title": dict(text=f"Backtest Report",
                                         x=0.05,
                                         y=0.98,
                                         yanchor="top",
                                         xanchor="left",
                                         font=dict(
                                             family="Courier New, monospace",
                                             size=25,
                                             # color="RebeccaPurple",
                                             variant="small-caps",
                                             weight="bold",
                                         ))})
    Options.layout.update({"legend": dict(orientation="h",
                                          yanchor="bottom",
                                          y=1.02,
                                          xanchor="right",
                                          x=1,
                                          bordercolor='lightgrey',  # 边框颜色
                                          borderwidth=1,  # 边框宽度
                                          bgcolor='rgba(255, 255, 255, 0.8)', )})
    layout = {k: v for k, v in Options.layout.items()}
    layout["height"] = rows * 350
    fig.update_layout(**layout, yaxis2=dict(overlaying="y",
                                            side="right",
                                            gridcolor='lightgrey',
                                            gridwidth=1,
                                            griddash='dot',
                                            range=[2 * mdd, 0]),)
    for row in [1, 2]:
        fig.update_xaxes(**_axis["xaxis"], row=row, col=1)
        fig.update_yaxes(**_axis["yaxis"], row=row, col=1)
    if data.index.name == "date":
        data.index = pd.to_datetime(data.index)
        fig.update_xaxes(hoverformat='%Y-%m-%d',
                         tickformat="%Y-%m-%d",
                         rangebreaks=guess_plotly_rangebreaks(data.index))
        # 添加滑块
        fig.update_xaxes(row=1, col=1,
                         rangeslider={'autorange': True, 'visible': True, 'thickness': 0.03})
    return fig, metric