import matplotlib.pylab as plt
import numpy as np
from numpy.fft import rfft, irfft, rfftfreq, rfft2, irfft2, fft2, ifft2
import astropy.io.fits as pyfits
import scipy.linalg

# A function to calculate the elements of a covariance matrix
def cov(df, dt, nch, ntime):
    if dt>ntime/2:
        dt-=ntime
    if df>nch/2:
        df-=nch
    if dt<=-ntime/2:
        dt+=ntime
    if df<=-nch/2:
        df+=nch
    if dt==0:
        dt=1
    if df==0:
        df=1
    return (abs(dt)/8)**-1*(abs(df)/4)**-1
    
# fill the content of a covariance matrix
def fill_cov_matrix(nch=16, ntime=32):
    result=np.zeros([nch*ntime, nch*ntime])
    for t1 in range(ntime):
        for f1 in range(nch):
            for t2 in range(ntime):
                for f2 in range(nch):
                    i=f1*ntime+t1
                    j=f2*ntime+t2
                    result[i,j]=cov(f2-f1, t2-t1, nch, ntime)
    return result

# generating the equivlant deconvolution kernel
def fill_cov_kernel(nch=16, ntime=32):
    result=np.zeros([ntime, nch])
    for t1 in range(ntime):
        for f1 in range(nch):
            result[t1, f1]=cov(f1, t1, nch, ntime)
    return result

# flattern a 2-D array to 1-D
def flattern(m):
    return m.reshape(m.size, order='F')

# unflattern a 1-D array to 2-D
def unflattern(m, nch):
    ntime=m.size//nch
    return m.reshape((ntime, nch), order='F')

ntime=32
nch=16
m=fill_cov_matrix(nch, ntime)
k=fill_cov_kernel(nch, ntime)
plt.imshow(m)
plt.colorbar()

plt.figure(figsize=(5,7))
plt.imshow(k)
plt.colorbar()

data2d=np.random.normal(size=(ntime, nch))
plt.imshow(data2d)

data1d=flattern(data2d)
convolved1d=unflattern(scipy.linalg.solve(m, data1d), nch)
plt.imshow(convolved1d)

convolved2d=irfft2(rfft2(data2d)/rfft2(k))
plt.imshow(convolved2d)

plt.imshow(convolved1d-convolved2d)
plt.colorbar()