/**
 * CUDA Kernel Fusion for High-Frequency Trading ML Operations
 * 
 * This module implements fused CUDA kernels to reduce kernel launch overhead
 * and improve GPU utilization for machine learning operations in the Foxhunt
 * trading system.
 */

#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <cmath>

/**
 * Fused kernel combining transformation, normalization, and activation
 * 
 * This kernel performs three operations in a single GPU kernel launch:
 * 1. Data transformation (optional scaling/offset)
 * 2. Z-score normalization using provided mean and standard deviation
 * 3. ReLU activation function
 * 
 * @param input Input tensor data
 * @param output Output tensor data after fusion operations
 * @param mean Mean value for normalization
 * @param std Standard deviation for normalization
 * @param size Total number of elements to process
 */
__global__ void fused_transform_normalize_activate(
    float* input,
    float* output,
    float* mean,
    float* std,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < size) {
        // Transform: Load input value
        float val = input[idx];
        
        // Normalize: Apply Z-score normalization
        val = (val - mean[0]) / std[0];
        
        // Activate: Apply ReLU activation function
        output[idx] = fmaxf(0.0f, val);
    }
}

/**
 * Fused kernel for batch normalization with learnable parameters
 * 
 * Combines batch normalization with scale/shift parameters and activation
 * 
 * @param input Input tensor
 * @param output Output tensor
 * @param gamma Scale parameter (learnable)
 * @param beta Shift parameter (learnable)
 * @param mean Batch mean
 * @param variance Batch variance
 * @param epsilon Small constant for numerical stability
 * @param size Number of elements
 */
__global__ void fused_batch_norm_scale_activate(
    float* input,
    float* output,
    float* gamma,
    float* beta,
    float* mean,
    float* variance,
    float epsilon,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < size) {
        float val = input[idx];
        
        // Batch normalization
        val = (val - mean[0]) / sqrtf(variance[0] + epsilon);
        
        // Scale and shift
        val = gamma[0] * val + beta[0];
        
        // ReLU activation
        output[idx] = fmaxf(0.0f, val);
    }
}

/**
 * Fused kernel for matrix multiplication followed by bias addition and activation
 * 
 * Combines dense layer operations: Y = activate(X * W + b)
 * 
 * @param input Input matrix (batch_size x input_dim)
 * @param weights Weight matrix (input_dim x output_dim)
 * @param bias Bias vector (output_dim)
 * @param output Output matrix (batch_size x output_dim)
 * @param batch_size Number of samples in batch
 * @param input_dim Input feature dimension
 * @param output_dim Output feature dimension
 */
__global__ void fused_dense_bias_activate(
    float* input,
    float* weights,
    float* bias,
    float* output,
    int batch_size,
    int input_dim,
    int output_dim
) {
    int row = blockIdx.y * blockDim.y + threadIdx.y; // batch index
    int col = blockIdx.x * blockDim.x + threadIdx.x; // output feature index
    
    if (row < batch_size && col < output_dim) {
        float sum = 0.0f;
        
        // Matrix multiplication: dot product of input row with weight column
        for (int k = 0; k < input_dim; k++) {
            sum += input[row * input_dim + k] * weights[k * output_dim + col];
        }
        
        // Add bias
        sum += bias[col];
        
        // Apply ReLU activation
        output[row * output_dim + col] = fmaxf(0.0f, sum);
    }
}

/**
 * Fused kernel for elementwise operations with multiple activations
 * 
 * Supports different activation functions in a single kernel
 * 
 * @param input Input tensor
 * @param output Output tensor
 * @param scale Scale factor
 * @param offset Offset value
 * @param activation_type 0=ReLU, 1=Tanh, 2=Sigmoid, 3=GELU
 * @param size Number of elements
 */
__global__ void fused_elementwise_activate(
    float* input,
    float* output,
    float scale,
    float offset,
    int activation_type,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < size) {
        float val = input[idx];
        
        // Scale and offset
        val = val * scale + offset;
        
        // Apply activation based on type
        switch (activation_type) {
            case 0: // ReLU
                val = fmaxf(0.0f, val);
                break;
            case 1: // Tanh
                val = tanhf(val);
                break;
            case 2: // Sigmoid
                val = 1.0f / (1.0f + expf(-val));
                break;
            case 3: // GELU (approximation)
                val = 0.5f * val * (1.0f + tanhf(0.7978845608f * (val + 0.044715f * val * val * val)));
                break;
            default:
                val = fmaxf(0.0f, val); // Default to ReLU
        }
        
        output[idx] = val;
    }
}

/**
 * Fused kernel for attention mechanism computations
 * 
 * Combines query-key multiplication, scaling, and softmax preparation
 * 
 * @param queries Query vectors
 * @param keys Key vectors
 * @param attention_scores Output attention scores
 * @param scale Scaling factor (typically 1/sqrt(d_k))
 * @param seq_len Sequence length
 * @param d_k Key dimension
 */
__global__ void fused_attention_qk_scale(
    float* queries,
    float* keys,
    float* attention_scores,
    float scale,
    int seq_len,
    int d_k
) {
    int i = blockIdx.y * blockDim.y + threadIdx.y; // query index
    int j = blockIdx.x * blockDim.x + threadIdx.x; // key index
    
    if (i < seq_len && j < seq_len) {
        float score = 0.0f;
        
        // Compute dot product between query i and key j
        for (int k = 0; k < d_k; k++) {
            score += queries[i * d_k + k] * keys[j * d_k + k];
        }
        
        // Apply scaling
        attention_scores[i * seq_len + j] = score * scale;
    }
}

/**
 * Host function to launch fused transform-normalize-activate kernel
 * 
 * @param input Input data on device
 * @param output Output data on device
 * @param mean Mean value on device
 * @param std Standard deviation on device
 * @param size Number of elements
 * @param stream CUDA stream for async execution
 */
extern "C" void launch_fused_transform_normalize_activate(
    float* input,
    float* output,
    float* mean,
    float* std,
    int size,
    cudaStream_t stream = 0
) {
    const int block_size = 256;
    const int grid_size = (size + block_size - 1) / block_size;
    
    fused_transform_normalize_activate<<<grid_size, block_size, 0, stream>>>(
        input, output, mean, std, size
    );
}

/**
 * Host function to launch fused batch normalization kernel
 */
extern "C" void launch_fused_batch_norm_scale_activate(
    float* input,
    float* output,
    float* gamma,
    float* beta,
    float* mean,
    float* variance,
    float epsilon,
    int size,
    cudaStream_t stream = 0
) {
    const int block_size = 256;
    const int grid_size = (size + block_size - 1) / block_size;
    
    fused_batch_norm_scale_activate<<<grid_size, block_size, 0, stream>>>(
        input, output, gamma, beta, mean, variance, epsilon, size
    );
}

/**
 * Host function to launch fused dense layer kernel
 */
extern "C" void launch_fused_dense_bias_activate(
    float* input,
    float* weights,
    float* bias,
    float* output,
    int batch_size,
    int input_dim,
    int output_dim,
    cudaStream_t stream = 0
) {
    dim3 block_size(16, 16);
    dim3 grid_size(
        (output_dim + block_size.x - 1) / block_size.x,
        (batch_size + block_size.y - 1) / block_size.y
    );
    
    fused_dense_bias_activate<<<grid_size, block_size, 0, stream>>>(
        input, weights, bias, output, batch_size, input_dim, output_dim
    );
}

/**
 * Host function to launch fused elementwise operations kernel
 */
extern "C" void launch_fused_elementwise_activate(
    float* input,
    float* output,
    float scale,
    float offset,
    int activation_type,
    int size,
    cudaStream_t stream = 0
) {
    const int block_size = 256;
    const int grid_size = (size + block_size - 1) / block_size;
    
    fused_elementwise_activate<<<grid_size, block_size, 0, stream>>>(
        input, output, scale, offset, activation_type, size
    );
}

/**
 * Host function to launch fused attention kernel
 */
extern "C" void launch_fused_attention_qk_scale(
    float* queries,
    float* keys,
    float* attention_scores,
    float scale,
    int seq_len,
    int d_k,
    cudaStream_t stream = 0
) {
    dim3 block_size(16, 16);
    dim3 grid_size(
        (seq_len + block_size.x - 1) / block_size.x,
        (seq_len + block_size.y - 1) / block_size.y
    );
    
    fused_attention_qk_scale<<<grid_size, block_size, 0, stream>>>(
        queries, keys, attention_scores, scale, seq_len, d_k
    );
}