# ===================================================================
# MODULE 6, LESSON 4: ALGORITHMIC TRADING AND REINFORCEMENT LEARNING
# ===================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from collections import deque
import random
import warnings
warnings.filterwarnings('ignore')

# Set style
sns.set_style("whitegrid")
np.random.seed(42)
random.seed(42)

print("="*70)
print("ALGORITHMIC TRADING WITH REINFORCEMENT LEARNING")
print("="*70)

# ----------------------------------------------------------------
# PART A: GENERATE SYNTHETIC PRICE DATA
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART A: Generating Synthetic Price Data")
print("-"*60)

def generate_price_data(n_days=1000, volatility=0.02, trend=0.0002):
    """
    Generate synthetic daily price data.
    """
    # Random walk with drift
    returns = np.random.normal(trend, volatility, n_days)
    prices = 100 * np.exp(np.cumsum(returns))
    return prices

# Generate data
prices = generate_price_data(n_days=2000, volatility=0.015, trend=0.0001)
returns = np.diff(prices) / prices[:-1]

# Create DataFrame
df = pd.DataFrame({
    'price': prices,
    'return': np.append(0, returns)
})

print(f"Data points: {len(df)}")
print(f"Price range: ${df['price'].min():.2f} – ${df['price'].max():.2f}")
print(f"Mean daily return: {df['return'].mean()*100:.4f}%")
print(f"Volatility: {df['return'].std()*100:.4f}%")

# Split into train (80%) and test (20%)
train_size = int(0.8 * len(df))
train_df = df.iloc[:train_size].reset_index(drop=True)
test_df = df.iloc[train_size:].reset_index(drop=True)

print(f"Training data: {len(train_df)} days")
print(f"Test data: {len(test_df)} days")

# Visualise
fig, ax = plt.subplots(figsize=(14, 5))
ax.plot(train_df.index, train_df['price'], label='Training Data', color='blue')
ax.plot(range(len(train_df), len(train_df)+len(test_df)), test_df['price'], label='Test Data', color='red')
ax.axvline(x=len(train_df), color='black', linestyle='--', label='Train/Test Split')
ax.set_xlabel('Day')
ax.set_ylabel('Price ($)')
ax.set_title('Synthetic Price Data')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('price_data.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART B: TRADING ENVIRONMENT
# ----------------------------------------------------------------

class TradingEnv:
    """
    A simple trading environment for RL.
    - State: price, returns, technical indicators (SMA, volatility).
    - Actions: 0 = hold, 1 = buy, 2 = sell.
    - Reward: profit from the trade (or 0 if holding).
    """
    
    def __init__(self, data, initial_balance=10000, transaction_cost=0.001):
        self.data = data
        self.initial_balance = initial_balance
        self.transaction_cost = transaction_cost
        
        # Trading parameters
        self.position = 0  # Number of shares held
        self.cash = initial_balance
        self.balance = initial_balance
        self.trades = []
        
        # Technical indicators
        self.window_sma = 20
        self.window_vol = 10
        
        # State: we'll use price, return, SMA, volatility
        self.state_dim = 4
        self.n_actions = 3  # hold, buy, sell
        
        # Reset
        self.reset()
    
    def reset(self):
        """Reset the environment to the initial state."""
        self.position = 0
        self.cash = self.initial_balance
        self.balance = self.initial_balance
        self.trades = []
        self.current_step = self.window_vol  # Start after warm-up
        return self._get_state()
    
    def _get_state(self):
        """Get the current state from the data."""
        if self.current_step < len(self.data):
            price = self.data.iloc[self.current_step]['price']
            ret = self.data.iloc[self.current_step]['return']
            
            # Simple Moving Average (SMA)
            sma = self.data.iloc[max(0, self.current_step-self.window_sma):self.current_step]['price'].mean()
            
            # Volatility (rolling std)
            vol = self.data.iloc[max(0, self.current_step-self.window_vol):self.current_step]['return'].std()
            
            # Normalise features
            state = np.array([
                price / 100,  # Normalise price
                ret * 100,    # Normalise return (percentage)
                (price - sma) / sma,  # Price relative to SMA
                vol * 100     # Normalise volatility
            ])
            return state
        else:
            return np.zeros(self.state_dim)
    
    def step(self, action):
        """
        Take an action in the environment.
        Returns: next_state, reward, done, info
        """
        # Get current price
        price = self.data.iloc[self.current_step]['price']
        
        # Execute action
        reward = 0
        done = False
        
        # Action: 0 = hold, 1 = buy, 2 = sell
        if action == 1:  # Buy
            if self.cash > 0:
                # Buy as many shares as possible
                shares_to_buy = int(self.cash * 0.95 / price)  # Keep some cash
                if shares_to_buy > 0:
                    # Transaction cost
                    cost = shares_to_buy * price * (1 + self.transaction_cost)
                    self.cash -= cost
                    self.position += shares_to_buy
                    self.trades.append({'step': self.current_step, 'action': 'buy', 
                                       'price': price, 'shares': shares_to_buy})
        
        elif action == 2:  # Sell
            if self.position > 0:
                # Sell all shares
                proceeds = self.position * price * (1 - self.transaction_cost)
                self.cash += proceeds
                self.trades.append({'step': self.current_step, 'action': 'sell', 
                                   'price': price, 'shares': self.position})
                self.position = 0
        
        # Update balance (cash + position value)
        self.balance = self.cash + self.position * price
        
        # Calculate reward: change in balance
        # But we only want to reward profitable trades, not small fluctuations
        # We'll use the change in portfolio value
        old_balance = self.balance - self.position * price * 0  # We'll track prev balance
        # Actually, use the previous balance tracked
        if not hasattr(self, 'prev_balance'):
            self.prev_balance = self.initial_balance
        
        # Reward = log change in balance (to handle scaling)
        # reward = np.log(self.balance / self.prev_balance)
        # Simpler: reward = change in balance / 100 (scaled)
        reward = (self.balance - self.prev_balance) / 100
        self.prev_balance = self.balance
        
        # Move to next step
        self.current_step += 1
        
        # Check if done
        if self.current_step >= len(self.data) - 1:
            done = True
        
        # Get next state
        next_state = self._get_state()
        
        return next_state, reward, done, {'balance': self.balance}
    
    def render(self):
        """Print the current state."""
        price = self.data.iloc[self.current_step]['price']
        print(f"Step: {self.current_step}, Price: ${price:.2f}, Cash: ${self.cash:.2f}, "
              f"Position: {self.position}, Balance: ${self.balance:.2f}")

# ----------------------------------------------------------------
# PART C: SIMPLE Q-LEARNING AGENT
# ----------------------------------------------------------------

class QLearningAgent:
    """
    A simple Q-learning agent for trading.
    Uses discretised state space for simplicity.
    """
    
    def __init__(self, state_dims, n_actions, learning_rate=0.01, 
                 discount_factor=0.9, epsilon=0.1):
        self.state_dims = state_dims
        self.n_actions = n_actions
        self.lr = learning_rate
        self.gamma = discount_factor
        self.epsilon = epsilon
        
        # Discretise state space
        self.bins = [10, 10, 10, 10]  # Bins per state dimension
        self.q_table = {}
    
    def _discretise_state(self, state):
        """Convert continuous state to discrete bins."""
        # Normalise each dimension to [0, 1] and bin it
        # We'll use simple quantisation based on typical ranges
        ranges = [
            (0.8, 1.2),   # price (0.8-1.2 of initial)
            (-0.05, 0.05), # return (-5% to 5%)
            (-0.1, 0.1),  # price/SMA deviation
            (0, 0.05)     # volatility
        ]
        
        discretised = []
        for i, (val, (low, high)) in enumerate(zip(state, ranges)):
            if val < low:
                bin_idx = 0
            elif val > high:
                bin_idx = self.bins[i] - 1
            else:
                bin_idx = int((val - low) / (high - low) * (self.bins[i] - 1))
            discretised.append(min(bin_idx, self.bins[i] - 1))
        
        return tuple(discretised)
    
    def get_q_value(self, state, action):
        """Get Q-value from q-table, initialise if not present."""
        state_key = self._discretise_state(state)
        if state_key not in self.q_table:
            self.q_table[state_key] = np.zeros(self.n_actions)
        return self.q_table[state_key][action]
    
    def act(self, state):
        """Select an action using epsilon-greedy."""
        if np.random.random() < self.epsilon:
            return np.random.randint(self.n_actions)
        else:
            state_key = self._discretise_state(state)
            if state_key not in self.q_table:
                self.q_table[state_key] = np.zeros(self.n_actions)
            return np.argmax(self.q_table[state_key])
    
    def learn(self, state, action, reward, next_state, done):
        """Update Q-values using Q-learning."""
        state_key = self._discretise_state(state)
        next_state_key = self._discretise_state(next_state)
        
        if state_key not in self.q_table:
            self.q_table[state_key] = np.zeros(self.n_actions)
        if next_state_key not in self.q_table:
            self.q_table[next_state_key] = np.zeros(self.n_actions)
        
        # Current Q-value
        q_current = self.q_table[state_key][action]
        
        # Target Q-value
        q_next_max = np.max(self.q_table[next_state_key]) if not done else 0
        q_target = reward + self.gamma * q_next_max
        
        # Update
        self.q_table[state_key][action] += self.lr * (q_target - q_current)

# ----------------------------------------------------------------
# PART D: TRAINING THE AGENT
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Training Q-Learning Agent")
print("-"*60)

# Create environment and agent
env = TradingEnv(train_df)
agent = QLearningAgent(
    state_dims=4,
    n_actions=3,
    learning_rate=0.01,
    discount_factor=0.9,
    epsilon=0.2
)

# Training parameters
n_episodes = 100
episode_rewards = []
episode_balances = []

print("Training in progress...")
for episode in range(n_episodes):
    state = env.reset()
    done = False
    total_reward = 0
    
    # Reset environment for each episode
    # We need to re-initialise the environment to reset the data pointer
    # Since we're using a fixed dataset, we can't truly reset to the start each episode.
    # Instead, we train on sequential data.
    
    # For this demonstration, we'll run through the training data
    steps = 0
    while not done and steps < len(train_df) - 10:
        # Select action
        action = agent.act(state)
        
        # Take action
        next_state, reward, done, info = env.step(action)
        
        # Learn
        agent.learn(state, action, reward, next_state, done)
        
        state = next_state
        total_reward += reward
        steps += 1
    
    episode_rewards.append(total_reward)
    episode_balances.append(env.balance)
    
    # Decay epsilon
    agent.epsilon = max(0.01, 0.2 * (1 - episode / n_episodes))
    
    if (episode + 1) % 20 == 0:
        print(f"Episode {episode+1}/{n_episodes}, Reward: {total_reward:.2f}, Balance: ${env.balance:.2f}")

# Plot training progress
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

ax = axes[0]
ax.plot(episode_rewards, 'b-', linewidth=0.5, alpha=0.5)
ax.set_xlabel('Episode')
ax.set_ylabel('Total Reward')
ax.set_title('Training Episode Rewards')
ax.grid(True, alpha=0.3)

ax = axes[1]
ax.plot(episode_balances, 'g-', linewidth=0.5, alpha=0.5)
ax.axhline(y=env.initial_balance, color='red', linestyle='--', label='Initial Balance')
ax.set_xlabel('Episode')
ax.set_ylabel('Final Balance ($)')
ax.set_title('Final Balance per Episode')
ax.legend()
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('training_progress.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART E: EVALUATION ON TEST DATA
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Evaluating on Test Data")
print("-"*60)

# Create a new environment for testing
test_env = TradingEnv(test_df, initial_balance=env.balance)

# Test the trained agent (exploration off)
agent.epsilon = 0.0
state = test_env.reset()
done = False
total_reward = 0
test_steps = 0
test_balances = [test_env.balance]
test_prices = [test_df.iloc[0]['price']]
test_actions = []

while not done and test_steps < len(test_df) - 1:
    action = agent.act(state)
    next_state, reward, done, info = test_env.step(action)
    
    state = next_state
    total_reward += reward
    test_steps += 1
    
    test_balances.append(test_env.balance)
    test_prices.append(test_df.iloc[test_steps]['price'])
    test_actions.append(action)

print(f"Test completed! {test_steps} steps")
print(f"Final balance: ${test_env.balance:.2f}")
print(f"Total return: {(test_env.balance - test_env.initial_balance) / test_env.initial_balance * 100:.2f}%")

# Calculate performance metrics
returns_series = np.diff(test_balances) / test_balances[:-1]
sharpe_ratio = np.mean(returns_series) / np.std(returns_series) * np.sqrt(252) if np.std(returns_series) > 0 else 0

# Maximum drawdown
peak = np.maximum.accumulate(test_balances)
drawdown = (peak - test_balances) / peak
max_drawdown = np.max(drawdown)

print(f"Sharpe ratio (annualised): {sharpe_ratio:.2f}")
print(f"Maximum drawdown: {max_drawdown*100:.2f}%")

# Visualise test results
fig, axes = plt.subplots(2, 1, figsize=(14, 10))

# Price and balance
ax = axes[0]
ax.plot(test_prices, label='Price', color='blue', linewidth=1.5)
ax2 = ax.twinx()
ax2.plot(test_balances, label='Portfolio Balance', color='green', linewidth=2)
ax2.axhline(y=test_env.initial_balance, color='red', linestyle='--', linewidth=1, label='Initial Balance')
ax.set_xlabel('Day')
ax.set_ylabel('Price ($)', color='blue')
ax2.set_ylabel('Balance ($)', color='green')
ax.set_title('Price vs Portfolio Balance (Test Data)')
ax.legend(loc='upper left')
ax2.legend(loc='upper right')
ax.grid(True, alpha=0.3)

# Actions
ax = axes[1]
action_names = ['Hold', 'Buy', 'Sell']
action_colors = ['gray', 'green', 'red']
for action in range(3):
    action_indices = [i for i, a in enumerate(test_actions) if a == action]
    if action_indices:
        ax.scatter(action_indices, [test_prices[i] for i in action_indices], 
                  label=action_names[action], color=action_colors[action], alpha=0.7, s=20)
ax.plot(test_prices, color='blue', alpha=0.3, linewidth=0.5)
ax.set_xlabel('Day')
ax.set_ylabel('Price ($)')
ax.set_title('Trading Actions on Test Data')
ax.legend()
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('test_results.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART F: PERFORMANCE COMPARISON WITH BUY-AND-HOLD
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Performance Comparison")
print("-"*60)

# Buy and hold strategy (buy at first day, sell at last day)
initial_investment = test_env.initial_balance
bh_shares = int(initial_investment / test_prices[0])
bh_final_value = bh_shares * test_prices[-1]
bh_return = (bh_final_value - initial_investment) / initial_investment * 100

# RL strategy return
rl_return = (test_env.balance - test_env.initial_balance) / test_env.initial_balance * 100

print(f"Buy and Hold Return: {bh_return:.2f}%")
print(f"RL Strategy Return: {rl_return:.2f}%")
print(f"Excess Return (RL - BH): {rl_return - bh_return:.2f}%")

# Comparison visualisation
fig, ax = plt.subplots(figsize=(10, 6))

# Compare cumulative returns
bh_cum = np.cumprod(1 + np.diff(test_prices) / test_prices[:-1]) - 1
rl_cum = (np.array(test_balances) - test_balances[0]) / test_balances[0]

ax.plot(bh_cum, label='Buy and Hold', color='blue', linewidth=2)
ax.plot(rl_cum, label='RL Strategy', color='green', linewidth=2)
ax.axhline(y=0, color='black', linestyle='-', alpha=0.3)
ax.set_xlabel('Day')
ax.set_ylabel('Cumulative Return')
ax.set_title('Cumulative Return: RL vs Buy and Hold')
ax.legend()
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('performance_comparison.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART G: REGULATORY AND RISK CONSIDERATIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART G: Regulatory and Risk Considerations")
print("-"*60)

print("""
Key Regulatory Requirements for Algorithmic Trading:

1. MiFID II (EU):
   - Algorithmic trading must be clearly identified.
   - Firms must have effective risk controls (circuit breakers, position limits).
   - Systems must be tested and monitored.

2. SEC Rule 15c3-5 (US):
   - Market Access Rule – requires controls on market access.
   - Pre-trade risk management.

3. FINRA Rule 3110:
   - Supervision of algorithmic trading strategies.
   - Written policies and procedures.

4. ESMA Guidelines:
   - Algorithmic trading must not create disorderly trading conditions.
   - Systems must be resilient and have fallback procedures.

Key Risks:
  - Market manipulation (spoofing, layering).
  - Flash crashes (e.g., May 2010).
  - Overfitting to historical data.
  - Model drift and degradation.
  - Systemic risk from similar strategies.

Best Practices:
  - Extensive backtesting and forward-testing (paper trading).
  - Risk limits (max position, max loss, max drawdown).
  - Human oversight (kill switches).
  - Regular model validation (per SR 11-7).
  - Compliance with market conduct rules.
  - Documentation of all algorithms and controls.
""")

# ----------------------------------------------------------------
# PART H: SUMMARY AND RECOMMENDATIONS
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART H: Summary and Recommendations")
print("="*70)

print("""
Reinforcement Learning for Trading – Key Takeaways:

1. RL is powerful for sequential decision-making problems like trading.
2. Q-learning is a good starting point for discrete action spaces.
3. DQN and PPO are state-of-the-art for more complex environments.
4. Challenges: non-stationarity, market impact, overfitting.
5. Always include transaction costs and realistic market conditions.
6. Extensive backtesting and validation are essential.
7. Regulatory compliance (MiFID II, SEC, FINRA) is non-negotiable.
8. Risk management (position limits, stop-loss) must be built in.
9. Human oversight is required – no fully autonomous trading without supervision.
10. Start with simulation and paper trading before live deployment.

Recommended Next Steps:
1. Implement DQN with experience replay for more robust learning.
2. Add technical indicators (RSI, MACD, Bollinger Bands) to state space.
3. Implement a multi-asset trading environment.
4. Incorporate transaction costs and slippage more realistically.
5. Perform walk-forward validation (rolling retraining).
6. Explore transfer learning from synthetic to real data.
7. Study market microstructure and order book dynamics.
""")

print("="*70)
print("END OF LESSON 4 – MODULE 6")
print("="*70)