"""
SAC Trainer Component.

Handles the offline training process for the SAC agent, including loading 
dependencies from a specified GRU run, preparing data for the environment, 
and executing the training loop.
"""

import tensorflow as tf
import json
import logging
import os
import sys
import yaml
import joblib
import pandas as pd
import numpy as np
from datetime import datetime
from tqdm import tqdm
from tensorflow.keras.callbacks import TensorBoard

# Import necessary components from the pipeline
# Use absolute imports assuming the package structure is correct
from gru_sac_predictor.src.data_loader import DataLoader
from gru_sac_predictor.src.feature_engineer import FeatureEngineer
from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.calibrator import Calibrator
from gru_sac_predictor.src.sac_agent import SACTradingAgent
from gru_sac_predictor.src.trading_env import TradingEnv
try:
    from gru_sac_predictor.src.features import minimal_whitelist # For FE fallback
except ImportError:
    minimal_whitelist = [] # Define empty if import fails

logger = logging.getLogger(__name__)

class SACTrainer:
    """Manages the offline SAC training workflow."""

    def __init__(self, config: dict, base_models_dir: str, base_logs_dir: str, base_results_dir: str):
        """
        Initialize the SACTrainer.

        Args:
            config (dict): The main pipeline configuration dictionary.
            base_models_dir (str): The base directory where all run models are stored (e.g., project_root/models).
            base_logs_dir (str): Base directory for logs.
            base_results_dir (str): Base directory for results.
        """
        self.config = config
        self.base_models_dir = base_models_dir
        self.sac_cfg = config['sac']
        self.env_cfg = config.get('environment', {})
        self.control_cfg = config.get('control', {})
        self.data_cfg = config['data']

        # Generate a specific run ID for this SAC training instance
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.sac_train_run_id = f"sac_train_{timestamp}"
        logger.info(f"Initializing SACTrainer with Run ID: {self.sac_train_run_id}")

        # Setup directories for this specific SAC training run
        self.sac_run_models_dir = os.path.join(self.base_models_dir, self.sac_train_run_id)
        self.sac_run_logs_dir = os.path.join(base_logs_dir, self.sac_train_run_id)
        self.sac_run_results_dir = os.path.join(base_results_dir, self.sac_train_run_id)
        self.sac_tb_log_dir = os.path.join(self.sac_run_logs_dir, 'tensorboard')

        os.makedirs(self.sac_run_models_dir, exist_ok=True)
        os.makedirs(self.sac_run_logs_dir, exist_ok=True)
        os.makedirs(self.sac_run_results_dir, exist_ok=True)
        os.makedirs(self.sac_tb_log_dir, exist_ok=True)

        # Configure logging specifically for this trainer instance if needed
        # For now, relies on the pipeline's logger setup
        logger.info(f"  SAC Models Dir: {self.sac_run_models_dir}")
        logger.info(f"  SAC Logs Dir:   {self.sac_run_logs_dir}")
        logger.info(f"  SAC Results Dir:{self.sac_run_results_dir}")
        logger.info(f"  SAC TB Dir:     {self.sac_tb_log_dir}")
        
        # Save config subset relevant to SAC training?
        # Or assume full config is saved by the main pipeline

    def _load_gru_dependencies(self, gru_run_id: str) -> dict | None:
        """
        Loads artifacts (whitelist, scaler, GRU model, T) from a completed GRU pipeline run.

        Args:
            gru_run_id (str): The run ID of the GRU pipeline run.

        Returns:
            dict | None: A dictionary containing the loaded dependencies 
                         ('whitelist', 'scaler', 'gru_model', 'optimal_T'), or None on failure.
        """
        logger.info(f"--- Loading Dependencies from GRU Run ID: {gru_run_id} ---")
        gru_run_models_dir = os.path.join(self.base_models_dir, f"run_{gru_run_id}")
        if not os.path.exists(gru_run_models_dir):
             logger.error(f"Models directory for GRU run {gru_run_id} not found at: {gru_run_models_dir}")
             return None
             
        dependencies = {}

        # 1. Load Whitelist
        whitelist_path = os.path.join(gru_run_models_dir, f"final_whitelist_{gru_run_id}.json")
        try:
            with open(whitelist_path, 'r') as f:
                dependencies['whitelist'] = json.load(f)
            logger.info(f"Loaded whitelist ({len(dependencies['whitelist'])} features) from {whitelist_path}")
            if not dependencies['whitelist']:
                 raise ValueError("Loaded whitelist is empty.")
        except Exception as e:
            logger.error(f"Failed to load whitelist from {whitelist_path}: {e}", exc_info=True)
            return None

        # 2. Load Scaler
        scaler_path = os.path.join(gru_run_models_dir, f"feature_scaler_{gru_run_id}.joblib")
        try:
            dependencies['scaler'] = joblib.load(scaler_path)
            logger.info(f"Loaded scaler from {scaler_path}")
        except Exception as e:
            logger.error(f"Failed to load scaler from {scaler_path}: {e}", exc_info=True)
            return None

        # 3. Load GRU Model
        model_path = os.path.join(gru_run_models_dir, f"gru_model_{gru_run_id}.keras")
        # Need a temporary GRU handler instance to use its load method
        temp_gru_handler = GRUModelHandler(run_id="temp_load", models_dir="temp_load")
        dependencies['gru_model'] = temp_gru_handler.load(model_path)
        if dependencies['gru_model'] is None:
             logger.error(f"Failed to load GRU model from {model_path}")
             return None
        logger.info(f"Loaded GRU model from {model_path}")

        # 4. Load Optimal Temperature
        temp_path = os.path.join(gru_run_models_dir, f"calibration_temp_{gru_run_id}.npy")
        try:
            dependencies['optimal_T'] = float(np.load(temp_path))
            logger.info(f"Loaded optimal temperature T={dependencies['optimal_T']:.4f} from {temp_path}")
        except Exception as e:
            logger.error(f"Failed to load optimal temperature from {temp_path}: {e}", exc_info=True)
            # Allow continuation without T? Or require it? Let's require it for now.
            return None
            
        logger.info("--- Successfully loaded all GRU dependencies ---")
        return dependencies

    def _prepare_data_for_sac(self, gru_dependencies: dict) -> tuple | None:
        """
        Replicates the necessary data loading and preparation steps using the loaded 
        GRU dependencies, specifically focusing on the VALIDATION dataset to 
        generate inputs for the TradingEnv.

        Args:
            gru_dependencies (dict): The dictionary returned by _load_gru_dependencies.

        Returns:
            tuple | None: A tuple containing (mu_val, sigma_val, p_cal_val, actual_ret_val)
                          for the validation set, or None on failure.
        """
        logger.info("--- Preparing Validation Data for SAC Environment --- ")
        try:
            # 1. Load Raw Data (using a temporary DataLoader)
            temp_data_loader = DataLoader(db_dir=self.data_cfg['db_dir'])
            df_raw = temp_data_loader.load_data(
                ticker=self.data_cfg['ticker'],
                exchange=self.data_cfg['exchange'],
                start_date=self.data_cfg['start_date'],
                end_date=self.data_cfg['end_date'],
                interval=self.data_cfg['interval']
            )
            if df_raw is None or df_raw.empty: raise ValueError("Raw data loading failed")
            df_raw.dropna(subset=['open', 'high', 'low', 'close', 'volume'], inplace=True)
            logger.info("Loaded raw data.")

            # 2. Engineer Base Features (using a temporary FeatureEngineer)
            # Pass the *minimal* whitelist as a fallback if the loaded one causes issues
            temp_feature_engineer = FeatureEngineer(minimal_whitelist=minimal_whitelist)
            df_engineered = temp_feature_engineer.add_base_features(df_raw)
            df_engineered.dropna(inplace=True) # Drop NaNs after feature eng
            if df_engineered.empty: raise ValueError("Dataframe empty after feature engineering.")
            logger.info("Engineered base features.")

            # 3. Prune Features using *loaded* whitelist
            loaded_whitelist = gru_dependencies['whitelist']
            missing_in_eng = [f for f in loaded_whitelist if f not in df_engineered.columns]
            if missing_in_eng:
                 raise ValueError(f"Features from loaded whitelist missing in engineered data: {missing_in_eng}")
            df_features = df_engineered[loaded_whitelist]
            logger.info("Pruned features using loaded whitelist.")

            # 4. Define Labels
            horizon = self.config['gru'].get('prediction_horizon', 5)
            target_ret_col = f'fwd_log_ret_{horizon}'
            target_dir_col = f'direction_label_{horizon}'
            df_engineered[target_ret_col] = np.log(df_engineered['close'].shift(-horizon) / df_engineered['close'])
            df_engineered[target_dir_col] = (df_engineered[target_ret_col] > 0).astype(int)
            # Align by dropping NaNs in targets AND ensuring indices match features
            df_engineered.dropna(subset=[target_ret_col, target_dir_col], inplace=True)
            common_index = df_features.index.intersection(df_engineered.index)
            if common_index.empty:
                raise ValueError("No common index between features and targets after label definition.")
            df_features = df_features.loc[common_index]
            df_targets = df_engineered.loc[common_index, [target_ret_col, target_dir_col]]
            logger.info("Defined labels and aligned features/targets.")
            
            # 5. Split Data (to get validation set indices)
            split_cfg = self.config['split_ratios']
            train_ratio, val_ratio = split_cfg['train'], split_cfg['validation']
            total_len = len(df_features)
            train_end_idx = int(total_len * train_ratio)
            val_end_idx = int(total_len * (train_ratio + val_ratio))
            val_indices = df_features.index[train_end_idx:val_end_idx]
            if val_indices.empty: raise ValueError("Validation split resulted in empty indices.")
            X_val_pruned = df_features.loc[val_indices]
            y_val = df_targets.loc[val_indices]
            logger.info(f"Isolated validation set data (Features: {X_val_pruned.shape}, Targets: {y_val.shape}).")

            # 6. Scale Validation Features using *loaded* scaler
            scaler = gru_dependencies['scaler']
            numeric_cols = X_val_pruned.select_dtypes(include=np.number).columns
            X_val_scaled = X_val_pruned.copy()
            if not numeric_cols.empty:
                 X_val_scaled[numeric_cols] = scaler.transform(X_val_pruned[numeric_cols])
            logger.info("Scaled validation features using loaded scaler.")

            # 7. Create Validation Sequences
            lookback = self.config['gru']['lookback']
            X_val_seq = []
            y_val_seq_targets = [] # Store corresponding targets
            val_seq_indices = []   # Store corresponding indices
            features_np = X_val_scaled.values
            targets_np = y_val.values # Contains both ret and dir
            for i in range(lookback, len(features_np)):
                X_val_seq.append(features_np[i-lookback : i])
                y_val_seq_targets.append(targets_np[i]) # Target corresponds to end of sequence
                val_seq_indices.append(y_val.index[i])
            
            if not X_val_seq:
                 raise ValueError("Validation sequence creation resulted in empty list.")
                 
            X_val_seq = np.array(X_val_seq)
            y_val_seq_targets = np.array(y_val_seq_targets)
            actual_ret_val_seq = y_val_seq_targets[:, 0] # First column is return
            y_dir_val_seq = y_val_seq_targets[:, 1] # Second column is direction
            logger.info(f"Created validation sequences (X shape: {X_val_seq.shape}).")

            # 8. Get GRU Predictions on Validation Sequences using *loaded* GRU model
            gru_model = gru_dependencies['gru_model']
            # Use a temporary handler instance with the loaded model
            temp_gru_handler = GRUModelHandler(run_id="temp_predict", models_dir="temp")
            temp_gru_handler.model = gru_model # Assign the loaded model
            predictions_val = temp_gru_handler.predict(X_val_seq)
            if predictions_val is None or len(predictions_val) < 3:
                 raise ValueError("GRU prediction on validation sequences failed.")
            mu_val_pred = predictions_val[0].flatten()
            log_sigma_val_pred = predictions_val[1][:, 1].flatten()
            p_raw_val_pred = predictions_val[2].flatten()
            sigma_val_pred = np.exp(log_sigma_val_pred)
            logger.info("Generated GRU predictions on validation sequences.")
            
            # Verify lengths
            n_seq = len(X_val_seq)
            if not (len(mu_val_pred) == n_seq and len(sigma_val_pred) == n_seq and \
                      len(p_raw_val_pred) == n_seq and len(actual_ret_val_seq) == n_seq):
                 raise ValueError(f"Length mismatch after validation predictions: Expected {n_seq}, got mu={len(mu_val_pred)}, sigma={len(sigma_val_pred)}, p_raw={len(p_raw_val_pred)}, ret={len(actual_ret_val_seq)}")

            # 9. Calibrate Predictions using *loaded* optimal_T
            optimal_T = gru_dependencies['optimal_T']
            # Use a temporary calibrator instance
            temp_calibrator = Calibrator(edge_threshold=0.5) # Edge threshold doesn't matter here
            temp_calibrator.optimal_T = optimal_T
            p_cal_val_pred = temp_calibrator.calibrate(p_raw_val_pred)
            logger.info(f"Calibrated validation predictions using loaded T={optimal_T:.4f}.")

            # 10. Return the necessary components for the TradingEnv
            logger.info("--- Successfully prepared validation data for SAC Environment ---")
            return mu_val_pred, sigma_val_pred, p_cal_val_pred, actual_ret_val_seq

        except Exception as e:
            logger.error(f"Error preparing data for SAC environment: {e}", exc_info=True)
            return None

    def _load_agent_for_resume(self, agent: SACTradingAgent) -> None:
        """Loads agent weights if resuming is specified in config."""
        load_run_id = self.control_cfg.get('sac_resume_run_id')
        load_step = self.control_cfg.get('sac_resume_step', 'final')
        current_edge_threshold = self.config.get('calibration', {}).get('edge_threshold', 0.55)

        if not load_run_id:
            logger.info("No SAC resume run ID specified. Starting training from scratch.")
            return

        # Construct path relative to base models dir
        if load_step == 'final':
            load_path = os.path.join(self.base_models_dir, f"{load_run_id}", 'sac_agent_final')
        else:
            try:
                step_num = int(load_step)
                load_path = os.path.join(self.base_models_dir, f"{load_run_id}", f'sac_agent_step_{step_num}')
            except ValueError:
                logger.error(f"Invalid sac_resume_step: {load_step}. Must be 'final' or an integer. Starting fresh.")
                return

        logger.info(f"Attempting to load SAC agent from {load_path} to resume training...")
        if os.path.exists(load_path):
            try:
                loaded_meta = agent.load(load_path)
                # Check for Buffer Purge on Load 
                saved_edge_thr = loaded_meta.get('edge_threshold_config')
                if saved_edge_thr is not None and abs(saved_edge_thr - current_edge_threshold) > 1e-6:
                    logger.warning(f'Edge threshold mismatch on load (Saved={saved_edge_thr:.3f}, Current={current_edge_threshold:.3f}). Clearing replay buffer before resuming.')
                    agent.clear_buffer()
                elif saved_edge_thr is None:
                    logger.warning("Loaded SAC agent metadata did not contain 'edge_threshold_config'. Cannot verify consistency.")
                else:
                    logger.info('Edge threshold consistent with loaded agent metadata.')
            except Exception as e:
                logger.error(f"Failed to load SAC agent for resume: {e}. Starting fresh.", exc_info=True)
        else:
            logger.warning(f"SAC agent path not found for resume: {load_path}. Starting fresh.")
            
    def _training_loop(self, agent: SACTradingAgent, env: TradingEnv) -> str | None:
        """Runs the main SAC training loop."""
        total_steps = self.sac_cfg.get('total_training_steps', 100000)
        batch_size = self.sac_cfg.get('batch_size', 256)
        log_interval = self.sac_cfg.get('log_interval', 1000)
        save_interval = self.sac_cfg.get('save_interval', 10000)
        generate_new_on_epoch = self.config.get('experience', {}).get('generate_new_on_epoch', False)
        epoch_len = env.n_steps

        # Initialize TensorBoard
        summary_writer = None
        try:
            summary_writer = tf.summary.create_file_writer(self.sac_tb_log_dir)
            log_to_tensorboard = True
            logger.info(f"TensorBoard logging initialized to: {self.sac_tb_log_dir}")
        except Exception as e:
            logger.error(f"Failed to initialize TensorBoard: {e}")
            log_to_tensorboard = False

        logger.info(f"Starting SAC training loop: {total_steps=}, {batch_size=}, {log_interval=}, {save_interval=}")
        state = env.reset()
        episode_reward = 0
        episode_steps = 0
        episode_num = 0
        all_rewards = []
        final_save_path = None # Track the last save path

        for step in tqdm(range(total_steps), desc="SAC Training Steps"):
            if generate_new_on_epoch and step > 0 and step % epoch_len == 0:
                logger.info(f"Start of epoch {step // epoch_len + 1}. Clearing replay buffer.")
                agent.clear_buffer()

            action = agent.get_action(state, deterministic=False)
            next_state, reward, done, info = env.step(action[0])
            agent.buffer.add(state, action, reward, next_state, float(done))
            state = next_state
            episode_reward += reward
            episode_steps += 1

            if len(agent.buffer) >= agent.min_buffer_size:
                metrics = agent.train(batch_size)
                if metrics and log_to_tensorboard and summary_writer:
                    with summary_writer.as_default(step=step):
                        for key, value in metrics.items():
                            tf.summary.scalar(f'SAC_Metrics/{key}', value)
            
            if done:
                if log_to_tensorboard and summary_writer:
                     with summary_writer.as_default(step=step):
                          tf.summary.scalar('SAC_Episode/Reward', episode_reward)
                          tf.summary.scalar('SAC_Episode/Steps', episode_steps)
                all_rewards.append(episode_reward)
                state = env.reset()
                episode_reward = 0
                episode_steps = 0
                episode_num += 1

            if (step + 1) % save_interval == 0:
                save_path = os.path.join(self.sac_run_models_dir, f'sac_agent_step_{step+1}')
                os.makedirs(save_path, exist_ok=True)
                agent.save(save_path)
                logger.info(f"SAC agent weights saved at step {step+1} to {save_path}")
                final_save_path = save_path # Update last saved path

        # Save final agent
        logger.info("Saving final SAC agent...")
        final_save_path = os.path.join(self.sac_run_models_dir, 'sac_agent_final')
        os.makedirs(final_save_path, exist_ok=True)
        agent.save(final_save_path)
        logger.info(f"Final SAC agent weights saved to {final_save_path}")

        # Save rewards 
        rewards_df = pd.DataFrame({'episode_reward': all_rewards})
        rewards_df.to_csv(os.path.join(self.sac_run_results_dir, f'sac_episode_rewards_{self.sac_train_run_id}.csv'))

        if summary_writer: summary_writer.close()
        env.close()
        return final_save_path

    def train(self, gru_run_id_for_sac: str) -> str | None:
        """
        Main entry point to start the SAC training process.

        Args:
            gru_run_id_for_sac (str): The run ID of the GRU pipeline run whose artifacts should be used.
        
        Returns:
            str | None: Path to the final saved SAC agent, or None if training failed.
        """
        logger.info(f"=== Starting SAC Training Process (SAC Run ID: {self.sac_train_run_id}) ===")
        logger.info(f"Using artifacts from GRU Run ID: {gru_run_id_for_sac}")

        # 1. Load GRU dependencies
        gru_dependencies = self._load_gru_dependencies(gru_run_id_for_sac)
        if gru_dependencies is None:
             logger.error("Failed to load GRU dependencies. Aborting SAC training.")
             return None

        # 2. Prepare data for SAC environment (using validation set)
        env_data = self._prepare_data_for_sac(gru_dependencies)
        if env_data is None:
             logger.error("Failed to prepare data for SAC environment. Aborting SAC training.")
             return None
        mu_val, sigma_val, p_cal_val, actual_ret_val = env_data

        # 3. Initialize Environment
        logger.info("Initializing Trading Environment...")
        env = TradingEnv(
            mu_predictions=mu_val,
            sigma_predictions=sigma_val,
            p_cal_predictions=p_cal_val,
            actual_returns=actual_ret_val,
            initial_capital=self.env_cfg.get('initial_capital', 10000.0),
            transaction_cost=self.env_cfg.get('transaction_cost', 0.0005)
        )
        logger.info(f"TradingEnv initialized with {env.n_steps} steps.")

        # 4. Initialize SAC Agent
        logger.info("Initializing SAC Agent...")
        current_edge_threshold = self.config.get('calibration', {}).get('edge_threshold', 0.55)
        agent = SACTradingAgent(
            state_dim=env.state_dim,
            action_dim=env.action_dim,
            gamma=self.sac_cfg.get('gamma', 0.99),
            tau=self.sac_cfg.get('tau', 0.005),
            initial_lr=self.sac_cfg.get('actor_lr', 3e-4),
            lr_decay_rate=self.sac_cfg.get('lr_decay_rate', 0.96),
            decay_steps=self.sac_cfg.get('decay_steps', 100000),
            buffer_capacity=self.sac_cfg.get('buffer_max_size', 100000),
            ou_noise_stddev=self.sac_cfg.get('ou_noise_stddev', 0.2),
            alpha=self.sac_cfg.get('alpha', 0.2),
            alpha_auto_tune=self.sac_cfg.get('alpha_auto_tune', True),
            target_entropy=self.sac_cfg.get('target_entropy', -1.0 * env.action_dim),
            min_buffer_size=self.sac_cfg.get('min_buffer_size', 1000),
            edge_threshold_config=current_edge_threshold # Pass edge threshold
        )
        logger.info("SAC Agent initialized.")

        # 5. Load agent weights if resuming
        self._load_agent_for_resume(agent)

        # 6. Run training loop
        final_agent_path = self._training_loop(agent, env)

        if final_agent_path:
            logger.info(f"=== SAC Training Process Completed Successfully ===")
        else:
            logger.error("=== SAC Training Process Failed ===")
            
        return final_agent_path 