import websocket
import json
import pandas as pd
import numpy as np
from ta.trend import SMAIndicator
from ta.momentum import RSIIndicator
from ta.volatility import AverageTrueRange
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.dates as mdates
import logging
import time
from datetime import timedelta, datetime

# Setup logging
logging.basicConfig(filename='trade_log.txt', level=logging.INFO, format='%(asctime)s - %(message)s')

# Strategy parameters
initial_capital = 100000
base_size = 0.01  # 1% of portfolio
risk_per_trade = 0.01  # 1% risk
simulation_duration = 600  # 10 minutes in seconds
buffer_threshold = 0.1  # 10% buffer for position changes
ticker = "BTC/USDT"

# TRADING COSTS - Realistic exchange fees
maker_fee = 0.001  # 0.1% maker fee (limit orders)
taker_fee = 0.001  # 0.1% taker fee (market orders) 
spread_cost = 0.0005  # 0.05% bid-ask spread cost
slippage_cost = 0.0002  # 0.02% slippage on market orders

# Data storage
data = []
position_dollars = 0  # Position size in dollars (positive = long, negative = short)
btc_units = 0  # Actual BTC units (positive = owned, negative = borrowed/short)
entry_price = 0
stop_loss = 0
trailing_stop = 0
highest_price = 0  # For long positions
lowest_price = 0   # For short positions
portfolio_value = initial_capital
cumulative_costs = 0  # Track all trading costs
trade_log = []
cost_log = []  # Track costs over time
entries = []
exits = []
df_global = pd.DataFrame()
start_time = None  # Will be set when first data arrives
current_price_base = None  # For price scaling

# Plot setup - WHITE BACKGROUND with 3 panels
plt.style.use('default')  # Use default (white) background
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(16, 12), height_ratios=[2, 1, 1])
fig.suptitle(f'Real-Time {ticker} Long/Short Strategy with Trading Costs', fontsize=16, fontweight='bold')

# Create right axis for real price
ax1_right = ax1.twinx()

# Main plot (price and SMAs)
price_line, = ax1.plot([], [], label='Price', color='blue', linewidth=2)
sma5_line, = ax1.plot([], [], label='SMA5', color='orange', linewidth=1.5)
sma15_line, = ax1.plot([], [], label='SMA15', color='green', linewidth=1.5)
stop_loss_line, = ax1.plot([], [], label='Stop Loss', color='red', linestyle='--', alpha=0.8)
trailing_stop_line, = ax1.plot([], [], label='Trailing Stop', color='purple', linestyle='--', alpha=0.8)

# Buy/Sell markers - larger and more visible
buy_scatter = ax1.scatter([], [], color='lime', marker='^', label='BUY (Long)', s=200, edgecolor='darkgreen', linewidth=2, zorder=10)
short_scatter = ax1.scatter([], [], color='orange', marker='v', label='SHORT', s=200, edgecolor='darkorange', linewidth=2, zorder=10)
close_long_scatter = ax1.scatter([], [], color='red', marker='x', label='Close Long', s=200, edgecolor='darkred', linewidth=2, zorder=10)
close_short_scatter = ax1.scatter([], [], color='cyan', marker='+', label='Close Short', s=200, edgecolor='darkcyan', linewidth=2, zorder=10)

ax1.set_ylabel('Price Change (%)', fontsize=12, fontweight='bold')
ax1_right.set_ylabel('Real Price (USDT)', fontsize=12, fontweight='bold', color='gray')
ax1.legend(loc='upper left', fontsize=10)
ax1.grid(True, alpha=0.3)
ax1.set_title('Price Chart with Long/Short Indicators', fontsize=12, fontweight='bold')

# Portfolio subplot
portfolio_line, = ax2.plot([], [], label='Portfolio Value (After Costs)', color='darkblue', linewidth=2)
portfolio_gross_line, = ax2.plot([], [], label='Portfolio Value (Before Costs)', color='lightblue', linewidth=1, linestyle='--', alpha=0.7)
ax2.axhline(y=initial_capital, color='gray', linestyle='-', alpha=0.5, label='Initial Capital')
ax2.axhline(y=initial_capital * 1.01, color='lightgreen', linestyle='--', alpha=0.7, label='+1%')
ax2.axhline(y=initial_capital * 0.99, color='lightcoral', linestyle='--', alpha=0.7, label='-1%')

ax2.set_ylabel('Portfolio Value ($)', fontsize=12, fontweight='bold')
ax2.legend(loc='upper left', fontsize=10)
ax2.grid(True, alpha=0.3)
ax2.set_title('Portfolio Performance (Net vs Gross)', fontsize=12, fontweight='bold')

# Costs subplot (NEW - 3rd panel)
costs_line, = ax3.plot([], [], label='Cumulative Costs', color='red', linewidth=2)
fees_line, = ax3.plot([], [], label='Trading Fees', color='orange', linewidth=1.5, alpha=0.8)
spread_line, = ax3.plot([], [], label='Spread Costs', color='purple', linewidth=1.5, alpha=0.8)
slippage_line, = ax3.plot([], [], label='Slippage Costs', color='brown', linewidth=1.5, alpha=0.8)

ax3.set_ylabel('Costs ($)', fontsize=12, fontweight='bold')
ax3.set_xlabel('Time (MM:SS)', fontsize=12, fontweight='bold')
ax3.legend(loc='upper left', fontsize=10)
ax3.grid(True, alpha=0.3)
ax3.set_title('Trading Costs Breakdown', fontsize=12, fontweight='bold')

def calculate_indicators(df):
    """Calculate technical indicators with proper error handling"""
    try:
        if len(df) >= 5:
            df['sma5'] = SMAIndicator(df['close'], window=5).sma_indicator()
        else:
            df['sma5'] = df['close']
            
        if len(df) >= 15:
            df['sma15'] = SMAIndicator(df['close'], window=15).sma_indicator()
        else:
            df['sma15'] = df['close']
            
        if len(df) >= 7:
            df['rsi'] = RSIIndicator(df['close'], window=7).rsi()
            df['atr'] = AverageTrueRange(df['high'], df['low'], df['close'], window=7).average_true_range()
        else:
            df['rsi'] = pd.Series([50] * len(df), index=df.index)
            df['atr'] = pd.Series([1.0] * len(df), index=df.index)
            
        if len(df) >= 10:
            df['roc'] = df['close'].pct_change(10) * 100
        else:
            df['roc'] = pd.Series([0.0] * len(df), index=df.index)
            
    except Exception as e:
        print(f"Error calculating indicators: {e}")
        df['sma5'] = df['close']
        df['sma15'] = df['close']
        df['rsi'] = pd.Series([50] * len(df), index=df.index)
        df['atr'] = pd.Series([1.0] * len(df), index=df.index)
        df['roc'] = pd.Series([0.0] * len(df), index=df.index)
    
    return df

def calculate_trading_costs(trade_value, trade_type="market"):
    """Calculate all trading costs for a given trade"""
    # Base trading fee (maker or taker)
    fee_rate = maker_fee if trade_type == "limit" else taker_fee
    trading_fee = trade_value * fee_rate
    
    # Spread cost (always applies)
    spread_cost_amount = trade_value * spread_cost
    
    # Slippage (only for market orders)
    slippage_cost_amount = trade_value * slippage_cost if trade_type == "market" else 0
    
    total_cost = trading_fee + spread_cost_amount + slippage_cost_amount
    
    return {
        'total_cost': total_cost,
        'trading_fee': trading_fee,
        'spread_cost': spread_cost_amount,
        'slippage_cost': slippage_cost_amount
    }

def get_time_in_seconds(timestamp):
    """Convert timestamp to seconds from start"""
    if start_time is None:
        return 0
    return (timestamp - start_time).total_seconds()

def format_time_axis(seconds):
    """Format seconds as MM:SS"""
    minutes = int(seconds // 60)
    secs = int(seconds % 60)
    return f"{minutes:02d}:{secs:02d}"

def get_position_type():
    """Return current position type"""
    if position_dollars > 0:
        return "LONG"
    elif position_dollars < 0:
        return "SHORT"
    else:
        return "FLAT"

# WebSocket callback functions
def on_message(ws, message):
    global data, position_dollars, btc_units, entry_price, stop_loss, trailing_stop
    global highest_price, lowest_price, portfolio_value, cumulative_costs, df_global, start_time, current_price_base
    
    try:
        msg = json.loads(message)
        if 'k' in msg:
            kline = msg['k']
            if kline['x']:  # Closed candle
                timestamp = pd.to_datetime(kline['t'], unit='ms')
                close = float(kline['c'])
                high = float(kline['h'])
                low = float(kline['l'])
                
                # Set start time on first data point
                if start_time is None:
                    start_time = timestamp
                    current_price_base = close
                    print(f"🕐 Long/Short Strategy with Costs started at {start_time}")
                    print(f"💰 Base price: ${current_price_base:.2f}")
                    print(f"💸 Trading Costs: Maker {maker_fee*100:.1f}%, Taker {taker_fee*100:.1f}%, Spread {spread_cost*100:.2f}%, Slippage {slippage_cost*100:.2f}%")
                
                # Calculate time in seconds from start
                time_seconds = get_time_in_seconds(timestamp)
                
                data.append({
                    'timestamp': timestamp, 
                    'time_seconds': time_seconds,
                    'close': close, 
                    'high': high, 
                    'low': low
                })
                
                # Keep only last 100 data points for efficiency
                data[:] = data[-100:] if len(data) > 100 else data
                
                if len(data) > 5:  # Minimum data needed
                    df = pd.DataFrame(data)
                    df = calculate_indicators(df)
                    df_global = df.copy()  # Store for plotting
                    
                    current = df.iloc[-1]
                    prev = df.iloc[-2] if len(df) > 1 else None
                    
                    # Risk management factors
                    if len(df) >= 5:
                        atr_avg = df['atr'].rolling(5).mean().iloc[-1]
                        vol_factor = 0.8 if current['atr'] > 2 * atr_avg else 1.0
                    else:
                        vol_factor = 1.0
                        
                    signal_factor = 1.2 if abs(current['rsi'] - 50) > 20 else 1.0  # Strong RSI signal
                    market_factor = 1.2 if abs(current['roc']) > 3 else (0.8 if abs(current['roc']) < 1 else 1.0)
                    
                    # ENTRY LOGIC - Only when no position
                    if position_dollars == 0 and prev is not None and len(df) >= 15:
                        
                        # LONG ENTRY: SMA5 crosses above SMA15 AND RSI > 50
                        if (current['sma5'] > current['sma15'] and prev['sma5'] <= prev['sma15'] and 
                            current['rsi'] > 50):
                            # Calculate position size in dollars (positive = long)
                            gross_position = (base_size * signal_factor * vol_factor * market_factor) * initial_capital
                            
                            # Calculate trading costs
                            costs = calculate_trading_costs(gross_position, "market")  # Market order for entries
                            cumulative_costs += costs['total_cost']
                            
                            # Net position after costs
                            position_dollars = gross_position
                            entry_price = current['close']
                            # Calculate actual BTC units (positive = owned)
                            btc_units = position_dollars / entry_price
                            stop_loss = entry_price - 1.5 * current['atr']
                            highest_price = entry_price
                            trailing_stop = stop_loss
                            entries.append((time_seconds, entry_price, 'BUY'))
                            
                            # Update portfolio with costs
                            portfolio_value -= costs['total_cost']
                            
                            print(f"🚀 LONG ENTRY: Time={format_time_axis(time_seconds)}, Price=${entry_price:.2f}")
                            print(f"   Investment: ${position_dollars:.2f}, BTC Units: {btc_units:.6f}, Stop: ${stop_loss:.2f}")
                            print(f"   💸 Entry Costs: ${costs['total_cost']:.2f} (Fee: ${costs['trading_fee']:.2f}, Spread: ${costs['spread_cost']:.2f}, Slippage: ${costs['slippage_cost']:.2f})")
                            logging.info(f"LONG: Time={format_time_axis(time_seconds)}, Price={entry_price}, Investment=${position_dollars:.2f}, BTC={btc_units:.6f}, Costs=${costs['total_cost']:.2f}")
                            
                            # Log costs
                            cost_log.append({
                                'timestamp': timestamp,
                                'time_seconds': time_seconds,
                                'cumulative_costs': cumulative_costs,
                                'trading_fees': costs['trading_fee'],
                                'spread_costs': costs['spread_cost'],
                                'slippage_costs': costs['slippage_cost']
                            })
                        
                        # SHORT ENTRY: SMA5 crosses below SMA15 AND RSI < 50
                        elif (current['sma5'] < current['sma15'] and prev['sma5'] >= prev['sma15'] and 
                              current['rsi'] < 50):
                            # Calculate position size in dollars (negative = short)
                            gross_position = (base_size * signal_factor * vol_factor * market_factor) * initial_capital
                            
                            # Calculate trading costs
                            costs = calculate_trading_costs(gross_position, "market")  # Market order for entries
                            cumulative_costs += costs['total_cost']
                            
                            # Net position after costs
                            position_dollars = -gross_position  # Negative for short
                            entry_price = current['close']
                            # Calculate actual BTC units (negative = borrowed/short)
                            btc_units = position_dollars / entry_price  # Will be negative
                            stop_loss = entry_price + 1.5 * current['atr']  # Stop above entry for shorts
                            lowest_price = entry_price
                            trailing_stop = stop_loss
                            entries.append((time_seconds, entry_price, 'SHORT'))
                            
                            # Update portfolio with costs
                            portfolio_value -= costs['total_cost']
                            
                            print(f"🔻 SHORT ENTRY: Time={format_time_axis(time_seconds)}, Price=${entry_price:.2f}")
                            print(f"   Short Size: ${abs(position_dollars):.2f}, BTC Units: {btc_units:.6f}, Stop: ${stop_loss:.2f}")
                            print(f"   💸 Entry Costs: ${costs['total_cost']:.2f} (Fee: ${costs['trading_fee']:.2f}, Spread: ${costs['spread_cost']:.2f}, Slippage: ${costs['slippage_cost']:.2f})")
                            logging.info(f"SHORT: Time={format_time_axis(time_seconds)}, Price={entry_price}, Size=${abs(position_dollars):.2f}, BTC={btc_units:.6f}, Costs=${costs['total_cost']:.2f}")
                            
                            # Log costs
                            cost_log.append({
                                'timestamp': timestamp,
                                'time_seconds': time_seconds,
                                'cumulative_costs': cumulative_costs,
                                'trading_fees': costs['trading_fee'],
                                'spread_costs': costs['spread_cost'],
                                'slippage_costs': costs['slippage_cost']
                            })
                    
                    # EXIT LOGIC
                    if position_dollars != 0 and len(df) >= 15:
                        exit_reason = ""
                        should_exit = False
                        
                        if position_dollars > 0:  # LONG EXIT
                            if current['sma5'] < current['sma15']:
                                exit_reason = "SMA Crossover (Long → Exit)"
                                should_exit = True
                            elif current['close'] < stop_loss:
                                exit_reason = "Stop Loss (Long)"
                                should_exit = True
                                
                        else:  # SHORT EXIT
                            if current['sma5'] > current['sma15']:
                                exit_reason = "SMA Crossover (Short → Exit)"
                                should_exit = True
                            elif current['close'] > stop_loss:
                                exit_reason = "Stop Loss (Short)"
                                should_exit = True
                                
                        if should_exit:
                            # Calculate exit costs
                            trade_value = abs(position_dollars)
                            costs = calculate_trading_costs(trade_value, "market")  # Market order for exits
                            cumulative_costs += costs['total_cost']
                            
                            # Profit calculation (before costs):
                            if position_dollars > 0:  # Closing long
                                gross_profit = (current['close'] - entry_price) * btc_units
                                exits.append((time_seconds, current['close'], 'CLOSE_LONG'))
                                print(f"🚪 CLOSE LONG: Time={format_time_axis(time_seconds)}, {exit_reason}, Price=${current['close']:.2f}")
                            else:  # Closing short
                                gross_profit = (entry_price - current['close']) * abs(btc_units)  # Profit when price drops
                                exits.append((time_seconds, current['close'], 'CLOSE_SHORT'))
                                print(f"🚪 CLOSE SHORT: Time={format_time_axis(time_seconds)}, {exit_reason}, Price=${current['close']:.2f}")
                            
                            # Net profit after exit costs
                            net_profit = gross_profit - costs['total_cost']
                            portfolio_value += gross_profit - costs['total_cost']  # Add net profit
                            
                            print(f"   💰 Gross P&L: ${gross_profit:+.2f}")
                            print(f"   💸 Exit Costs: ${costs['total_cost']:.2f} (Fee: ${costs['trading_fee']:.2f}, Spread: ${costs['spread_cost']:.2f}, Slippage: ${costs['slippage_cost']:.2f})")
                            print(f"   📊 Net P&L: ${net_profit:+.2f}, Portfolio: ${portfolio_value:,.2f}")
                            
                            logging.info(f"EXIT: Time={format_time_axis(time_seconds)}, Price={current['close']}, Reason={exit_reason}, GrossProfit=${gross_profit:.2f}, NetProfit=${net_profit:.2f}, Costs=${costs['total_cost']:.2f}")
                            
                            # Log costs
                            cost_log.append({
                                'timestamp': timestamp,
                                'time_seconds': time_seconds,
                                'cumulative_costs': cumulative_costs,
                                'trading_fees': costs['trading_fee'],
                                'spread_costs': costs['spread_cost'],
                                'slippage_costs': costs['slippage_cost']
                            })
                            
                            # Reset position variables
                            position_dollars = 0
                            btc_units = 0
                            entry_price = 0
                            stop_loss = 0
                            trailing_stop = 0
                            highest_price = 0
                            lowest_price = 0
                    
                    # Update portfolio tracking
                    if position_dollars != 0:
                        if position_dollars > 0:  # Long position
                            current_btc_value = btc_units * current['close']
                            unrealized_pnl = current_btc_value - position_dollars
                        else:  # Short position
                            unrealized_pnl = (entry_price - current['close']) * abs(btc_units)
                            
                        current_portfolio_value = portfolio_value + unrealized_pnl
                        trade_log.append({'timestamp': timestamp, 'time_seconds': time_seconds, 'portfolio_value': current_portfolio_value})
                    else:
                        trade_log.append({'timestamp': timestamp, 'time_seconds': time_seconds, 'portfolio_value': portfolio_value})
                    
                    # Always log current costs (even if no new trades)
                    if not cost_log or cost_log[-1]['time_seconds'] != time_seconds:
                        cost_log.append({
                            'timestamp': timestamp,
                            'time_seconds': time_seconds,
                            'cumulative_costs': cumulative_costs,
                            'trading_fees': 0,  # No new fees
                            'spread_costs': 0,
                            'slippage_costs': 0
                        })
                    
                    # Print status every 10 seconds
                    if len(data) % 10 == 0:
                        pos_type = get_position_type()
                        status_emoji = {"LONG": "📈", "SHORT": "📉", "FLAT": "⏸️"}[pos_type]
                        price_change = ((current['close'] - current_price_base) / current_price_base) * 100
                        portfolio_change = ((portfolio_value - initial_capital) / initial_capital) * 100
                        cost_percentage = (cumulative_costs / initial_capital) * 100
                        
                        if position_dollars != 0:
                            if position_dollars > 0:  # Long
                                current_btc_value = btc_units * current['close']
                                unrealized_pnl = current_btc_value - position_dollars
                                print(f"{status_emoji} {pos_type} | Time: {format_time_axis(time_seconds)} | Price: ${current['close']:.2f} ({price_change:+.2f}%)")
                                print(f"   Investment: ${position_dollars:.2f} | Current Value: ${current_btc_value:.2f} | Unrealized P&L: ${unrealized_pnl:+.2f}")
                                print(f"   💸 Total Costs: ${cumulative_costs:.2f} ({cost_percentage:.3f}% of capital)")
                            else:  # Short
                                unrealized_pnl = (entry_price - current['close']) * abs(btc_units)
                                print(f"{status_emoji} {pos_type} | Time: {format_time_axis(time_seconds)} | Price: ${current['close']:.2f} ({price_change:+.2f}%)")
                                print(f"   Short Size: ${abs(position_dollars):.2f} | Entry: ${entry_price:.2f} | Unrealized P&L: ${unrealized_pnl:+.2f}")
                                print(f"   💸 Total Costs: ${cumulative_costs:.2f} ({cost_percentage:.3f}% of capital)")
                        else:
                            print(f"{status_emoji} {pos_type} | Time: {format_time_axis(time_seconds)} | Price: ${current['close']:.2f} ({price_change:+.2f}%) | Portfolio: ${portfolio_value:,.2f} ({portfolio_change:+.2f}%)")
                            print(f"   💸 Total Costs: ${cumulative_costs:.2f} ({cost_percentage:.3f}% of capital)")
                    
    except Exception as e:
        print(f"❌ Error in on_message: {e}")

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

def on_close(ws, close_status_code, close_msg):
    print("🔌 WebSocket closed")
    if trade_log:
        pd.DataFrame(trade_log).to_csv('trade_log.csv', index=False)
        print("💾 Trade log saved")
    if cost_log:
        pd.DataFrame(cost_log).to_csv('cost_log.csv', index=False)
        print("💾 Cost log saved")

def on_open(ws):
    print("🔗 WebSocket opened - Starting long/short trend following strategy with costs")

# Animation update function
def update_plot(frame):
    global df_global, current_price_base
    
    try:
        if df_global.empty or len(df_global) == 0 or current_price_base is None:
            return (price_line, sma5_line, sma15_line, stop_loss_line, trailing_stop_line, 
                    portfolio_line, portfolio_gross_line, costs_line, fees_line, spread_line, slippage_line,
                    buy_scatter, short_scatter, close_long_scatter, close_short_scatter)
        
        df = df_global.copy()
        
        # Calculate price change percentage for left axis
        df['price_change_pct'] = ((df['close'] - current_price_base) / current_price_base) * 100
        df['sma5_change_pct'] = ((df['sma5'] - current_price_base) / current_price_base) * 100
        df['sma15_change_pct'] = ((df['sma15'] - current_price_base) / current_price_base) * 100
        
        # Update main plot - price change percentage
        price_line.set_data(df['time_seconds'], df['price_change_pct'])
        
        if 'sma5' in df.columns and not df['sma5'].isna().all():
            sma5_line.set_data(df['time_seconds'], df['sma5_change_pct'])
        else:
            sma5_line.set_data([], [])
            
        if 'sma15' in df.columns and not df['sma15'].isna().all():
            sma15_line.set_data(df['time_seconds'], df['sma15_change_pct'])
        else:
            sma15_line.set_data([], [])
        
        # Update markers - separate long and short entries/exits
        long_entries = [(t, p) for t, p, action in entries if action == 'BUY']
        short_entries = [(t, p) for t, p, action in entries if action == 'SHORT']
        long_exits = [(t, p) for t, p, action in exits if action == 'CLOSE_LONG']
        short_exits = [(t, p) for t, p, action in exits if action == 'CLOSE_SHORT']
        
        if long_entries:
            buy_times, buy_prices = zip(*long_entries)
            buy_prices_pct = [((price - current_price_base) / current_price_base) * 100 for price in buy_prices]
            buy_scatter.set_offsets(np.c_[buy_times, buy_prices_pct])
        else:
            buy_scatter.set_offsets(np.array([]).reshape(0, 2))
            
        if short_entries:
            short_times, short_prices = zip(*short_entries)
            short_prices_pct = [((price - current_price_base) / current_price_base) * 100 for price in short_prices]
            short_scatter.set_offsets(np.c_[short_times, short_prices_pct])
        else:
            short_scatter.set_offsets(np.array([]).reshape(0, 2))
            
        if long_exits:
            close_long_times, close_long_prices = zip(*long_exits)
            close_long_prices_pct = [((price - current_price_base) / current_price_base) * 100 for price in close_long_prices]
            close_long_scatter.set_offsets(np.c_[close_long_times, close_long_prices_pct])
        else:
            close_long_scatter.set_offsets(np.array([]).reshape(0, 2))
            
        if short_exits:
            close_short_times, close_short_prices = zip(*short_exits)
            close_short_prices_pct = [((price - current_price_base) / current_price_base) * 100 for price in close_short_prices]
            close_short_scatter.set_offsets(np.c_[close_short_times, close_short_prices_pct])
        else:
            close_short_scatter.set_offsets(np.array([]).reshape(0, 2))
        
        # Handle trade_log for portfolio (net and gross)
        if trade_log:
            trade_df = pd.DataFrame(trade_log)
            portfolio_line.set_data(trade_df['time_seconds'], trade_df['portfolio_value'])
            
            # Calculate gross portfolio (before costs)
            gross_portfolio = trade_df['portfolio_value'] + cumulative_costs
            portfolio_gross_line.set_data(trade_df['time_seconds'], gross_portfolio)
        
        # Handle cost_log for costs breakdown
        if cost_log:
            cost_df = pd.DataFrame(cost_log)
            costs_line.set_data(cost_df['time_seconds'], cost_df['cumulative_costs'])
            
            # Calculate cumulative costs by type
            cost_df['cum_fees'] = cost_df['trading_fees'].cumsum()
            cost_df['cum_spread'] = cost_df['spread_costs'].cumsum()
            cost_df['cum_slippage'] = cost_df['slippage_costs'].cumsum()
            
            fees_line.set_data(cost_df['time_seconds'], cost_df['cum_fees'])
            spread_line.set_data(cost_df['time_seconds'], cost_df['cum_spread'])
            slippage_line.set_data(cost_df['time_seconds'], cost_df['cum_slippage'])
        
        # Fixed time scale: 0 to 600 seconds (10 minutes)
        ax1.set_xlim(0, 600)
        ax2.set_xlim(0, 600)
        ax3.set_xlim(0, 600)
        
        # Price scale: ±2%
        ax1.set_ylim(-2, 2)
        
        # Right axis for real prices
        current_price = df['close'].iloc[-1]
        price_range = current_price * 0.02  # 2%
        ax1_right.set_ylim(current_price - price_range, current_price + price_range)
        
        # Portfolio scale: initial capital ±1%
        capital_range = initial_capital * 0.01  # 1%
        ax2.set_ylim(initial_capital - capital_range, initial_capital + capital_range)
        
        # Costs scale: 0 to max costs with some headroom
        if cumulative_costs > 0:
            ax3.set_ylim(0, cumulative_costs * 1.2)
        else:
            ax3.set_ylim(0, 100)  # Default range
        
        # Set tick marks for time (every minute = 60 seconds)
        time_ticks = range(0, 601, 60)  # 0, 60, 120, ..., 600
        time_labels = [format_time_axis(t) for t in time_ticks]
        for ax in [ax1, ax2, ax3]:
            ax.set_xticks(time_ticks)
            ax.set_xticklabels(time_labels)
        
        # Add grid lines for better readability
        ax1.grid(True, alpha=0.3)
        ax2.grid(True, alpha=0.3)
        ax3.grid(True, alpha=0.3)
        
    except Exception as e:
        print(f"Plot update error: {e}")
    
    return (price_line, sma5_line, sma15_line, stop_loss_line, trailing_stop_line, 
            portfolio_line, portfolio_gross_line, costs_line, fees_line, spread_line, slippage_line,
            buy_scatter, short_scatter, close_long_scatter, close_short_scatter)

# Run WebSocket and animation
if __name__ == "__main__":
    print("🚀 Starting Enhanced Long/Short Strategy with Trading Costs")
    print("📈 LONG: SMA5 crosses above SMA15 + RSI > 50")
    print("📉 SHORT: SMA5 crosses below SMA15 + RSI < 50")
    print("💸 Trading Costs Included:")
    print(f"   • Maker Fee: {maker_fee*100:.1f}%")
    print(f"   • Taker Fee: {taker_fee*100:.1f}%")
    print(f"   • Spread Cost: {spread_cost*100:.2f}%")
    print(f"   • Slippage Cost: {slippage_cost*100:.2f}%")
    print("📊 3-Panel View: Price/SMAs, Portfolio (Net vs Gross), Trading Costs")
    print("⚡ Realistic trading with all costs included!")
    
    ws_url = "wss://stream.binance.com:9443/ws/btcusdt@kline_1s"
    ws = websocket.WebSocketApp(ws_url, on_message=on_message, on_error=on_error, 
                                on_close=on_close, on_open=on_open)
    
    import threading
    ws_thread = threading.Thread(target=ws.run_forever)
    ws_thread.daemon = True
    ws_thread.start()
    
    try:
        ani = FuncAnimation(fig, update_plot, interval=1000, blit=True, cache_frame_data=False)
        plt.tight_layout()
        plt.show()
        
        time.sleep(simulation_duration)
    except KeyboardInterrupt:
        print("\n⚠️ Interrupted by user")
    except Exception as e:
        print(f"❌ Error: {e}")
    finally:
        print("🛑 Shutting down...")
        ws.close()
        
        # Print final summary with costs analysis
        if entries or exits:
            print(f"\n📊 LONG/SHORT TRADING SUMMARY WITH COSTS:")
            long_entries_count = len([e for e in entries if e[2] == 'BUY'])
            short_entries_count = len([e for e in entries if e[2] == 'SHORT'])
            long_exits_count = len([e for e in exits if e[2] == 'CLOSE_LONG'])
            short_exits_count = len([e for e in exits if e[2] == 'CLOSE_SHORT'])
            
            print(f"   Long Entries: {long_entries_count}")
            print(f"   Short Entries: {short_entries_count}")
            print(f"   Long Exits: {long_exits_count}")
            print(f"   Short Exits: {short_exits_count}")
            print(f"   Total Trades: {long_entries_count + short_entries_count}")
            
            # Cost analysis
            cost_percentage = (cumulative_costs / initial_capital) * 100
            print(f"\n💸 COST ANALYSIS:")
            print(f"   Total Trading Costs: ${cumulative_costs:.2f}")
            print(f"   Cost as % of Capital: {cost_percentage:.3f}%")
            print(f"   Average Cost per Trade: ${cumulative_costs/(long_entries_count + short_entries_count):.2f}" if (long_entries_count + short_entries_count) > 0 else "   No completed trades")
            
            # Performance analysis
            gross_return = ((portfolio_value + cumulative_costs - initial_capital) / initial_capital) * 100
            net_return = ((portfolio_value - initial_capital) / initial_capital) * 100
            print(f"\n📈 PERFORMANCE ANALYSIS:")
            print(f"   Gross Return (before costs): {gross_return:+.2f}%")
            print(f"   Net Return (after costs): {net_return:+.2f}%")
            print(f"   Cost Impact: {gross_return - net_return:.2f}%")
            print(f"   Final Portfolio: ${portfolio_value:,.2f}")
        
        print("✅ Long/Short Strategy with Costs completed!")