import numpy as np
import multiprocessing
from abc import ABC, abstractmethod
import torch
import triton
import triton.language as tl


class PrefixComputer(ABC):
    """Abstract base class for computing prefix length matrices."""
    
    @abstractmethod
    def compute_prefix_matrix(self, queries):
        pass


class CPUPrefixComputer(PrefixComputer):
    """CPU-based prefix matrix computation."""
    
    def __init__(self, use_parallel=False):
        self.use_parallel = use_parallel
    
    def compute_prefix_matrix(self, queries):
        """Compute prefix length matrix using CPU."""
        if self.use_parallel:
            return self._compute_parallel(queries)
        else:
            return self._compute_serial(queries)
    
    def _compute_serial(self, queries):
        """Serial CPU implementation."""
        n = len(queries)
        prefix_matrix = np.zeros((n, n), dtype=np.int32)
        
        for i in range(n):
            for j in range(i + 1, n):
                length = self._longest_common_prefix_length(queries[i], queries[j])
                prefix_matrix[i, j] = length
                prefix_matrix[j, i] = length
        
        return prefix_matrix
    
    def _compute_parallel(self, queries):
        """Parallel CPU implementation using multiprocessing."""
        n = len(queries)
        if n <= 1:
            return np.zeros((n, n), dtype=np.int32)
        
        def worker(args):
            i, j, queries = args
            return i, j, self._longest_common_prefix_length(queries[i], queries[j])
        
        pairs = [(i, j, queries) for i in range(n) for j in range(i + 1, n)]
        
        try:
            with multiprocessing.Pool() as pool:
                results = pool.map(worker, pairs)
            
            prefix_matrix = np.zeros((n, n), dtype=np.int32)
            for i, j, length in results:
                prefix_matrix[i, j] = length
                prefix_matrix[j, i] = length
            
            return prefix_matrix
        except Exception as e:
            print(f"Parallel processing failed: {e}, falling back to serial")
            return self._compute_serial(queries)
    
    @staticmethod
    def _longest_common_prefix_length(list1, list2):
        """Calculate the length of the longest common prefix between two lists."""
        length = 0
        for a, b in zip(list1, list2):
            if a == b:
                length += 1
            else:
                break
        return length


class TritonPrefixComputer(PrefixComputer):
    """GPU-based prefix matrix computation using Triton."""
    
    def __init__(self, fallback_to_cpu=True):
        self.fallback_to_cpu = fallback_to_cpu
        self.cpu_computer = CPUPrefixComputer() if fallback_to_cpu else None
    
    def compute_prefix_matrix(self, queries):
        """Compute prefix length matrix using GPU with Triton."""
        n = len(queries)
        if n <= 1:
            return np.zeros((n, n), dtype=np.int32)
        
        try:
            if not torch.cuda.is_available():
                if self.fallback_to_cpu:
                    print("CUDA not available, falling back to CPU")
                    return self.cpu_computer.compute_prefix_matrix(queries)
                else:
                    raise RuntimeError("CUDA not available")
            
            return self._compute_gpu(queries)
        
        except Exception as e:
            if self.fallback_to_cpu:
                print(f"GPU computation failed: {e}, falling back to CPU")
                return self.cpu_computer.compute_prefix_matrix(queries)
            else:
                raise e
    
    def _compute_gpu(self, queries):
        """GPU implementation using Triton."""
        n = len(queries)
        queries_flat, query_lengths, query_starts, max_len = self._prepare_gpu_data(queries, 'cuda')
        
        prefix_matrix = torch.empty((n, n), dtype=torch.int32, device='cuda')
        
        def grid(meta):
            return (triton.cdiv(n, meta['BLOCK_SIZE_M']), triton.cdiv(n, meta['BLOCK_SIZE_N']))
        
        self._compute_prefix_kernel[grid](
            queries_flat, query_lengths, query_starts, prefix_matrix,
            n_queries=n,
            stride_pm_row=prefix_matrix.stride(0),
            stride_pm_col=prefix_matrix.stride(1),
            BLOCK_SIZE_M=16,
            BLOCK_SIZE_N=16,
            MAX_QUERY_LEN=max_len
        )
        
        return prefix_matrix.cpu().numpy()

    @staticmethod
    @triton.jit
    def _compute_prefix_kernel(
        queries_flat_ptr, query_lengths_ptr, query_starts_ptr,
        prefix_matrix_ptr,
        n_queries,
        stride_pm_row, stride_pm_col,
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
        MAX_QUERY_LEN: tl.constexpr
    ):
        """Triton kernel to compute prefix length matrix with block processing."""
        pid_m = tl.program_id(axis=0)
        pid_n = tl.program_id(axis=1)
        
        offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        
        mask_m = offs_m < n_queries
        mask_n = offs_n < n_queries
        
        # Only compute upper triangle
        valid_pair_mask = (mask_m[:, None] & mask_n[None, :]) & (offs_m[:, None] < offs_n[None, :])
        
        # Load query metadata
        starts_m = tl.load(query_starts_ptr + offs_m, mask=mask_m, other=0)
        lengths_m = tl.load(query_lengths_ptr + offs_m, mask=mask_m, other=0)
        starts_n = tl.load(query_starts_ptr + offs_n, mask=mask_n, other=0)
        lengths_n = tl.load(query_lengths_ptr + offs_n, mask=mask_n, other=0)
        
        # Initialize prefix lengths
        prefix_len = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32)
        
        # Compute minimum lengths for each pair
        min_len = tl.minimum(lengths_m[:, None], lengths_n[None, :])
        
        # Process each position
        for k in range(MAX_QUERY_LEN):
            # Check if k is within bounds for each pair
            k_valid = (k < min_len) & valid_pair_mask
            
            # Load elements for each query in the block
            elem_m = tl.load(queries_flat_ptr + starts_m + k, mask=mask_m & (k < lengths_m), other=-1)
            elem_n = tl.load(queries_flat_ptr + starts_n + k, mask=mask_n & (k < lengths_n), other=-2)
            
            # Compare elements
            matches = (elem_m[:, None] == elem_n[None, :]) & k_valid
            
            # Update prefix lengths only for matching elements
            prefix_len = tl.where(matches, prefix_len + 1, prefix_len)
            
            # Stop checking pairs that don't match
            valid_pair_mask = valid_pair_mask & matches
        
        # Store results
        prefix_matrix_ptrs = prefix_matrix_ptr + offs_m[:, None] * stride_pm_row + offs_n[None, :] * stride_pm_col
        tl.store(prefix_matrix_ptrs, prefix_len, mask=(mask_m[:, None] & mask_n[None, :]) & (offs_m[:, None] < offs_n[None, :]))
        
        # Store symmetric results
        sym_prefix_matrix_ptrs = prefix_matrix_ptr + offs_n[:, None] * stride_pm_row + offs_m[None, :] * stride_pm_col
        tl.store(sym_prefix_matrix_ptrs, tl.trans(prefix_len), mask=tl.trans((mask_m[:, None] & mask_n[None, :]) & (offs_m[:, None] < offs_n[None, :])))

    def _prepare_gpu_data(self, queries, device):
        """Convert list of queries to flat tensors for GPU processing."""
        if not queries:
            return (torch.tensor([], dtype=torch.int32, device=device),
                    torch.tensor([], dtype=torch.int32, device=device),
                    torch.tensor([], dtype=torch.int32, device=device),
                    0)
        
        query_lengths_np = np.array([len(q) for q in queries], dtype=np.int32)
        max_len = int(np.max(query_lengths_np)) if len(query_lengths_np) > 0 else 0

        query_starts_np = np.zeros(len(queries), dtype=np.int32)
        if len(queries) > 1:
            query_starts_np[1:] = np.cumsum(query_lengths_np)[:-1]
        
        queries_flat_list = []
        for query in queries:
            if not query:
                continue
            if query and isinstance(query[0], (int, float)):
                queries_flat_list.extend([int(item) for item in query])
            else:
                queries_flat_list.extend([hash(str(item)) % (2**31 - 1) for item in query])
        
        queries_flat_np = np.array(queries_flat_list, dtype=np.int32)
        
        queries_flat = torch.tensor(queries_flat_np, device=device)
        query_lengths = torch.tensor(query_lengths_np, device=device)
        query_starts = torch.tensor(query_starts_np, device=device)
        
        return queries_flat, query_lengths, query_starts, max_len