import numpy as np
from scipy.linalg import lu, solve_triangular

def rand(shape):
    if not isinstance(shape, tuple):
        raise TypeError("Shape must be a tuple, e.g. (3, 3) or (2, 2, 2)")
    return array(np.random.rand(*shape))

def eye(n):
    return array(np.eye(n))
    
class array:
    data = None #np.array object
    err = 1e-12
    
    def __init__(self, data):
        self.data = np.array(data)

    @property
    def norm(self):
        return float(np.linalg.norm(self.data))

    def __getitem__(self, key): #supports A[:,0] in the numpy way
        return array(self.data[key])

    def __truediv__(self, other):
        if isinstance(other, (int, float, complex)):
            return array(self.data / other)
    
        if isinstance(other, array):
            if np.prod(other.data.shape) == 1:
                #return array(self.data / float(other.data))
                return array(self.data / other.data.item())
    
        raise TypeError("Division only supported with scalar values or scalar-like array")

    def __mul__(self, other):  # overloads * operator (C = A*B)
        if isinstance(other, array):
            result = self.data @ other.data
        else:
            result = self.data * other
    
        # If result is scalar-like (e.g. array([[6.32]])), return float
        if result.size == 1:
            return result.item()
        else:
            return array(result)

    def __rmul__(self, other):
        return self * other

    def __add__(self, other):        
        if isinstance(other, array):
            result = self.data + other.data
        else:
            result = self.data + other
        return array(result)

    def __radd__(self, other):
        return self + other


    def __rshift__(self, other): #overloads >> operator (x = b >> A)
        A = other.data
        b = self.data
        
        if A.shape[0] == A.shape[1]: #if square
            if  np.all(np.abs(np.tril(A, -1)) < self.err): # Upper triangular
                if not np.any(np.abs(np.diag(A)) < self.err): # Full rank
                    x = solve_triangular(A, b, lower=False)
                    return array(x)
                
            elif np.all(np.abs(np.triu(A, 1)) < self.err): # Lower triangular
                if not np.any(np.abs(np.diag(A)) < self.err): # Full rank
                    x = solve_triangular(A, b, lower=True)
                    return array(x)

        x, _, rank, _ = np.linalg.lstsq(A, b)

        if rank < min(A.shape):
            print("Matrix is rank deficient")

        return array(x)

    def __eq__(self, other):
        if not isinstance(other, array):
            return False
        elif self.data.shape != other.data.shape:
            return False
        else:
            return bool(np.all(np.abs(self.data - other.data) < self.err))

    @property
    def T(self): #Transpose property (A.T)
        if self.data.ndim == 1:  # for 1D arrays
            return array(self.data.reshape(-1, 1))
        else:
            return array(self.data.T)
    
    @property
    def LU(self): #LU decomposition property(P, L, U = A.LU)
        A = self.data
        P, L, U = lu(A)
        if A.shape[0] != A.shape[1]: #if square
            print("LU: Matrix is not square")
        return array(P), array(L), array(U)

    @property
    def R(self):
        A = self.data
        R = np.linalg.qr(A, mode='r')

        if np.any(np.abs(np.diag(R)) < self.err): # not Full rank
            print("R: A is rank deficient")
        return array(R)

    @property
    def QR(self):
        A = self.data
        Q, R = np.linalg.qr(A)
    
        if np.any(np.abs(np.diag(R)) < self.err):  # not Full rank
            print("QR: A is rank deficient")
    
        return array(Q), array(R)


    @property
    def shape(self):
        return self.data.shape

    # def __repr__(self): #string representation used by print(A)
    #     return repr(self.data);

    def __repr__(self):
        masked = np.where(np.abs(self.data) < self.err, 0, self.data) #replace small cvalues with 0
        return repr(masked).replace("0.00000000e+00", "0             ")