import triton
import triton.language as tl
import torch
import numpy as np

@triton.jit
def intersection_kernel_triton(
    prompts_flat_ptr,
    prompt_offsets_ptr,
    prompt_lengths_ptr,
    result_matrix_ptr,
    n_prompts,
    BLOCK_SIZE: tl.constexpr,
    MAX_LEN: tl.constexpr,
):
    """
    Corrected Triton kernel for computing intersection counts using a block-based approach.
    This version uses a vectorized method within the block to avoid indexing errors and improve performance.
    """
    # Get program IDs for the 2D grid
    pid_i = tl.program_id(0)
    pid_j = tl.program_id(1)

    # Calculate base indices for the block
    base_i = pid_i * BLOCK_SIZE
    base_j = pid_j * BLOCK_SIZE

    # The outer loops over the block's rows (i) and columns (j) are unrolled by the compiler.
    for i_offset in range(BLOCK_SIZE):
        global_i_idx = base_i + i_offset
        # Skip any out-of-bounds work for this row
        if global_i_idx < n_prompts:
            # Load metadata and data for the single 'i' prompt
            offset_i = tl.load(prompt_offsets_ptr + global_i_idx)
            length_i = tl.load(prompt_lengths_ptr + global_i_idx)
            
            i_range = tl.arange(0, MAX_LEN)
            mask_i_elements = i_range < length_i
            prompt_i_data = tl.load(prompts_flat_ptr + offset_i + i_range, mask=mask_i_elements, other=-1)

            for j_offset in range(BLOCK_SIZE):
                global_j_idx = base_j + j_offset
                # Skip any out-of-bounds work for this column
                if global_j_idx < n_prompts:
                    # Load metadata and data for the single 'j' prompt
                    offset_j = tl.load(prompt_offsets_ptr + global_j_idx)
                    length_j = tl.load(prompt_lengths_ptr + global_j_idx)
                    
                    j_range = tl.arange(0, MAX_LEN)
                    mask_j_elements = j_range < length_j
                    prompt_j_data = tl.load(prompts_flat_ptr + offset_j + j_range, mask=mask_j_elements, other=-2)

                    # --- Vectorized Intersection Calculation (for one i, j pair) ---
                    # Use broadcasting to create a comparison matrix
                    elements_i = prompt_i_data[:, None]  # Shape: [MAX_LEN, 1]
                    elements_j = prompt_j_data[None, :]  # Shape: [1, MAX_LEN]
                    
                    match_matrix = (elements_i == elements_j)
                    
                    # Apply masks to only consider valid elements of each prompt
                    valid_matches = match_matrix & mask_i_elements[:, None] & mask_j_elements[None, :]
                    
                    # For each element in prompt_i, check if it has at least one match in prompt_j
                    found_in_j = tl.sum(valid_matches, axis=1) > 0
                    
                    # The intersection count is the sum of unique elements from i found in j
                    count = tl.sum(found_in_j)

                    # Store the final count for the (i, j) pair
                    result_offset = global_i_idx * n_prompts + global_j_idx
                    tl.store(result_matrix_ptr + result_offset, count)


@triton.jit
def distance_kernel_triton(
    intersection_matrix_ptr,
    lengths_ptr,
    result_matrix_ptr,
    method_id,
    n_prompts,
    BLOCK_SIZE: tl.constexpr,
):
    """Triton kernel for computing distance matrix from intersections."""
    # Get program IDs
    pid_i = tl.program_id(0)
    pid_j = tl.program_id(1)

    # Create 1D ranges for row and column indices for this block
    i_indices = pid_i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    j_indices = pid_j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

    # Create 1D masks for these indices
    i_mask = i_indices < n_prompts
    j_mask = j_indices < n_prompts

    # Load 1D vectors of lengths, then broadcast to 2D
    len_i_vec = tl.load(lengths_ptr + i_indices, mask=i_mask, other=1)
    len_j_vec = tl.load(lengths_ptr + j_indices, mask=j_mask, other=1)
    len_i = len_i_vec[:, None]
    len_j = len_j_vec[None, :]

    # Create 2D index grids for loading the intersection matrix
    i_grid = i_indices[:, None]
    j_grid = j_indices[None, :]
    
    # Create the 2D mask for the block
    mask = (i_grid < n_prompts) & (j_grid < n_prompts)

    # Load intersection counts
    intersection_offset = i_grid * n_prompts + j_grid
    intersections = tl.load(intersection_matrix_ptr + intersection_offset, mask=mask, other=0)

    # Handle zero-length prompts
    zero_length_mask = (len_i == 0) | (len_j == 0)

    # Compute similarities
    max_len = tl.maximum(len_i, len_j)
    # Avoid division by zero
    max_len = tl.where(max_len == 0, 1, max_len)
    similarity = intersections.to(tl.float32) / max_len.to(tl.float32)

    # Apply similarity methods
    if method_id == 1:  # square
        similarity = similarity * similarity
    elif method_id == 2:  # sharp
        a = 2.0
        exp_a = tl.exp(a)
        exp_term = tl.exp(a * similarity * 2.0)
        similarity = (exp_term - 1.0) / (exp_a - 1.0)
        similarity = tl.minimum(tl.maximum(similarity, 0.0), 1.0)

    # Convert to distance
    distance = 1.0 - similarity

    # Handle zero-length prompts by setting their distance to 1.0
    distance = tl.where(zero_length_mask, 1.0, distance)

    # Store results
    tl.store(result_matrix_ptr + intersection_offset, distance, mask=mask)


@triton.jit
def intersection_kernel_vectorized(
    prompts_flat_ptr,
    prompt_offsets_ptr,
    prompt_lengths_ptr,
    result_matrix_ptr,
    n_prompts,
    MAX_LEN: tl.constexpr,
):
    """
    Vectorized intersection kernel where each program handles one row.
    This version avoids the explicit inner python loop to prevent ValueError.
    """
    # Each program instance computes intersections for one row
    row_idx = tl.program_id(0)

    # Early exit if row index is out of bounds
    if row_idx >= n_prompts:
        return

    # Load metadata for the row prompt
    offset_i = tl.load(prompt_offsets_ptr + row_idx)
    length_i = tl.load(prompt_lengths_ptr + row_idx)

    # Load the entire row prompt into SRAM
    prompt_indices_i = offset_i + tl.arange(0, MAX_LEN)
    prompt_mask_i = tl.arange(0, MAX_LEN) < length_i
    prompt_data_i = tl.load(prompts_flat_ptr + prompt_indices_i, mask=prompt_mask_i, other=-1)

    # Iterate through all columns to compute intersections with the current row
    for col_idx in range(n_prompts):
        # Load metadata for the column prompt
        offset_j = tl.load(prompt_offsets_ptr + col_idx)
        length_j = tl.load(prompt_lengths_ptr + col_idx)

        # Load the entire column prompt into SRAM
        prompt_indices_j = offset_j + tl.arange(0, MAX_LEN)
        prompt_mask_j = tl.arange(0, MAX_LEN) < length_j
        prompt_data_j = tl.load(prompts_flat_ptr + prompt_indices_j, mask=prompt_mask_j, other=-1)

        # --- Vectorized Intersection Calculation ---
        # Create broadcastable tensors (col_vec vs row_vec)
        elements_i = prompt_data_i[:, None]
        elements_j = prompt_data_j[None, :]
        
        # Create a 2D matrix of matches
        match_matrix = (elements_i == elements_j)
        
        # Apply masks to only consider valid elements of each prompt
        valid_matches = match_matrix & prompt_mask_i[:, None] & prompt_mask_j[None, :]
        
        # For each element in prompt_i, see if it has at least one match in prompt_j
        # Sum along axis 1 to collapse the columns (prompt_j elements)
        found_in_j = tl.sum(valid_matches, axis=1) > 0
        
        # The intersection count is the sum of unique elements from i found in j
        intersection_count = tl.sum(found_in_j)

        # Store the final count
        result_offset = row_idx * n_prompts + col_idx
        tl.store(result_matrix_ptr + result_offset, intersection_count)


class TritonOperations:
    """Triton implementation for clustering operations."""

    def __init__(self, device="cuda"):
        self.device = device
        self.method_map = {"direct": 0, "square": 1, "sharp": 2}

    def prepare_prompts(self, prompts):
        """Prepare prompts for Triton operations."""
        if not prompts:
            empty_tensor = torch.tensor([], dtype=torch.int32, device=self.device)
            return empty_tensor, empty_tensor, empty_tensor

        # Sort and deduplicate each prompt
        sorted_prompts = []
        for q in prompts:
            if len(q) > 0:
                unique_sorted = np.unique(np.array(q, dtype=np.int32))
                sorted_prompts.append(unique_sorted)
            else:
                sorted_prompts.append(np.array([], dtype=np.int32))

        # Flatten all prompts into a single array
        non_empty_prompts = [q for q in sorted_prompts if len(q) > 0]
        if not non_empty_prompts:
            # Handle case where all prompts are empty
            prompt_lengths = np.array([0] * len(prompts), dtype=np.int32)
            prompt_offsets = np.zeros(len(prompts), dtype=np.int32)
            prompts_flat = np.array([], dtype=np.int32)
        else:
            prompts_flat = np.concatenate(non_empty_prompts)
            prompt_lengths = np.array([len(q) for q in sorted_prompts], dtype=np.int32)
            prompt_offsets = np.zeros(len(sorted_prompts), dtype=np.int32)
            if len(sorted_prompts) > 1:
                prompt_offsets[1:] = np.cumsum(prompt_lengths[:-1])

        # Convert to PyTorch tensors and move to the specified device
        prompts_flat_tensor = torch.from_numpy(prompts_flat).to(self.device)
        prompt_lengths_tensor = torch.from_numpy(prompt_lengths).to(self.device)
        prompt_offsets_tensor = torch.from_numpy(prompt_offsets).to(self.device)

        return prompts_flat_tensor, prompt_offsets_tensor, prompt_lengths_tensor

    def compute_distance_matrix(self, prompts, method="sharp", block_size=16, use_block_kernel=True):
        """Compute distance matrix using Triton kernels."""
        n = len(prompts)
        if n < 2:
            return np.array([])
        
        prompts_flat, prompt_offsets, prompt_lengths = self.prepare_prompts(prompts)
        if n > 0 and prompts_flat.numel() == 0:
            return np.ones((n, n), dtype=np.float32)
        
        # Allocate result tensors on the GPU
        intersection_matrix = torch.zeros((n, n), dtype=torch.int32, device=self.device)
        distance_matrix = torch.zeros((n, n), dtype=torch.float32, device=self.device)
        
        # Determine max prompt length for kernels that need it
        max_len = int(prompt_lengths.max().item()) if prompt_lengths.numel() > 0 else 0
        
        # Cap max length to a reasonable power of 2 for Triton performance
        if max_len > 512: MAX_LEN = 1024
        elif max_len > 256: MAX_LEN = 512
        elif max_len > 128: MAX_LEN = 256
        elif max_len > 64: MAX_LEN = 128
        elif max_len > 32: MAX_LEN = 64
        else: MAX_LEN = 32

        # --- Launch Intersection Kernel ---
        if use_block_kernel:
            # Use block-based kernel
            cdiv = lambda a, b: (a + b - 1) // b
            grid = (cdiv(n, block_size), cdiv(n, block_size))
            
            intersection_kernel_triton[grid](
                prompts_flat,
                prompt_offsets,
                prompt_lengths,
                intersection_matrix,
                n,
                BLOCK_SIZE=block_size,
                MAX_LEN=MAX_LEN,
                num_warps=4,
            )
        else:
            # Use row-wise vectorized kernel
            grid_size = (n,)  # One program per row
            intersection_kernel_vectorized[grid_size](
                prompts_flat,
                prompt_offsets,
                prompt_lengths,
                intersection_matrix,
                n,
                MAX_LEN=MAX_LEN,
                num_warps=4,
            )
        
        # --- Launch Distance Kernel ---
        method_id = self.method_map.get(method, 0)
        cdiv = lambda a, b: (a + b - 1) // b
        grid = (cdiv(n, block_size), cdiv(n, block_size))
        
        distance_kernel_triton[grid](
            intersection_matrix,
            prompt_lengths,
            distance_matrix,
            method_id,
            n,
            BLOCK_SIZE=block_size,
            num_warps=4
        )
        
        # Copy the result back to CPU and convert to a NumPy array
        return distance_matrix.cpu().numpy()