import argparse

import numpy as np
import scipy.sparse as sp
import struct
import scipy.sparse.linalg as spla # For matrix norm
import time
import os

def loadSparseMatrixBinary(f):
    """
    Loads a sparse matrix from the custom binary format (.csrbin).

    Args:
        filename (str): The path to the .csrbin file.

    Returns:
        scipy.sparse.csr_matrix: The loaded sparse matrix.

    Raises:
        ValueError: If the file format is incorrect or sizes don't match.
        IOError: If the file cannot be read.
    """
    INT_SIZE = 8  # Expecting int64_t from the C++ writer
    FLT_SIZE = 8  # Expecting double from the C++ writer
    EXPECTED_MAGIC = b'CSRB'
    EXPECTED_VERSION = 1

    if isinstance(f, str):
        f = open(f, 'rb')

    try:
        # --- Read Header ---
        magic = f.read(4)
        if magic != EXPECTED_MAGIC:
            raise ValueError(f"Invalid magic number. Expected {EXPECTED_MAGIC}, got {magic}")

        version, int_size_file, flt_size_file, reserved = struct.unpack('<BBBB', f.read(4))
        # '<' means little-endian, 'B' means unsigned char (1 byte)

        if version != EXPECTED_VERSION:
            print(f"Warning: File version {version} differs from expected {EXPECTED_VERSION}.")
        if int_size_file != INT_SIZE:
            raise ValueError(f"Integer size mismatch. Expected {INT_SIZE}, file has {int_size_file}")
        if flt_size_file != FLT_SIZE:
            raise ValueError(f"Float size mismatch. Expected {FLT_SIZE}, file has {flt_size_file}")

        height, width, nnz = struct.unpack('<QQQ', f.read(24))
        # '<' means little-endian, 'Q' means unsigned long long (8 bytes)

        i_count = height + 1
        j_count = nnz
        data_count = nnz

        if nnz == 0: # Handle empty matrix case
            print("Warning: Matrix file contains zero non-zero elements.")
            # Return an empty matrix with correct shape
            return sp.csr_matrix((height, width), dtype=np.float64)


        # --- Read Arrays ---

        # Read I array (Row Pointers)
        expected_i_bytes = i_count * INT_SIZE
        I_array = np.fromfile(f, dtype=np.int64, count=i_count) # Read as int64
        if I_array.size != i_count:
            raise ValueError(f"Error reading I array. Expected {i_count} elements, read {I_array.size}. File truncated or corrupt?")

        # Read J array (Column Indices)
        expected_j_bytes = j_count * INT_SIZE
        J_array = np.fromfile(f, dtype=np.int64, count=j_count) # Read as int64
        if J_array.size != j_count:
            raise ValueError(f"Error reading J array. Expected {j_count} elements, read {J_array.size}. File truncated or corrupt?")

        # Read Data array (Values)
        expected_data_bytes = data_count * FLT_SIZE
        Data_array = np.fromfile(f, dtype=np.float64, count=data_count) # Read as float64
        if Data_array.size != data_count:
            raise ValueError(f"Error reading Data array. Expected {data_count} elements, read {Data_array.size}. File truncated or corrupt?")


        # --- Construct SciPy CSR Matrix ---
        sparse_matrix = sp.csr_matrix((Data_array, J_array, I_array), shape=(height, width))

        if sparse_matrix.nnz != nnz:
            print(f"Warning: NNZ mismatch after loading. Header NNZ: {nnz}, Scipy NNZ: {sparse_matrix.nnz}")


        return sparse_matrix

    except Exception as e:
        raise RuntimeError(f"An error occurred while reading: {e}")

def loadSparseMatrix(filename):
    """
    Loads a sparse matrix from the custom binary format (.csrbin).

    Args:
        filename (str): The path to the .csrbin file.

    Returns:
        scipy.sparse.csr_matrix: The loaded sparse matrix.

    Raises:
        ValueError: If the file format is incorrect or sizes don't match.
        IOError: If the file cannot be read.
    """
    with open(filename, 'rb') as f:
        # Check magic number
        magic = f.read(4)
        if magic != b'CSRB':
            raise ValueError(f"Invalid magic number. Expected 'CSRB', got {magic}")

        # Read the rest of the file
        f.seek(0, 0)
        sm = loadSparseMatrixBinary(f)

    return sm
def loadBlockMatrix(filename):
    smList = list()
    with open(filename, 'rb') as f:
        f.seek(0, 2)
        fileSize = f.tell()
        f.seek(0, 0)
        magic = f.read(4)
        if magic != b'BLCK':
            raise ValueError(f"Invalid magic number. Expected 'BLCK'. got {magic}")
        size = struct.unpack('<B', f.read(1))[0]
        print(f"Size: {size}")
        while f.tell() < fileSize:
            dataStartCard = f.read(9)
            if dataStartCard != b'DATASTART':
                raise ValueError(f"Invalid data start card. Expected 'DATASTART' Got {dataStartCard}.")
            blockId = struct.unpack(f'<ii', f.read(8))
            sm = loadSparseMatrixBinary(f)
            smList.append((sm, blockId))
            # unpack 2 ints as the block id
            dataEndCard = f.read(7)
            if dataEndCard != b'DATAEND':
                raise ValueError(f"Invalid data end card. Expected 'DATAEND'. Got {dataEndCard}.")
    outArray = np.empty(shape=(size, size), dtype=np.object_)

    for sm, blockId in smList:
        if blockId[0] >= size or blockId[1] >= size:
            raise ValueError(f"Block ID {blockId} out of range. Size: {size}")
        outArray[blockId[0], blockId[1]] = sm

    # Check if all blocks are filled
    return sp.bmat(outArray, format='csr')


def analyze_sparse_matrix(sp_mat):
    """
    Analyzes a SciPy sparse matrix and prints various statistics.

    Args:
        sp_mat (scipy.sparse.spmatrix): The sparse matrix to analyze.
                                         (e.g., csr_matrix, csc_matrix).
    """
    print("-" * 50)
    print("Sparse Matrix Analysis Report")
    print("-" * 50)

    rows, cols = sp_mat.shape
    print(f"Size (Shape):           {rows} rows x {cols} columns")

    if rows == 0 or cols == 0:
        print("\nMatrix is empty. No further analysis possible.")
        print("-" * 50)
        return

    nnz = sp_mat.nnz
    total_elements = rows * cols
    sparsity = 0.0
    if total_elements > 0:
        sparsity = 1.0 - (nnz / total_elements)
    else:
        sparsity = 1.0

    print(f"Non-zero elements (NNZ): {nnz} (~{nnz*8/(1024**2):.2f} MB)")
    print(f"Total elements:         {total_elements} (~{total_elements*8/(1024**3):.2f} GB)")
    print(f"Sparsity:               {sparsity:.6%} (percentage of zeros)")

    if nnz == 0:
        print("\nMatrix contains only zero elements.")
        diag_elements = sp_mat.diagonal()
        print(f"\nDiagonal Mean:          {np.mean(diag_elements):.6e}")
        print(f"Diagonal Max:           {np.max(diag_elements):.6e}")
        print(f"Diagonal Min:           {np.min(diag_elements):.6e}")
        print(f"Value Range (Min):      N/A (no non-zero values)")
        print(f"Value Range (Max):      N/A (no non-zero values)")
        print(f"Mean Non-Zero Value:    N/A (no non-zero values)")
        print(f"Relative Diagonal Norm: N/A (matrix norm is zero)")
        print("-" * 50)
        return

    all_values = sp_mat.data # Access non-zero values directly
    min_val = np.min(all_values)
    max_val = np.max(all_values)
    mean_val = np.mean(all_values)
    print(f"\nValue Range (Min):      {min_val:.6e}")
    print(f"Value Range (Max):      {max_val:.6e}")
    print(f"Mean Non-Zero Value:    {mean_val:.6e}")


    print("\n--- Diagonal Properties ---")
    start_diag = time.time()
    diag_elements = sp_mat.diagonal()
    end_diag = time.time()
    print(f"(Diagonal extraction time: {end_diag - start_diag:.4f}s)")

    if diag_elements.size > 0: # Should always be true unless rows=0 (handled above)
        mean_diag = np.mean(diag_elements)
        max_diag = np.max(diag_elements)
        min_diag = np.min(diag_elements)

        diag_nonzero = diag_elements[diag_elements != 0]
        if diag_nonzero.size > 0:
            mean_diag_nz = np.mean(diag_nonzero)
            print(f"Mean Diagonal (all):    {mean_diag:.6e}")
            print(f"Mean Diagonal (non-zero):{mean_diag_nz:.6e} ({diag_nonzero.size} elements)")
        else:
            print(f"Mean Diagonal (all):    {mean_diag:.6e}")
            print(f"Mean Diagonal (non-zero): N/A (all diagonal elements are zero)")

        print(f"Max Diagonal:           {max_diag:.6e}")
        print(f"Min Diagonal:           {min_diag:.6e}")

        # 5. "Diagonality" - Relative Diagonal Norm (using Frobenius norm)
        # The Frobenius norm is sqrt(sum(abs(A_ij)^2))
        start_norm = time.time()
        norm_diag = np.linalg.norm(diag_elements)
        norm_matrix = spla.norm(sp_mat, ord='fro')
        end_norm = time.time()
        print(f"(Norm calculation time: {end_norm - start_norm:.4f}s)")

        if norm_matrix > 1e-15: # Avoid division by zero
            diagonality_ratio = norm_diag / norm_matrix
            print(f"\nRelative Diagonal Norm (Frobenius): {diagonality_ratio:.6f}")
            print(f"  (Ratio of ||diag(A)||_F / ||A||_F)")
            print(f"  (Diagonal Norm = {norm_diag:.6e}, Matrix Norm = {norm_matrix:.6e})")
            if diagonality_ratio > 0.99:
                print("  -> Matrix is strongly diagonal dominant by norm.")
            elif diagonality_ratio < 0.1:
                print("  -> Matrix norm is dominated by off-diagonal elements.")
        else:
            print("\nRelative Diagonal Norm: N/A (matrix Frobenius norm is zero)")

    else: # Should not happen if rows > 0
        print("\nCould not extract diagonal (matrix has zero rows?).")


    # 6. Other Useful Stats
    print("\n--- Other Properties ---")
    is_square = rows == cols
    print(f"Is Square:              {is_square}")
    if is_square:
        try:
            diff_norm = spla.norm(sp_mat - sp_mat.T, ord='fro')
            if diff_norm < 1e-10 * norm_matrix : # Check relative difference norm
                print(f"Is Symmetric (approx):  True (||A - A.T||_F / ||A||_F < 1e-10)")
            else:
                print(f"Is Symmetric (approx):  False (||A - A.T||_F = {diff_norm:.2e})")
        except Exception as e:
            print(f"Is Symmetric (approx):  Check failed ({e})")
    else:
        print(f"Is Symmetric (approx):  False (not square)")


    print("-" * 50)

def load_and_analyze_sparse_matrix(filename: str):
    sm = loadSparseMatrixBinary(filename)
    analyze_sparse_matrix(sm)

def compute_frobenius_distance(sparseMat):
    identityMat = sp.eye(sparseMat.shape[0], sparseMat.shape[1], format='csr')
    diffMat = sparseMat - identityMat
    normDistance = np.sqrt(diffMat.data.dot(diffMat.data))
    frobNormIdentity = np.sqrt(identityMat.shape[0])

    return normDistance/frobNormIdentity

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple tool to get some statistics about a sparse matrix from mfem")
    parser.add_argument("path", help="path to the output file", type=str)

    args = parser.parse_args()

    load_and_analyze_sparse_matrix(args.filename)
