import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

def _build_calibrated_decoder(dim, stride):
    """
    Builds a B-Spline decoder that is NUMERICALLY CALIBRATED to ensure
    unity gain (Input 1.0 -> Output 1.0).
    Cardinal Cubic B-Spline. - uses 4 control points (degree 3). knots are spaced equally in time.
    The Basis Function B(t) is a bell-shaped curve defined piecewise over the range [−2,2]
    S(t)=∑_i P_i⋅ B(t - i)
    "fixed" basis guarantees that the output is a smooth function of the control points.
    """
    # 1. Use Symmetric Kernel Size (Odd number ensures perfect center)
    kernel_size = stride * 4 + 1
    
    # 2. Generate Theoretical Kernel
    t = torch.linspace(-2.0, 2.0, kernel_size)
    t_abs = torch.abs(t)
    y = torch.zeros_like(t)
    
    mask1 = t_abs <= 1
    y[mask1] = 2/3 - t_abs[mask1]**2 + 0.5 * t_abs[mask1]**3
    mask2 = (t_abs > 1) & (t_abs <= 2)
    y[mask2] = (2 - t_abs[mask2])**3 / 6
    
    # 3. Normalize for Discrete Partition of Unity
    # In a Transpose Conv, we need the sum of overlapping kernels to be 1.0.
    # The easiest way to fix this is to scale the kernel until the output magnitude matches.
    # We assume 'y * stride' is the rough scale, but we correct it below.
    y = y * stride 
    
    # CALIBRATION STEP:
    # Create a dummy layer and run 1.0 through it to measure the gain.
    dummy_layer = nn.ConvTranspose1d(1, 1, kernel_size, stride=stride, padding=2*stride, bias=False)
    dummy_layer.weight.data = y.view(1, 1, -1)
    
    with torch.no_grad():
        test_input = torch.ones(1, 1, 100) # Constant 1.0
        test_output = dummy_layer(test_input)
        # Measure the middle (steady state) value
        measured_gain = test_output[..., 50:80].mean().item()
    
    # Correction Factor
    correction = 1.0 / measured_gain
    y_calibrated = y * correction
    
    # 4. Create Real Layer
    # Padding = 2*stride ensures the output size is roughly Input * Stride
    layer = nn.ConvTranspose1d(dim, dim, kernel_size, stride=stride, 
                                groups=dim, padding=2*stride, bias=False)
    
    layer.weight.data = y_calibrated.view(1, 1, -1).repeat(dim, 1, 1)
    layer.weight.requires_grad = False
    return layer


class NeuralSplineFilter(nn.Module):
    def __init__(self, input_dim, spacing=4, hidden_dim=64, k1=7, k2=5, constraints=None, delta_scale=None):
        super().__init__()
        """
        constraints : D x 2 tensor
        """
        self.kwargs_summary = {'input_dim': input_dim, 'spacing': spacing, 'hidden_dim': hidden_dim, 'k1': k1, 'k2': k2, 'constraints': constraints, 'delta_scale': delta_scale}
        # Ensure hidden capacity
        hidden_dim = max(hidden_dim, input_dim)
        
        self.spacing = spacing
        self.input_dim = input_dim
        
        # --- 1. The Learnable Encoder (Residual Correction) ---
        # This branch learns the "Delta" to add to the rough downsampling.
        # We initialize it to output ZERO to start.
        k3 = spacing
        self.encoder = nn.Sequential(
            nn.Conv1d(input_dim, hidden_dim, kernel_size=k1, padding=k1//2, padding_mode='reflect'),
            nn.LeakyReLU(0.2),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=k2, padding=k2//2, padding_mode='reflect'),
            nn.LeakyReLU(0.2),
            nn.Conv1d(hidden_dim, input_dim, kernel_size=k3, stride=spacing)
        )
        # Force the encoder to start silent (Zero Output)
        self.apply(self._init_weights_to_zero)

        # --- 2. The Fixed Decoder (Self-Calibrated) ---
        self.decoder = _build_calibrated_decoder(input_dim, spacing)

        # --- 3. Calculate & Print Recommended T ---
        # Formula: RF = k1 + (k2-1) + (k3-1) -> Simple stack approximation
        receptive_field = k1 + (k2 - 1) + (k3 - 1)

        # Safety Factor 8: Ensures roughly 90% valid data, 10% padding context
        ideal_T = receptive_field * 8 
        
        # Round to nearest power of 2 (e.g., 96 -> 128) for GPU efficiency
        recommended_T = 2 ** int(np.ceil(np.log2(ideal_T)))
        
        # Clamp to a reasonable minimum (e.g. 64)
        recommended_T = max(64, recommended_T)
        
        # print(f"Receptive Field: {receptive_field} frames. Recommended minimum temporal length (T): {recommended_T}")
        self.with_constraints = constraints is not None
        if constraints is not None:
            assert constraints.shape[0] == input_dim and constraints.shape[1] == 2, "Constraints must be of shape (D, 2)"
            self.register_buffer('constraints_diff', constraints[:,1] - constraints[:,0])
            self.register_buffer('constraints_min', constraints[:,0])
            self.register_buffer('constraints_max', constraints[:,1])

        self.delta_scale = delta_scale

    def _init_weights_to_zero(self, m):
        """Initializes the last layer of the encoder to zero."""
        if isinstance(m, nn.Conv1d):
            # Standard init for earlier layers
            nn.init.kaiming_normal_(m.weight, nonlinearity='leaky_relu')
            if m.bias is not None: nn.init.constant_(m.bias, 0.0)
            
            # If this is the downsampling layer (stride > 1), zero it out.
            if m.stride[0] == self.spacing:
                nn.init.constant_(m.weight, 0.0)
                if m.bias is not None: nn.init.constant_(m.bias, 0.0)

    def forward(self, x, get_loss=False):
        # --- A. Shape Handling ---
        orig_shape = x.shape
        if x.dim() == 2: x = x.unsqueeze(0).permute(0, 2, 1) # (T, D) -> (1, D, T)
        elif x.dim() == 3: x = x.permute(0, 2, 1)            # (B, T, D) -> (B, D, T)
        
        input_length = x.shape[-1]
        
        # --- B. Padding Input (Reflect to avoid edge jumps) ---
        pad_len = (self.spacing - (input_length % self.spacing)) % self.spacing
        if pad_len > 0:
            x = F.pad(x, (0, pad_len), mode='reflect')

        # --- C. Control Points (Residual Connection) ---
        # 1. The Rough Guess: Deterministic Downsample (Preserves Amplitude!)
        #    We calculate target length exactly
        target_k = x.shape[-1] // self.spacing
        cp_initial = F.interpolate(x, size=target_k, mode='linear', align_corners=True)
        
        # 2. The Correction: Learned Encoder (Starts at 0.0)
        #    We perform center cropping on the encoder output to match interpolation size if needed
        cp_delta = self.encoder(x)
        
        # 1. Apply Tanh limit in Value Space (Physical Units)
        if self.delta_scale is not None:
            # cp_delta is now in range [-delta_scale, +delta_scale] (e.g., +/- 0.05 meters)
            cp_delta = torch.tanh(cp_delta) * self.delta_scale

        # ... [Shape handling logic] ...

        # 3. Combine
        if self.with_constraints:
            min_v = self.constraints_min.view(1, -1, 1)
            diff_v = self.constraints_diff.view(1, -1, 1)
            eps = 1e-6

            # A. Calculate Initial Logits (Your existing code)
            cp_normalized = (cp_initial - min_v) / (diff_v + eps)
            cp_normalized = torch.clamp(cp_normalized, eps, 1.0 - eps)
            cp_initial_logits = torch.log(cp_normalized / (1.0 - cp_normalized))
            
            # B. CONVERT DELTA: Value Space -> Logit Space
            # Step 1: Normalize delta to [0-1] ratio relative to the range
            # e.g., if range is 100m and delta is 1m, ratio is 0.01
            delta_ratio = cp_delta / (diff_v + eps)
            
            # Step 2: Scale for Sigmoid Derivative
            # The derivative of Sigmoid at the center (0.5) is 0.25.
            # To achieve a change of Y, you need a change of 4*Y in X.
            # Factor 4.0 ensures that near the center, the physical limit matches delta_scale.
            cp_delta_logits = delta_ratio * 4.0
            
            # C. Add in Logit Space
            control_points = cp_initial_logits + cp_delta_logits
            
        else:
            # Standard additive residual in Value Space
            control_points = cp_initial + cp_delta

        # --- D. Decode (Robust Boundary Handling) ---
        # We pad the CONTROL POINTS to ensure the spline reaches the edge.
        # Padding 2 points allows the kernel to "see" past the edge.
        cp_pad = 2
        control_points_padded = F.pad(control_points, (cp_pad, cp_pad), mode='replicate')
        
        if self.with_constraints:
            # Apply constraints via affine transform
            control_points_padded = torch.sigmoid(control_points_padded) * self.constraints_diff.view(1, -1, 1) + self.constraints_min.view(1, -1, 1)

        smoothed = self.decoder(control_points_padded)

        # --- E. Center Crop ---
        # The decoder output is larger than input due to padding.
        # We must align the center exactly. 
        # With symmetric kernel K = 4*S + 1 and padding P = 2*S, the shift is perfectly 0.
        # But we added cp_pad on the left.
        # Shift = cp_pad * stride
        start = cp_pad * self.spacing
        smoothed = smoothed[..., start : start + x.shape[-1]]

        # --- F. Cleanup ---
        # Remove input padding
        if pad_len > 0:
            smoothed = smoothed[..., :input_length]

        ##### test ####
        if self.with_constraints:
            offset = 1e-5
            assert torch.all(control_points_padded >= self.constraints_min[None,:,None] - offset) and torch.all(control_points_padded <= self.constraints_max[None,:,None] + offset), "Constraints violated on control points!"
            assert torch.all(smoothed >= self.constraints_min[None,:,None] - offset) and torch.all(smoothed <= self.constraints_max[None,:,None] + offset), "Constraints violated!"
            # print('constraints checked!')
        ###############

        loss_out = None
        if get_loss:
            # Regularization
            reg_loss = torch.mean((x[..., :input_length] - smoothed)**2)
            acc_loss = torch.mean((control_points[:,:,:-2] - 
                                   2*control_points[:,:,1:-1] + 
                                   control_points[:,:,2:])**2)
            loss_out = {'reg_loss': reg_loss, 'control_acc_loss': acc_loss}

        # Restore shapes
        smoothed = smoothed.permute(0, 2, 1)
        control_points = control_points.permute(0, 2, 1)
        if len(orig_shape) == 2:
            smoothed = smoothed.squeeze(0)
            control_points = control_points.squeeze(0)
            
        return smoothed, control_points, loss_out


class NeuralSplineFilterDecoder(nn.Module):
    def __init__(self, input_length, input_dim, chunk_size=10000, spacing=4, constraints=None, 
                 init_cp_params=None, init_params=None, apply_constraints_before_decode=True):
        super().__init__()
        
        self.input_length = input_length
        self.input_dim = input_dim
        self.spacing = spacing
        self.chunk_size = chunk_size
        self.apply_constraints_before_decode = apply_constraints_before_decode
        if not apply_constraints_before_decode:
            print("SplineDecoder: Applying constraints AFTER decoding.")
        
        # 1. Calculation (Trust the Math, not the User arg)
        # We calculate exactly how many chunks cover the input_length
        self.total_knots = int(np.ceil(input_length / spacing)) + 4 
        self.knots_per_chunk = int(np.ceil(chunk_size / spacing))
        self.num_chunks = int(np.ceil(self.total_knots / self.knots_per_chunk))
        
        print(f"SplineDecoder: Initialized {self.num_chunks} chunks for {self.total_knots} total knots.")

        # --- PREPARE INITIALIZATION DATA ---
        init_tensor = None
        
        if init_cp_params is not None:
            # User gave sparse knots
            init_tensor = init_cp_params
            if init_tensor.dim() == 2: init_tensor = init_tensor.permute(1, 0).unsqueeze(0)
                
        elif init_params is not None:
            # User gave dense frames (T, D) -> We downsample to knots
            # Ensure shape (1, D, T)
            if init_params.dim() == 2: 
                x = init_params.permute(1, 0).unsqueeze(0)
            else:
                x = init_params

            # --- THE FIX ---
            # 1. Create a grid of time indices: 0, 4, 8, ...
            # Shape: (1, 1, K)
            indices = torch.arange(self.total_knots, device=x.device) * self.spacing
            
            # 2. Handle the "Padding Knots" at the end
            # The indices will eventually go past the end of the video (e.g. index 104 for a 100-frame video).
            # We CLAMP them to the last valid frame (Replication Padding).
            indices = torch.clamp(indices, max=self.input_length - 1)
            
            # 3. Gather the values
            # x is (1, D, T). We select specific columns.
            # Result: (1, D, K)
            init_tensor = x[..., indices.long()]
            
            print(f"SplineDecoder: Sampled init params with stride {self.spacing}. Fixed time alignment.")

        # 2. HANDLE CONSTRAINTS (Inverse Sigmoid)
        self.with_constraints = constraints is not None
        if self.with_constraints:
            self.register_buffer('constraints_min', constraints[:,0].view(1, -1, 1))
            self.register_buffer('constraints_diff', (constraints[:,1] - constraints[:,0]).view(1, -1, 1))
            
            if init_tensor is not None:
                # Value -> Logit conversion logic is CORRECT.
                # It prevents the initialization from drifting when sigmoid is applied later.
                min_v = self.constraints_min
                diff_v = self.constraints_diff
                eps = 1e-6
                
                init_normalized = (init_tensor - min_v) / (diff_v + eps)
                init_normalized = torch.clamp(init_normalized, eps, 1.0 - eps)
                init_tensor = torch.log(init_normalized / (1.0 - init_normalized))
                print("SplineDecoder: Converted initialization data to Logit Space.")

        # 3. CREATE CHUNKS
        self.chunks = nn.ParameterList()
        
        for i in range(self.num_chunks):
            start_k = i * self.knots_per_chunk
            end_k = min((i + 1) * self.knots_per_chunk, self.total_knots)
            size = end_k - start_k
            
            if size > 0:
                if init_tensor is not None:
                    # Clone is critical so gradients don't leak back to init_tensor
                    chunk_data = init_tensor[..., start_k:end_k].clone()
                    p = nn.Parameter(chunk_data)
                else:
                    p = nn.Parameter(torch.zeros(1, input_dim, size))
                self.chunks.append(p)

        self.decoder = _build_calibrated_decoder(input_dim, spacing)

    def forward(self, chunk_idx, **kwargs):
        if chunk_idx < 0 or chunk_idx >= len(self.chunks):
            raise ValueError(f"Chunk index {chunk_idx} out of bounds")

        # --- 1. Fetch & Stitch Knots ---
        core_knots = self.chunks[chunk_idx]
        
        # Why pad_k=2? 
        # Kernel size is 4*spacing. Radius is 2*spacing.
        # We need 2 knots of context to cover that radius.
        pad_k = 2 
        
        knots_to_decode = [core_knots]
        
        # Padding Logic (Correct B-Spline Boundary conditions)
        if chunk_idx > 0:
            knots_to_decode.insert(0, self.chunks[chunk_idx - 1][..., -pad_k:])
        else:
            # knots_to_decode.insert(0, core_knots[..., :1].repeat(1, 1, pad_k)) # Clamp Start
            if core_knots.shape[-1] >= 2:
                p_0 = core_knots[..., 0:1]      # (1, D, 1)
                p_1 = core_knots[..., 1:2]      # (1, D, 1)
                slope = p_1 - p_0               # Velocity
                
                # Extrapolation Ramp: [1, 2] -> We want P_-1, P_-2
                ramp = torch.arange(1, pad_k + 1, device=core_knots.device).view(1, 1, -1)
                
                # Backwards: P_pad = P_0 - slope * ramp
                # Note: ramp is 1, 2. So we get P0-v, P0-2v. Correct.
                # We need to flip it so it connects: [P_-2, P_-1] -> [P_0...]
                padding = p_0 - (slope * ramp)
                padding = torch.flip(padding, dims=[-1]) 
                
                knots_to_decode.insert(0, padding)
            else:
                knots_to_decode.insert(0, core_knots[..., :1].repeat(1, 1, pad_k))

        if chunk_idx < len(self.chunks) - 1:
            knots_to_decode.append(self.chunks[chunk_idx + 1][..., :pad_k])
        else:
            # knots_to_decode.append(core_knots[..., -1:].repeat(1, 1, pad_k)) # Clamp End
            # We must Extrapolate velocity forwards to avoid "Zero Velocity Wall"
            if core_knots.shape[-1] >= 2:
                p_last = core_knots[..., -1:]      # (1, D, 1)
                p_prev = core_knots[..., -2:-1]    # (1, D, 1)
                slope = p_last - p_prev            # Velocity
                
                # Extrapolation Ramp: [1, 2]
                ramp = torch.arange(1, pad_k + 1, device=core_knots.device).view(1, 1, -1)
                
                # Forwards: P_pad = P_last + slope * ramp
                padding = p_last + (slope * ramp)
                
                knots_to_decode.append(padding)
            else:
                knots_to_decode.append(core_knots[..., -1:].repeat(1, 1, pad_k))
            

        knots_window = torch.cat(knots_to_decode, dim=-1)

        # --- 2. Constraints ---
        if self.with_constraints and self.apply_constraints_before_decode:
            knots_window = torch.sigmoid(knots_window) * self.constraints_diff + self.constraints_min

        # --- 3. Decode ---
        # Output shape is approx: (Knots * Spacing)
        decoded_window = self.decoder(knots_window)
        
        # --- 4. Crop (The most important math) ---
        # We added 'pad_k' knots to the left. 
        # Since decoder is centered, we must trim exactly 'pad_k * spacing' frames.
        frames_to_crop_left = pad_k * self.spacing
        valid_length = core_knots.shape[-1] * self.spacing
        
        decoded_frames = decoded_window[..., frames_to_crop_left : frames_to_crop_left + valid_length]
        
        # --- 5. Global Boundary Safety ---
        # If the last chunk overshoots the original input_length, trim it.
        global_start_frame = chunk_idx * self.knots_per_chunk * self.spacing
        
        if global_start_frame + decoded_frames.shape[-1] > self.input_length:
             valid_len = self.input_length - global_start_frame
             decoded_frames = decoded_frames[..., :valid_len]

        assert len(decoded_frames) == 1
        if self.with_constraints and (not self.apply_constraints_before_decode):
            decoded_frames = torch.sigmoid(decoded_frames) * self.constraints_diff + self.constraints_min

        return decoded_frames.permute(0, 2, 1).squeeze(0), {}

class StreamSplineProcessor:
    def __init__(self, model, chunk_size=2000, overlap=200):
        """
        Args:
            model: Your NeuralSplineFilter instance.
            chunk_size: How many frames to process per batch (memory limit).
            overlap: Safety buffer size. Must be > model receptive field.
                     For your spline model, 100-200 is plenty safe.
        """
        self.model = model
        self.chunk_size = chunk_size
        self.overlap = overlap
        
    def process_sequence(self, huge_input_tensor):
        """
        Args:
            huge_input_tensor: (Total_Frames, D) - e.g., (100000, 45)
                               Can be on CPU to save GPU memory.
        Returns:
            smoothed_output: (Total_Frames, D) - Seamlessly stitched.
        """
        total_frames = huge_input_tensor.shape[0]
        result_list = []
        
        # We step forward by 'step_size' (the valid region size)
        step_size = self.chunk_size - (2 * self.overlap)
        
        if step_size <= 0:
            raise ValueError("Chunk size must be larger than 2 * overlap!")

        current_idx = 0
        
        print(f"Processing {total_frames} frames in overlapping chunks...")
        
        while current_idx < total_frames:
            # 1. Determine Window Indices (Global)
            # We want to produce valid output for [current_idx, current_idx + step_size]
            # To do that, we need context: [start_pad, end_pad]
            
            start_idx = max(0, current_idx - self.overlap)
            end_idx = min(total_frames, current_idx + step_size + self.overlap)
            
            # 2. Extract Chunk (Load to GPU)
            # chunk shape: (B, D, T_chunk)
            chunk = huge_input_tensor[start_idx:end_idx].cuda()
            
            # 3. Run Inference
            # We assume your model accepts (T, D) or we reshape it
            with torch.no_grad():
                # Model returns (Smoothed, Knots) -> we only need smoothed here
                smoothed_chunk, _ = self.model(chunk)
            
            # 4. Crop to Valid Region
            # We need to figure out how much padding we actually added on the left/right
            pad_left = current_idx - start_idx
            
            # We want to take 'step_size' frames, but clamp if we are at the end of video
            valid_length = min(step_size, total_frames - current_idx)
            
            # Crop the "dirty" edges
            # valid_chunk is strictly the clean center
            valid_chunk = smoothed_chunk[pad_left : pad_left + valid_length]
            
            # 5. Store result (Move to CPU to save GPU RAM)
            result_list.append(valid_chunk.cpu())
            
            # Move pointer
            current_idx += valid_length
            
        # 6. Concatenate all clean chunks
        full_output = torch.cat(result_list, dim=0)
        
        return full_output
