from dash import Input, Output, State
from datetime import timedelta
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from models.backtest import backtest

def register_backtest_callbacks(app):
    @app.callback(
        [Output('backtest-results', 'data'),
         Output('combined-heatmap', 'figure'),
         Output('forwardtest-results', 'data')],
        Input('run-backtest', 'n_clicks'),
        [State('bt-end-date', 'date'),
         State('window-start', 'value'),
         State('window-end', 'value'),
         State('window-step', 'value'),
         State('threshold-start', 'value'),
         State('threshold-end', 'value'),
         State('threshold-step', 'value'),
         State('time-factor', 'value'),
         State('commission', 'value'),
         State('heatmap-title', 'value')],  # Add new input for title
        prevent_initial_call=True
    )
    def run_backtest_and_create_heatmaps(n_clicks, bt_end_date, window_start, window_end, window_step,
                                         threshold_start, threshold_end, threshold_step, time_factor,
                                         commission, heatmap_title):
        if n_clicks == 0:
            return {}, {}, {}

        try:
            # Get data from the global scope
            from app import full_data

            # Parse inputs
            bt_end_date = pd.to_datetime(bt_end_date)
            ft_start_date = bt_end_date + timedelta(days=1)
            commission = commission / 100

            # Create window and threshold lists
            window_list = list(range(window_start, window_end, window_step))
            threshold_steps = int((threshold_end - threshold_start) / threshold_step)
            threshold_list = [round(threshold_start + i * threshold_step, 2)
                             for i in range(threshold_steps)]

            # Split data
            bt_df = full_data[:bt_end_date]
            ft_df = full_data[ft_start_date:]

            # Run backtests
            bt_results = [backtest(bt_df, w, t, time_factor, commission, 'heatmap')
                          for w in window_list for t in threshold_list]
            ft_results = [backtest(ft_df, w, t, time_factor, commission, 'heatmap')
                          for w in window_list for t in threshold_list]

            # Reverse threshold list for display
            threshold_list.reverse()

            # Process results
            bt_results_df = pd.DataFrame(bt_results)
            ft_results_df = pd.DataFrame(ft_results)

            bt_data_table = bt_results_df.pivot(index='threshold', columns='window', values='sharpe')
            ft_data_table = ft_results_df.pivot(index='threshold', columns='window', values='sharpe')

            # Sort tables
            bt_data_table = bt_data_table.sort_index(ascending=False)
            ft_data_table = ft_data_table.sort_index(ascending=False)

            # Use custom title or default if none provided
            bt_title = f"{heatmap_title}_BT_sharpe" if heatmap_title else "Backtest Sharpe Ratio"
            ft_title = f"{heatmap_title}_FT_sharpe" if heatmap_title else "Forward Test Sharpe Ratio"

            # Create heatmap figure
            fig = make_subplots(
                rows=1, cols=2,
                subplot_titles=(bt_title, ft_title),
                horizontal_spacing=0.17,
                column_widths=[0.5, 0.5]
            )

            # Custom colorscale
            custom_colorscale = [
                [0.0, '#FFFFBF'],  # Light yellow
                [0.2, '#E0F3F8'],  # Pale blue
                [0.4, '#ABD9E9'],  # Very light blue
                [0.6, '#74ADD1'],  # Light blue
                [0.8, '#4575B4'],  # Blue
                [1.0, '#313695']   # Deep blue
            ]

            # Add backtest heatmap
            fig.add_trace(
                go.Heatmap(
                    z=bt_data_table.values,
                    x=bt_data_table.columns,
                    y=np.round(bt_data_table.index, 2),
                    text=np.round(bt_data_table.values, 2),
                    texttemplate='%{text}',
                    textfont={"size": 9},
                    colorscale=custom_colorscale,
                    showscale=True,
                    colorbar=dict(
                        x=0.45,
                        thickness=15,
                        title="Sharpe Ratio",
                        titleside="right",
                        len=0.9
                    )
                ),
                row=1, col=1
            )

            # Add forward test heatmap
            fig.add_trace(
                go.Heatmap(
                    z=ft_data_table.values,
                    x=ft_data_table.columns,
                    y=np.round(ft_data_table.index, 2),
                    text=np.round(ft_data_table.values, 2),
                    texttemplate='%{text}',
                    textfont={"size": 9},
                    colorscale=custom_colorscale,
                    showscale=True,
                    colorbar=dict(
                        x=1.02,
                        thickness=15,
                        title="Sharpe Ratio",
                        titleside="right",
                        len=0.9
                    )
                ),
                row=1, col=2
            )

            # Calculate dimensions
            cell_height = 25
            cell_width = 15
            height = max(cell_height * len(threshold_list), 500)
            width = max(cell_width * len(window_list) * 3.2, 1800)

            # Update layout
            fig.update_layout(
                height=height,
                width=width,
                showlegend=False,
                title_x=0.5,
                margin=dict(l=50, r=120, t=50, b=50),
                paper_bgcolor='white',
                plot_bgcolor='white',
                font=dict(size=12)
            )

            # Update axes
            for i in range(1, 3):
                fig.update_xaxes(
                    title_text="Window",
                    tickmode='array',
                    tickvals=window_list,
                    ticktext=window_list,
                    tickangle=45,
                    row=1, col=i
                )
                fig.update_yaxes(
                    title_text="Threshold",
                    tickmode='array',
                    tickvals=threshold_list,
                    ticktext=[f"{x:.2f}" for x in threshold_list],
                    row=1, col=i
                )

            return bt_results_df.to_dict('records'), fig, ft_results_df.to_dict('records')

        except Exception as e:
            print(f"Error in heatmap generation: {str(e)}")
            return {}, {}, {}