import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# From NumPy
np_array = np.random.randn(1000, 10)
tensor = torch.from_numpy(np_array).float()

# Zeros, ones, identity
zeros = torch.zeros(3, 4)
ones = torch.ones(3, 4)
eye = torch.eye(5)

# Random tensors
uniform = torch.rand(1000, 10)  # Uniform [0, 1]
normal = torch.randn(1000, 10)  # Standard normal

# With specific dtype
double_tensor = torch.randn(1000, 10, dtype=torch.float64)
int_tensor = torch.randint(0, 10, (1000,))

# Element-wise operations (vectorised)
a = torch.randn(1000, 10)
b = torch.randn(1000, 10)

c = a + b          # Addition
d = a * b          # Element-wise multiplication
e = a @ b.T        # Matrix multiplication
f = torch.matmul(a, b.T)  # Same as @

# Mathematical functions
mean = a.mean(dim=0)          # Mean along columns
std = a.std(dim=0, unbiased=True)  # Sample standard deviation
log_returns = torch.log(a / a.shift(1))  # Log returns

# Reduction operations
sum_all = a.sum()
max_val = a.max()
min_val = a.min()
argmax = a.argmax(dim=1)      # Index of max along rows

# Check GPU availability
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# Move tensors to GPU
tensor_gpu = tensor.to(device)

# Multi-GPU setup (if available)
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

# Pin memory for faster data transfer
pin_memory = True

# Use in-place operations to save memory
x = torch.randn(10000, 10000)
x.add_(1)           # In-place addition (modifies x)
x.mul_(2)           # In-place multiplication

# Clear GPU cache when needed
torch.cuda.empty_cache()

# Use float16 for memory-constrained environments
x_half = x.half()   # Converts to float16

# Create a tensor with gradient tracking
x = torch.tensor([2.0, 3.0], requires_grad=True)

# Define a function
y = x[0]**2 + x[1]**2 + 3*x[0]*x[1]

# Compute gradients
y.backward()  # Computes dy/dx

print(x.grad)  # tensor([13., 12.])
# Manual: ∂y/∂x0 = 2*x0 + 3*x1 = 4 + 9 = 13
# Manual: ∂y/∂x1 = 2*x1 + 3*x0 = 6 + 6 = 12

def sharpe_ratio(returns, risk_free=0.0):
    """
    Sharpe Ratio = (E[R] - r_f) / std(R)
    """
    mean_return = returns.mean()
    std_return = returns.std(unbiased=False)
    return (mean_return - risk_free) / (std_return + 1e-8)

# Portfolio returns (linear combination of asset returns)
weights = torch.randn(10, requires_grad=True)
weights = weights / weights.sum()  # Normalise to sum to 1

asset_returns = torch.randn(1000, 10)  # 1000 days, 10 assets
portfolio_returns = asset_returns @ weights

# Maximise Sharpe ratio (minimise negative)
loss = -sharpe_ratio(portfolio_returns)

# Compute gradient of loss w.r.t weights
loss.backward()
gradient = weights.grad
# This gradient tells us how to adjust weights to improve Sharpe ratio

# Detach to stop gradient tracking (for validation)
with torch.no_grad():
    val_loss = model(val_data)

# Reset gradients before each backward pass
optimizer.zero_grad()  # Or model.zero_grad()

# In-place operations break autograd
x = torch.tensor([2.0], requires_grad=True)
y = x * 2
y.add_(1)  # This is allowed but prevents gradient computation for the add

class SharpeLoss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, returns, weights):
        """
        Forward pass: compute Sharpe Ratio.
        """
        portfolio_returns = returns @ weights
        mean_r = portfolio_returns.mean()
        std_r = portfolio_returns.std(unbiased=False)
        sharpe = mean_r / (std_r + 1e-8)
        ctx.save_for_backward(returns, weights, portfolio_returns, mean_r, std_r)
        return sharpe

    @staticmethod
    def backward(ctx, grad_output):
        """
        Backward pass: compute gradient of Sharpe w.r.t weights.
        """
        returns, weights, portfolio_returns, mean_r, std_r = ctx.saved_tensors

        # Gradient: ∂Sharpe/∂w
        # Derivation: ∂/∂w [(w^T μ) / sqrt(w^T Σ w)]
        # = (1/σ) [ μ - (SR) * (Σ w) / σ ]
        mu = returns.mean(dim=0)
        sigma = std_r
        sigma_w = returns.T @ weights / len(returns)

        grad_weights = (1/sigma) * (mu - (mean_r/sigma) * sigma_w)
        return None, grad_weights * grad_output  # None for returns

class FinancialNN(nn.Module):
    def __init__(self, input_dim, hidden_dims, output_dim, dropout_rate=0.3):
        super(FinancialNN, self).__init__()

        layers = []
        prev_dim = input_dim

        # Build hidden layers
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))
            prev_dim = hidden_dim

        # Output layer
        layers.append(nn.Linear(prev_dim, output_dim))

        # For classification: nn.Sigmoid() or nn.Softmax(dim=1)
        # For regression: no activation

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

# Initialise model
input_dim = 50  # Number of features
hidden_dims = [128, 64, 32]
output_dim = 1  # Regression for return prediction

model = FinancialNN(input_dim, hidden_dims, output_dim)
model = model.to(device)

def init_weights(m):
    if isinstance(m, nn.Linear):
        # Xavier/Glorot init for ReLU
        nn.init.xavier_uniform_(m.weight)
        nn.init.constant_(m.bias, 0.0)

model.apply(init_weights)

# Alternative initialisation schemes
# nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')
# nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
# nn.init.xavier_normal_(m.weight)
# nn.init.orthogonal_(m.weight)

def model_summary(model):
    """
    Print the number of parameters in each layer.
    """
    total_params = 0
    for name, param in model.named_parameters():
        if param.requires_grad:
            num_params = param.numel()
            print(f"{name}: {num_params:,} parameters")
            total_params += num_params
    print(f"\nTotal trainable parameters: {total_params:,}")
    return total_params

from torch.utils.data import DataLoader, TensorDataset

def prepare_data(X, y, batch_size=128, shuffle=True):
    """
    Create DataLoader for training/validation.
    """
    dataset = TensorDataset(
        torch.from_numpy(X.values).float(),
        torch.from_numpy(y.values).float()
    )
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
    return dataloader

# Example usage
X_train, X_val, X_test, y_train, y_val, y_test = split_data()
train_loader = prepare_data(X_train, y_train, batch_size=128, shuffle=True)
val_loader = prepare_data(X_val, y_val, batch_size=128, shuffle=False)
test_loader = prepare_data(X_test, y_test, batch_size=128, shuffle=False)

# Mean Squared Error (for regression)
criterion_mse = nn.MSELoss()

# Mean Absolute Error (robust to outliers)
criterion_mae = nn.L1Loss()

# Huber Loss (combines MSE and MAE)
criterion_huber = nn.HuberLoss(delta=1.0)

# Cross-Entropy Loss (for classification)
criterion_ce = nn.BCEWithLogitsLoss()  # For binary classification
criterion_ce_multi = nn.CrossEntropyLoss()  # For multi-class

# Custom Quantile Loss (for VaR prediction)
def quantile_loss(y_true, y_pred, quantile=0.95):
    """
    Pinball loss for quantile regression.
    """
    error = y_true - y_pred
    return torch.max(quantile * error, (quantile - 1) * error).mean()

# Adam (default for most financial AI)
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5, betas=(0.9, 0.999))

# SGD with momentum (can outperform Adam on some problems)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-5)

# Learning rate schedulers
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=10, verbose=True
)

# Cosine annealing
scheduler = optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100, eta_min=1e-6
)

# Step decay
scheduler = optim.lr_scheduler.StepLR(
    optimizer, step_size=30, gamma=0.1
)

def train_model(model, train_loader, val_loader, optimizer, criterion,
                epochs=100, scheduler=None, patience=20, checkpoint_path='best_model.pth'):
    """
    Complete training loop with early stopping and model checkpointing.
    """
    best_val_loss = float('inf')
    patience_counter = 0
    train_losses = []
    val_losses = []

    for epoch in range(epochs):
        # Training phase
        model.train()
        train_loss = 0.0
        for X_batch, y_batch in train_loader:
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)

            # Forward pass
            y_pred = model(X_batch)
            loss = criterion(y_pred.squeeze(), y_batch)

            # Backward pass
            optimizer.zero_grad()
            loss.backward()

            # Gradient clipping (prevents exploding gradients)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()

            train_loss += loss.item() * X_batch.size(0)

        train_loss /= len(train_loader.dataset)
        train_losses.append(train_loss)

        # Validation phase
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for X_batch, y_batch in val_loader:
                X_batch = X_batch.to(device)
                y_batch = y_batch.to(device)

                y_pred = model(X_batch)
                loss = criterion(y_pred.squeeze(), y_batch)
                val_loss += loss.item() * X_batch.size(0)

        val_loss /= len(val_loader.dataset)
        val_losses.append(val_loss)

        # Update learning rate
        if scheduler is not None:
            if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
                scheduler.step(val_loss)
            else:
                scheduler.step()

        # Early stopping and checkpointing
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            torch.save(model.state_dict(), checkpoint_path)
            print(f"Epoch {epoch}: New best model saved (val_loss: {val_loss:.6f})")
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(f"Early stopping at epoch {epoch}")
                break

        if epoch % 10 == 0:
            print(f"Epoch {epoch}: train_loss={train_loss:.6f}, val_loss={val_loss:.6f}")

    # Load best model
    model.load_state_dict(torch.load(checkpoint_path))
    return model, train_losses, val_losses

# L2 regularisation (weight decay)
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)

# L1 regularisation (requires manual implementation)
def l1_regularisation(model, lambda_l1=1e-5):
    l1_norm = sum(p.abs().sum() for p in model.parameters())
    return lambda_l1 * l1_norm

# In training loop:
loss = criterion(y_pred, y_batch) + l1_regularisation(model)

class DropoutNN(nn.Module):
    def __init__(self, input_dim, hidden_dims, output_dim, dropout_rate=0.5):
        super(DropoutNN, self).__init__()
        layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))  # Dropout after each hidden layer
            prev_dim = hidden_dim
        layers.append(nn.Linear(prev_dim, output_dim))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

class BatchNormNN(nn.Module):
    def __init__(self, input_dim, hidden_dims, output_dim):
        super(BatchNormNN, self).__init__()
        layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.BatchNorm1d(hidden_dim))  # Batch norm after linear
            layers.append(nn.ReLU())
            prev_dim = hidden_dim
        layers.append(nn.Linear(prev_dim, output_dim))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

# Correct order for batch norm + dropout
layers.append(nn.Linear(prev_dim, hidden_dim))
layers.append(nn.BatchNorm1d(hidden_dim))
layers.append(nn.ReLU())
layers.append(nn.Dropout(dropout_rate))
# Note: Batch norm before ReLU, dropout after ReLU

def evaluate_regression(y_true, y_pred):
    """
    Compute financial regression metrics.
    """
    y_true = y_true.numpy()
    y_pred = y_pred.numpy()

    mse = np.mean((y_true - y_pred)**2)
    mae = np.mean(np.abs(y_true - y_pred))
    rmse = np.sqrt(mse)
    r2 = 1 - np.sum((y_true - y_pred)**2) / np.sum((y_true - y_true.mean())**2)

    # Directional accuracy
    direction_true = np.sign(y_true)
    direction_pred = np.sign(y_pred)
    direction_accuracy = np.mean(direction_true == direction_pred)

    return {
        'MSE': mse, 'MAE': mae, 'RMSE': rmse, 'R2': r2,
        'Direction_Accuracy': direction_accuracy
    }

def evaluate_classification(y_true, y_pred_prob, threshold=0.5):
    """
    Compute financial classification metrics.
    """
    y_pred = (y_pred_prob > threshold).float()
    y_true = y_true.float()

    # Confusion matrix
    tp = ((y_true == 1) & (y_pred == 1)).sum().item()
    tn = ((y_true == 0) & (y_pred == 0)).sum().item()
    fp = ((y_true == 0) & (y_pred == 1)).sum().item()
    fn = ((y_true == 1) & (y_pred == 0)).sum().item()

    accuracy = (tp + tn) / (tp + tn + fp + fn + 1e-8)
    precision = tp / (tp + fp + 1e-8)
    recall = tp / (tp + fn + 1e-8)
    f1 = 2 * precision * recall / (precision + recall + 1e-8)

    # Confusion matrix
    conf_matrix = np.array([[tn, fp], [fn, tp]])

    return {
        'Accuracy': accuracy,
        'Precision': precision,
        'Recall': recall,
        'F1': f1,
        'Confusion_Matrix': conf_matrix
    }

import random
import numpy as np
import torch

def set_seed(seed=42):
    """
    Set all random seeds for reproducibility.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # For deterministic algorithms
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)