dqn_base

import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from collections import deque
import random
from sklearn.preprocessing import RobustScaler
from tensorflow.keras.layers import LSTM, Dense, Dropout, BatchNormalization, Conv1D, MaxPooling1D, Activation, Add, \
    Layer, Concatenate
from tensorflow.keras.regularizers import l2


class TradingEnvironment:

    def __init__(
            self,
            data: pd.DataFrame,
            window_size: int = 30,
            commission: float = 0.001,
            spread: float = 0.0002,
            risk_free_rate: float = 0.02,
            train_split: float = 0.8
    ):
        self.data = data.copy()
        self.window_size = window_size
        self.commission = commission
        self.spread = spread
        self.risk_free_rate = risk_free_rate

        self.max_position_duration = 20
        self.position_penalty_threshold = 10
        self.inactivity_penalty = 0.001
        self.max_position_size = 0.2
        self.stop_loss = 0.02
        self.trailing_stop = 0.03

        self.base_features = ['Open', 'High', 'Low', 'Close', 'Volume']
        self.tech_features = [
            'SMA_10', 'SMA_30', 'EMA_10', 'EMA_30', 'RSI',
            'MACD', 'MACD_Signal', 'BB_middle', 'BB_upper', 'BB_lower',
            'ATR', 'OBV', 'ADX', 'Stochastic_K', 'Stochastic_D',
            'MFI', 'Volume_SMA'
        ]
        self.all_features = self.base_features + self.tech_features

        self._initialize_indicator_buffers()

        train_idx = int(len(data) * train_split)
        self.train_data = data.iloc[:train_idx]
        self.scaler = self._fit_scaler(self.train_data)

        self.n_features = len(self.all_features)
        self.state_size = self.window_size * self.n_features
        self.action_size = 3

        self.reset()

    def _initialize_indicator_buffers(self):
        buffer_size = max(30, 14 + 3)
        self.price_buffer = deque(maxlen=buffer_size)
        self.volume_buffer = deque(maxlen=buffer_size)
        self.high_buffer = deque(maxlen=buffer_size)
        self.low_buffer = deque(maxlen=buffer_size)

        self.ema_10_alpha = 2.0 / (10 + 1)
        self.ema_30_alpha = 2.0 / (30 + 1)

        self.ema_10 = None
        self.ema_30 = None
        self.macd_12 = None
        self.macd_26 = None
        self.macd_signal = None

        self.prev_obv = 0
        self.prev_high = None
        self.prev_low = None
        self.prev_close = None

    def _fit_scaler(self, train_data: pd.DataFrame) -> RobustScaler:
        scaler = RobustScaler()
        base_features_df = train_data[self.base_features]
        scaler.fit(base_features_df)

        return scaler

    def _calculate_ema(self, price: float, prev_ema: float, alpha: float) -> float:
        if prev_ema is None:
            return price

        return alpha * price + (1 - alpha) * prev_ema

    def _calculate_rolling_metrics(self, step_idx: int) -> dict:
        current_data = self.data.iloc[step_idx]
        price = current_data['Close']
        volume = current_data['Volume']
        high = current_data['High']
        low = current_data['Low']

        self.price_buffer.append(price)
        self.volume_buffer.append(volume)
        self.high_buffer.append(high)
        self.low_buffer.append(low)

        metrics = {}

        if len(self.price_buffer) >= 10:
            metrics['SMA_10'] = np.mean(list(self.price_buffer)[-10:])
        else:
            metrics['SMA_10'] = price

        if len(self.price_buffer) >= 30:
            metrics['SMA_30'] = np.mean(list(self.price_buffer)[-30:])
        else:
            metrics['SMA_30'] = price

        self.ema_10 = self._calculate_ema(price, self.ema_10, self.ema_10_alpha)
        self.ema_30 = self._calculate_ema(price, self.ema_30, self.ema_30_alpha)
        metrics['EMA_10'] = self.ema_10
        metrics['EMA_30'] = self.ema_30

        alpha_12 = 2.0 / (12 + 1)
        alpha_26 = 2.0 / (26 + 1)
        self.macd_12 = self._calculate_ema(price, self.macd_12, alpha_12)
        self.macd_26 = self._calculate_ema(price, self.macd_26, alpha_26)

        macd = self.macd_12 - self.macd_26
        metrics['MACD'] = macd

        self.macd_signal = self._calculate_ema(macd, self.macd_signal, 2.0 / (9 + 1))
        metrics['MACD_Signal'] = self.macd_signal

        if len(self.price_buffer) >= 20:
            rolling_mean = np.mean(list(self.price_buffer)[-20:])
            rolling_std = np.std(list(self.price_buffer)[-20:])
            metrics['BB_middle'] = rolling_mean
            metrics['BB_upper'] = rolling_mean + (2 * rolling_std)
            metrics['BB_lower'] = rolling_mean - (2 * rolling_std)
        else:
            metrics['BB_middle'] = price
            metrics['BB_upper'] = price
            metrics['BB_lower'] = price

        if len(self.price_buffer) >= 2:
            if len(self.price_buffer) >= 14:
                avg_gain = np.mean([max(price - self.price_buffer[-i - 2], 0)
                                    for i in range(13)])
                avg_loss = np.mean([max(self.price_buffer[-i - 2] - price, 0)
                                    for i in range(13)])

                if avg_loss == 0:
                    metrics['RSI'] = 100
                else:
                    rs = avg_gain / avg_loss
                    metrics['RSI'] = 100 - (100 / (1 + rs))
            else:
                metrics['RSI'] = 50
        else:
            metrics['RSI'] = 50

        if self.prev_high is not None:
            tr = max([
                high - low,
                abs(high - self.prev_close),
                abs(low - self.prev_close)
            ])

            if len(self.price_buffer) >= 14:
                metrics['ATR'] = np.mean([tr] + [max(
                    self.high_buffer[-i - 1] - self.low_buffer[-i - 1],
                    abs(self.high_buffer[-i - 1] - self.price_buffer[-i - 2]),
                    abs(self.low_buffer[-i - 1] - self.price_buffer[-i - 2])
                ) for i in range(13)])
            else:
                metrics['ATR'] = tr
        else:
            metrics['ATR'] = high - low

        if self.prev_close is not None:
            if price > self.prev_close:
                self.prev_obv += volume

            elif price < self.prev_close:
                self.prev_obv -= volume

        metrics['OBV'] = self.prev_obv

        if len(self.price_buffer) >= 14:
            low_14 = min(list(self.low_buffer)[-14:])
            high_14 = max(list(self.high_buffer)[-14:])

            if high_14 - low_14 > 0:
                k = 100 * (price - low_14) / (high_14 - low_14)
            else:
                k = 50

            metrics['Stochastic_K'] = k

            if len(self.price_buffer) >= 16:
                k_values = []

                for i in range(3):
                    if i + 14 <= len(self.price_buffer):
                        period_low = min(list(self.low_buffer)[-(i + 14):-i] if i > 0 else list(self.low_buffer)[-14:])
                        period_high = max(
                            list(self.high_buffer)[-(i + 14):-i] if i > 0 else list(self.high_buffer)[-14:])

                        if period_high - period_low > 0:
                            period_k = 100 * (self.price_buffer[-(i + 1)] - period_low) / (period_high - period_low)
                        else:
                            period_k = 50

                        k_values.append(period_k)

                metrics['Stochastic_D'] = np.mean(k_values) if k_values else k
            else:
                metrics['Stochastic_D'] = k
        else:
            metrics['Stochastic_K'] = 50
            metrics['Stochastic_D'] = 50

        if len(self.volume_buffer) >= 20:
            metrics['Volume_SMA'] = np.mean(list(self.volume_buffer)[-20:])
        else:
            metrics['Volume_SMA'] = volume

        if len(self.price_buffer) >= 14:
            prev_typical_prices = [
                (self.high_buffer[-i - 1] + self.low_buffer[-i - 1] + self.price_buffer[-i - 1]) / 3
                for i in range(13)
            ]

            positive_flows = sum(
                (self.high_buffer[-i - 1] + self.low_buffer[-i - 1] + self.price_buffer[-i - 1]) / 3 *
                self.volume_buffer[-i - 1]
                for i in range(14)
                if i < len(prev_typical_prices) and
                (self.high_buffer[-i - 1] + self.low_buffer[-i - 1] + self.price_buffer[-i - 1]) / 3 >
                (self.high_buffer[-i - 2] + self.low_buffer[-i - 2] + self.price_buffer[-i - 2]) / 3
            )

            negative_flows = sum(
                (self.high_buffer[-i - 1] + self.low_buffer[-i - 1] + self.price_buffer[-i - 1]) / 3 *
                self.volume_buffer[-i - 1]
                for i in range(14)
                if i < len(prev_typical_prices) and
                (self.high_buffer[-i - 1] + self.low_buffer[-i - 1] + self.price_buffer[-i - 1]) / 3 <
                (self.high_buffer[-i - 2] + self.low_buffer[-i - 2] + self.price_buffer[-i - 2]) / 3
            )

            if negative_flows == 0:
                metrics['MFI'] = 100
            else:
                money_ratio = positive_flows / negative_flows
                metrics['MFI'] = 100 - (100 / (1 + money_ratio))
        else:
            metrics['MFI'] = 50

        if len(self.price_buffer) >= 2:
            if self.prev_close is not None:
                tr = max([
                    high - low,
                    abs(high - self.prev_close),
                    abs(low - self.prev_close)
                ])
            else:
                tr = high - low

            if self.prev_high is not None and self.prev_low is not None:
                up_move = high - self.prev_high
                down_move = self.prev_low - low

                plus_dm = max(0, up_move) if up_move > down_move else 0
                minus_dm = max(0, down_move) if down_move > up_move else 0
            else:
                plus_dm = 0
                minus_dm = 0

            if not hasattr(self, 'smoothed_tr'):
                self.smoothed_tr = tr
                self.smoothed_plus_dm = plus_dm
                self.smoothed_minus_dm = minus_dm
            else:
                alpha = 1 / 14
                self.smoothed_tr = (self.smoothed_tr * (1 - alpha)) + (tr * alpha)
                self.smoothed_plus_dm = (self.smoothed_plus_dm * (1 - alpha)) + (plus_dm * alpha)
                self.smoothed_minus_dm = (self.smoothed_minus_dm * (1 - alpha)) + (minus_dm * alpha)

            if self.smoothed_tr > 0:
                plus_di = 100 * (self.smoothed_plus_dm / self.smoothed_tr)
                minus_di = 100 * (self.smoothed_minus_dm / self.smoothed_tr)

                di_sum = abs(plus_di) + abs(minus_di)
                if di_sum > 0:
                    dx = 100 * (abs(plus_di - minus_di) / di_sum)
                else:
                    dx = 0

                if not hasattr(self, 'adx'):
                    self.adx = dx
                else:
                    self.adx = (self.adx * 13 + dx) / 14

                metrics['ADX'] = self.adx
            else:
                metrics['ADX'] = 0
        else:
            metrics['ADX'] = 0

        self.prev_high = high
        self.prev_low = low
        self.prev_close = price

        return metrics

    def get_state(self) -> np.ndarray:
        if self.current_step < self.window_size:
            padding_size = self.window_size - self.current_step - 1
            state_matrix = np.zeros((self.window_size, len(self.all_features)))

            for i in range(self.current_step + 1):
                current_features = self._get_features(i)
                state_matrix[padding_size + i] = current_features

        else:
            state_matrix = np.array([
                self._get_features(i)
                for i in range(
                    self.current_step - self.window_size + 1,
                    self.current_step + 1
                )
            ])

        return state_matrix.flatten()

    def _get_features(self, idx: int) -> np.ndarray:
        base_data = self.data.iloc[idx:idx + 1][self.base_features]

        tech_data = self._calculate_rolling_metrics(idx)
        tech_values = [tech_data[feature] for feature in self.tech_features]

        scaled_base = self.scaler.transform(base_data)

        all_features = np.concatenate([
            scaled_base.flatten(),
            tech_values
        ])

        return all_features

    def step(self, action: int) -> tuple:
        current_price = self.data.iloc[self.current_step + self.window_size]['Close']

        if self.shares > 0:
            self.current_holding_period += 1
            self.consecutive_holds += 1
            self.highest_price = max(self.highest_price, current_price)
        else:
            self.current_holding_period = 0
            self.consecutive_holds = 0

        position_reward = 0
        holding_penalty = 0
        trading_cost = 0

        if self.current_holding_period >= self.max_position_duration:
            action = 1

        if self.current_holding_period > self.position_penalty_threshold:
            holding_penalty = -0.001 * (self.current_holding_period - self.position_penalty_threshold)

        if self.consecutive_holds > self.position_penalty_threshold:
            holding_penalty -= self.inactivity_penalty * self.consecutive_holds

        if self._check_stop_loss(current_price):
            action = 1

        if action == 0 and self.shares == 0:
            shares_to_buy = self._calculate_position_size(current_price)
            if shares_to_buy > 0:
                transaction_price = self._calculate_transaction_price(current_price, True)
                commission_cost = shares_to_buy * transaction_price * self.commission
                total_cost = (shares_to_buy * transaction_price) + commission_cost

                if total_cost <= self.cash:
                    self.shares = shares_to_buy
                    self.cash -= total_cost
                    self.entry_price = transaction_price
                    self.highest_price = transaction_price
                    self.trades.append(('buy', self.current_step + self.window_size, transaction_price))
                    trading_cost -= commission_cost
                    self.current_position_start = self.current_step
                    self.consecutive_holds = 0

        elif action == 1 and self.shares > 0:
            transaction_price = self._calculate_transaction_price(current_price, False)
            commission_cost = self.shares * transaction_price * self.commission
            total_revenue = (self.shares * transaction_price) - commission_cost

            self.cash += total_revenue
            trading_cost -= commission_cost
            self.trades.append(('sell', self.current_step + self.window_size, transaction_price))

            if self.current_position_start is not None:
                position_duration = self.current_step - self.current_position_start
                self.total_holding_periods.append(position_duration)

            self.shares = 0
            self.entry_price = None
            self.highest_price = 0
            self.current_position_start = None
            self.current_holding_period = 0
            self.consecutive_holds = 0

        portfolio_value = self.cash + (self.shares * current_price)
        self.total_value_history.append(portfolio_value)

        reward = 0
        if len(self.total_value_history) > 1:
            returns = (portfolio_value - self.total_value_history[-2]) / self.total_value_history[-2]
            risk_adjusted_return = returns - (self.risk_free_rate / 252)

            reward = (risk_adjusted_return * 100)
            reward += position_reward
            reward += trading_cost
            reward += holding_penalty

            if len(self.trades) == 0:
                reward -= 0.1

            if len(self.trades) > 2:
                recent_trades = self.trades[-3:]

                if len(set([t[0] for t in recent_trades])) == 1:
                    reward -= 0.5

        self.current_step += 1
        done = self.current_step >= len(self.data) - self.window_size

        if done:
            self._calculate_performance_metrics()

            if len(self.trades) < 2:
                reward -= 2.0

            if self.total_holding_periods:
                avg_holding_period = sum(self.total_holding_periods) / len(self.total_holding_periods)

                if avg_holding_period > self.position_penalty_threshold:
                    reward -= 0.5 * (avg_holding_period - self.position_penalty_threshold)

        next_state = self.get_state() if not done else None

        return next_state, reward, done

    def _calculate_position_size(self, price: float) -> int:
        portfolio_value = self.cash + (self.shares * price)
        max_position = portfolio_value * self.max_position_size

        return min(int(self.cash // price), int(max_position // price))

    def _check_stop_loss(self, current_price: float) -> bool:
        if self.shares > 0 and self.entry_price is not None:
            loss_percentage = (current_price - self.entry_price) / self.entry_price
            trailing_stop_price = self.highest_price * (1 - self.trailing_stop)
            return loss_percentage <= -self.stop_loss or current_price <= trailing_stop_price

        return False

    def _calculate_transaction_price(self, base_price: float, is_buy: bool) -> float:
        if is_buy:
            return base_price * (1 + self.spread / 2)

        return base_price * (1 - self.spread / 2)

    def _calculate_performance_metrics(self):
        returns = np.diff(self.total_value_history) / self.total_value_history[:-1]
        excess_returns = returns - (self.risk_free_rate / 252)

        self.sharpe_ratio = np.sqrt(252) * (np.mean(excess_returns) / np.std(excess_returns)) if len(returns) > 1 else 0

        cumulative_returns = np.cumprod(1 + returns)
        running_max = np.maximum.accumulate(cumulative_returns)
        drawdowns = (running_max - cumulative_returns) / running_max
        self.max_drawdown = np.max(drawdowns) if len(drawdowns) > 0 else 0

        if len(self.trades) > 1:
            profitable_trades = sum(1 for i in range(1, len(self.trades), 2)
                                    if self.trades[i][2] > self.trades[i - 1][2])
            self.win_rate = profitable_trades / (len(self.trades) // 2)
        else:
            self.win_rate = 0

    def reset(self) -> np.ndarray:
        self.current_step = 0
        self.cash = 10000
        self.shares = 0
        self.entry_price = None
        self.highest_price = 0
        self.total_value_history = [self.cash]
        self.trades = []
        self.positions = []
        self.current_position_start = None
        self.current_holding_period = 0
        self.consecutive_holds = 0
        self.total_holding_periods = []

        if hasattr(self, 'smoothed_tr'):
            del self.smoothed_tr
        if hasattr(self, 'smoothed_plus_dm'):
            del self.smoothed_plus_dm
        if hasattr(self, 'smoothed_minus_dm'):
            del self.smoothed_minus_dm
        if hasattr(self, 'adx'):
            del self.adx

        self._initialize_indicator_buffers()

        return self.get_state()

    def render(self, mode='human'):
        portfolio_value = self.total_value_history[-1]
        position = "Long" if self.shares > 0 else "None"
        print(f"Step: {self.current_step}")
        print(f"Portfolio Value: ${portfolio_value:.2f}")
        print(f"Cash: ${self.cash:.2f}")
        print(f"Position: {position}")
        print(f"Number of Trades: {len(self.trades)}")
        if len(self.trades) > 0:
            print(f"Last Trade: {self.trades[-1]}")


class SumTree:

    def __init__(self, capacity):
        self.capacity = capacity
        self.tree = np.zeros(2 * capacity - 1)
        self.data = np.zeros(capacity, dtype=object)
        self.write = 0
        self.n_entries = 0

    def __len__(self):
        return self.n_entries

    def _propagate(self, idx, change):
        parent = (idx - 1) // 2
        self.tree[parent] += change
        if parent != 0:
            self._propagate(parent, change)

    def _retrieve(self, idx, s):
        left = 2 * idx + 1
        right = left + 1

        if left >= len(self.tree):
            return idx

        if s <= self.tree[left]:
            return self._retrieve(left, s)
        else:
            return self._retrieve(right, s - self.tree[left])

    def total(self):
        return self.tree[0]

    def add(self, priority, data):
        idx = self.write + self.capacity - 1
        self.data[self.write] = data
        self.update(idx, priority)
        self.write = (self.write + 1) % self.capacity
        self.n_entries = min(self.n_entries + 1, self.capacity)

    def update(self, idx, priority):
        change = priority - self.tree[idx]
        self.tree[idx] = priority
        self._propagate(idx, change)

    def get(self, s):
        idx = self._retrieve(0, s)
        dataIdx = idx - self.capacity + 1
        return (idx, self.tree[idx], self.data[dataIdx])


class MultiHeadAttention(Layer):

    def __init__(self, head_size=32, num_heads=4):
        super(MultiHeadAttention, self).__init__()
        self.head_size = head_size
        self.num_heads = num_heads

    def build(self, input_shape):
        self.query_dense = [Dense(self.head_size) for _ in range(self.num_heads)]
        self.key_dense = [Dense(self.head_size) for _ in range(self.num_heads)]
        self.value_dense = [Dense(self.head_size) for _ in range(self.num_heads)]
        self.combine_heads = Dense(input_shape[-1])

    def call(self, x):
        attention_outputs = []

        for i in range(self.num_heads):
            query = self.query_dense[i](x)
            key = self.key_dense[i](x)
            value = self.value_dense[i](x)

            score = tf.matmul(query, key, transpose_b=True)
            score = score / tf.math.sqrt(tf.cast(self.head_size, tf.float32))
            attention_weights = tf.nn.softmax(score, axis=-1)
            head_output = tf.matmul(attention_weights, value)
            attention_outputs.append(head_output)

        multi_head_output = Concatenate()(attention_outputs)
        return self.combine_heads(multi_head_output)


class AdvancedDQNAgent:

    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        self.memory_capacity = 50000
        self.memory = SumTree(self.memory_capacity)
        self.gamma = 0.99
        self.epsilon = 1.0
        self.epsilon_min = 0.01
        self.epsilon_decay = 0.995
        self.learning_rate = 0.0001
        self.n_features = 22
        self.grad_clip_norm = 1.0
        self.action_noise_sigma = 0.1
        self.model = self._build_model()
        self.target_model = self._build_model()

        self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.learning_rate)
        self.optimizer.clipnorm = self.grad_clip_norm

        self.batch_size = 128
        self.min_replay_memory = 1000
        self.update_target_frequency = 5
        self.n_step = 3
        self.n_step_buffer = deque(maxlen=self.n_step)
        self.per_alpha = 0.6
        self.per_beta = 0.4
        self.per_beta_increment = 0.001
        self.per_epsilon = 1e-6

    @tf.function(reduce_retracing=True)
    def train_step(self, states, targets, sample_weights):
        with tf.GradientTape() as tape:
            predictions = self.model(states, training=True)
            element_wise_loss = self._huber_loss(targets, predictions)
            per_sample_loss = tf.reduce_mean(element_wise_loss, axis=1)
            weighted_loss = tf.reduce_mean(sample_weights * per_sample_loss)

        gradients = tape.gradient(weighted_loss, self.model.trainable_variables)
        gradients, _ = tf.clip_by_global_norm(gradients, self.grad_clip_norm)
        self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
        return weighted_loss

    def _build_model(self):
        sequence_length = self.state_size // self.n_features
        inputs = tf.keras.Input(shape=(self.state_size,))
        reshaped = tf.keras.layers.Reshape((sequence_length, self.n_features))(inputs)

        conv1 = Conv1D(128, 3, padding='same')(reshaped)
        conv1 = BatchNormalization()(conv1)
        conv1 = Activation('relu')(conv1)
        conv1 = Dropout(0.2)(conv1)
        conv1_residual = conv1

        conv2 = Conv1D(128, 3, padding='same')(conv1)
        conv2 = BatchNormalization()(conv2)
        conv2 = Activation('relu')(conv2)
        conv2 = Dropout(0.2)(conv2)

        conv2 = Add()([conv2, conv1_residual])
        conv2 = MaxPooling1D(2)(conv2)

        conv3 = Conv1D(256, 3, padding='same')(conv2)
        conv3 = BatchNormalization()(conv3)
        conv3 = Activation('relu')(conv3)
        conv3 = Dropout(0.2)(conv3)
        conv3_residual = conv3

        conv4 = Conv1D(256, 3, padding='same')(conv3)
        conv4 = BatchNormalization()(conv4)
        conv4 = Activation('relu')(conv4)
        conv4 = Dropout(0.2)(conv4)

        conv4 = Add()([conv4, conv3_residual])
        conv4 = MaxPooling1D(2)(conv4)

        attention_output = MultiHeadAttention(head_size=32, num_heads=4)(conv4)

        lstm1 = LSTM(128, return_sequences=True, recurrent_regularizer=l2(0.01))(attention_output)
        lstm1 = Dropout(0.2)(lstm1)
        lstm2 = LSTM(64, recurrent_regularizer=l2(0.01))(lstm1)
        lstm2 = BatchNormalization()(lstm2)

        class NoisyDense(Layer):
            def __init__(self, units, activation=None):
                super(NoisyDense, self).__init__()
                self.units = units
                self.activation = tf.keras.activations.get(activation)

            def build(self, input_shape):
                self.w_mu = self.add_weight(
                    shape=(input_shape[-1], self.units),
                    initializer='glorot_uniform',
                    trainable=True)
                self.w_sigma = self.add_weight(
                    shape=(input_shape[-1], self.units),
                    initializer='glorot_uniform',
                    trainable=True)
                self.b_mu = self.add_weight(
                    shape=(self.units,),
                    initializer='zeros',
                    trainable=True)
                self.b_sigma = self.add_weight(
                    shape=(self.units,),
                    initializer='zeros',
                    trainable=True)

            def call(self, inputs, training=None):
                if training:
                    w_epsilon = tf.random.normal(self.w_mu.shape)
                    b_epsilon = tf.random.normal(self.b_mu.shape)
                    w = self.w_mu + self.w_sigma * w_epsilon
                    b = self.b_mu + self.b_sigma * b_epsilon
                else:
                    w = self.w_mu
                    b = self.b_mu

                output = tf.matmul(inputs, w) + b
                if self.activation is not None:
                    output = self.activation(output)
                return output

        value_stream = NoisyDense(32, activation='relu')(lstm2)
        value_stream = BatchNormalization()(value_stream)
        value_stream = NoisyDense(1)(value_stream)

        advantage_stream = NoisyDense(32, activation='relu')(lstm2)
        advantage_stream = BatchNormalization()(advantage_stream)
        advantage_stream = NoisyDense(self.action_size)(advantage_stream)

        outputs = tf.keras.layers.Add()([
            value_stream,
            tf.keras.layers.Subtract()([
                advantage_stream,
                tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1, keepdims=True))(advantage_stream)
            ])
        ])

        model = tf.keras.Model(inputs=inputs, outputs=outputs)

        optimizer = tf.keras.optimizers.Adam(learning_rate=self.learning_rate)
        optimizer.clipnorm = self.grad_clip_norm

        model.compile(
            optimizer=optimizer,
            loss=self._huber_loss
        )

        return model

    def _compute_td_error(self, state, action, reward, next_state, done):
        current_q = self.model.predict(state.reshape(1, -1), verbose=0)[0][action]
        if done:
            target_q = reward
        else:
            next_q_values = self.target_model.predict(next_state.reshape(1, -1), verbose=0)[0]
            target_q = reward + self.gamma * np.max(next_q_values)
        return abs(target_q - current_q)

    def remember(self, state, action, reward, next_state, done):
        self.n_step_buffer.append((state, action, reward, next_state, done))

        if len(self.n_step_buffer) < self.n_step:
            return

        state, action, n_step_reward, next_state, done = self.n_step_buffer[0]
        for i in range(1, self.n_step):
            if self.n_step_buffer[i][4]:
                break
            n_step_reward += self.gamma ** i * self.n_step_buffer[i][2]

        td_error = self._compute_td_error(state, action, n_step_reward, next_state, done)
        priority = (td_error + self.per_epsilon) ** self.per_alpha

        self.memory.add(priority, (state, action, n_step_reward, next_state, done))

    def _huber_loss(self, y_true, y_pred, delta=1.0):
        error = y_true - y_pred
        abs_error = tf.abs(error)
        quadratic = tf.minimum(abs_error, delta)
        linear = abs_error - quadratic
        return 0.5 * tf.square(quadratic) + delta * linear

    def act(self, state, training=True):
        if training and np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)

        state = np.array(state).reshape(1, -1)
        act_values = self.model.predict(state, verbose=0)[0]

        if training:
            noise = np.random.normal(0, self.action_noise_sigma, self.action_size)
            act_values += noise

        return np.argmax(act_values)

    def evaluate(self, state):
        state = np.array(state).reshape(1, -1)
        act_values = self.model.predict(state, verbose=0)[0]
        return np.argmax(act_values)

    def replay(self):
        if len(self.memory) < self.min_replay_memory:
            return

        minibatch = []
        idxs = []
        segment = self.memory.total() / self.batch_size
        priorities = []

        self.per_beta = min(1.0, self.per_beta + self.per_beta_increment)

        for i in range(self.batch_size):
            a = segment * i
            b = segment * (i + 1)
            s = random.uniform(a, b)
            idx, priority, data = self.memory.get(s)
            priorities.append(priority)
            minibatch.append(data)
            idxs.append(idx)

        states = np.vstack([x[0] for x in minibatch])
        actions = np.array([x[1] for x in minibatch])
        rewards = np.array([x[2] for x in minibatch])
        next_states = np.vstack([x[3] for x in minibatch])
        dones = np.array([x[4] for x in minibatch])

        priorities = np.array(priorities, dtype=np.float32)
        probs = priorities / self.memory.total()
        weights = (self.memory.n_entries * probs) ** (-self.per_beta)
        weights = weights / weights.max()

        with tf.device('/CPU:0'):
            current_q_values = self.model.predict(states, batch_size=self.batch_size, verbose=0)
            next_q_values_main = self.model.predict(next_states, batch_size=self.batch_size, verbose=0)
            next_q_values_target = self.target_model.predict(next_states, batch_size=self.batch_size, verbose=0)

        targets = current_q_values.copy()

        for i in range(self.batch_size):
            if dones[i]:
                targets[i][actions[i]] = rewards[i]
            else:
                best_action = np.argmax(next_q_values_main[i])
                targets[i][actions[i]] = rewards[i] + self.gamma ** self.n_step * next_q_values_target[i][best_action]

            td_error = abs(targets[i][actions[i]] - current_q_values[i][actions[i]])
            priority = (td_error + self.per_epsilon) ** self.per_alpha
            self.memory.update(idxs[i], priority)

        states = tf.convert_to_tensor(states, dtype=tf.float32)
        targets = tf.convert_to_tensor(targets, dtype=tf.float32)
        weights = tf.convert_to_tensor(weights, dtype=tf.float32)

        loss = self.train_step(states, targets, weights)

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

        return loss

    def update_target_model(self):
        tau = 0.001
        weights = self.model.get_weights()
        target_weights = self.target_model.get_weights()

        for i in range(len(weights)):
            target_weights[i] = tau * weights[i] + (1 - tau) * target_weights[i]

        self.target_model.set_weights(target_weights)


def create_realistic_market_data(n_points=1000, volatility=0.02):
    np.random.seed(42)
    dates = pd.date_range(start='2020-01-01', periods=n_points)

    returns = np.random.normal(0.0002, volatility, n_points)
    price = 100 * np.exp(np.cumsum(returns))

    trend = np.linspace(0, 20, n_points)
    seasonal = 5 * np.sin(np.linspace(0, 8 * np.pi, n_points))
    price = price + trend + seasonal

    data = pd.DataFrame({
        'Open': price * (1 + np.random.normal(0, 0.001, n_points)),
        'High': price * (1 + np.abs(np.random.normal(0, 0.002, n_points))),
        'Low': price * (1 - np.abs(np.random.normal(0, 0.002, n_points))),
        'Close': price * (1 + np.random.normal(0, 0.001, n_points)),
        'Volume': np.random.gamma(2, 100000, n_points) * (1 + np.abs(returns))
    }, index=dates)

    data['High'] = data[['Open', 'High', 'Close']].max(axis=1)
    data['Low'] = data[['Open', 'Low', 'Close']].min(axis=1)

    return data


def visualize_training(env, episode_rewards, total_value_history, episode):
    plt.figure(figsize=(15, 12))

    plt.subplot(3, 1, 1)
    plt.plot(episode_rewards)
    plt.title(f'Episode Rewards Over Time (Episode {episode})')
    plt.xlabel('Episode')
    plt.ylabel('Total Reward')

    plt.subplot(3, 1, 2)
    plt.plot(total_value_history)
    plt.title('Portfolio Value Over Time')
    plt.xlabel('Step')
    plt.ylabel('Total Value ($)')

    plt.subplot(3, 1, 3)
    plt.plot(env.data.index[env.window_size:], env.data['Close'][env.window_size:], label='Close Price')
    plt.title('Trading Actions')
    plt.xlabel('Date')
    plt.ylabel('Price')

    for trade in env.trades:
        action, step, price = trade
        color = 'g' if action == 'buy' else 'r'
        marker = '^' if action == 'buy' else 'v'
        plt.plot(env.data.index[step], price, marker=marker, color=color, markersize=10)

    plt.legend()
    plt.tight_layout()
    plt.show()


class CustomEarlyStopping:

    def __init__(self, patience=10, min_delta: float = 0, mode='max'):
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.best = float('inf') if mode == 'min' else float('-inf')
        self.patience_counter = 0
        self.stop_training = False

    def __call__(self, current_value):
        if self.mode == 'min':
            if current_value < self.best - self.min_delta:
                self.best = current_value
                self.patience_counter = 0
            else:
                self.patience_counter += 1
        else:
            if current_value > self.best + self.min_delta:
                self.best = current_value
                self.patience_counter = 0
            else:
                self.patience_counter += 1

        if self.patience_counter >= self.patience:
            self.stop_training = True
            return True

        return False


def train_trading_agent(episodes=500, evaluation_interval=10):
    data = create_realistic_market_data(n_points=1000)
    env = TradingEnvironment(data)
    agent = AdvancedDQNAgent(env.state_size, env.action_size)

    episode_rewards = []
    best_sharpe = float('-inf')
    reward_window = deque(maxlen=10)

    early_stopping = CustomEarlyStopping(
        patience=20,
        min_delta=0.01,
        mode='max'
    )

    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        done = False

        while not done:
            action = agent.act(state)
            next_state, reward, done = env.step(action)
            total_reward += reward

            if next_state is not None:
                agent.remember(state, action, reward, next_state, done)
                state = next_state

            if len(agent.memory) >= agent.min_replay_memory:
                agent.replay()

        if episode % agent.update_target_frequency == 0:
            agent.update_target_model()

        episode_rewards.append(total_reward)
        reward_window.append(total_reward)
        avg_reward = np.mean(reward_window)

        if env.sharpe_ratio > best_sharpe:
            best_sharpe = env.sharpe_ratio
            agent.model.save('best_trading_model.keras')

        print(f"\nEpisode: {episode + 1}/{episodes}")
        print(f"Total Reward: {total_reward:.2f}")
        print(f"Final Portfolio Value: ${env.total_value_history[-1]:.2f}")
        print(f"Sharpe Ratio: {env.sharpe_ratio:.2f}")
        print(f"Max Drawdown: {env.max_drawdown:.2%}")
        print(f"Win Rate: {env.win_rate:.2%}")
        print(f"Epsilon: {agent.epsilon:.3f}")
        print(f"Average Reward (last 10): {avg_reward:.2f}")
        print("----------------------------------------")

        # if (episode + 1) % evaluation_interval == 0:
        #     visualize_training(env, episode_rewards, env.total_value_history, episode + 1)

        if early_stopping(avg_reward):
            print("Early stopping triggered!")
            break

    return agent, env, episode_rewards


def evaluate_agent(agent, data, episodes=10):
    env = TradingEnvironment(data)
    evaluation_results = []

    for episode in range(episodes):
        state = env.reset()
        done = False

        while not done:
            action = agent.act(state, training=False)
            next_state, reward, done = env.step(action)
            state = next_state if not done else None

        result = {
            'episode': episode + 1,
            'final_value': env.total_value_history[-1],
            'returns': (env.total_value_history[-1] - 10000) / 10000 * 100,
            'num_trades': len(env.trades),
            'sharpe_ratio': env.sharpe_ratio,
            'max_drawdown': env.max_drawdown,
            'win_rate': env.win_rate
        }
        evaluation_results.append(result)

        print(f"\nEvaluation Episode {episode + 1}:")
        print(f"Final Portfolio Value: ${result['final_value']:.2f}")
        print(f"Return: {result['returns']:.2f}%")
        print(f"Sharpe Ratio: {result['sharpe_ratio']:.2f}")
        print(f"Maximum Drawdown: {result['max_drawdown']:.2%}")
        print(f"Win Rate: {result['win_rate']:.2%}")
        print(f"Number of Trades: {result['num_trades']}")
        print("----------------------------------------")

    return evaluation_results


if __name__ == "__main__":
    trained_agent, final_env, rewards = train_trading_agent()

    bull_market_data = create_realistic_market_data(n_points=500, volatility=0.015)
    bear_market_data = create_realistic_market_data(n_points=500, volatility=0.025)

    print("\nBull Market Evaluation:")
    bull_results = evaluate_agent(trained_agent, bull_market_data)

    print("\nBear Market Evaluation:")
    bear_results = evaluate_agent(trained_agent, bear_market_data)

    print("\nOverall Performance Summary:")
    for market_type, results in [("Bull Market", bull_results), ("Bear Market", bear_results)]:
        returns = [r['returns'] for r in results]
        sharpe_ratios = [r['sharpe_ratio'] for r in results]
        drawdowns = [r['max_drawdown'] for r in results]
        win_rates = [r['win_rate'] for r in results]

        print(f"\n{market_type} Results:")
        print(f"Average Return: {np.mean(returns):.2f}%")
        print(f"Return Std Dev: {np.std(returns):.2f}%")
        print(f"Average Sharpe Ratio: {np.mean(sharpe_ratios):.2f}")
        print(f"Average Max Drawdown: {np.mean(drawdowns):.2%}")
        print(f"Average Win Rate: {np.mean(win_rates):.2%}")
        print(f"Best Return: {max(returns):.2f}%")
        print(f"Worst Return: {min(returns):.2f}%")