PCA

import numpy as np
import matplotlib.pylab as plt
import h5py
from scipy.ndimage.filters import median_filter
from scipy.ndimage import rotate

from sdp.correction.function import subtract_dark_frame, subtract_flat_field, copy_metadata, get_dark_frame, get_flat_field, make_scan
from sdp.correction.corrector import Corrector

from sdp.mathlib.spectrum  import power_spectrum_1d
from scipy.linalg import svd
from numpy.fft import rfft2, irfft2,fftfreq

#/////////////////////////////////////////////////////////////

class pca2d:
    def __init__(self, data_d):
        self.cov_matrix = None
        self.data = data_d
        self.num_frames = self.data.shape[0]
        self.y_size = self.data.shape[1]
        self.x_size = self.data.shape[2]
        self.u_matrix = None
        self.s_matrix = None
        self.vh_matrix = None

    def calc_cov_matrix(self, init=0, fin=-1, step=1, excl_list=[], summing='row', filename=None):
        sum = np.zeros((self.y_size, self.x_size), dtype='float64')
        if fin is -1:
            fin = self.num_frames
        overall = (fin-init)
        for j in range(init,fin,step):
            if j in excl_list:
                overall = overall - 1
                continue
            sum =sum + self.data[j]
        if overall < 1:
            print('number of frame summed < 1')
        aver_frame = sum / overall
        if summing is 'column':
            self.cov_matrix = np.zeros((self.y_size, self.y_size), dtype='float64')
            for j in range(init, fin, step):
                if j in excl_list:
                    continue
                self.cov_matrix = self.cov_matrix + np.dot((self.data[j]-aver_frame), np.transpose(self.data[j]-aver_frame))
        elif summing is 'row':
            self.cov_matrix = np.zeros((self.x_size, self.x_size), dtype='float64')
            for j in range(init,fin, step):
                if j in excl_list:
                    continue
                self.cov_matrix = self.cov_matrix + np.dot(np.transpose(self.data[j]-aver_frame),(self.data[j]-aver_frame))
        else:
            pass#TODO
        self.cov_matrix = self.cov_matrix / overall
        if filename is not None:
            dsc = h5py.File(filename, 'w')
            ds = dsc.create_dataset('cov_matrix', self.cov_matrix.shape,
                                    dtype=self.cov_matrix.dtype)
            ds[:,:] = self.cov_matrix[:,:]
            dsc.close()
        self.svd()

    def load_cov_matrix(self, filename,dpath):
        dsc = h5py.File(filename, 'r')
        ds = dsc[dpath]
        self.cov_matrix = ds[:,:]
        self.svd()

    def set_cov_matrix(self, matrix):
        self.cov_matrix = matrix
        self.svd()

    def get_cov_matrix(self):
        if cov_matrix is not None:
            return self.cov_matrix

    def svd(self):
        if self.cov_matrix is not None:
            u,s,vh = svd(self.cov_matrix)
            self.u_matrix = u
            self.s_matrix = s
            self.vh_matrix = vh
        else:
            pass#TODO

    def get_pca_vector(self, frame_number):
        if self.u_matrix is not None:
            v = np.dot(self.data[frame_number], self.u_matrix)
        else:
            pass#TODO
        return v

    def get_original_recon(self, frame_number):
        return np.dot(self.get_pca_vector(frame_number), self.u_matrix.T)

    def get_frame_recon(self, frame_number,max_mode_number, excl_list = []):
        v = self.get_pca_vector(frame_number)
        frame = np.zeros((v.shape[0], self.u_matrix.shape[1]), dtype='float32')
        ut=self.u_matrix.T
        for j in range(max_mode_number):
            if  j in excl_list:
                continue
            vt = v[:,j].reshape(v.shape[0],1)
            utt = ut[j].reshape(1,ut[j].shape[0])
            frame = frame + np.dot(vt,utt)
        return frame