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

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

    def __mul__(self, other): #overloads * operator (C = A*B)
        result = np.matmul(self.data, other.data)
        return array(result)

    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)

    @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 shape(self):
        return self.data.shape

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