"""
This module demonstrates a Fourier Neural Network applied to synthetic financial time series data.

A Fourier Neural Network uses sinusoidal activation functions to model periodic patterns in data,
which can be particularly useful for capturing cyclical behaviors in financial time series.

In this example, we generate synthetic stock price data with underlying periodic components,
train the network to predict future stock prices, and visualize the results.
"""

import numpy as np
import matplotlib.pyplot as plt

def generate_synthetic_stock_data(num_points: int, noise_level: float = 0.1) -> tuple:
    """
    Generates synthetic stock price data with periodic components, trends, and noise.

    Args:
        num_points (int): Number of data points to generate.
        noise_level (float): Standard deviation of the Gaussian noise.

    Returns:
        tuple: Tuple containing time steps and corresponding stock prices.
    """
    np.random.seed(0)  # For reproducibility
    t: np.ndarray = np.arange(num_points)
    trend: np.ndarray = t * 0.05  # Linear trend
    seasonal: np.ndarray = 2 * np.sin(0.02 * t) + 0.5 * np.sin(0.1 * t)  # Seasonal components
    noise: np.ndarray = np.random.normal(0, noise_level, num_points)  # Random noise
    stock_prices: np.ndarray = 50 + trend + seasonal + noise  # Synthetic stock prices
    return t, stock_prices

class FourierNeuralNetwork:
    """
    A Fourier Neural Network that uses sinusoidal activation functions to model periodic data.

    Attributes:
        W1 (np.ndarray): Weights of the first (input to hidden) layer.
        b1 (np.ndarray): Biases of the first layer.
        W2 (np.ndarray): Weights of the second (hidden to output) layer.
        b2 (np.ndarray): Biases of the second layer.
    """

    def __init__(self, input_size: int, hidden_size: int, output_size: int):
        """
        Initializes the Fourier Neural Network with random weights and biases.

        Args:
            input_size (int): Number of input neurons.
            hidden_size (int): Number of neurons in the hidden layer.
            output_size (int): Number of output neurons.
        """
        # Scale down the weights to prevent large activations
        self.W1: np.ndarray = np.random.randn(input_size, hidden_size) * 0.1
        self.b1: np.ndarray = np.zeros((1, hidden_size))
        self.W2: np.ndarray = np.random.randn(hidden_size, output_size) * 0.1
        self.b2: np.ndarray = np.zeros((1, output_size))

    def forward(self, X: np.ndarray) -> np.ndarray:
        """
        Performs a forward pass through the network.

        Args:
            X (np.ndarray): Input data of shape (n_samples, input_size).

        Returns:
            np.ndarray: Output of the network of shape (n_samples, output_size).
        """
        self.Z1: np.ndarray = np.dot(X, self.W1) + self.b1  # Linear combination
        self.A1: np.ndarray = np.sin(self.Z1)               # Sinusoidal activation
        Z2: np.ndarray = np.dot(self.A1, self.W2) + self.b2 # Output layer
        return Z2

    def compute_loss(self, Y_pred: np.ndarray, Y_true: np.ndarray) -> float:
        """
        Computes the Mean Squared Error loss.

        Args:
            Y_pred (np.ndarray): Predicted outputs.
            Y_true (np.ndarray): True outputs.

        Returns:
            float: The Mean Squared Error loss.
        """
        loss: float = np.mean((Y_pred - Y_true) ** 2)
        return loss

    def train(self, X: np.ndarray, Y: np.ndarray, learning_rate: float, epochs: int):
        """
        Trains the network using gradient descent optimization.

        Args:
            X (np.ndarray): Input data of shape (n_samples, input_size).
            Y (np.ndarray): Target data of shape (n_samples, output_size).
            learning_rate (float): Learning rate for weight updates.
            epochs (int): Number of times to iterate over the training data.
        """
        n_samples: int = X.shape[0]
        for epoch in range(epochs):
            # Forward pass
            Y_pred: np.ndarray = self.forward(X)

            # Compute loss
            loss: float = self.compute_loss(Y_pred, Y)

            # Backward pass
            dY_pred: np.ndarray = (2 / n_samples) * (Y_pred - Y)        # Loss derivative w.r.t Y_pred
            dW2: np.ndarray = np.dot(self.A1.T, dY_pred)                # Gradient for W2
            db2: np.ndarray = np.sum(dY_pred, axis=0, keepdims=True)    # Gradient for b2

            dA1: np.ndarray = np.dot(dY_pred, self.W2.T)                # Gradient w.r.t activation
            dZ1: np.ndarray = dA1 * np.cos(self.Z1)                     # Gradient w.r.t pre-activation
            dW1: np.ndarray = np.dot(X.T, dZ1)                          # Gradient for W1
            db1: np.ndarray = np.sum(dZ1, axis=0, keepdims=True)        # Gradient for b1

            # Update weights and biases
            self.W1 -= learning_rate * dW1
            self.b1 -= learning_rate * db1
            self.W2 -= learning_rate * dW2
            self.b2 -= learning_rate * db2

            # Print the loss every 10% of the total epochs
            if (epoch + 1) % max(1, epochs // 10) == 0 or epoch == 0:
                print(f"Epoch {epoch+1}/{epochs}, Loss: {loss:.6f}")