import numpy as np
import pandas as pd
from ta.volatility import AverageTrueRange, BollingerBands
from ta.momentum import RSIIndicator
from ta.trend import MACD, ADXIndicator
from sklearn.preprocessing import LabelEncoder, StandardScaler
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, Dense, Dropout, MultiHeadAttention, LayerNormalization, GlobalAveragePooling1D
from tensorflow.keras.models import Model
import tensorflow as tf
from binance.client import Client
import pickle
import seaborn as sns
def load_data(client, symbol, interval='1s', limit=1000):
    """
    Fetch historical kline data for a given symbol.
    """
    try:
        historicals = client.get_klines(symbol=symbol, interval=interval, limit=limit)
        data = []
        for kl in historicals:
            open_time = pd.to_datetime(kl[0], unit='ms')
            open_price = float(kl[1])
            high = float(kl[2])
            low = float(kl[3])
            close = float(kl[4])
            volume = float(kl[5])
            data.append([open_time, open_price, high, low, close, volume])
        df = pd.DataFrame(data, columns=['open_time', 'open', 'high', 'low', 'close', 'volume'])
        df.set_index('open_time', inplace=True)
        return df
    except Exception as e:
        print(f"Error fetching data for {symbol}: {e}")
        return pd.DataFrame()

def add_technical_indicators(df):
    """
    Add technical indicators to the DataFrame.
    """
    # RSI
    rsi = RSIIndicator(close=df['close'], window=14)
    df['RSI'] = rsi.rsi()
    
    # 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()
    
    # ATR
    atr = AverageTrueRange(high=df['high'], low=df['low'], close=df['close'], window=14)
    df['ATR'] = atr.average_true_range()
    
    # 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']
    
    # ADX
    adx = ADXIndicator(high=df['high'], low=df['low'], close=df['close'], window=14)
    df['ADX'] = adx.adx()
    
    # Drop initial rows with NaN values
    df.dropna(inplace=True)
    return df

from sklearn.preprocessing import LabelEncoder
import numpy as np



def classify_market_environments(df):
    """
    Classify market environments into numerical categories for model training.
    """
    # Volatility Classification
    df['ATR_mavg'] = df['ATR'].rolling(window=14).mean()
    df['ATR_vol'] = np.where(df['ATR'] > 1.2 * df['ATR_mavg'], 'High',
                             np.where(df['ATR'] < 0.8 * df['ATR_mavg'], 'Low', 'Medium'))
    
    df['BB_mavg'] = df['BB_width'].rolling(window=20).mean()
    df['BB_vol'] = np.where(df['BB_width'] > 1.2 * df['BB_mavg'], 'High',
                            np.where(df['BB_width'] < 0.8 * df['BB_mavg'], 'Low', 'Medium'))
    
    df['daily_return'] = df['close'].pct_change()
    df['RV'] = df['daily_return'].rolling(window=20).std() * np.sqrt(252)
    rv_80 = df['RV'].quantile(0.8)
    rv_20 = df['RV'].quantile(0.2)
    df['RV_vol'] = np.where(df['RV'] > rv_80, 'High',
                            np.where(df['RV'] < rv_20, 'Low', 'Medium'))
    
    df['Volatility'] = df[['ATR_vol', 'BB_vol', 'RV_vol']].mode(axis=1)[0]
    
    df['Trend'] = np.where((df['MACD'] > df['MACD_signal']) & (df['MACD'] > 0), 'Bullish',
                           np.where((df['MACD'] < df['MACD_signal']) & (df['MACD'] < 0), 'Bearish', 'Neutral'))
    
    df['Trend_strength'] = np.where(df['ADX'] > 25, 'Strong', 'Weak')
    
    df['Market_Environment'] = df.apply(lambda row: f"{row['Volatility']} Vol/{row['Trend']}" if row['Trend_strength'] == 'Strong' else f"{row['Volatility']} Vol/Neutral", axis=1)

    # Encode categorical data to numeric
    label_encoders = {}
    for column in ['ATR_vol', 'BB_vol', 'RV_vol', 'Volatility', 'Trend', 'Trend_strength', 'Market_Environment']:
        le = LabelEncoder()
        df[column] = le.fit_transform(df[column].astype(str))
        label_encoders[column] = le  # Save encoder for each column if you need to decode later

    return df # left encoder  # Return encoders for possible reverse transformation if needed


def calculate_leg_data(df):
    """
    Calculate leg data: current and previous leg percent change and length.
    """
    df['percent_delta'] = df['close'].pct_change()
    df = df.reset_index(drop=True)
    
    previous_leg_change = 0
    previous_leg_length = 0
    current_leg_change = 0
    current_leg_length = 0
    
    previous_changes = []
    previous_lengths = []
    current_changes = []
    current_lengths = []
    
    for i in range(len(df)):
        if i == 0:
            current_leg_change = 0
            current_leg_length = 0
        else:
            percent_delta = df.at[i, 'percent_delta']
            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):
                    current_leg_change += percent_delta
                    current_leg_length += 1
                else:
                    previous_leg_change = current_leg_change
                    previous_leg_length = current_leg_length
                    current_leg_change = percent_delta
                    current_leg_length = 1
        previous_changes.append(previous_leg_change)
        previous_lengths.append(previous_leg_length)
        current_changes.append(current_leg_change)
        current_lengths.append(current_leg_length)
    
    df['previous_leg_change'] = previous_changes
    df['previous_leg_length'] = previous_lengths
    df['current_leg_change'] = current_changes
    df['current_leg_length'] = current_lengths
    
    df.drop(columns=['percent_delta'], inplace=True)
    return df

def classify_percent_change(df):
    """
    Classify percent changes into 7 categories.
    """
    df['percent_change'] = df['close'].pct_change()
    df.dropna(inplace=True)
    
    percentiles = df['percent_change'].quantile([0.05, 0.20, 0.40, 0.60, 0.80, 0.95]).to_dict()
    
    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'
    
    df['percent_change_classification'] = df['percent_change'].apply(lambda x: classify(x, percentiles))
    return df

def encode_and_scale(df, label_encoders=None, scalers=None):
    """
    Encode categorical features and scale numerical features.
    """
    if label_encoders is None:
        label_encoders = {}
    if scalers is None:
        scalers = {}
    
    # Encode categorical features
    categorical_features = ['percent_change_classification', 'Market_Environment']
    for col in categorical_features:
        le = LabelEncoder()
        df[f'{col}_encoded'] = le.fit_transform(df[col])
        label_encoders[col] = le
    
    # Scale numerical features
    numerical_features = ['RSI', 'MACD', 'MACD_signal', 'ATR', 'BB_width',
                          'previous_leg_change', 'previous_leg_length',
                          'current_leg_change', 'current_leg_length']
    scaler = StandardScaler()
    df[numerical_features] = scaler.fit_transform(df[numerical_features])
    scalers['numerical'] = scaler
    
    return df, label_encoders, scalers

def prepare_sequences(df, seq_length=60, target_steps=1):
    """
    Prepare input sequences and corresponding targets for the Transformer.
    """
    FEATURES = ['RSI', 'MACD', 'MACD_signal', 'ATR', 'BB_width',
               'previous_leg_change', 'previous_leg_length',
               'current_leg_change', 'current_leg_length',
               'Market_Environment']
    
    TARGETS = ['percent_change_classification']
    
    # Assuming 'Market_Environment' is already encoded
    X = []
    Y = []
    
    for i in range(len(df) - seq_length - target_steps + 1):
        seq = df.iloc[i:i + seq_length][FEATURES + ['percent_change_classification_encoded']].values
        target = df.iloc[i + seq_length]['percent_change_classification_encoded']
        X.append(seq)
        Y.append(target)
    
    X = np.array(X)
    Y = np.array(Y)
    
    return X, Y
    
    


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.
    """
    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):
        # Layer Normalization 1
        x_norm = LayerNormalization(epsilon=1e-6)(x)
        
        # Multi-Head Attention
        attention = MultiHeadAttention(key_dim=head_size, num_heads=num_heads, dropout=dropout)(x_norm, x_norm)
        attention = Dropout(dropout)(attention)
        x = x + attention  # Residual connection

        # Layer Normalization 2
        x_norm = LayerNormalization(epsilon=1e-6)(x)
        
        # Feed-Forward Network
        ffn = Dense(ff_dim, activation="relu")(x_norm)
        ffn = Dropout(dropout)(ffn)
        ffn = Dense(input_shape[-1])(ffn)
        ffn = Dropout(dropout)(ffn)
        x = x + ffn  # Residual connection

    # 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 (Binary Classification)
    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.
    """
    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

create the main