import numpy as np
import cv2
import matplotlib.pyplot as plt
import jovian

class ImageCompressionSVD:
    
    def __init__(self,img,p_dim = False):
        if len(img.shape) > 2:
            self.img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        else:
            self.img = img
        self.h , self.w = self.img.shape
        
        self.U, self.S, self.VT = np.linalg.svd(self.img, full_matrices=False)
        
        self.S_D = np.diag(self.S)
        self.S_n = self.S/np.sum(self.S_D)
        self.cumulativeSV = np.cumsum(self.S_n)
       
        if p_dim == True:
            print("Image Size : {} * {}". format(self.h, self.w))
            print("Dimention of U matrix : {}".format(self.U.shape))
            print("Dimention of S matrix : {}".format(self.S.shape))
            print("Dimention of Vt matrix : {}".format(self.VT.shape))
    
    def getImg(self, r = None):
        self.rr = r
        if self.rr == None:
            print("Error!! Enter the rank")
        else:
            Img = self.U[:,:self.rr]@self.S_D[:self.rr,:self.rr]@self.VT[:self.rr, :]
        return Img
    
    def plotSV(self): 
       

        plt.figure(figsize = (14,5))
        plt.subplot(1,2,1)
        plt.plot(self.cumulativeSV)
        plt.xlabel('singular values index')
        plt.ylabel('cumulative sum')

        plt.subplot(1,2,2)
        plt.plot(self.S)
        plt.xlabel('singular values index')
        plt.ylabel('singular values')
        plt.show()
    
    def rforCompress(self,percent = 50):
        self.percent = percent
        self.c_array = np.floor(self.cumulativeSV*100)
        self.result_array = np.where(self.c_array == self.percent)
        
        while len(self.result_array[0]) == 0:
            self.percent = self.percent + 1
            self.result_array = np.where(self.c_array == self.percent)
        else:
            self.result = np.amin(self.result_array[0])
            return print("For approx {} % energy retainment, select rank = {}".format(self.percent,self.result))
        return self.result_array
    
    def showLR(self, r = (5,10,50)):
        self.r1 = r[0]
        self.r2 = r[1]
        self.r3 = r[2]
        j = 1
        plt.figure(figsize = (15,15))
        for r in (0,self.r1, self.r2, self.r3):
            plt.subplot(2,2,j)
            if r == 0:
                img = plt.imshow(self.img)
                img.set_cmap('gray')
                plt.axis('off')
                plt.title('Original Image') 
                j = j+1
            else:
                Xapprox = self.U[:,:r]@self.S_D[:r,:r]@self.VT[:r, :]
                
                self.a,self.b = self.U[:,:r].shape
                self.c,self.d = self.S_D[:r,:r].shape
                self.e,self.f = self.VT[:r, :].shape
               
                self.new_size = self.a*self.b + self.c*self.d + self.e*self.f
                self.ori_size = self.h * self.w
                
                self.comp = ((self.ori_size - self.new_size)/self.ori_size) * 100
               
                img = plt.imshow(Xapprox)
                img.set_cmap('gray')
                #plt.axis('off')
                plt.title('r = {}'.format(r))
                plt.xlabel('Percentage compressed = {:.3f}'.format(self.comp))
               
                j = j+1
        plt.tight_layout()

img = cv2.imread("bird.jpg")

SVD = ImageCompressionSVD(img)

SVD.plotSV()