import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
from tensorflow.keras.layers import Input, LayerNormalization, MultiHeadAttention, Dropout, Dense, GlobalAveragePooling1D
from tensorflow.keras.models import Model
import tensorflow as tf
from ta.volatility import AverageTrueRange, BollingerBands
from ta.momentum import RSIIndicator
from ta.trend import MACD, ADXIndicator
from binance.client import Client  # Uncomment if using Binance API
import websocket  # Uncomment if using WebSockets
import json
import datetime
from collections import deque
import pickle
import threading
import time

# ----------------------------
# Configuration and Global Variables
# ----------------------------

# Binance API credentials (replace with your actual API key and secret)
BINANCE_API_KEY = 'YOUR_API_KEY'
BINANCE_API_SECRET = 'YOUR_API_SECRET'

client = Client(BINANCE_API_KEY, BINANCE_API_SECRET)

# Data collection settings
SYMBOLS = ['BTCUSDT']  # Add more symbols as needed
INTERVAL = '1s'  # Binance supports intervals like '1m', '5m', etc.
SEQ_LENGTH = 201
TARGET_STEPS = 1
LIMIT = 1000

# Initialize data_frames dictionary to store DataFrames for each symbol
data_frames = {symbol: pd.DataFrame(columns=[
    'open_time', 'open', 'high', 'low', 'close', 'volume',
    'RSI', 'MACD', 'MACD_signal', 'MACD_diff', 'ATR',
    'BB_upper', 'BB_lower', 'BB_width', 'ADX',
    'daily_return', 'RV',
    'MA_short', 'MA_long', 'MA_Trend',
    'RSI_Trend', 'MACD_Trend', 'ADX_Strength',
    'Market_Trend', 'ATR_Vol', 'BB_Vol', 'RV_Vol',
    'Volatility', 'Market_Environment',
    'previous_leg_change', 'previous_leg_length',
    'current_leg_change', 'current_leg_length',
    'percent_change', 'normalized_close_price',
    'percent_change_classification',
    'close_price',  # Target
    'leg_direction',  # Target
    'leg_length'  # Target
]) for symbol in SYMBOLS}

# Lock for thread-safe operations on data_frames
data_lock = threading.Lock()

# ----------------------------
# Data Loading Functions
# ----------------------------

def load_data(symbols, interval=INTERVAL, sequence_length=SEQ_LENGTH):
    """
    Connect to Binance's Combined WebSocket Stream and collect live kline data.
    """
    import websocket  # Ensure websocket is imported
    import json
    import datetime

    # Create the combined stream name
    streams = '/'.join([f"{symbol.lower()}@kline_{interval}" for symbol in symbols])
    ws_url = f"wss://stream.binance.com:9443/stream?streams={streams}"
    ws = websocket.WebSocketApp(ws_url, on_message=on_message, on_error=on_error, on_close=on_close)
    ws.on_open = on_open

    # Run WebSocket in a separate thread
    wst = threading.Thread(target=ws.run_forever)
    wst.daemon = True
    wst.start()

    # Keep the main thread alive to allow data collection
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        ws.close()

def on_message(ws, message):
    global data_frames
    msg = json.loads(message)

    # Check if the message contains 'data' and it is a kline event
    if 'data' in msg and msg['data']['e'] == 'kline':
        symbol = msg['data']['s']
        k = msg['data']['k']
        is_final = k['x']
        open_time = datetime.datetime.fromtimestamp(k['t'] / 1000)
        o = float(k['o'])
        h = float(k['h'])
        l = float(k['l'])
        c = float(k['c'])
        v = float(k['v'])

        new_row = {
            'open_time': open_time,
            'open': o,
            'high': h,
            'low': l,
            'close': c,
            'volume': v
        }

        with data_lock:
            # Replace append with pd.concat
            new_df = pd.DataFrame([new_row])
            data_frames[symbol] = pd.concat([data_frames[symbol], new_df], ignore_index=True)

            # Keep only the latest necessary data to save memory
            if len(data_frames[symbol]) > SEQ_LENGTH + TARGET_STEPS + 200:
                data_frames[symbol] = data_frames[symbol].iloc[-(SEQ_LENGTH + TARGET_STEPS + 200):].reset_index(drop=True)

        print(f"Received data for {symbol}: {new_row}")
        print(f"Data length for {symbol}: {len(data_frames[symbol])}")

def on_error(ws, error):
    print(f"WebSocket error: {error}")

def on_close(ws, close_status_code, close_msg):
    print("WebSocket connection closed")

def on_open(ws):
    print("WebSocket connection opened")

# ----------------------------
# Pre-Loading Historical Data Function
# ----------------------------

def pre_load_data(symbols, interval = INTERVAL, limit):
    """
    Fetches historical kline data for the given symbols and populates data_frames.

    :param symbols: List of symbols to fetch data for.
    :param interval: Kline interval (e.g., '1m', '5m').
    :param limit: Number of historical data points to fetch.
    :return: None
    """
    for symbol in symbols:
        try:
            historicals = client.get_klines(symbol=symbol, interval=interval, limit=limit)
            for kl in historicals:
                open_time = datetime.datetime.fromtimestamp(kl[0] / 1000)
                o = float(kl[1])
                h = float(kl[2])
                l = float(kl[3])
                c = float(kl[4])
                v = float(kl[5])

                new_row = {
                    'open_time': open_time,
                    'open': o,
                    'high': h,
                    'low': l,
                    'close': c,
                    'volume': v
                }

                # Replace append with pd.concat
                new_df = pd.DataFrame([new_row])
                data_frames[symbol] = pd.concat([data_frames[symbol], new_df], ignore_index=True)
            print(f"Pre-loaded {len(historicals)} data points for {symbol}.")
        except Exception as e:
            print(f"Error fetching historical data for {symbol}: {e}")

# ----------------------------
# Data Processing and Feature Engineering
# ----------------------------

def calculate_technical_indicators(df):
    """
    Calculate technical indicators and add them to the DataFrame.
    """
    if len(df) < 14:
        # Not enough data to calculate indicators
        return df

    # Calculate RSI
    rsi = RSIIndicator(close=df['close'], window=14)
    df['RSI'] = rsi.rsi()

    # Calculate MACD
    macd = MACD(close=df['close'], window_slow=26, window_fast=12, window_sign=9)
    df['MACD'] = macd.macd()
    df['MACD_signal'] = macd.macd_signal()
    df['MACD_diff'] = macd.macd_diff()

    # Calculate ATR
    atr = AverageTrueRange(high=df['high'], low=df['low'], close=df['close'], window=14)
    df['ATR'] = atr.average_true_range()

    # Calculate Bollinger Bands Width
    bollinger = BollingerBands(close=df['close'], window=20, window_dev=2)
    df['BB_upper'] = bollinger.bollinger_hband()
    df['BB_lower'] = bollinger.bollinger_lband()
    df['BB_width'] = (df['BB_upper'] - df['BB_lower']) / df['close']

    # Calculate ADX
    adx = ADXIndicator(high=df['high'], low=df['low'], close=df['close'], window=14)
    df['ADX'] = adx.adx()

    # Calculate Realized Volatility (Rolling 20-window standard deviation of returns)
    df['daily_return'] = df['close'].pct_change()
    df['RV'] = df['daily_return'].rolling(window=20).std() * np.sqrt(252)  # Annualized

    return df

def classify_trend(df):
    """
    Classify market trend as Bullish, Neutral, or Bearish based on moving averages, RSI, MACD, and ADX.
    """
    # Calculate Moving Averages
    df['MA_short'] = df['close'].rolling(window=50).mean()
    df['MA_long'] = df['close'].rolling(window=200).mean()

    # Trend based on Moving Averages
    def ma_trend(row):
        if pd.isna(row['MA_short']) or pd.isna(row['MA_long']):
            return 'Neutral'
        if row['MA_short'] > row['MA_long']:
            return 'Bullish'
        elif row['MA_short'] < row['MA_long']:
            return 'Bearish'
        else:
            return 'Neutral'

    df['MA_Trend'] = df.apply(ma_trend, axis=1)

    # Trend based on RSI
    def rsi_trend(row):
        if pd.isna(row['RSI']):
            return 'Neutral'
        if row['RSI'] > 60:
            return 'Bullish'
        elif row['RSI'] < 40:
            return 'Bearish'
        else:
            return 'Neutral'

    df['RSI_Trend'] = df.apply(rsi_trend, axis=1)

    # Trend based on MACD
    def macd_trend(row):
        if pd.isna(row['MACD']) or pd.isna(row['MACD_signal']):
            return 'Neutral'
        if row['MACD'] > row['MACD_signal'] and row['MACD'] > 0:
            return 'Bullish'
        elif row['MACD'] < row['MACD_signal'] and row['MACD'] < 0:
            return 'Bearish'
        else:
            return 'Neutral'

    df['MACD_Trend'] = df.apply(macd_trend, axis=1)

    # Trend Strength based on ADX
    def adx_strength(row):
        if pd.isna(row['ADX']):
            return 'Weak'
        if row['ADX'] > 25:
            return 'Strong'
        else:
            return 'Weak'

    df['ADX_Strength'] = df.apply(adx_strength, axis=1)

    # Combine Trends
    def combined_trend(row):
        trends = [row['MA_Trend'], row['RSI_Trend'], row['MACD_Trend']]
        bullish = trends.count('Bullish')
        bearish = trends.count('Bearish')
        neutral = trends.count('Neutral')

        if bullish > bearish and bullish > neutral:
            base_trend = 'Bullish'
        elif bearish > bullish and bearish > neutral:
            base_trend = 'Bearish'
        else:
            base_trend = 'Neutral'

        if row['ADX_Strength'] == 'Strong':
            return base_trend
        else:
            return 'Neutral'  # Treat weak trends as neutral

    df['Market_Trend'] = df.apply(combined_trend, axis=1)

    return df

def classify_volatility(df):
    """
    Classify volatility as High, Medium, or Low based on ATR, Bollinger Bands Width, and Realized Volatility.
    """
    if len(df) < 20:
        df['Volatility'] = 'Medium'
        return df

    # ATR Volatility Classification
    atr_mavg = df['ATR'].rolling(window=14).mean().shift(1)  # Shift to avoid look-ahead bias
    df['ATR_Vol'] = df.apply(
        lambda row: 'High' if row['ATR'] > 1.2 * atr_mavg.get(row.name, np.nan)
                    else ('Medium' if 0.9 * atr_mavg.get(row.name, np.nan) <= row['ATR'] <= 1.1 * atr_mavg.get(row.name, np.nan)
                          else 'Low'),
        axis=1
    )

    # Bollinger Bands Width Classification
    bb_mavg = df['BB_width'].rolling(window=20).mean().shift(1)
    df['BB_Vol'] = df.apply(
        lambda row: 'High' if row['BB_width'] > 1.2 * bb_mavg.get(row.name, np.nan)
                    else ('Medium' if 0.9 * bb_mavg.get(row.name, np.nan) <= row['BB_width'] <= 1.1 * bb_mavg.get(row.name, np.nan)
                          else 'Low'),
        axis=1
    )

    # Realized Volatility Classification
    rv_80 = df['RV'].quantile(0.8)
    rv_20 = df['RV'].quantile(0.2)
    df['RV_Vol'] = df['RV'].apply(
        lambda x: 'High' if x > rv_80
                  else ('Medium' if rv_20 <= x <= rv_80
                        else 'Low')
    )

    # Voting Mechanism for Volatility
    def volatility_vote(row):
        votes = [row['ATR_Vol'], row['BB_Vol'], row['RV_Vol']]
        high = votes.count('High')
        medium = votes.count('Medium')
        low = votes.count('Low')
        if high > medium and high > low:
            return 'High'
        elif medium > high and medium > low:
            return 'Medium'
        else:
            return 'Low'

    df['Volatility'] = df.apply(volatility_vote, axis=1)

    return df

def define_market_environment(df):
    """
    Combine Volatility and Market Trend to define Market Environment.
    """
    df['Market_Environment'] = df.apply(lambda row: f"{row['Volatility']} Vol/{row['Market_Trend']}", axis=1)
    return df

def calculate_leg_data(df):
    """
    Track the magnitude and duration of each leg (price movement direction).
    """
    previous_leg_change = 0
    previous_leg_length = 0
    current_leg_change = 0
    current_leg_length = 0

    legs = []

    for index, row in df.iterrows():
        if index == 0:
            # Initialize
            legs.append({
                'previous_leg_change': previous_leg_change,
                'previous_leg_length': previous_leg_length,
                'current_leg_change': current_leg_change,
                'current_leg_length': current_leg_length
            })
            continue

        percent_delta = (row['close'] - df.at[index-1, 'close']) / df.at[index-1, 'close']

        if current_leg_length == 0:
            current_leg_change = percent_delta
            current_leg_length = 1
        else:
            if (current_leg_change > 0 and percent_delta > 0) or (current_leg_change < 0 and percent_delta < 0):
                # Continue current leg
                current_leg_change += percent_delta
                current_leg_length += 1
            else:
                # Direction changed
                previous_leg_change = current_leg_change
                previous_leg_length = current_leg_length
                current_leg_change = percent_delta
                current_leg_length = 1

        legs.append({
            'previous_leg_change': previous_leg_change,
            'previous_leg_length': previous_leg_length,
            'current_leg_change': current_leg_change,
            'current_leg_length': current_leg_length
        })

    legs_df = pd.DataFrame(legs)
    df = pd.concat([df, legs_df], axis=1)
    return df

def classify_percent_change(df):
    """
    Discretize continuous percent changes into 7 classification buckets.
    """
    # Calculate percent changes
    df['percent_change'] = df['close'].pct_change()

    # Normalize Close Price for Regression Target
    df['normalized_close_price'] = (df['close'] - df['close'].mean()) / df['close'].std()

    # Remove NaNs
    df = df.dropna(subset=['percent_change', 'normalized_close_price']).reset_index(drop=True)

    # Determine percentiles
    percentiles = df['percent_change'].quantile([0.05, 0.20, 0.40, 0.60, 0.80, 0.95]).to_dict()

    # Define classification function
    def classify(x, p):
        if x < p[0.05]:
            return 'Down a Lot'
        elif x < p[0.20]:
            return 'Down Moderate'
        elif x < p[0.40]:
            return 'Down a Little'
        elif x < p[0.60]:
            return 'No Change'
        elif x < p[0.80]:
            return 'Up a Little'
        elif x < p[0.95]:
            return 'Up Moderate'
        else:
            return 'Up a Lot'

    # Apply classification
    df['percent_change_classification'] = df['percent_change'].apply(lambda x: classify(x, percentiles))

    return df

def update_features(symbol):
    """
    Update technical indicators and classifications for a given symbol.
    """
    with data_lock:
        df = data_frames[symbol].copy()

    if len(df) < 200:
        # Not enough data to calculate indicators and classify
        return

    # Calculate technical indicators
    df = calculate_technical_indicators(df)

    # Classify trend
    df = classify_trend(df)

    # Classify volatility
    df = classify_volatility(df)

    # Define market environment
    df = define_market_environment(df)

    # Calculate leg data
    df = calculate_leg_data(df)

    # Classify percent change
    df = classify_percent_change(df)

    with data_lock:
        data_frames[symbol] = df

# ----------------------------
# Label Encoding and Scaling
# ----------------------------

def encode_and_scale(df, label_encoders=None, scalers=None):
    """
    Encodes categorical features and scales numerical features.

    Args:
        df (pd.DataFrame): DataFrame to encode and scale.
        label_encoders (dict): Existing label encoders; creates new ones if None.
        scalers (dict): Existing scalers; creates new ones if None.

    Returns:
        df (pd.DataFrame): Encoded and scaled DataFrame.
        label_encoders (dict): Updated label encoders.
        scalers (dict): Updated scalers.
    """
    if label_encoders is None:
        label_encoders = {}

    if scalers is None:
        scalers = {}

    # Encode categorical features
    categorical_features = ['percent_change_classification', 'Market_Trend', 'Volatility']
    for col in categorical_features:
        if col not in label_encoders:
            le = LabelEncoder()
            df[col] = le.fit_transform(df[col])
            label_encoders[col] = le
        else:
            le = label_encoders[col]
            df[col] = le.transform(df[col])

    # Encode 'Market_Environment' if needed
    # Assuming 'Market_Environment' is a combined string and not directly used as a feature
    # If it's to be used as a feature, encode it similarly
    # For this script, it's used as a feature directly

    # Scale numerical features
    numerical_features = [
        'RSI', 'MACD', 'MACD_signal', 'ATR', 'BB_width', 'ADX', 
        'previous_leg_change', 'previous_leg_length',
        'current_leg_change', 'current_leg_length'
    ]
    for col in numerical_features:
        if col not in scalers:
            scaler = StandardScaler()
            df[col] = scaler.fit_transform(df[[col]])
            scalers[col] = scaler
        else:
            scaler = scalers[col]
            df[col] = scaler.transform(df[[col]])

    return df, label_encoders, scalers

# ----------------------------
# Sequence Preparation
# ----------------------------

# Define your features and targets
FEATURES = [
    'RSI', 'MACD', 'MACD_signal', 'ATR', 'BB_width', 'ADX', 
    'previous_leg_change', 'previous_leg_length',
    'current_leg_change', 'current_leg_length',
    'Volatility', 'Market_Environment'
]

TARGETS = [
    'close_price', 
    'percent_change_classification', 
    'leg_direction', 
    'leg_length'
]

def prepare_sequences(df, seq_length=SEQ_LENGTH, target_steps=TARGET_STEPS):
    """
    Prepare sequences of data for the Transformer model.

    Args:
        df (pd.DataFrame): DataFrame containing all features and targets.
        seq_length (int): Length of each input sequence.
        target_steps (int): Number of future steps to predict.

    Returns:
        X (np.ndarray): Array of input sequences.
        Y (dict): Dictionary of target arrays for each output.
    """
    sequences = []
    targets_close = []
    targets_class = []
    targets_leg_dir = []
    targets_leg_length = []

    for i in range(len(df) - seq_length - target_steps):
        seq = df.iloc[i:i+seq_length][FEATURES].values
        target_close = df.iloc[i+seq_length]['normalized_close_price']
        target_class = df.iloc[i+seq_length]['percent_change_classification']
        # Assuming 'leg_direction' is already in the DataFrame
        # If not, you need to define it based on 'current_leg_change'
        # For this example, let's define 'leg_direction' as 1 if 'current_leg_change' > 0 else 0
        leg_dir = 1 if df.iloc[i+seq_length]['current_leg_change'] > 0 else 0
        target_leg_dir = leg_dir
        target_leg_length = df.iloc[i+seq_length]['current_leg_length']

        sequences.append(seq)
        targets_close.append(target_close)
        targets_class.append(target_class)
        targets_leg_dir.append(target_leg_dir)
        targets_leg_length.append(target_leg_length)

    X = np.array(sequences)
    Y = {
        'close_price': np.array(targets_close),
        'percent_change_classification': np.array(targets_class),
        'leg_direction': np.array(targets_leg_dir),
        'leg_length': np.array(targets_leg_length)
    }

    return X, Y

# ----------------------------
# Model Architecture: Multi-Output Transformer
# ----------------------------

def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0.1):
    """
    Transformer Encoder Block.
    """
    # Layer Normalization 1
    x = LayerNormalization(epsilon=1e-6)(inputs)
    
    # Multi-Head Attention
    x = MultiHeadAttention(key_dim=head_size, num_heads=num_heads, dropout=dropout)(x, x)
    
    # Dropout
    x = Dropout(dropout)(x)
    
    # Residual Connection
    res = x + inputs
    
    # Layer Normalization 2
    x = LayerNormalization(epsilon=1e-6)(res)
    
    # Feed-Forward Network
    x = Dense(ff_dim, activation="relu")(x)
    x = Dropout(dropout)(x)
    x = Dense(inputs.shape[-1])(x)
    
    # Second Residual Connection
    return x + res

def build_multi_output_transformer_model(
    input_shape, 
    num_classes_classification, 
    head_size=256, 
    num_heads=4, 
    ff_dim=128, 
    num_transformer_blocks=2, 
    dropout=0.1
):
    """
    Builds a Transformer model with multiple outputs.
    
    Parameters:
    - input_shape: Tuple, shape of the input data (sequence_length, num_features)
    - num_classes_classification: Int, number of classes for classification tasks
    - head_size: Int, dimensionality of each attention head
    - num_heads: Int, number of attention heads
    - ff_dim: Int, dimensionality of the feed-forward network
    - num_transformer_blocks: Int, number of Transformer encoder blocks
    - dropout: Float, dropout rate
    
    Returns:
    - Keras Model object
    """
    inputs = Input(shape=input_shape, name='input')

    # Positional Encoding (Learnable)
    position_indices = tf.range(start=0, limit=input_shape[0], delta=1)
    positional_embedding = Dense(input_shape[1])(tf.expand_dims(position_indices, axis=1))
    x = inputs + positional_embedding  # Broadcasting addition

    # Transformer Encoder Blocks
    for _ in range(num_transformer_blocks):
        x = transformer_encoder(
            x, 
            head_size=head_size, 
            num_heads=num_heads, 
            ff_dim=ff_dim, 
            dropout=dropout
        )

    # Global Average Pooling
    x = GlobalAveragePooling1D()(x)
    x = Dropout(dropout)(x)
    x = Dense(128, activation="relu")(x)
    x = Dropout(dropout)(x)

    # Output: Predicted Close Price (Regression)
    close_price_output = Dense(1, activation='linear', name='close_price')(x)

    # Output: Percent Change Classification (7 classes)
    percent_change_output = Dense(num_classes_classification, activation='softmax', name='percent_change_classification')(x)

    # Output: Leg Direction (Assuming Binary Classification: Continue or Reverse)
    leg_direction_output = Dense(2, activation='softmax', name='leg_direction')(x)

    # Output: Leg Length (Regression)
    leg_length_output = Dense(1, activation='linear', name='leg_length')(x)

    # Define the Model
    model = Model(
        inputs=inputs, 
        outputs=[
            close_price_output, 
            percent_change_output, 
            leg_direction_output, 
            leg_length_output
        ]
    )
    
    return model

def compile_multi_output_model(model):
    """
    Compiles the multi-output Transformer model with appropriate loss functions and metrics.
    
    Parameters:
    - model: Keras Model object
    
    Returns:
    - Compiled Keras Model
    """
    model.compile(
        optimizer='adam',
        loss={
            'close_price': 'mse',
            'percent_change_classification': 'sparse_categorical_crossentropy',
            'leg_direction': 'sparse_categorical_crossentropy',
            'leg_length': 'mse'
        },
        metrics={
            'close_price': ['mae'],
            'percent_change_classification': ['accuracy'],
            'leg_direction': ['accuracy'],
            'leg_length': ['mae']
        }
    )
    return model

# ----------------------------
# Training Strategies: Walk-Forward Training
# ----------------------------

def walk_forward_training(
    df, 
    initial_days=20, 
    test_days=3, 
    epochs=10, 
    batch_size=32,
    label_encoders=None,
    scalers=None
):
    """
    Implements walk-forward training.
    
    Parameters:
    - df (pd.DataFrame): DataFrame containing all data.
    - initial_days (int): Number of days for initial training.
    - test_days (int): Number of days to test before retraining.
    - epochs (int): Number of training epochs.
    - batch_size (int): Training batch size.
    - label_encoders (dict): Existing label encoders.
    - scalers (dict): Existing scalers.
    
    Returns:
    - model: Trained Keras Model
    """
    # Initialize the model
    model = build_multi_output_transformer_model(
        input_shape=(SEQ_LENGTH, len(FEATURES)),
        num_classes_classification=7,  # 7-level classification
        head_size=256,
        num_heads=4,
        ff_dim=128,
        num_transformer_blocks=2,
        dropout=0.1
    )
    model = compile_multi_output_model(model)
    model.summary()

    samples_per_day = 1440  # Adjust based on your data interval (e.g., 1m intervals)
    initial_samples = initial_days * samples_per_day
    test_samples = test_days * samples_per_day

    total_samples = len(df)
    start = 0
    end = initial_samples

    while end + test_samples <= total_samples:
        train_data = df.iloc[start:end]
        test_data = df.iloc[end:end + test_samples]

        # Prepare sequences
        X_train, Y_train = prepare_sequences(train_data, seq_length=SEQ_LENGTH, target_steps=TARGET_STEPS)
        X_test, Y_test = prepare_sequences(test_data, seq_length=SEQ_LENGTH, target_steps=TARGET_STEPS)

        # Encode and scale
        X_train_df = pd.DataFrame(X_train.reshape(-1, X_train.shape[-1]), columns=FEATURES)
        X_test_df = pd.DataFrame(X_test.reshape(-1, X_test.shape[-1]), columns=FEATURES)

        X_train_df, label_encoders, scalers = encode_and_scale(X_train_df, label_encoders, scalers)
        X_test_df, _, _ = encode_and_scale(X_test_df, label_encoders, scalers)

        X_train_processed = X_train_df.values.reshape(X_train.shape)
        X_test_processed = X_test_df.values.reshape(X_test.shape)

        # Train the model
        print(f"\nTraining on data samples {start} to {end}...")
        model.fit(
            X_train_processed, Y_train,
            validation_data=(X_test_processed, Y_test),
            epochs=epochs,
            batch_size=batch_size
        )

        # Fine-tune with the most recent day
        fine_tune_start = end
        fine_tune_end = end + samples_per_day
        if fine_tune_end > total_samples:
            fine_tune_end = total_samples

        fine_tune_data = df.iloc[fine_tune_start:fine_tune_end]
        X_finetune, Y_finetune = prepare_sequences(fine_tune_data, seq_length=SEQ_LENGTH, target_steps=TARGET_STEPS)

        # Encode and scale fine-tune data
        X_finetune_df = pd.DataFrame(X_finetune.reshape(-1, X_finetune.shape[-1]), columns=FEATURES)
        X_finetune_df, _, _ = encode_and_scale(X_finetune_df, label_encoders, scalers)
        X_finetune_processed = X_finetune_df.values.reshape(X_finetune.shape)

        print(f"\nFine-tuning on data samples {fine_tune_start} to {fine_tune_end}...")
        model.fit(
            X_finetune_processed, Y_finetune,
            epochs=1,
            batch_size=batch_size
        )

        # Move the window forward
        start += test_samples
        end += test_samples

    return model

# ----------------------------
# Main Execution: Complete Script
# ----------------------------

def main():
    # ----------------------------
    # Step 1: Pre-load Historical Data
    # ----------------------------
    print("Pre-loading historical data...")
    pre_load_data(SYMBOLS, INTERVAL, limit=LIMIT)  # Adjust limit as needed

    # ----------------------------
    # Step 2: Start Live Data Collection
    # ----------------------------
    print("Starting live data collection...")
    data_thread = threading.Thread(target=load_data, args=(SYMBOLS, INTERVAL, SEQ_LENGTH))
    data_thread.daemon = True
    data_thread.start()

    # Allow some time to collect initial data
    print("Collecting initial data...")
    time.sleep(60)  # Wait for 1 minute; adjust based on your needs

    # ----------------------------
    # Step 3: Update Features for Each Symbol
    # ----------------------------
    print("Updating features for each symbol...")
    for symbol in SYMBOLS:
        update_features(symbol)

    # ----------------------------
    # Step 4: Combine Data from All Symbols
    # ----------------------------
    # For simplicity, we'll concatenate data from all symbols
    combined_df = pd.concat([data_frames[symbol] for symbol in SYMBOLS], axis=0).reset_index(drop=True)

    # ----------------------------
    # Step 5: Encode Categorical Features and Scale Numerical Features
    # ----------------------------
    print("Encoding categorical features and scaling numerical features...")
    combined_df, label_encoders, scalers = encode_and_scale(combined_df)

    # ----------------------------
    # Step 6: Prepare Sequences and Targets
    # ----------------------------
    print("Preparing input sequences and targets...")
    X, Y = prepare_sequences(combined_df, seq_length=SEQ_LENGTH, target_steps=TARGET_STEPS)

    print(f"Input shape: {X.shape}")
    print(f"Targets: { {k: v.shape for k, v in Y.items()} }")

    # ----------------------------
    # Step 7: Split Data into Training and Validation Sets
    # ----------------------------
    # Here, we use an 80-20 split. Adjust as needed.
    X_train, X_val, Y_train, Y_val = train_test_split(
        X, Y, test_size=0.2, shuffle=False
    )

    print(f"Training set shape: {X_train.shape}")
    print(f"Validation set shape: {X_val.shape}")

    # ----------------------------
    # Step 8: Build and Compile the Model
    # ----------------------------
    print("Building and compiling the model...")
    input_shape = X_train.shape[1:]  # (sequence_length, num_features)
    num_classes_classification = len(combined_df['percent_change_classification'].unique())
    model = build_multi_output_transformer_model(
        input_shape=input_shape,
        num_classes_classification=num_classes_classification,
        head_size=256,
        num_heads=4,
        ff_dim=128,
        num_transformer_blocks=2,
        dropout=0.1
    )
    model = compile_multi_output_model(model)

    # ----------------------------
    # Step 9: Train the Model
    # ----------------------------
    print("Training the model...")
    history = model.fit(
        X_train, Y_train,
        validation_data=(X_val, Y_val),
        epochs=20,  # Adjust the number of epochs as needed
        batch_size=32
    )

    # ----------------------------
    # Step 10: Evaluate the Model
    # ----------------------------
    print("Evaluating the model on validation set...")
    evaluation = model.evaluate(X_val, Y_val)
    print(f"Evaluation Results: {evaluation}")

    # ----------------------------
    # Step 11: Save the Model and Encoders/Scalers
    # ----------------------------
    print("Saving the model and encoders/scalers...")
    model.save("multi_output_transformer_model.h5")

    # Save Label Encoders
    with open("label_encoders.pkl", "wb") as f:
        pickle.dump(label_encoders, f)

    # Save Scalers
    with open("scalers.pkl", "wb") as f:
        pickle.dump(scalers, f)

    print("Model and encoders/scalers saved successfully.")

    # ----------------------------
    # Step 12: Visualization (Optional)
    # ----------------------------
    # Plot training and validation losses and accuracies for each output
    print("Plotting training history...")
    outputs = ['close_price', 'percent_change_classification', 'leg_direction', 'leg_length']
    for output in outputs:
        plt.figure(figsize=(12, 4))
        plt.plot(history.history['loss'], label='Train Loss')
        plt.plot(history.history['val_loss'], label='Validation Loss')
        plt.title(f'Loss for {output}')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.show()

        if 'classification' in output or 'direction' in output:
            plt.figure(figsize=(12, 4))
            plt.plot(history.history[f'{output}_accuracy'], label='Train Accuracy')
            plt.plot(history.history[f'val_{output}_accuracy'], label='Validation Accuracy')
            plt.title(f'Accuracy for {output}')
            plt.xlabel('Epoch')
            plt.ylabel('Accuracy')
            plt.legend()
            plt.show()

    # ----------------------------
    # Step 13: Implement Training Strategies (Walk-Forward)
    # ----------------------------
    # Uncomment and adjust as needed.
    # print("Starting Walk-Forward Training...")
    # trained_model_walk_forward = walk_forward_training(
    #     df=combined_df,
    #     initial_days=20,
    #     test_days=3,
    #     epochs=10,
    #     batch_size=32,
    #     label_encoders=label_encoders,
    #     scalers=scalers
    # )
    # print("Walk-Forward Training Complete.")

    # ----------------------------
    # Step 14: Save Training History
    # ----------------------------
    # Optionally, save the training history for later analysis
    with open("training_history.pkl", "wb") as f:
        pickle.dump(history.history, f)
    
    print("Training history saved successfully.")

    # ----------------------------
    # Step 15: Making Predictions and Extracting Confidence Scores
    # ----------------------------
    # Example function to make predictions
    def make_prediction(model, sequence, label_encoders, scalers):
        """
        Makes a prediction for a single input sequence.

        Args:
            model (tf.keras.Model): Trained Keras model.
            sequence (np.ndarray): Input sequence of shape (1, seq_length, num_features).
            label_encoders (dict): Dictionary of label encoders.
            scalers (dict): Dictionary of scalers.

        Returns:
            dict: Predictions for each output with decoded labels and confidence scores.
        """
        predictions = model.predict(sequence)
        
        # Decode Close Price (Assuming it's normalized, you might need to inverse transform)
        predicted_close_price = predictions[0][0][0]  # Regression output

        # Percent Change Classification
        percent_change_probs = predictions[1][0]
        percent_change_class = np.argmax(percent_change_probs)
        percent_change_confidence = np.max(percent_change_probs)
        percent_change_label = label_encoders['percent_change_classification'].inverse_transform([percent_change_class])[0]

        # Leg Direction
        leg_direction_probs = predictions[2][0]
        leg_direction_class = np.argmax(leg_direction_probs)
        leg_direction_confidence = np.max(leg_direction_probs)
        leg_direction_label = 'Up' if leg_direction_class == 1 else 'Down'  # Assuming 0: Down, 1: Up

        # Leg Length
        predicted_leg_length = predictions[3][0][0]  # Regression output

        return {
            'predicted_close_price': predicted_close_price,
            'percent_change_classification': {
                'label': percent_change_label,
                'confidence': percent_change_confidence
            },
            'leg_direction': {
                'label': leg_direction_label,
                'confidence': leg_direction_confidence
            },
            'predicted_leg_length': predicted_leg_length
        }

    # Example usage:
    if len(X_val) > 0:
        sample_index = 0  # Change as needed
        sample_sequence = X_val[sample_index].reshape(1, SEQ_LENGTH, len(FEATURES))  # Reshape for the model

        prediction = make_prediction(model, sample_sequence, label_encoders, scalers)
        print("\nPrediction Results:")
        print(f"Predicted Close Price: {prediction['predicted_close_price']}")
        print(f"Percent Change Classification: {prediction['percent_change_classification']['label']} "
              f"with confidence {prediction['percent_change_classification']['confidence']:.2f}")
        print(f"Leg Direction: {prediction['leg_direction']['label']} "
              f"with confidence {prediction['leg_direction']['confidence']:.2f}")
        print(f"Predicted Leg Length: {prediction['predicted_leg_length']}")  length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 9), 'open': 75000.0, 'high': 75000.0, 'low': 74999.99, 'close': 74999.99, 'volume': 0.10279}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 10), 'open': 74999.99, 'high': 75000.0, 'low': 74999.99, 'close': 75000.0, 'volume': 0.0671}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 11), 'open': 74999.99, 'high': 75000.0, 'low': 74999.99, 'close': 74999.99, 'volume': 0.03098}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 12), 'open': 74999.99, 'high': 75000.0, 'low': 74996.54, 'close': 74996.54, 'volume': 1.10308}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 13), 'open': 74996.53, 'high': 74996.54, 'low': 74992.0, 'close': 74992.01, 'volume': 1.01328}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 14), 'open': 74992.01, 'high': 74996.53, 'low': 74992.0, 'close': 74992.01, 'volume': 5.87681}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 15), 'open': 74992.0, 'high': 74992.01, 'low': 74992.0, 'close': 74992.01, 'volume': 0.00934}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 16), 'open': 74992.01, 'high': 74999.8, 'low': 74992.01, 'close': 74999.79, 'volume': 1.71529}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 17), 'open': 74999.8, 'high': 74999.8, 'low': 74999.79, 'close': 74999.8, 'volume': 0.55518}
Data length for BTCUSDT: 402
Received data for BTCUSDT: {'open_time': datetime.datetime(2024, 11, 6, 23, 49, 18), 'open': 74999.79, 'high': 74999.79, 'low': 74999.79, 'close': 74999.79, 'volume': 0.15345}
Data length for BTCUSDT: 402
Updating features for each symbol...
Encoding categorical features and scaling numerical features...
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1137: RuntimeWarning: invalid value encountered in divide
  updated_mean = (last_sum + new_sum) / updated_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1142: RuntimeWarning: invalid value encountered in divide
  T = new_sum / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1162: RuntimeWarning: invalid value encountered in divide
  new_unnormalized_variance -= correction**2 / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1137: RuntimeWarning: invalid value encountered in divide
  updated_mean = (last_sum + new_sum) / updated_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1142: RuntimeWarning: invalid value encountered in divide
  T = new_sum / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1162: RuntimeWarning: invalid value encountered in divide
  new_unnormalized_variance -= correction**2 / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1137: RuntimeWarning: invalid value encountered in divide
  updated_mean = (last_sum + new_sum) / updated_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1142: RuntimeWarning: invalid value encountered in divide
  T = new_sum / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1162: RuntimeWarning: invalid value encountered in divide
  new_unnormalized_variance -= correction**2 / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1137: RuntimeWarning: invalid value encountered in divide
  updated_mean = (last_sum + new_sum) / updated_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1142: RuntimeWarning: invalid value encountered in divide
  T = new_sum / new_sample_count
/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/sklearn/utils/extmath.py:1162: RuntimeWarning: invalid value encountered in divide
  new_unnormalized_variance -= correction**2 / new_sample_count
Preparing input sequences and targets...
Traceback (most recent call last):
  File "/home/elliot/Downloads/quant work/transformers/script7.py", line 1061, in <module>
    main()
  File "/home/elliot/Downloads/quant work/transformers/script7.py", line 878, in main
    X, Y = prepare_sequences(combined_df, seq_length=SEQ_LENGTH, target_steps=TARGET_STEPS)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/elliot/Downloads/quant work/transformers/script7.py", line 579, in prepare_sequences
    leg_dir = 1 if df.iloc[i+seq_length]['current_leg_change'] > 0 else 0
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/elliot/Downloads/quant work/transformers/trans/lib/python3.11/site-packages/pandas/core/generic.py", line 1577, in __nonzero__
    raise ValueError(
ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().


if __name__ == "__main__":
    main()