"""
Stock Market Time Series Dataset - Advanced ML Models

Author: Molla Samser
Organization: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in/
Email: help@rskworld.in
Phone: +91 93305 39277
Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from typing import Tuple, Optional, List
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import warnings
warnings.filterwarnings('ignore')


class StockMLModels:
    """
    Advanced Machine Learning models for stock price prediction.
    """
    
    def __init__(self, data_path: str):
        """
        Initialize the ML Models.
        
        Args:
            data_path: Path to the stock CSV file
        """
        self.data_path = Path(data_path)
        self.symbol = self.data_path.stem
        self.df = self._load_data()
        self.scaler = MinMaxScaler()
        self.model = None
        
        sns.set_style('whitegrid')
        plt.rcParams['figure.figsize'] = (12, 6)
        
    def _load_data(self) -> pd.DataFrame:
        """Load and prepare the data."""
        df = pd.read_csv(self.data_path, parse_dates=['Date'], index_col='Date')
        df.sort_index(inplace=True)
        return df
    
    def prepare_lstm_data(self, feature_columns: List[str] = None, 
                         sequence_length: int = 60,
                         test_size: float = 0.2) -> Tuple:
        """
        Prepare data for LSTM model.
        
        Args:
            feature_columns: List of feature columns to use
            sequence_length: Number of time steps to look back
            test_size: Proportion of data for testing
            
        Returns:
            Tuple of (X_train, X_test, y_train, y_test, scaler)
        """
        if feature_columns is None:
            feature_columns = ['Close']
        
        # Extract features
        data = self.df[feature_columns].values
        
        # Scale the data
        scaled_data = self.scaler.fit_transform(data)
        
        # Create sequences
        X, y = [], []
        for i in range(sequence_length, len(scaled_data)):
            X.append(scaled_data[i-sequence_length:i])
            y.append(scaled_data[i, 0])  # Predict Close price
        
        X, y = np.array(X), np.array(y)
        
        # Split the data
        split_idx = int(len(X) * (1 - test_size))
        X_train, X_test = X[:split_idx], X[split_idx:]
        y_train, y_test = y[:split_idx], y[split_idx:]
        
        return X_train, X_test, y_train, y_test, self.scaler
    
    def build_lstm_model(self, input_shape: Tuple, units: int = 50) -> object:
        """
        Build LSTM model.
        
        Note: Requires TensorFlow/Keras installation
        pip install tensorflow
        
        Args:
            input_shape: Shape of input data (sequence_length, n_features)
            units: Number of LSTM units
            
        Returns:
            Compiled LSTM model
        """
        try:
            from tensorflow.keras.models import Sequential
            from tensorflow.keras.layers import LSTM, Dense, Dropout
            from tensorflow.keras.optimizers import Adam
            
            model = Sequential([
                LSTM(units=units, return_sequences=True, input_shape=input_shape),
                Dropout(0.2),
                LSTM(units=units, return_sequences=True),
                Dropout(0.2),
                LSTM(units=units),
                Dropout(0.2),
                Dense(units=1)
            ])
            
            model.compile(optimizer=Adam(learning_rate=0.001),
                         loss='mean_squared_error',
                         metrics=['mae'])
            
            return model
            
        except ImportError:
            print("TensorFlow not installed. Install with: pip install tensorflow")
            print("Using simpler model instead...")
            return None
    
    def train_lstm(self, epochs: int = 50, batch_size: int = 32,
                  sequence_length: int = 60, verbose: int = 1) -> dict:
        """
        Train LSTM model.
        
        Args:
            epochs: Number of training epochs
            batch_size: Batch size for training
            sequence_length: Sequence length for time series
            verbose: Verbosity level
            
        Returns:
            Dictionary with training history and predictions
        """
        print(f"\n{'='*70}")
        print(f"Training LSTM Model for {self.symbol}")
        print(f"{'='*70}\n")
        
        # Prepare data
        X_train, X_test, y_train, y_test, scaler = self.prepare_lstm_data(
            feature_columns=['Close', 'Volume', 'MA_20', 'MA_50', 'RSI', 'MACD'],
            sequence_length=sequence_length
        )
        
        print(f"Training data shape: {X_train.shape}")
        print(f"Testing data shape: {X_test.shape}\n")
        
        # Build model
        model = self.build_lstm_model(input_shape=(X_train.shape[1], X_train.shape[2]))
        
        if model is None:
            return self._simple_ml_fallback(X_train, X_test, y_train, y_test)
        
        # Train model
        history = model.fit(
            X_train, y_train,
            epochs=epochs,
            batch_size=batch_size,
            validation_split=0.1,
            verbose=verbose
        )
        
        # Make predictions
        train_predictions = model.predict(X_train, verbose=0)
        test_predictions = model.predict(X_test, verbose=0)
        
        # Inverse transform predictions
        train_predictions = scaler.inverse_transform(
            np.concatenate([train_predictions, np.zeros((len(train_predictions), 5))], axis=1)
        )[:, 0]
        test_predictions = scaler.inverse_transform(
            np.concatenate([test_predictions, np.zeros((len(test_predictions), 5))], axis=1)
        )[:, 0]
        y_train_inv = scaler.inverse_transform(
            np.concatenate([y_train.reshape(-1, 1), np.zeros((len(y_train), 5))], axis=1)
        )[:, 0]
        y_test_inv = scaler.inverse_transform(
            np.concatenate([y_test.reshape(-1, 1), np.zeros((len(y_test), 5))], axis=1)
        )[:, 0]
        
        # Calculate metrics
        train_rmse = np.sqrt(mean_squared_error(y_train_inv, train_predictions))
        test_rmse = np.sqrt(mean_squared_error(y_test_inv, test_predictions))
        train_mae = mean_absolute_error(y_train_inv, train_predictions)
        test_mae = mean_absolute_error(y_test_inv, test_predictions)
        train_r2 = r2_score(y_train_inv, train_predictions)
        test_r2 = r2_score(y_test_inv, test_predictions)
        
        print(f"\n{'='*70}")
        print("Model Performance Metrics")
        print(f"{'='*70}")
        print(f"Training RMSE: ${train_rmse:.2f}")
        print(f"Testing RMSE:  ${test_rmse:.2f}")
        print(f"Training MAE:  ${train_mae:.2f}")
        print(f"Testing MAE:   ${test_mae:.2f}")
        print(f"Training R²:   {train_r2:.4f}")
        print(f"Testing R²:    {test_r2:.4f}")
        print(f"{'='*70}\n")
        
        self.model = model
        
        return {
            'model': model,
            'history': history.history,
            'train_predictions': train_predictions,
            'test_predictions': test_predictions,
            'y_train': y_train_inv,
            'y_test': y_test_inv,
            'metrics': {
                'train_rmse': train_rmse,
                'test_rmse': test_rmse,
                'train_mae': train_mae,
                'test_mae': test_mae,
                'train_r2': train_r2,
                'test_r2': test_r2
            }
        }
    
    def _simple_ml_fallback(self, X_train, X_test, y_train, y_test) -> dict:
        """Fallback to simple ML model if LSTM unavailable."""
        from sklearn.ensemble import RandomForestRegressor
        from sklearn.linear_model import LinearRegression
        
        print("Using Random Forest Regressor as fallback...")
        
        # Reshape data for sklearn
        X_train_2d = X_train.reshape(X_train.shape[0], -1)
        X_test_2d = X_test.reshape(X_test.shape[0], -1)
        
        # Train model
        model = RandomForestRegressor(n_estimators=100, random_state=42)
        model.fit(X_train_2d, y_train)
        
        # Predictions
        train_predictions = model.predict(X_train_2d)
        test_predictions = model.predict(X_test_2d)
        
        # Metrics
        train_rmse = np.sqrt(mean_squared_error(y_train, train_predictions))
        test_rmse = np.sqrt(mean_squared_error(y_test, test_predictions))
        
        print(f"\nRandom Forest RMSE - Train: {train_rmse:.4f}, Test: {test_rmse:.4f}\n")
        
        return {
            'model': model,
            'train_predictions': train_predictions,
            'test_predictions': test_predictions,
            'y_train': y_train,
            'y_test': y_test
        }
    
    def plot_predictions(self, results: dict, save_path: Optional[str] = None):
        """
        Plot actual vs predicted prices.
        
        Args:
            results: Results dictionary from training
            save_path: Path to save the plot
        """
        fig, axes = plt.subplots(2, 1, figsize=(14, 10))
        
        # Training predictions
        axes[0].plot(results['y_train'], label='Actual', linewidth=2, alpha=0.8)
        axes[0].plot(results['train_predictions'], label='Predicted', 
                    linewidth=2, alpha=0.8, linestyle='--')
        axes[0].set_title(f'{self.symbol} - Training Set Predictions', 
                         fontsize=14, fontweight='bold')
        axes[0].set_ylabel('Price ($)', fontsize=12)
        axes[0].legend(loc='best')
        axes[0].grid(True, alpha=0.3)
        
        # Testing predictions
        axes[1].plot(results['y_test'], label='Actual', linewidth=2, alpha=0.8)
        axes[1].plot(results['test_predictions'], label='Predicted', 
                    linewidth=2, alpha=0.8, linestyle='--')
        axes[1].set_title(f'{self.symbol} - Testing Set Predictions', 
                         fontsize=14, fontweight='bold')
        axes[1].set_xlabel('Time Steps', fontsize=12)
        axes[1].set_ylabel('Price ($)', fontsize=12)
        axes[1].legend(loc='best')
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Predictions plot saved to: {save_path}")
        
        plt.show()
    
    def plot_training_history(self, history: dict, save_path: Optional[str] = None):
        """
        Plot training history.
        
        Args:
            history: Training history dictionary
            save_path: Path to save the plot
        """
        if 'loss' not in history:
            print("No training history available to plot.")
            return
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        
        # Loss
        axes[0].plot(history['loss'], label='Training Loss', linewidth=2)
        axes[0].plot(history['val_loss'], label='Validation Loss', linewidth=2)
        axes[0].set_title('Model Loss', fontsize=14, fontweight='bold')
        axes[0].set_xlabel('Epoch', fontsize=12)
        axes[0].set_ylabel('Loss', fontsize=12)
        axes[0].legend(loc='best')
        axes[0].grid(True, alpha=0.3)
        
        # MAE
        if 'mae' in history:
            axes[1].plot(history['mae'], label='Training MAE', linewidth=2)
            axes[1].plot(history['val_mae'], label='Validation MAE', linewidth=2)
            axes[1].set_title('Model MAE', fontsize=14, fontweight='bold')
            axes[1].set_xlabel('Epoch', fontsize=12)
            axes[1].set_ylabel('MAE', fontsize=12)
            axes[1].legend(loc='best')
            axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Training history plot saved to: {save_path}")
        
        plt.show()
    
    def predict_future(self, days: int = 30, sequence_length: int = 60) -> pd.DataFrame:
        """
        Predict future stock prices.
        
        Args:
            days: Number of days to predict
            sequence_length: Sequence length used in training
            
        Returns:
            DataFrame with predictions
        """
        if self.model is None:
            print("Model not trained yet. Please train the model first.")
            return None
        
        # Get last sequence
        data = self.df[['Close', 'Volume', 'MA_20', 'MA_50', 'RSI', 'MACD']].values
        scaled_data = self.scaler.transform(data)
        last_sequence = scaled_data[-sequence_length:]
        
        predictions = []
        current_sequence = last_sequence.copy()
        
        for _ in range(days):
            # Reshape for prediction
            pred_input = current_sequence.reshape(1, sequence_length, 6)
            
            # Predict
            pred = self.model.predict(pred_input, verbose=0)[0, 0]
            predictions.append(pred)
            
            # Update sequence (simplified - uses predicted value)
            new_row = np.array([pred, current_sequence[-1, 1], current_sequence[-1, 2],
                               current_sequence[-1, 3], current_sequence[-1, 4],
                               current_sequence[-1, 5]])
            current_sequence = np.vstack([current_sequence[1:], new_row])
        
        # Inverse transform predictions
        predictions = np.array(predictions).reshape(-1, 1)
        predictions_inv = self.scaler.inverse_transform(
            np.concatenate([predictions, np.zeros((len(predictions), 5))], axis=1)
        )[:, 0]
        
        # Create DataFrame
        last_date = self.df.index[-1]
        future_dates = pd.date_range(start=last_date + pd.Timedelta(days=1),
                                     periods=days, freq='D')
        
        forecast_df = pd.DataFrame({
            'Date': future_dates,
            'Predicted_Price': predictions_inv
        })
        forecast_df.set_index('Date', inplace=True)
        
        return forecast_df


class FeatureEngineering:
    """
    Advanced feature engineering for stock data.
    """
    
    @staticmethod
    def add_technical_features(df: pd.DataFrame) -> pd.DataFrame:
        """
        Add advanced technical indicators.
        
        Args:
            df: DataFrame with OHLCV data
            
        Returns:
            DataFrame with additional features
        """
        df = df.copy()
        
        # Returns
        df['Returns'] = df['Close'].pct_change()
        df['Log_Returns'] = np.log(df['Close'] / df['Close'].shift(1))
        
        # Volatility
        df['Volatility_20'] = df['Returns'].rolling(window=20).std()
        df['Volatility_50'] = df['Returns'].rolling(window=50).std()
        
        # Momentum
        df['Momentum_10'] = df['Close'] - df['Close'].shift(10)
        df['Momentum_20'] = df['Close'] - df['Close'].shift(20)
        
        # Rate of Change
        df['ROC_10'] = ((df['Close'] - df['Close'].shift(10)) / df['Close'].shift(10)) * 100
        df['ROC_20'] = ((df['Close'] - df['Close'].shift(20)) / df['Close'].shift(20)) * 100
        
        # Bollinger Bands
        df['BB_Middle'] = df['Close'].rolling(window=20).mean()
        df['BB_Upper'] = df['BB_Middle'] + 2 * df['Close'].rolling(window=20).std()
        df['BB_Lower'] = df['BB_Middle'] - 2 * df['Close'].rolling(window=20).std()
        df['BB_Width'] = df['BB_Upper'] - df['BB_Lower']
        
        # Average True Range (ATR)
        high_low = df['High'] - df['Low']
        high_close = np.abs(df['High'] - df['Close'].shift())
        low_close = np.abs(df['Low'] - df['Close'].shift())
        ranges = pd.concat([high_low, high_close, low_close], axis=1)
        true_range = np.max(ranges, axis=1)
        df['ATR'] = true_range.rolling(14).mean()
        
        # On-Balance Volume (OBV)
        df['OBV'] = (np.sign(df['Close'].diff()) * df['Volume']).fillna(0).cumsum()
        
        # Price Position
        df['Price_Position'] = (df['Close'] - df['Low']) / (df['High'] - df['Low'])
        
        # Lag features
        for lag in [1, 3, 7, 14]:
            df[f'Close_Lag_{lag}'] = df['Close'].shift(lag)
            df[f'Volume_Lag_{lag}'] = df['Volume'].shift(lag)
        
        return df


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Advanced ML Models")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Example usage
    print("\nNote: Install TensorFlow for LSTM: pip install tensorflow")
    print("\nInitializing ML Models for AAPL...")
    
    ml_model = StockMLModels('data/AAPL.csv')
    
    # Train LSTM
    print("\nTraining LSTM model (this may take a few minutes)...")
    results = ml_model.train_lstm(epochs=20, batch_size=32)
    
    # Plot results
    if results:
        ml_model.plot_predictions(results)
        if 'history' in results:
            ml_model.plot_training_history(results['history'])
    
    # Feature engineering example
    print("\nAdding advanced technical features...")
    df_enhanced = FeatureEngineering.add_technical_features(ml_model.df)
    print(f"\nOriginal features: {len(ml_model.df.columns)}")
    print(f"Enhanced features: {len(df_enhanced.columns)}")
    print(f"\nNew features added: {list(set(df_enhanced.columns) - set(ml_model.df.columns))}")
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)