#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from astropy.io import fits
import sys

infile = sys.argv[1]


with fits.open(infile) as fitsfile:
    #print(fitsfile[1].header['CDSOUT'])
    
    fig,axs = plt.subplots(4,1)
    print(axs)


    hdunum=0
    for hdu in fitsfile:
        print(hdu)
        if isinstance(hdu, fits.ImageHDU) or isinstance(hdu, fits.CompImageHDU):
            ncol = int(hdu.header['NCOL'])
            nsamp = int(hdu.header['NSAMP'])
            if nsamp==1:
                nsamp = 100
            PERIOD_X = 1
            PERIOD_Y = 1
            ax = axs[hdunum]
            window_x = nsamp*PERIOD_X
            usable_x = (hdu.data.shape[1]//window_x)*window_x

            if hdu.data.shape[0]>1:
                usable_y = ((hdu.data.shape[0]-1)//PERIOD_Y)*PERIOD_Y
                trimmed = hdu.data[1:usable_y+1,:usable_x]
            else:
                trimmed = hdu.data[:,:usable_x]
            print(trimmed.shape)
            for dy in range(PERIOD_Y):
                data = trimmed[dy::PERIOD_Y,:]
                wrapped = data.reshape(-1,window_x)
                ax.plot(wrapped.mean(axis=0), label="dy={0}".format(dy))
                #ax.plot(wrapped.mean(axis=0).reshape(-1,5).mean(axis=1), label="dy={0}".format(dy))




            ax.legend(loc='upper right')

            hdunum+=1

##ax.set_yscale('log')
#ax.set_xscale('log')
#if spectrogram_hdu>=0:
#    if ds!=1.0:
#        axs[1].set_xlabel("time [s]")
#        axs[1].set_ylabel("frequency [Hz]")
#    else:
#        axs[1].set_xlabel("time [sample]")
#        axs[1].set_ylabel("frequency [1/sample]")
#
#ax.legend(loc='upper right')
#ax.grid()
#ax.set_title(infile)
#if ds!=1.0:
#    ax.set_xlabel("frequency [Hz]")
#    ax.set_ylabel("PSD [ADU^2/Hz]")
#else:
#    ax.set_xlabel("frequency [1/sample]")
#    ax.set_ylabel("PSD [ADU^2*sample]")
plt.show()