import numpy as np
from . import maxwell_core

class Linear:
    def __init__(self, in_features, out_features, bias=False):
        self.in_features = in_features
        self.out_features = out_features
        
        # Initialize latent weights (float32)
        # Using He Initialization for better convergence
        limit = np.sqrt(2.0 / in_features)
        self.weight_latent = np.random.randn(in_features, out_features).astype(np.float32) * limit
        
        self.grad_weight = np.zeros_like(self.weight_latent)
        
        # Buffers for Forward/Backward
        self.last_input = None
        self.last_output = None
        
    def forward(self, x):
        """
        x: [Batch, In_Features]
        """
        self.last_input = x
        batch_size = x.shape[0]
        
        # 1. Quantize Weights (Ternary: -1, 0, 1)
        # STE: Forward uses quantized, Backward updates latent
        w_quant = self._quantize(self.weight_latent)
        
        # 2. Pack Weights
        w_packed = maxwell_core.create_packed_weights(w_quant)
        
        # 3. CUDA Forward
        y = np.zeros((batch_size, self.out_features), dtype=np.float32)
        maxwell_core.run_maxwell_cuda(x, w_packed, y)
        
        self.last_output = y
        return y
    
    def backward(self, grad_output):
        """
        grad_output: [Batch, Out_Features]
        Returns: grad_input [Batch, In_Features]
        """
        # 1. Compute Weight Gradients (grad_W = Input^T * grad_Output)
        # We use our optimized CUDA Kernel for this
        maxwell_core.run_maxwell_backward_weights_cuda(
            self.last_input, 
            grad_output, 
            self.grad_weight
        )
        
        # 2. Compute Input Gradients (grad_Input = grad_Output * W^T)
        # NOTE: For this demo, we assume this is the first layer or we don't need to backprop to input.
        # If needed, we would implement run_maxwell_backward_data_cuda
        return None 
    
    def _quantize(self, w):
        # Deterministic Ternary Quantization with Threshold 0.05 for stability
        # In 'Extreme SOTA', we might use stochastic quantization.
        threshold = 0.05
        return np.where(w > threshold, 1, np.where(w < -threshold, -1, 0)).astype(np.int8)

