import heapq
class SparseMatrixCSR:
    def __init__(self, data, indices, indptr, shape):
        self.data = data
        self.indices = indices
        self.indptr = indptr
        self.shape = shape
        
    def __repr__(self):
        f = f"Matrix of size {self.shape}"
        d = f"Data : {' '.join(map(str, self.data))}"
        i = f"Row Indices : {' '.join(map(str, self.indices))}" 
        p = f"Column Indices : {' '.join(map(str, self.indptr))}"  
        return '\n'.join((f, d, i, p))
        
    @staticmethod
    def from_dense(matrix):
        # Create a CSR matrix from a dense matrix
        data, indices, indptr = [], [], [0]
        for row in matrix:
            for j, value in enumerate(row):
                if value != 0:
                    data.append(value)
                    indices.append(j)
            indptr.append(len(data))
        return SparseMatrixCSR(data, indices, indptr, (len(matrix), len(matrix[0])))

    def to_dense(self):
        dense = [[0] * self.shape[1] for _ in range(self.shape[0])]
        for row in range(self.shape[0]):
            for idx in range(self.indptr[row], self.indptr[row + 1]):
                col = self.indices[idx]
                dense[row][col] = self.data[idx]
        return dense
        
    def permute(self, ordering):
        """
        Permute the matrix rows and columns according to a given ordering.
        """
        import numpy as np

        n = len(ordering)
        inverse_order = np.argsort(ordering)

        # Permute rows
        data, indices, indptr = [], [], [0]
        for i in ordering:
            for idx in range(self.indptr[i], self.indptr[i + 1]):
                data.append(self.data[idx])
                indices.append(inverse_order[self.indices[idx]])
            indptr.append(len(data))

        # Permute columns
        indices = [ordering[idx] for idx in indices]
        return SparseMatrixCSR(data, indices, indptr, self.shape)

    def __getitem__(self, key):
        """
        Retrieve the value at position (i, j) in the matrix.
        """
        i, j = key
        for idx in range(self.indptr[i], self.indptr[i + 1]):
            if self.indices[idx] == j:
                return self.data[idx]
        return 0  # Element is zero if not explicitly stored

def sparse_lu(A):
    """
    Perform LU decomposition on a sparse matrix A in CSR format.
    :param A: SparseMatrixCSR, input sparse matrix
    :return: (L, U), lower and upper triangular matrices in CSR format
    """
    from copy import deepcopy

    n = A.shape[0]
    L_data, L_indices, L_indptr = [], [], [0]
    U_data, U_indices, U_indptr = [], [], [0]

    # Initialize working copies of A for factorization
    U = deepcopy(A)

    for k in range(n):
        # Select the pivot (diagonal element)
        pivot = None
        # Store pivot explicitly in U
        for idx in range(U.indptr[k], U.indptr[k + 1]):
            col = U.indices[idx]
            if col == k:
                U.data[idx] = pivot
                break
        else:
            # If pivot is missing, add it
            U.data.append(pivot)
            U.indices.append(k)


        # Update U (row elimination)
        for idx in range(U.indptr[k], U.indptr[k + 1]):
            col = U.indices[idx]
            if col > k:
                U.data[idx] /= pivot

        # Append to L
        for idx in range(U.indptr[k], U.indptr[k + 1]):
            col = U.indices[idx]
            if col < k:
                L_data.append(U.data[idx])
                L_indices.append(col)

        L_data.append(1.0)  # Diagonal element is always 1 for L
        L_indices.append(k)

        # Update U (subtract row multiples)
        for i in range(k + 1, n):
            factor = U[i, k] / pivot
            if factor != 0:
                for idx in range(U.indptr[k], U.indptr[k + 1]):
                    j = U.indices[idx]
                    U_ij = U[i, j] - factor * U.data[idx]

                    # Update U only if non-zero
                    if abs(U_ij) > 1e-10:  # Threshold to avoid precision issues
                        for u_idx in range(U.indptr[i], U.indptr[i + 1]):
                            if U.indices[u_idx] == j:
                                U.data[u_idx] = U_ij
                                break
                        else:  # New non-zero
                            U.data.append(U_ij)
                            U.indices.append(j)

        L_indptr.append(len(L_data))
        U_indptr.append(len(U_data))

    # Create L and U matrices
    L = SparseMatrixCSR(L_data, L_indices, L_indptr, (n, n))
    U = SparseMatrixCSR(U_data, U_indices, U_indptr, (n, n))

    return L, U



def minimum_degree_ordering(A):
    """
    Compute a reordering of rows/columns to reduce fill-in using a minimum degree strategy.
    :param A: SparseMatrixCSR, input matrix
    :return: list, reordering of rows/columns
    """


    n = A.shape[0]
    adjacency_list = {i: set() for i in range(n)}

    # Build adjacency list (symmetric for undirected graph)
    for i in range(n):
        for idx in range(A.indptr[i], A.indptr[i + 1]):
            j = A.indices[idx]
            if i != j:
                adjacency_list[i].add(j)
                adjacency_list[j].add(i)

    # Priority queue for degrees
    degree = {i: len(neighbors) for i, neighbors in adjacency_list.items()}
    min_heap = [(deg, node) for node, deg in degree.items()]
    heapq.heapify(min_heap)

    # Elimination ordering
    ordering = []
    eliminated = set()

    while min_heap:
        _, node = heapq.heappop(min_heap)
        if node in eliminated:
            continue
        ordering.append(node)
        eliminated.add(node)

        # Update neighbors
        for neighbor in adjacency_list[node]:
            adjacency_list[neighbor].remove(node)
            degree[neighbor] -= 1

        # Update heap
        new_heap = [(degree[n], n) for n in adjacency_list if n not in eliminated]
        heapq.heapify(new_heap)
        min_heap = new_heap

    return ordering


def solve(A, b):
    """
    Solve Ax = b using sparse LU decomposition with AMD ordering.
    :param A: SparseMatrixCSR, the matrix A in CSR format
    :param b: list, the dense vector b
    :return: list, the solution vector x
    """
    # Step 1: Compute the ordering
    ordering = minimum_degree_ordering(A)

    # Step 2: Permute the matrix and vector
    A_permuted = A.permute(ordering)
    b_permuted = [b[i] for i in ordering]

    # Step 3: Decompose permuted A into L and U
    L, U = sparse_lu(A_permuted)

    # Step 4: Forward substitution to solve Ly = b_permuted
    y = [0] * A.shape[0]
    for i in range(A.shape[0]):
        sum_terms = 0
        for idx in range(L.indptr[i], L.indptr[i + 1]):
            col = L.indices[idx]
            sum_terms += L.data[idx] * y[col]
        y[i] = (b_permuted[i] - sum_terms)  # Divide by L[i, i] which is 1 for unit lower triangular

    # Step 5: Backward substitution to solve Ux = y
    x = [0] * A.shape[0]
    for i in range(A.shape[0] - 1, -1, -1):
        sum_terms = 0
        for idx in range(U.indptr[i], U.indptr[i + 1]):
            col = U.indices[idx]
            if col > i:
                sum_terms += U.data[idx] * x[col]
        diag_idx = next(
            (idx for idx in range(U.indptr[i], U.indptr[i + 1]) if U.indices[idx] == i), None
        )
        if diag_idx is None:
            raise ValueError(f"Diagonal element at row {i} is missing in U. Matrix may be singular.")

        x[i] = (y[i] - sum_terms) / U.data[diag_idx]

    # Step 6: Reverse the permutation on the solution vector
    inverse_ordering = [0] * A.shape[0]
    for i, p in enumerate(ordering):
        inverse_ordering[p] = i

    x_original_order = [x[inverse_ordering[i]] for i in range(A.shape[0])]

    return x_original_order

# Example matrix and vector
dense_matrix = [
    [4, 0, 0, 0],
    [3, 5, 0, 0],
    [0, 1, 6, 0],
    [0, 0, 2, 7]
]
b = [8, 7, 10, 13]

# Convert to CSR
A = SparseMatrixCSR.from_dense(dense_matrix)
print(A)
print(A.to_dense())
# Solve Ax = b
# x = solve(A, b)

# print("Solution vector x:", x)