#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 15 17:05:37 2017

@author: Nikolaus Zen Prusinski
Copyright (c) 2017 NZP Productions
"""
# import needed extensions
from numpy import *
import matplotlib.pyplot as plt # plotting package
import pyfits
import urllib, cStringIO
import os
from astropy.io.fits import getval
from scipy import stats
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as smf
import pandas as pd

def LinearityTest(): #Collect Signal Intensity from the images and create graph with Linear Regression
    flatField0125 = 'http://www.nzp.guru/ccd/Flat Fields/0.125s Flat Field.fit'
    flatField025 = 'http://www.nzp.guru/ccd/Flat Fields/0.25s Flat Field.fit'
    flatField05 = 'http://www.nzp.guru/ccd/Flat Fields/0.5s Flat Field.fit'
    flatField075 = 'http://www.nzp.guru/ccd/Flat Fields/0.75s Flat Field.fit'
    flatField10 = 'http://www.nzp.guru/ccd/Flat Fields/1s Flat Field.fit'
    flatField15 = 'http://www.nzp.guru/ccd/Flat Fields/1.5s Flat Field.fit'
    flatField20 = 'http://www.nzp.guru/ccd/Flat Fields/2s Flat Field.fit'
    flatField25 = 'http://www.nzp.guru/ccd/Flat Fields/2.5s Flat Field.fit'
    
    flatFieldLow1 = 'http://www.nzp.guru/ccd/Low/0.125s.fit'
    flatFieldLow2 = 'http://www.nzp.guru/ccd/Low/0.25s.fit'
    
    flatFieldMed1 = 'http://www.nzp.guru/ccd/Medium/0.125s.fit'
    flatFieldMed2 = 'http://www.nzp.guru/ccd/Medium/0.25s.fit'
    flatFieldMed3 = 'http://www.nzp.guru/ccd/Medium/0.5s.fit'
    
    
    time = [0.125, 0.25, 0.5, 0.75, 1, 1.5, 2, 2.5]; #Establish list of exposure times
    timeLow = [0.125, 0.25];
    timeMed = [0.125, 0.25, 0.5];
    
    
    flatField_list = [flatField0125, flatField025, flatField05, flatField075, flatField10, flatField15, flatField20, flatField25]; # List of flatFields
    flatField_list_Low = [flatFieldLow1, flatFieldLow2];
    flatField_list_Med = [flatFieldMed1, flatFieldMed2, flatFieldMed3];
    
    counts = [] #empty list
    countsLow = []
    countsMed = []


    for flatField in flatField_list:
        val = getval(cStringIO.StringIO(urllib.urlopen(flatField).read()), 'CWHITE')
        print val,"ADU"
        counts.append(val)
    print counts
    
    for flatField in flatField_list_Low:
        val = getval(cStringIO.StringIO(urllib.urlopen(flatField).read()), 'CWHITE')
        print val,"ADU"
        countsLow.append(val)
    print countsLow

    for flatField in flatField_list_Med:
        val = getval(cStringIO.StringIO(urllib.urlopen(flatField).read()), 'CWHITE')
        print val,"ADU"
        countsMed.append(val)
    print countsMed
    
    #for each image - find CWHITE ADU value, and add it to counts list
    
    ADSatLevel = 68918.918 # Full Well Capacity/Gain (theoretcial) -- 25,500 e-/0.37 e-/ADU = 68918.918 ADU
    
    slope, intercept, r_value, prob2, see = stats.mstats.linregress([time],[counts]) #gather linear regression data
    slopeLow, interceptLow, r_valueLow, prob2Low, seeLow = stats.mstats.linregress([timeLow],[countsLow]) #gather linear regression data    
    slopeMed, interceptMed, r_valueMed, prob2Med, seeMed = stats.mstats.linregress([timeMed],[countsMed]) #gather linear regression data    


#    y = np.array([counts])
#    X = np.array([time])
#    
#    time = sm.add_constant(X)
    #df = pd.DataFrame({"counts": [counts], "time": [time]})
  #  results = smf.OLS(counts, time).fit()
  
    df = pd.DataFrame({'Y':counts, 'X':time})
  
    results = smf.ols('Y ~ 1 + X',data=df).fit()
    print results.summary()
    
    
    x = np.linspace(0,3.5,100) #establish space for plotting the line
    
    plt.plot([time], [counts], 'C1o') # plot points
    plt.plot(x, slope*x + intercept, 'C1', label = 'Linear Regression High Resolution') # plot regression line

    plt.plot([timeMed], [countsMed], 'C2o') # plot points
    plt.plot(x, slopeMed*x + interceptMed, 'C2', label = 'Linear Regression Medium Resolution') # plot regression line
    
    plt.plot([timeLow], [countsLow], 'C3o') # plot points
    plt.plot(x, slopeLow*x + interceptLow, 'C3', label = 'Linear Regression Low Resolution') # plot regression line
    
    
#    plt.plot(x, 23923.962*x + 1061.854, c = 'm', label = 'Regression without 3s Point') # plot regression line without 3s Point
    plt.axhline(y=ADSatLevel, xmin = 0,xmax = 3.5, c ='C4', dashes = [6, 4], label = 'A/D Saturation Level') #plot saturation level
    plt.axis([0,3.5,0,80000])
    plt.xlabel('Exposure Time (s)')
    plt.ylabel('Signal Intensity (ADU)')
    plt.title('SBIG STT-8300 Linearity Test')
    plt.legend()
    plt.show()
    
    #print data
    
    print "Slope = ",slope
    print "Y-Intercept = ",intercept
    print "Coefficient of Determination (R^2) = ",(r_value)**2
    print "Correlation Coefficient (R) = ",r_value


def mpv(fileName, low, high, returnMean = True, returnDev = False): #find mean pixel value of an image
    # read the file
    h = pyfits.open(fileName)
    
    # copy the image data into a numpy array
    img = h[0].data
    
    plt.ion() # do plots in interactive mode    
    colmap = plt.get_cmap('gray') # load gray colormap to match color
    
    # plot the image
    plt.figure(1)
    plt.imshow(img, cmap=colmap) # plot image using gray colorbar
    
    # img is a 2-d array, need to change to 1-d to make a histogram
    #imgh = 1.0*img # make a copy
    nx, ny = img.shape # find the size of the array
    imgh = reshape(img, nx*ny) # change the shape to be 1d
    
    # print statistics about the image
    print 'Image minimum = ', min(imgh)
    print 'Image maximum = ', max(imgh)
    print 'Image mean = ', mean(imgh)
    print 'Image standard deviation = ', std(imgh)
    
    # plot a histogram of the image values
    plt.figure(2)
    plt.hist(imgh, bins = 100, histtype='stepfilled') #This should be a Gaussian Distribution
    #Flats have a wider area under the curve, while bias images are narrow
    
    plt.show() # display the plots
    
    plow = low #establish ranges to cut the images - so that the data is more accurate
    phi = high
    q = where((imgh >= plow) & (imgh <= phi))
    imghcut = imgh[q] #new image is within the range
    
    print 'Image minimum = ', min(imghcut) #collect stats about the image
    print 'Image maximum = ', max(imghcut)
    print 'Image mean = ', mean(imghcut)
    print 'Image standard deviation = ', std(imghcut)
        
    if (returnMean == True): #choice if we want mean or std. dev.
        return mean(imghcut)
    if (returnDev == True):
        return std(imghcut);
    

def stdDev(fileName1, fileName2, newFileName, lowBound, highBound): #find std deviation of two images when subtracted
    # This is a similar procedure to above - we are loading the images into arrays and displaying them 
    #read the files
    h1 = pyfits.open(fileName1)
    h2 = pyfits.open(fileName2)
    
    # copy the image data into a numpy array
    pic1 = h1[0].data
    pic2 = h2[0].data
    
    plt.ion() # do plots in interactive mode
    colmap = plt.get_cmap('gray') # load gray colormap
    
    # plot the first image
    plt.figure(1)  
    plt.imshow(pic1, cmap=colmap) # plot image using gray colorbar
    
    # plot the second image in another window
    plt.figure(2)  
    plt.imshow(pic2, cmap=colmap) # plot image using gray colorbar
    
    # find the difference in the images using subtraction
    diff = pic2 - pic1
    
    # plot the difference image
    plt.figure(3)
    plt.imshow(diff, cmap=colmap) # plot image using gray colorbar
    
    plt.show() # display the images
    
    pyfits.writeto(newFileName, diff, clobber = True) #create a new file to gather data from
    deviation = mpv(newFileName, lowBound, highBound, returnMean = False, returnDev = True) #use the mean pixel value fx to return standard deviation of the image
    os.remove(newFileName) #remove the file from directory
    return deviation;



#urls of the images
flatField1 = 'http://www.nzp.guru/ccd/Flat Fields/2.5s Flat Field.fit'
flatField2 = 'http://www.nzp.guru/ccd/Flat Fields/2.5s Flat Field (2).fit'
Bias1 = 'http://www.nzp.guru/ccd/Bias Images/Bias 1.fit'
Bias2 = 'http://www.nzp.guru/ccd/Bias Images/Bias 2.fit'   
    
#gather images from server to be used and calculate mean pixel value
f1 = mpv(cStringIO.StringIO(urllib.urlopen(flatField1).read()), 0, 70000, returnMean = True, returnDev = False)
f2 = mpv(cStringIO.StringIO(urllib.urlopen(flatField2).read()), 0, 70000, returnMean = True, returnDev = False)
b1 = mpv(cStringIO.StringIO(urllib.urlopen(Bias1).read()), 0, 70000, returnMean = True, returnDev = False)
b2 = mpv(cStringIO.StringIO(urllib.urlopen(Bias2).read()), 0, 70000, returnMean = True, returnDev = False)

#gather images and calculate std. Deviation
sigmaF = stdDev(cStringIO.StringIO(urllib.urlopen(flatField1).read()), cStringIO.StringIO(urllib.urlopen(flatField2).read()), '2.5s Flat Diff.fit', 0, 10000)
sigmaB = stdDev(cStringIO.StringIO(urllib.urlopen(Bias1).read()), cStringIO.StringIO(urllib.urlopen(Bias2).read()), 'Bias Diff.fit', 0, 14000)

#Use equations on pgs. 72-73 of Handbook of CCD Astronomy to calculate gain, readnoise, and width of distribution
gain = ((f1+f2) - (b1+b2))/(((sigmaF)**2)-((sigmaB)**2))
readNoise = (gain * sigmaB)/sqrt(2)
sigmaADU = (sqrt(f1*gain))/gain
percentError = abs(((gain - 0.37)/0.37)*100)

LinearityTest() #calculate linearity of the data and print results

print "Gain = ",gain," e-/ADU"
print "Percent Error on Gain = ",percentError,"%"
print "Read Noise = ",readNoise," e- rms"
print "Distribution Width = ", sigmaADU