from scipy.sparse import coo_matrix
import struct
import numpy as np

def load(fname):
    MATRIX_MAGIC = 1211216
    VECTOR_MAGIC = 1211214
    
    with open(fname, "rb") as f:
        magic = struct.unpack(">i", f.read(4))[0]

        if magic == MATRIX_MAGIC:
            m = struct.unpack(">i", f.read(4))[0]
            n = struct.unpack(">i", f.read(4))[0]
            nonzero = struct.unpack(">i", f.read(4))[0]
            nonzero_row = struct.unpack(">%di" %(m), f.read(m*4))
            ks = struct.unpack(">%di" %(nonzero), f.read(nonzero*4))
            xs = np.array(struct.unpack(">%dd" %(nonzero), f.read(nonzero*8)))

            ii = np.zeros(nonzero, dtype=int)
            jj = np.zeros(nonzero, dtype=int)
            a = 0
            for i in range(0, m):
                for j in range(0, nonzero_row[i]):
                    ii[a] = i
                    jj[a] = ks[a]
                    a += 1

            return coo_matrix((xs, (ii, jj)), shape=(m, n))

        elif magic == VECTOR_MAGIC:
            m = struct.unpack(">i", f.read(4))[0]
            xs = np.array(struct.unpack(">%dd" %(m), f.read(m*8)))
            return xs
            
        else:
            sys.stderr.write("Incorrect PETSc magic number (read %d).\n" %(magic))
            exit(1)