# Import necessary libraries
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple

# Set a random seed for reproducibility
np.random.seed(42)

# Module Level Docstring:
"""
This Jupyter Notebook demonstrates how perturbation theory can be applied to neural networks in financial markets.
It covers two examples:
1. Stock Price Prediction with Perturbations.
2. Portfolio Optimization with Perturbed Expected Returns.
Neural networks are built from scratch using only NumPy.
"""

# -------------------------------------------------------------
# Utility Functions for Synthetic Data Generation, Perturbations, and Normalization
# -------------------------------------------------------------

def generate_synthetic_stock_data(n_samples: int = 1000) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate synthetic stock price data for training the neural network.
    
    :param n_samples: The number of data points to generate.
    :return: Tuple of stock prices (features) and the next-day prices (targets).
    """
    # Simulate daily stock price data with a random walk
    stock_prices = np.cumsum(np.random.normal(loc=0.001, scale=0.02, size=n_samples))
    next_day_prices = np.roll(stock_prices, -1)  # Target is the price of the next day
    stock_prices = stock_prices[:-1]  # Remove the last element to align lengths
    next_day_prices = next_day_prices[:-1]
    
    return stock_prices.reshape(-1, 1), next_day_prices.reshape(-1, 1)

def generate_synthetic_portfolio_data(n_assets: int = 10, n_samples: int = 100) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate synthetic expected returns and actual returns for a portfolio optimization task.
    
    :param n_assets: Number of assets in the portfolio.
    :param n_samples: Number of samples (historical periods).
    :return: Tuple of expected returns (features) and actual returns (targets).
    """
    expected_returns = np.random.normal(loc=0.001, scale=0.02, size=(n_samples, n_assets))
    actual_returns = expected_returns + np.random.normal(loc=0, scale=0.01, size=(n_samples, n_assets))
    
    return expected_returns, actual_returns

def apply_perturbation(data: np.ndarray, perturbation_factor: float = 0.01) -> np.ndarray:
    """
    Apply a small perturbation to the input data to simulate market noise.
    
    :param data: The input stock price data.
    :param perturbation_factor: The magnitude of the perturbation.
    :return: The perturbed stock price data.
    """
    perturbation = np.random.normal(loc=0, scale=perturbation_factor, size=data.shape)
    return data + perturbation

def normalize_data(data: np.ndarray) -> np.ndarray:
    """
    Normalize the input data to have zero mean and unit variance.
    
    :param data: The input data to normalize.
    :return: Normalized data.
    """
    return (data - np.mean(data, axis=0)) / np.std(data, axis=0)

# -------------------------------------------------------------
# Basic Neural Network Class (with Normalization)
# -------------------------------------------------------------

class SimpleNeuralNetwork:
    """
    A simple feedforward neural network implemented from scratch using NumPy.
    
    This network contains:
    - One input layer, one hidden layer with ReLU activation, and one output layer.
    """
    
    def __init__(self, input_size: int, hidden_size: int, output_size: int, learning_rate: float = 0.001):
        """
        Initialize the neural network with random weights and biases.
        
        :param input_size: Number of input features.
        :param hidden_size: Number of neurons in the hidden layer.
        :param output_size: Number of output neurons.
        :param learning_rate: The learning rate for training the network.
        """
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.learning_rate = learning_rate
        
        # Initialize weights and biases
        self.W1 = np.random.randn(input_size, hidden_size)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size)
        self.b2 = np.zeros((1, output_size))
    
    def relu(self, z: np.ndarray) -> np.ndarray:
        """ReLU activation function."""
        return np.maximum(0, z)
    
    def relu_derivative(self, z: np.ndarray) -> np.ndarray:
        """Derivative of the ReLU function."""
        return np.where(z > 0, 1, 0)
    
    def forward(self, X: np.ndarray) -> np.ndarray:
        """
        Perform the forward pass through the network.
        
        :param X: Input data (features).
        :return: Output of the network.
        """
        # Hidden layer
        self.Z1 = np.dot(X, self.W1) + self.b1
        self.A1 = self.relu(self.Z1)
        
        # Output layer (no activation for regression)
        self.Z2 = np.dot(self.A1, self.W2) + self.b2
        return self.Z2
    
    def backward(self, X: np.ndarray, y: np.ndarray, output: np.ndarray):
        """
        Perform the backward pass and update weights using gradient descent.
        
        :param X: Input data (features).
        :param y: True output (targets).
        :param output: Predicted output.
        """
        # Loss derivative (mean squared error)
        dZ2 = 2 * (output - y) / y.size
        
        # Gradients for output layer
        dW2 = np.dot(self.A1.T, dZ2)
        db2 = np.sum(dZ2, axis=0, keepdims=True)
        
        # Backpropagation through the hidden layer
        dA1 = np.dot(dZ2, self.W2.T)
        dZ1 = dA1 * self.relu_derivative(self.Z1)
        
        # Gradients for hidden layer
        dW1 = np.dot(X.T, dZ1)
        db1 = np.sum(dZ1, axis=0, keepdims=True)
        
        # Update weights and biases
        self.W1 -= self.learning_rate * dW1
        self.b1 -= self.learning_rate * db1
        self.W2 -= self.learning_rate * dW2
        self.b2 -= self.learning_rate * db2
    
    def train(self, X: np.ndarray, y: np.ndarray, epochs: int = 1000):
        """
        Train the neural network using gradient descent.
        
        :param X: Training data (features).
        :param y: Training targets (true output).
        :param epochs: Number of training iterations.
        """
        for epoch in range(epochs):
            # Forward pass
            output = self.forward(X)
            
            # Backward pass and weight update
            self.backward(X, y, output)
            
            # Print loss every 100 epochs
            if epoch % 100 == 0:
                loss = np.mean((output - y) ** 2)
                print(f"Epoch {epoch}, Loss: {loss:.4f}")
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict output for given input data.
        
        :param X: Input data.
        :return: Predicted output.
        """
        return self.forward(X)

# -------------------------------------------------------------
# Example 1: Stock Price Prediction with Perturbations
# -------------------------------------------------------------

# Generate synthetic stock price data
stock_data, stock_target = generate_synthetic_stock_data()

# Create and train the neural network for stock price prediction
nn_stock = SimpleNeuralNetwork(input_size=1, hidden_size=64, output_size=1, learning_rate=0.01)
nn_stock.train(stock_data, stock_target, epochs=1000)

# Apply perturbations to the stock data
perturbed_stock_data = apply_perturbation(stock_data, perturbation_factor=0.02)

# Predict using original and perturbed data
predictions_original = nn_stock.predict(stock_data)
predictions_perturbed = nn_stock.predict(perturbed_stock_data)

# Visualization: Original vs. Perturbed Stock Price Predictions
plt.figure(figsize=(10, 6))
plt.plot(stock_target, label='True Prices', color='blue')
plt.plot(predictions_original, label='Predicted Prices (Original)', linestyle='--', color='green')
plt.plot(predictions_perturbed, label='Predicted Prices (Perturbed)', linestyle='--', color='red')
plt.title('Stock Price Prediction: Original vs. Perturbed Data')
plt.xlabel('Days')
plt.ylabel('Price')
plt.legend()
plt.show()

# -------------------------------------------------------------
# Example 2: Portfolio Optimization with Perturbed Expected Returns
# -------------------------------------------------------------

# Generate and normalize synthetic portfolio data
portfolio_data, portfolio_target = generate_synthetic_portfolio_data()
portfolio_data = normalize_data(portfolio_data)
portfolio_target = normalize_data(portfolio_target)

# Create and train a neural network for portfolio optimization
nn_portfolio = SimpleNeuralNetwork(input_size=10, hidden_size=64, output_size=10, learning_rate=0.01)
nn_portfolio.train(portfolio_data, portfolio_target, epochs=1000)

# Apply perturbations to the expected returns
perturbed_portfolio_data = apply_perturbation(portfolio_data, perturbation_factor=0.01)

# Predict using original and perturbed data
portfolio_predictions_original = nn_portfolio.predict(portfolio_data)
portfolio_predictions_perturbed = nn_portfolio.predict(perturbed_portfolio_data)

# Visualization: Original vs. Perturbed Portfolio Returns
plt.figure(figsize=(10, 6))
plt.plot(np.mean(portfolio_target, axis=1), label='True Average Returns', color='blue')
plt.plot(np.mean(portfolio_predictions_original, axis=1), label='Predicted Returns (Original)', linestyle='--', color='green')
plt.plot(np.mean(portfolio_predictions_perturbed, axis=1), label='Predicted Returns (Perturbed)', linestyle='--', color='red')
plt.title('Portfolio Optimization: Original vs. Perturbed Returns (With Normalization)')
plt.xlabel('Samples')
plt.ylabel('Average Return')
plt.legend()
plt.show()