import numpy as np 
import matplotlib.pyplot as plt
%matplotlib inline 
import scipy.io
import glob
import os
import time
from astropy.io import fits
from time import localtime,strftime
tstart = strftime("%a, %d %b %Y %H:%M:%S", localtime())
print "Starting Time:", tstart
tstart = time.time()

imnames = sorted(glob.glob('Fedop/*Fedop.fits'))# sort list of images. Fedop: Doppler images taken with an iron (Fe) filter.

N = len(imnames)

imlist = np.zeros((N,1024,1024))

for i in range(N):
    imlist[i] = fits.open(imnames[i])[0].data # Read in data, and save as Python list.


plt.imshow(imlist[100], cmap='cubehelix') # cmap: Specify color map to make for easier visual understanding.
plt.show()

plt.plot(imlist[:,500,500]) #Plot slice through Python list corresponding to one spatial location in all images.
plt.xlabel('Image (45s cadence)')
plt.ylabel('Signal')
plt.show()

#Create a power spectrum.

#First, take n-dimensional Fourier Transform of image data.
fftspec = np.fft.fftn(imlist, s=(512,1024,1024)) # FFT
fftspec = np.fft.fftshift(fftspec) # Shift FFT such that lowest frequencies are at the left of the array, highest at the right.

#Power Spectrum = FFT * conj(FFT)
powspec = fftspec * np.conj(fftspec)
t = time.time() - tstart
print "FFT runtime:", t, "s"

#Analysis: run separately from FFT
tstart = time.time()

#combine positive and negative frequencies
newpowspec = np.zeros((256,1024,1024), dtype=complex) #Define new array of complex datatype.
for i in range(len(powspec)/2):
    newpowspec[i] = powspec[i] + powspec[-(i+1)] # combine powspec halves
    
newpowspec = newpowspec[::-1] #flip over frequency axis so frequencies run from lowest to highest.

plt.imshow(np.log(abs(newpowspec[90]))) #We plot the logarithm of the power to make the features more visible.
plt.show()

az = np.zeros((256,1024),dtype=complex)
#compute azimuthal avg
for i in range(1024):
    for j in range(256): # loop over all index locations in kaz
        kx = np.array(newpowspec[j,i,:]) # relevant slice along x axis
        ky = np.array(newpowspec[j,:,i]) # slice along y axis
        az[j,i] = np.sqrt(np.mean(kx)**2 + np.mean(ky)**2) #azimuthal average = x and y averages added in quadrature.

plt.imshow(np.log(abs(az))) #again, plot in log space to make all features more visible
plt.show()

newaz = np.zeros((256,512), dtype=complex)
for i in range(len(az[0,:])/2):
    for j in range(256):
        newaz[j,i] = az[j,i] + az[j,-(i+1)] #combine positive and negative spatial values for each temporal frequency. 
newaz[:,:] = newaz[:,::-1] # flip over spatial frequency axis.

Nt = 512 #steps in temporal direction
Nx = 1024 # steps in spatial x direction
Ny = 1024 # steps in spatial y direction
deltat = 45 #s
deltax = 1.41 #Mm/pixel. Mm: megameter = 1 million meters. 
deltay= 1.41 #Mm/pixel
dt = round(float(1. / (Nt*deltat)),6) # step size in temporal direction, rounded to the nearest hundredth
dnux = round(float(2*np.pi / (Nx*deltax)),6) #step size in spatial direction. multiply by 2pi factor due to 
                                             # wavenumber definition: k=2pi/lambda, lambda=Nx*deltax. round to make plot labels legible.
dnuy = round(float(2*np.pi / (Ny*deltay)),6)

print round(dt,6), dnux, dnuy
dnu = np.sqrt(dnux**2+dnuy**2) #adding x and y in quadrature
dnu = round(dnu,4)
print dnu

plt.imshow(np.log(abs(newaz[5:150,5:200])), cmap='cubehelix', interpolation='none') # trim shortest and longest frequencies, 
                                                                           # and don't interpolate pixels for easier viewing.
plt.ylabel('Frequency (mHz)')
ylocs = [0,25,55,85,115] # locations of axis labels, in pixels
dt = dt *1000 #milliHertz (mHz) instead of Hertz. 1 Hertz = 1 oscillation per second.
yvals = [0,30*dt,60*dt,90*dt,120*dt] # values of axis labels
plt.yticks(ylocs,yvals) #places axis labels equal to yvals at ylocs pixel locations in plot.
plt.xlabel('Wavenumber (Mm^-1)')
xlocs = [0,45,95,145,195]
xvals = [0,50*dnu, 100*dnu, 150*dnu, 200*dnu]
plt.xticks(xlocs,xvals)
plt.show()

plt.plot(np.log(abs(az[:,0])))
plt.plot(np.log(abs(az[:,10])))
plt.plot(np.log(abs(az[:,100]))) # these correspond to different wavenumbers: waves of different spatial sizes.
plt.legend(('az0','az10', 'az100'))
plt.show()
print "Max value and location:",np.max(az), "at", np.unravel_index(np.argmax(az), np.shape(az))