"""
Backtesting Engine.

Simulates trading strategy execution on historical test data, calculates 
performance metrics, and generates reports and plots.
"""

import os
import logging
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
from typing import Dict, Any, Tuple, Optional

# Import required components (use absolute paths)
from gru_sac_predictor.src.sac_agent import SACTradingAgent
from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.calibrator import Calibrator

logger = logging.getLogger(__name__)


def calculate_sharpe_ratio(returns, periods_per_year=252*24*60): # Default for 1-min data
    """Calculate annualized Sharpe ratio from a series of returns."""
    returns = pd.Series(returns)
    if returns.std() == 0:
        return 0.0
    # Assuming risk-free rate is 0 for simplicity
    return np.sqrt(periods_per_year) * returns.mean() / returns.std()

def calculate_max_drawdown(equity_curve):
    """Calculate the maximum drawdown from an equity curve series."""
    equity_curve = pd.Series(equity_curve)
    rolling_max = equity_curve.cummax()
    drawdown = (equity_curve - rolling_max) / rolling_max
    max_drawdown = drawdown.min()
    return abs(max_drawdown) # Return positive value

class Backtester:
    """Runs the backtest simulation and generates results."""

    def __init__(self, config: dict):
        """
        Initialize the Backtester.
        Args:
            config (dict): Pipeline configuration dictionary, expected to contain 
                           'backtest' and potentially 'calibration', 'sac' sections.
        """
        self.config = config
        self.backtest_cfg = config.get('backtest', {})
        self.cal_cfg = config.get('calibration', {})
        self.sac_cfg = config.get('sac', {})
        
        self.initial_capital = self.backtest_cfg.get('initial_capital', 10000.0)
        self.transaction_cost = self.backtest_cfg.get('transaction_cost', 0.0005)
        self.edge_threshold = self.cal_cfg.get('edge_threshold', 0.55)
        
        self.results_df: Optional[pd.DataFrame] = None
        self.metrics: Optional[Dict[str, Any]] = None
        
        logger.info("Backtester initialized.")
        logger.info(f"  Initial Capital: {self.initial_capital:.2f}")
        logger.info(f"  Transaction Cost: {self.transaction_cost*100:.4f}%")
        logger.info(f"  Edge Threshold: {self.edge_threshold:.3f}")

    def run_backtest(
        self, 
        sac_agent_load_path: Optional[str],
        X_test_seq: np.ndarray,
        y_test_seq_dict: Dict[str, np.ndarray],
        test_indices: pd.Index, 
        gru_handler: GRUModelHandler,
        calibrator: Calibrator,
        original_prices: Optional[pd.Series] = None # Added for plotting
    ) -> Tuple[Optional[pd.DataFrame], Optional[Dict[str, Any]]]:
        """
        Executes the backtesting simulation loop.

        Args:
            sac_agent_load_path (Optional[str]): Path to load the SAC agent from. If None, uses untrained agent.
            X_test_seq (np.ndarray): Test set feature sequences.
            y_test_seq_dict (Dict[str, np.ndarray]): Dict of test set targets (e.g., 'ret', 'dir').
            test_indices (pd.Index): Timestamps corresponding to the test sequences' targets.
            gru_handler (GRUModelHandler): Instance to get GRU predictions.
            calibrator (Calibrator): Instance to calibrate probabilities.
            original_prices (Optional[pd.Series]): Original close prices aligned with test_indices (for plotting).

        Returns:
            Tuple[Optional[pd.DataFrame], Optional[Dict[str, Any]]]: 
                - DataFrame with detailed backtest results per step.
                - Dictionary containing calculated performance metrics.
                Returns (None, None) if backtest cannot run.
        """
        logger.info("--- Starting Backtest Simulation ---")

        if X_test_seq is None or y_test_seq_dict is None or test_indices is None:
            logger.error("Test sequence data (X, y, indices) is missing. Cannot run backtest.")
            return None, None
        if gru_handler.model is None:
             logger.error("GRU model is not loaded in the handler. Cannot run backtest.")
             return None
             
        # 1. Initialize SAC Agent & Load Weights
        agent_state_dim = self.sac_cfg.get('state_dim', 5)
        agent = SACTradingAgent(
            state_dim=agent_state_dim, # Ensure this matches the state construction below
            action_dim=1, # Typically -1 to 1 for position
            gamma=self.sac_cfg.get('gamma', 0.99), # These params might not matter much if just loading weights
            tau=self.sac_cfg.get('tau', 0.005),
            initial_lr=self.sac_cfg.get('actor_lr', 3e-4),
            alpha=self.sac_cfg.get('alpha', 0.2),
            alpha_auto_tune=self.sac_cfg.get('alpha_auto_tune', True),
            # Other SAC params...
            edge_threshold_config=self.edge_threshold # Store config threshold
        )
        if sac_agent_load_path and os.path.exists(sac_agent_load_path):
            logger.info(f"Loading SAC agent weights from: {sac_agent_load_path}")
            try:
                agent.load(sac_agent_load_path)
            except Exception as e:
                logger.error(f"Failed to load SAC agent weights from {sac_agent_load_path}: {e}. Proceeding with untrained agent.", exc_info=True)
        else:
            logger.warning(f"SAC agent load path not found or not specified ({sac_agent_load_path}). Proceeding with untrained agent.")
            
        # 2. Get GRU Predictions on Test Set
        logger.info(f"Generating GRU predictions for {len(X_test_seq)} test sequences...")
        predictions_test = gru_handler.predict(X_test_seq)
        if predictions_test is None or len(predictions_test) < 3:
            logger.error("Failed to get GRU predictions on test set. Cannot run backtest.")
            return None, None
        
        mu_test = predictions_test[0].flatten()
        log_sigma_test = predictions_test[1][:, 1].flatten()
        p_raw_test = predictions_test[2].flatten()
        sigma_test = np.exp(log_sigma_test)
        
        # Extract actual returns and directions
        actual_ret_test = y_test_seq_dict.get('ret')
        actual_dir_test = y_test_seq_dict.get('dir')
        if actual_ret_test is None or actual_dir_test is None:
             logger.error("Actual return ('ret') or direction ('dir') missing from y_test_seq_dict. Cannot run backtest.")
             return None
             
        # Verify prediction lengths
        n_test = len(X_test_seq)
        if not (len(mu_test) == n_test and len(sigma_test) == n_test and \
                  len(p_raw_test) == n_test and len(actual_ret_test) == n_test and \
                  len(actual_dir_test) == n_test and len(test_indices) == n_test):
             logger.error(f"Length mismatch in test predictions/targets/indices: Expected {n_test}, got mu={len(mu_test)}, sigma={len(sigma_test)}, p_raw={len(p_raw_test)}, ret={len(actual_ret_test)}, dir={len(actual_dir_test)}, indices={len(test_indices)}")
             return None, None

        # 3. Calibrate Test Probabilities
        logger.info(f"Calibrating test predictions using T={calibrator.optimal_T:.4f}")
        p_cal_test = calibrator.calibrate(p_raw_test)
        edge_test = 2 * p_cal_test - 1 # Edge = P(up) - P(down)
        z_score_test = np.abs(mu_test) / (sigma_test + 1e-9) # abs(mu)/sigma
        # Generate GRU-based signals for confusion matrix
        gru_signal_test = calibrator.action_signal(p_cal_test)
        
        # 4. Simulation Loop
        capital = self.initial_capital
        current_position = 0.0 # Starts neutral (-1 to 1)
        equity_curve = [capital]
        positions = [current_position]
        actions_taken = [0.0] # SAC agent's desired fractional position
        pnl_steps = []
        trades_executed = [] # Store details of trades

        logger.info(f"Starting backtest simulation loop: {n_test} steps...")
        for i in range(n_test):
            # Construct current state for SAC agent
            # state = [mu, sigma, edge, |mu|/sigma, position]
            state = np.array([
                mu_test[i],
                sigma_test[i],
                edge_test[i],
                z_score_test[i],
                current_position # Position from start of the period
            ], dtype=np.float32)

            # Get action from SAC agent (deterministic for backtest)
            # Action output is assumed to be the target position [-1, 1]
            sac_action = agent.get_action(state, deterministic=True)[0] # Get scalar action
            target_position = np.clip(sac_action, -1.0, 1.0) # Ensure action is within bounds
            
            # Calculate PnL for the step based on position held *during* the step
            # Return is for the period ending at the *current* index i
            step_actual_return = actual_ret_test[i]
            gross_pnl = current_position * capital * (np.exp(step_actual_return) - 1)
            
            # Calculate trade size and cost
            # Trade happens *after* observing return for period i, to establish position for period i+1
            trade = target_position - current_position
            cost = abs(trade) * capital * self.transaction_cost
            
            # Calculate net PnL for the step
            net_pnl = gross_pnl - cost
            
            # Update capital
            capital += net_pnl
            
            # Store results for this step
            equity_curve.append(capital)
            positions.append(target_position) # Store the position held for the *next* step
            actions_taken.append(sac_action)
            pnl_steps.append(net_pnl)
            if abs(trade) > 1e-6: # Record non-zero trades
                 trades_executed.append({
                      'timestamp': test_indices[i],
                      'trade_size': trade,
                      'cost': cost,
                      'position_before': current_position,
                      'position_after': target_position
                 })
                 
            # Update position for the next iteration
            current_position = target_position 
            
            if capital <= 0:
                logger.warning(f"Capital depleted at step {i+1}. Stopping backtest.")
                n_test = i + 1 # Adjust length to current step
                break
                
        logger.info("Backtest simulation loop finished.")
        logger.info(f"Final Equity: {capital:.2f}")

        # 5. Prepare Results DataFrame
        if n_test == 0:
             logger.warning("Backtest executed 0 steps.")
             return pd.DataFrame(), {}
             
        results_data = {
            'equity': equity_curve[1:], # Exclude initial capital
            'position': positions[1:],   # Position held for the period starting at this index
            'action': actions_taken[1:],   # Agent's target position output
            'pnl': pnl_steps,           
            'actual_return': actual_ret_test[:n_test],
            'mu_pred': mu_test[:n_test],
            'sigma_pred': sigma_test[:n_test],
            'p_cal_pred': p_cal_test[:n_test],
            'edge_pred': edge_test[:n_test],
            'gru_signal': gru_signal_test[:n_test],
            'actual_dir': actual_dir_test[:n_test]
        }
        # Add original prices if available
        if original_prices is not None:
            # Ensure alignment with test_indices before adding
            aligned_prices = original_prices.loc[test_indices[:n_test]]
            results_data['close_price'] = aligned_prices.values
            
        self.results_df = pd.DataFrame(results_data, index=test_indices[:n_test])
        self.results_df['returns'] = self.results_df['equity'].pct_change().fillna(0.0)
        self.results_df['cumulative_return'] = (1 + self.results_df['returns']).cumprod() - 1
        
        # Calculate Buy & Hold Benchmark
        if 'close_price' in self.results_df.columns and not self.results_df.empty:
            bh_returns = self.results_df['close_price'].pct_change().fillna(0.0)
            self.results_df['bh_cumulative_return'] = (1 + bh_returns).cumprod() - 1
            bh_sharpe = calculate_sharpe_ratio(bh_returns)
        else:
            bh_sharpe = 0.0
            self.results_df['bh_cumulative_return'] = 0.0
            logger.warning("Could not calculate Buy & Hold benchmark due to missing price data.")

        # 6. Calculate Performance Metrics
        logger.info("Calculating performance metrics...")
        total_pnl = self.results_df['pnl'].sum()
        final_equity = self.results_df['equity'].iloc[-1]
        total_return_pct = (final_equity / self.initial_capital - 1) * 100
        sharpe = calculate_sharpe_ratio(self.results_df['returns'])
        max_dd = calculate_max_drawdown(self.results_df['equity'])
        
        wins = self.results_df[self.results_df['pnl'] > 0]['pnl']
        losses = self.results_df[self.results_df['pnl'] < 0]['pnl']
        profit_factor = wins.sum() / abs(losses.sum()) if losses.sum() != 0 else np.inf
        
        num_trades = len(trades_executed)
        # Simple win rate based on pnl > 0 per step holding position
        win_rate_steps = (self.results_df['pnl'] > 0).mean() 
        
        # Confusion Matrix for GRU signals
        conf_matrix = confusion_matrix(actual_dir_test[:n_test], gru_signal_test[:n_test], labels=[-1, 0, 1])
        class_report = classification_report(actual_dir_test[:n_test], gru_signal_test[:n_test], labels=[-1, 0, 1], target_names=['Short', 'Neutral', 'Long'], zero_division=0)

        self.metrics = {
            "Run ID": self.config.get('run_id_template', 'N/A').format(timestamp="..."), # Use actual run ID later
            "Test Period Start": test_indices[0].strftime('%Y-%m-%d %H:%M'),
            "Test Period End": test_indices[n_test-1].strftime('%Y-%m-%d %H:%M'),
            "Initial Capital": self.initial_capital,
            "Final Equity": final_equity,
            "Total Net PnL": total_pnl,
            "Total Return (%)": total_return_pct,
            "Annualized Sharpe Ratio": sharpe,
            "Max Drawdown (%)": max_dd * 100,
            "Profit Factor": profit_factor,
            "Number of Trades": num_trades,
            "Step Win Rate (%)": win_rate_steps * 100,
            "Transaction Cost (% per trade)": self.transaction_cost * 100,
            "Edge Threshold": self.edge_threshold,
            "Calibration Temperature (Optimal T)": calibrator.optimal_T,
            "Buy & Hold Sharpe Ratio": bh_sharpe,
            "Confusion Matrix (GRU Signal vs Actual Dir)": conf_matrix.tolist(), # Convert to list for saving
            "Classification Report (GRU Signal)": class_report
        }
        logger.info("--- Backtest Simulation Finished ---")
        return self.results_df, self.metrics

    def save_results(
        self,
        results_df: pd.DataFrame,
        metrics: Dict[str, Any],
        results_dir: str,
        run_id: str
    ):
        """
        Saves the backtest results, metrics report, and plots.
        
        Args:
            results_df (pd.DataFrame): DataFrame from run_backtest.
            metrics (Dict[str, Any]): Metrics dictionary from run_backtest.
            results_dir (str): Directory to save the results.
            run_id (str): The pipeline run ID for filenames.
        """
        logger.info("--- Saving Backtest Results --- ")
        if results_df is None or metrics is None:
             logger.warning("No results DataFrame or metrics to save.")
             return
             
        os.makedirs(results_dir, exist_ok=True)

        # 1. Save Metrics Report
        metrics_path = os.path.join(results_dir, f'performance_metrics_{run_id}.txt')
        try:
            with open(metrics_path, 'w') as f:
                f.write(f"--- Performance Metrics (Run ID: {run_id}) ---\n\n")
                for key, value in metrics.items():
                    if key == "Confusion Matrix (GRU Signal vs Actual Dir)":
                         f.write(f"{key}:\n{np.array(value)}\n\n") # Nicer print for matrix
                    elif key == "Classification Report (GRU Signal)":
                         f.write(f"{key}:\n{value}\n\n")
                    elif isinstance(value, float):
                         f.write(f"{key}: {value:.4f}\n")
                    else:
                         f.write(f"{key}: {value}\n")
            logger.info(f"Performance metrics saved to {metrics_path}")
        except Exception as e:
            logger.error(f"Failed to save metrics report: {e}", exc_info=True)

        # 2. Save Results DataFrame
        results_csv_path = os.path.join(results_dir, f'backtest_results_{run_id}.csv')
        try:
            results_df.to_csv(results_csv_path)
            logger.info(f"Detailed backtest results saved to {results_csv_path}")
        except Exception as e:
            logger.error(f"Failed to save results DataFrame: {e}", exc_info=True)
        
        # 3. Generate and Save Plots
        if self.config.get('control', {}).get('generate_plots', True):
            logger.info("Generating backtest plots...")
            try:
                # Plot 1: Multi-subplot (Price/Pred, Action, Equity/BH)
                fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)
                
                # Subplot 1: Price vs Prediction
                if 'close_price' in results_df.columns:
                    ax = axes[0]
                    ax.plot(results_df.index, results_df['close_price'], label='Actual Price', color='black', alpha=0.8)
                    # Reconstruct predicted price from mu (log return prediction)
                    # pred_price = results_df['close_price'].shift(1) * np.exp(results_df['mu_pred'])
                    # ax.plot(results_df.index, pred_price, label='Predicted Price (from mu)', color='blue', alpha=0.6)
                    # Plot mu +/- sigma directly (interpreting mu as predicted return)
                    ax.plot(results_df.index, results_df['mu_pred'], label='Predicted Log Return (mu)', color='blue', alpha=0.7)
                    ax.fill_between(results_df.index, 
                                     results_df['mu_pred'] - results_df['sigma_pred'], 
                                     results_df['mu_pred'] + results_df['sigma_pred'], 
                                     color='blue', alpha=0.2, label='Predicted Sigma')
                    ax.set_ylabel("Price / Log Return")
                    ax.set_title(f'Price, Predictions & Uncertainty (Run: {run_id})')
                    ax.legend()
                    ax.grid(True)
                else:
                    axes[0].text(0.5, 0.5, 'Price data not available for plotting', ha='center', va='center')
                    axes[0].set_title('Price vs Prediction (Data Missing)')

                # Subplot 2: SAC Agent Action
                ax = axes[1]
                ax.plot(results_df.index, results_df['action'], label='SAC Agent Target Position', color='red')
                ax.set_ylabel("Target Position (-1 to 1)")
                ax.set_ylim(-1.1, 1.1)
                ax.set_title('SAC Agent Action')
                ax.legend()
                ax.grid(True)

                # Subplot 3: Equity Curve vs Buy&Hold
                ax = axes[2]
                equity_norm = results_df['equity'] / self.initial_capital
                ax.plot(results_df.index, equity_norm, label='SAC Strategy Equity', color='green')
                if 'bh_cumulative_return' in results_df.columns:
                     bh_equity_norm = 1 + results_df['bh_cumulative_return']
                     ax.plot(results_df.index, bh_equity_norm, label='Buy & Hold Equity', color='gray', linestyle='--')
                ax.set_ylabel("Normalized Equity")
                ax.set_xlabel("Time")
                ax.set_title('Portfolio Equity vs Buy & Hold')
                ax.legend()
                ax.grid(True)
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M'))
                plt.xticks(rotation=45)
                
                plt.tight_layout()
                plot1_path = os.path.join(results_dir, f'backtest_summary_{run_id}.png')
                plt.savefig(plot1_path)
                plt.close(fig)
                logger.info(f"Summary plot saved to {plot1_path}")

                # Plot 2: Confusion Matrix Heatmap
                if "Confusion Matrix (GRU Signal vs Actual Dir)" in metrics:
                     cm = np.array(metrics["Confusion Matrix (GRU Signal vs Actual Dir)"])
                     fig_cm, ax_cm = plt.subplots(figsize=(6, 5))
                     sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", 
                                 xticklabels=['Pred Short', 'Pred Neutral', 'Pred Long'], 
                                 yticklabels=['Actual Short', 'Actual Neutral', 'Actual Long'],
                                 ax=ax_cm)
                     ax_cm.set_xlabel("Predicted Signal")
                     ax_cm.set_ylabel("Actual Direction")
                     ax_cm.set_title(f"GRU Signal Confusion Matrix (Run: {run_id})")
                     plt.tight_layout()
                     cm_plot_path = os.path.join(results_dir, f'confusion_matrix_{run_id}.png')
                     plt.savefig(cm_plot_path)
                     plt.close(fig_cm)
                     logger.info(f"Confusion matrix plot saved to {cm_plot_path}")

            except Exception as e:
                 logger.error(f"Failed to generate plots: {e}", exc_info=True)
        else:
            logger.info("Skipping plot generation as per config.")
            
        logger.info("--- Finished Saving Backtest Results ---") 