import pandas as pd
import numpy as np
import yfinance as yf
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import mean_squared_error, r2_score
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping
import tensorflow as tf
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
from tqdm.notebook import tqdm
import warnings
import pandas_ta as ta  # Ensure consistent usage

warnings.filterwarnings('ignore')
tf.random.set_seed(42)
np.random.seed(42)

# Function to get list of Indian stocks
def get_indian_stocks():
    nse_url = "https://archives.nseindia.com/content/equities/EQUITY_L.csv"
    try:
        nse_df = pd.read_csv(nse_url)
        return [f"{symbol}.NS" for symbol in nse_df['SYMBOL'].tolist()]
    except Exception as e:
        print(f"Error fetching Indian stocks list: {str(e)}")
        return []

# Function to fetch stock data with retries and multiple periods
def get_stock_data(symbol, periods=["5y", "max"], retries=3):
    """
    Fetches historical stock data for a given symbol using available periods.

    Parameters:
    - symbol (str): Stock symbol.
    - periods (list): List of periods to attempt fetching.
    - retries (int): Number of retry attempts per period.

    Returns:
    - pd.DataFrame or None: Historical stock data or None if fetching fails.
    """
    for period in periods:
        for attempt in range(retries):
            try:
                stock = yf.Ticker(symbol)
                df = stock.history(period=period)
                if df.empty:
                    raise ValueError("Empty DataFrame returned")
                if len(df) < 400:  # Ensure we have enough data
                    print(f"Insufficient data for {symbol} with period '{period}'. Trying next period...")
                    break  # Move to the next period
                return df
            except Exception as e:
                if attempt == retries - 1:
                    print(f"Error fetching data for {symbol} with period '{period}': {str(e)}")
                time.sleep(1)  # Wait before retrying
    return None

# Function to create features
def create_features(df):
    # Calculate Percentage Returns
    df['Returns'] = df['Close'].pct_change()
    
    # Simple Moving Averages
    df['SMA_50'] = df['Close'].rolling(window=50).mean()
    df['SMA_200'] = df['Close'].rolling(window=200).mean()
    
    # Volatility
    df['Volatility'] = df['Returns'].rolling(window=20).std() * np.sqrt(252)
    
    # **Corrected Technical Indicators Using pandas_ta**
    # RSI
    df['RSI'] = ta.rsi(df['Close'], length=14)  # Default length is 14
    
    # MACD
    macd = ta.macd(df['Close'], fast=12, slow=26, signal=9)
    df['MACD'] = macd['MACD_12_26_9']
    df['MACD_Signal'] = macd['MACDs_12_26_9']
    
    # Bollinger Bands
    bb = ta.bbands(df['Close'], length=20, std=2)
    df['BB_High'] = bb['BBU_20_2.0']
    df['BB_Low'] = bb['BBL_20_2.0']
    
    # On-Balance Volume
    df['OBV'] = ta.obv(df['Close'], df['Volume'])
    
    # Drop Rows with NaN Values
    df = df.dropna()
    
    return df

# Function to prepare data for LSTM
def prepare_data_for_lstm(data, look_back=60):
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data)
    
    X, y = [], []
    for i in range(look_back, len(scaled_data)):
        X.append(scaled_data[i-look_back:i])
        y.append(scaled_data[i, 3])  # 'Close' is the 4th column (0-based indexing)
    
    return np.array(X), np.array(y), scaler

# Function to create and train LSTM model
def create_lstm_model(input_shape):
    model = Sequential([
        LSTM(100, return_sequences=True, input_shape=input_shape),
        Dropout(0.2),
        LSTM(100, return_sequences=False),
        Dropout(0.2),
        Dense(50, activation='relu'),
        Dense(1)
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss='mse')
    return model

# Function to train model and make prediction
def analyze_stock(symbol):
    df = get_stock_data(symbol)
    if df is None or len(df) < 400:
        return None

    df = create_features(df)
    
    features = ['Open', 'High', 'Low', 'Close', 'Volume', 'Returns', 'SMA_50', 'SMA_200', 
                'Volatility', 'RSI', 'MACD', 'MACD_Signal', 'BB_High', 'BB_Low', 'OBV']
    X = df[features]
    y = df['Close']
    
    # Prepare data for LSTM
    X_lstm, y_lstm, scaler = prepare_data_for_lstm(df[features].values)
    X_train_lstm, X_test_lstm, y_train_lstm, y_test_lstm = train_test_split(
        X_lstm, y_lstm, test_size=0.2, random_state=42
    )
    
    # Train LSTM model
    lstm_model = create_lstm_model((X_train_lstm.shape[1], X_train_lstm.shape[2]))
    early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
    lstm_model.fit(
        X_train_lstm, y_train_lstm, epochs=100, batch_size=32, validation_split=0.2, 
        callbacks=[early_stopping], verbose=0
    )
    
    # Train Random Forest and Gradient Boosting models
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    
    rf_model = RandomForestRegressor(n_estimators=200, random_state=42, n_jobs=-1)
    rf_model.fit(X_train, y_train)
    
    gb_model = GradientBoostingRegressor(n_estimators=200, random_state=42)
    gb_model.fit(X_train, y_train)
    
    # Make predictions
    lstm_pred = lstm_model.predict(X_test_lstm)
    # Inverse transform only the 'Close' price
    lstm_pred_full = np.zeros((len(lstm_pred), X.shape[1]))
    lstm_pred_full[:, 3] = lstm_pred  # 'Close' is the 4th column
    lstm_pred_inv = scaler.inverse_transform(lstm_pred_full)[:, 3]
    
    rf_pred = rf_model.predict(X_test)
    gb_pred = gb_model.predict(X_test)
    
    # Ensemble prediction
    ensemble_pred = (lstm_pred_inv + rf_pred + gb_pred) / 3
    
    # Calculate metrics
    mse = mean_squared_error(y_test, ensemble_pred)
    r2 = r2_score(y_test, ensemble_pred)
    
    # Make future prediction
    last_60_days = df[features].values[-60:]
    last_60_days_scaled = scaler.transform(last_60_days)
    last_60_days_reshaped = last_60_days_scaled.reshape(1, 60, X.shape[1])
    lstm_future = lstm_model.predict(last_60_days_reshaped)
    lstm_future_full = np.zeros((1, X.shape[1]))
    lstm_future_full[0, 3] = lstm_future[0, 0]  # 'Close' is the 4th column
    lstm_future_inv = scaler.inverse_transform(lstm_future_full)[0, 3]
    
    rf_future = rf_model.predict(df[features].iloc[-1:].values)
    gb_future = gb_model.predict(df[features].iloc[-1:].values)
    
    ensemble_future = (lstm_future_inv + rf_future[0] + gb_future[0]) / 3
    
    current_price = df['Close'].iloc[-1]
    potential_return = (ensemble_future - current_price) / current_price
    
    avg_volume = df['Volume'].mean()
    volatility = df['Volatility'].iloc[-1]
    
    return {
        'symbol': symbol,
        'current_price': current_price,
        'predicted_price': ensemble_future,
        'potential_return': potential_return,
        'mse': mse,
        'r2': r2,
        'avg_volume': avg_volume,
        'volatility': volatility
    }

# Function to calculate confidence score
def calculate_confidence(stock_data):
    r2_weight = 0.4
    volume_weight = 0.3
    volatility_weight = 0.3
    
    r2_score_val = max(min(stock_data['r2'], 1), 0)  # Clip R2 score between 0 and 1
    volume_score = min(stock_data['avg_volume'] / 1e6, 1)  # Normalize volume, cap at 1
    volatility_score = 1 - min(stock_data['volatility'], 1)  # Lower volatility is better, cap at 1
    
    confidence = (r2_weight * r2_score_val +
                  volume_weight * volume_score +
                  volatility_weight * volatility_score)
    
    return confidence * 100  # Convert to percentage

# Function to scan all stocks
def scan_stocks(stocks):
    results = []
    
    with ThreadPoolExecutor(max_workers=5) as executor:
        future_to_stock = {executor.submit(analyze_stock, symbol): symbol for symbol in stocks}
        for future in tqdm(as_completed(future_to_stock), total=len(stocks), desc="Analyzing Stocks"):
            result = future.result()
            if result:
                confidence = calculate_confidence(result)
                result['confidence'] = confidence
                results.append(result)
    
    return sorted(results, key=lambda x: x['confidence'], reverse=True)

# Main function
def main():
    print("Fetching list of Indian stocks...")
    stocks = get_indian_stocks()
    print(f"Found {len(stocks)} stocks. Starting analysis...")
    
    if not stocks:
        print("No stocks to analyze. Exiting.")
        return
    
    # Analyze top 100 stocks for demonstration; adjust as needed
    stocks_to_analyze = stocks[:100]
    results = scan_stocks(stocks_to_analyze)
    
    if not results:
        print("No stocks were successfully analyzed. Please check the error logs for details.")
        return
    
    print("\nTop 10 suggested stocks for investment:")
    for i, stock in enumerate(results[:10], 1):
        print(f"{i}. {stock['symbol']}:")
        print(f"   Current Price: ₹{stock['current_price']:.2f}")
        print(f"   Predicted Price: ₹{stock['predicted_price']:.2f}")
        print(f"   Potential Return: {stock['potential_return']:.2%}")
        print(f"   Confidence: {stock['confidence']:.2f}%")
        print(f"   R2 Score: {stock['r2']:.4f}")
        print(f"   MSE: {stock['mse']:.4f}")
        print(f"   Average Daily Volume: {stock['avg_volume']:,.0f}")
        print(f"   Volatility: {stock['volatility']:.2%}")
        print()
    
    print(f"Total stocks analyzed: {len(results)}")
    print("Note: This analysis is based on historical data and should not be the sole basis for investment decisions.")
    print("Always conduct thorough research and consult with a qualified financial advisor before investing.")

if __name__ == "__main__":
    main()

This is the code I am working on. Problem is,
1. I don't know if this has proper logic, the reason being
2. It is extremely slow in compilation

Fetching list of Indian stocks...
Found 2022 stocks. Starting analysis...
Analyzing Stocks:   0%
 0/100 [03:17<?, ?it/s]
WARNING:tensorflow:5 out of the last 9 calls to <function TensorFlowTrainer.make_train_function.<locals>.one_step_on_iterator at 0x7c77f32aa0e0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
7/7 ━━━━━━━━━━━━━━━━━━━━ 3s 380ms/step
7/7 ━━━━━━━━━━━━━━━━━━━━ 4s 357ms/step
WARNING:tensorflow:5 out of the last 15 calls to <function TensorFlowTrainer.make_predict_function.<locals>.one_step_on_data_distributed at 0x7c77f0d88af0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
7/7 ━━━━━━━━━━━━━━━━━━━━ 2s 222ms/step
WARNING:tensorflow:5 out of the last 15 calls to <function TensorFlowTrainer.make_predict_function.<locals>.one_step_on_data_distributed at 0x7c77e8129e10> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
7/7 ━━━━━━━━━━━━━━━━━━━━ 5s 528ms/step
7/7 ━━━━━━━━━━━━━━━━━━━━ 2s 225ms/step

This has been output so far in 10 minutes. We need fast stock analysis as stock prices change instantly.