import pyfits
import pandas
import statsmodels
import mdp
import numpy as np
# pylab's already imported...
import agpy # need PCA tools
import aplpy

matplotlib.rcParams['figure.figsize'] = (12,7)
matplotlib.rcParams['savefig.dpi'] = 120

D = pyfits.getdata('/Users/adam/work/fastphot/TEST.fits')
RF = pyfits.getdata('/Users/adam/work/fastphot/Robberto_Found.fits')

D = D.astype('float')
DF = pandas.DataFrame(D)
cov = DF.cov()

imshow(cov)

imshow(DF)

# above shows that the first couple observations were wrong - maybe focusing?  Maybe off-source?  Probably off source, and just by chance overlapped enough to get included
D = D[:,30:] # exclude first 30 stars
# first observing block - correlations are more clear here...
firstchunk = D[18:516,:] # first dim is time

# filter out sections with too few data points
OKfirst = np.isnan(firstchunk).sum(axis=0) < 20
firstchunk_filtered = firstchunk[:,OKfirst]
FCF = pandas.DataFrame(firstchunk_filtered)
FCFcov = FCF.cov()
# filter the RA/Dec data (etc) als0
RFfirst = RF[30:][OKfirst]

imshow(log10(FCFcov))

# PCA subtraction - remove 0'th component
efuncarr,covmat,evals,evects = agpy.PCA_tools.efuncs(np.nan_to_num(firstchunk_filtered),return_others=True)
efuncarr_blanked = efuncarr.copy()
efuncarr_blanked[:,0] = 0
firstchunk_subd = np.inner(efuncarr_blanked,evects)

figure(); imshow(log10(firstchunk_subd.real))
# there is a clear difference between first 300ish and the rest!
figure(); imshow(evects.real);
figure(); semilogy(evals.real)

plot(efuncarr[:,:2])

# find out which are most correlated with which component (presumably first 2
# account for most of what we see above, in the PCA-sub'd waterfall plot)
plot(evects[:,0]) 

# select based on correlation with evect 1
set1 = evects[:,0] > -0.05
set1.shape

# redo above with set 1
efuncarr2,covmat2,evals2,evects2 = agpy.PCA_tools.efuncs(np.nan_to_num(firstchunk_filtered[:,set1]),return_others=True)
efuncarr_blanked2 = efuncarr2.copy()
efuncarr_blanked2[:,0] = 0
firstchunk_set1_subd = np.inner(efuncarr_blanked2,evects2)
print firstchunk_set1_subd.shape,firstchunk_filtered.shape

figure(figsize=[12,7],dpi=300)
imshow(log10(firstchunk_set1_subd.real+1000),aspect=1/3.,vmin=2.5)

plot(firstchunk_set1_subd[:,100])

subplot(311); plot(efuncarr2[:,0])
subplot(312); plot(efuncarr2[:,1])
subplot(313); 
plot(firstchunk_set1_subd[:,47]) # this one probably belongs in the other set

plot(firstchunk_set1_subd[:,47:51])

# redo above with anti-set 1
efuncarr3,covmat3,evals3,evects3 = agpy.PCA_tools.efuncs(np.nan_to_num(firstchunk_filtered[:,True-set1]),return_others=True)
efuncarr_blanked3 = efuncarr3.copy()
efuncarr_blanked3[:,0] = 0
firstchunk_set2_subd = np.inner(efuncarr_blanked3,evects3)
print firstchunk_set2_subd.shape,firstchunk_filtered.shape

imshow(log10(firstchunk_set2_subd.real+1000),aspect=1/5.) # add 1000 so it shows up...
figure()
plot(firstchunk_set2_subd[:,5])

# the subd image shows another subset/breakdown...
plot(efuncarr2[:,:5]*-1)
figure()
plot(efuncarr3[:,:5])

import warnings
warnings.filterwarnings('ignore')
import pywcs
FF = aplpy.FITSFigure('/Users/adam/work/orion/orion_finder_mosaic.fits')
FF.show_grayscale(pmin=5,pmax=95)
FF.show_markers(RFfirst[set1]['RAJ2000'],RFfirst[set1]['DEJ2000'])
FF.show_markers(RFfirst[True-set1]['RAJ2000'],RFfirst[True-set1]['DEJ2000'],edgecolor='b')

# What if we just get rid of ALL of the correlated components?  
# PCA subtraction - remove first 200 components
efuncarr4,covmat4,evals4,evects4 = agpy.PCA_tools.efuncs(np.nan_to_num(firstchunk_filtered),return_others=True)
efuncarr_blanked4 = efuncarr4.copy()
efuncarr_blanked4[:,:200] = 0
firstchunk_subd_huge = np.inner(efuncarr_blanked4,evects4)

LS = np.log10(firstchunk_subd_huge.real+1000)  #logscale, make all above zero
imshow(LS,vmin=np.percentile(LS,1),vmax=np.percentile(LS,99))

for ii in xrange(6):
    subplot(6,1,ii+1)
    plot(firstchunk_subd_huge[:,ii])

for ii in xrange(6):
    subplot(6,1,ii+1)
    plot(firstchunk_subd_huge[:,105+ii])

FTFC = np.fft.fft(firstchunk_subd_huge,axis=0)
for ii in xrange(6):
    subplot(6,1,ii+1)
    plot(np.fft.fftshift(FTFC[:,ii])[:FTFC.shape[0]/2])

for ii in xrange(6):
    subplot(6,1,ii+1)
    plot(np.fft.fftshift(FTFC[:,105+ii])[:FTFC.shape[0]/2])

for ii in xrange(6):
    subplot(6,1,ii+1)
    plot(firstchunk_subd_huge[:,105+ii])
    plot(np.fft.fft(firstchunk_subd_huge[:,105+ii]))

LSFT = np.log10(FTFC.real+1000)  #logscale, make all above zero
imshow(LSFT,vmin=np.percentile(LSFT,1),vmax=np.percentile(LSFT,99))

errorbar(arange(LSFT.shape[1]),LSFT.max(axis=0),yerr=LSFT.std(axis=0),linestyle='none')

plot(LSFT.max(axis=0)/LSFT.std(axis=0))
possibleperiod = LSFT.max(axis=0)/LSFT.std(axis=0) > 5e5

for ii in xrange(6):
    subplot(6,1,ii+1)
    plot(np.fft.fftshift(FTFC[:,possibleperiod][:,ii])[:FTFC.shape[0]/2])

# second observing block - correlations are more clear here...
secondchunk = D[516:1020,:] # second dim is time
# filter out sections with too few data points
OKsecond = np.isnan(secondchunk).sum(axis=0) < 10
secondchunk_filtered = secondchunk[:,OKsecond]
SCF = pandas.DataFrame(secondchunk_filtered)
SCFcov = SCF.cov()
# filter the RA/Dec data (etc) als0
RFsecond = RF[30:][OKsecond]

imshow(SCFcov)

# redo above with anti-set 1
efuncarr5,covmat5,evals5,evects5 = agpy.PCA_tools.efuncs(np.nan_to_num(secondchunk_filtered),return_others=True)
efuncarr_blanked5 = efuncarr5.copy()
efuncarr_blanked5[:,0] = 0
secondchunk_subd = np.inner(efuncarr_blanked5,evects5)
print secondchunk_subd.shape,secondchunk_filtered.shape

semilogy(evals5)

imshow(log10(secondchunk_subd.real+1000),aspect=1/3.)

efuncarr_blanked5[:,:7] = 0
secondchunk_subd = np.inner(efuncarr_blanked5,evects5)
print secondchunk_subd.shape,secondchunk_filtered.shape
LS = (log10(secondchunk_subd.real+1000))
imshow(LS,vmin=np.percentile(LS,1),vmax=np.percentile(LS,99),aspect=1/3.)

print firstchunk_filtered.shape,efuncarr.shape,FCF.shape,FCF.mean(axis=0).shape


atmonorm = efuncarr[:,0]/efuncarr[:,0].mean()
firstchunk_normed = firstchunk_filtered / array(FCF.mean(axis=0)) / atmonorm[:,newaxis]

nplots=10
for ii in xrange(nplots):
    subplot(nplots,1,ii+1)
    plot(firstchunk_normed[:,ii])

imshow(firstchunk_normed.real,vmin=0.5,vmax=2,aspect=0.5)

spiky = where(firstchunk_normed.max(axis=0) > 2.5)[0]
dippy = where(firstchunk_normed.min(axis=0) < 0.1)[0]
print dippy

#help(scatter)

# lots of overplots
import glob
fitsfiles = glob.glob('/Users/adam/observations/tcam/UT121105/proc-t0[0-9]50.corr.fits')
FFlist = []
for fn in fitsfiles:
    FF = aplpy.FITSFigure(fn)
    FF.show_grayscale(pmin=1,pmax=99)
    FF.show_markers(RFfirst[set1]['RAJ2000'],RFfirst[set1]['DEJ2000'])
    FF.show_markers(RFfirst[True-set1]['RAJ2000'],RFfirst[True-set1]['DEJ2000'],edgecolor='b')
    FF.show_markers(RFfirst[spiky]['RAJ2000'],RFfirst[spiky]['DEJ2000'],edgecolor=(0,1,0),marker='x')
    FF.show_markers(RFfirst[dippy]['RAJ2000'],RFfirst[dippy]['DEJ2000'],edgecolor=(1,0,1),marker='x')
    FFlist.append(FF)