#FINAL PROJECT - JOSHUA COUPE

import numpy as np
import matplotlib.pylab as plt
import urllib2
import netCDF4 as nc


def timeSeries(month,year_input,data, title):  
    """
    Given the month and year that the user inputs, this function takes in data and returns
    it in a format so that it's similar to the NAO. Also plots the NAO month by month for the year.
    """
    if year_input>1900:   #for user data
        index = year_input-1950
        place_holder = data[index].split()
        data_series = place_holder[1:]
        data_wanted = place_holder[month]
        months = np.arange(1,13,1)
 
        plt.figure()
        
        plt.xlabel('Months')
        plt.ylabel('Index')
        plt.title('NAO Index and Monthly '+title+' Standard Deviation for %s' % year_input)
        plt.plot(months,data_series)
        plt.plot(months,months*0,'b--')
        plt.plot(month,data_wanted,'rD')
       
        
        return data_series, data_wanted
        
    else:               #when no the function is called with no need to plot
        index = year_input
        place_holder= data[index].split()
        data_series = place_holder[1:]
        data_wanted = 0     #statistical analysis does not care about monthly NAO
        return data_series,data_wanted
        
    
    
    
def monthly_average(user_input_month,data,num_of_years): 
    """
    Calculates the average value (precip or temperature) for the given month from 1950-2014 
    or 2010.
    """
    start = user_input_month + 23
    i=0
    total=0
   
    while i < num_of_years:
        time = start+12*i
        total = total + data[time][laty][lony]
        i=i+1
    average = total/i
    return average
    
def deviation_from_normal(user_input_month,avg,data,num_of_years):   
    """
    Computes the value for one standard deviation away of the mean for precip or temperature
    """    
    deviation = 0
    start = user_input_month + 23
    i=0
    while i <num_of_years:
        time = start+12*i
        deviation = deviation + (data[time][laty][lony] - avg)**2
        i = i+1
    standard_deviation = np.sqrt(deviation/(i-1))
    return standard_deviation
    
def statisticalAnalysis(month, years,sd):
    """
    Computes the percentage of months that correspond as following:
    NAO+ with +temperature
    NAO- with -temperature
    NAO+ with +precipitation
    NAO- with -precipitation
    """
    if len(years) <64:     #if dealing with precipitation data
        monthly_NAO = theNAO[0:60,month-1]
        monthly_sd = sd[0:60,month-1]
        title= 'precipitation'
    else:                 #if dealing with temperature data
        monthly_NAO = theNAO[:,month-1]
        monthly_sd = sd[:,month-1]
        title = 'temperature'
    
    neg_avg_sd = 0
    pos_avg_sd = 0
    total_neg_NAO = 0
    total_pos_NAO = 0
    significance_of_neg_NAO = 0
    significance_of_pos_NAO = 0
    #goes through every NAO value and determines adds to significance if a +/-NAO month
    #corresponds with a +/- .5 sd temperature/precip deviation
    for i in range(len(monthly_NAO)):
        if monthly_NAO[i] < 0:
            neg_avg_sd = neg_avg_sd + monthly_sd[i]
            total_neg_NAO = total_neg_NAO + 1
            if monthly_sd[i] <= -.5:
                significance_of_neg_NAO = significance_of_neg_NAO + 1
        else:
            pos_avg_sd = pos_avg_sd + monthly_sd[i]
            total_pos_NAO = total_pos_NAO + 1
            if monthly_sd[i] >= .5:
                significance_of_pos_NAO = significance_of_pos_NAO + 1
            

    negative_correlation= significance_of_neg_NAO / float(total_neg_NAO)*100
    print '--------------------------------------------'
    print round(negative_correlation,2),'% of negative NAO months for',dates[month-1],'corresponds to a '+title+' decline greater than .5 standard deviation.'
    
    positive_correlation = significance_of_pos_NAO / float(total_pos_NAO)*100
    print round(positive_correlation,2),'% of positive NAO months for',dates[month-1],'corresponds to a '+title+' increase greater than .5 standard deviation.'
    print '--------------------------------------------'

    plt.plot(years,monthly_NAO)
    plt.plot(years,monthly_sd,'r--')
    
    return negative_correlation,positive_correlation
    
def average_deviation(data,num_of_years):
    """
    Creates an array of the monthly deviation from normal of temperature or precipitation that
    is formatted similar to the NAO index data
    """
    average = [0]*12
    standard_deviation = [0]*12
    month_counter = 1
    while month_counter <13:
        average[month_counter-1] = monthly_average(month_counter, data,num_of_years)
        month_counter = month_counter + 1

    month_counter=1
    while month_counter < 13:
        standard_deviation[month_counter-1] = deviation_from_normal(month_counter,average[month_counter-1],data,num_of_years)
        month_counter = month_counter + 1
    return average,standard_deviation
    
def sd_computation(average,data,standard_deviation,num_of_years):
    """
    The final part of the standard dev. computation with the end result: The standard deviation
    is returned as an array with the standard deviation of precipitation or temperature 
    for each month.
    """
    the_sd = np.zeros((num_of_years,12))
    y=0
    while y < num_of_years:
        m = 1
        while m < 13:
            time = 23 + m + y*12
            the_diff = data[time][laty][lony] - average[m-1]
            the_sd[y,m-1] = the_diff / standard_deviation[m-1]
            m = m+1
        y = y+1
    return the_sd    

def anomalyPlot(monthly_negative_correlation,monthly_positive_correlation, title):
    """
    Analysis of every month from 1950 to 2014 by percentage of how many months correlate
    with the NAO cycle.
    """
    plt.subplot(2,1,1)
    plt.title('Percentage of NAO- months corresponding to -.5 SD '+title+ ' Anomalies')
    plt.plot(months,monthly_negative_correlation[:])

    plt.subplot(2,1,2)
    plt.title('Percentage of NAO+ months corresponding to +.5 SD '+title+' Anomalies')
    plt.plot(months,monthly_positive_correlation[:])
    plt.show()
 
def plotPanelCompilation(years,sd, title):
    j=1    
    while j < 13:
        k=1
        
        while k<4:
            plt.subplot(3,1,k)
            monthly_negative_correlation[j-1], monthly_positive_correlation[j-1] = statisticalAnalysis(j,years,sd)
            if k == 1:
                plt.title('(1)'+dates[j-1]+', (2)'+dates[j]+' (3)'+dates[j+1]+' NAO(blue) with '+title+'(red)')
            j=j+1
            k = k+1

        plt.show()
   


###########################################################################
#END OF FUNCTIONS, BEGINNING MAIN PROGRAM
###########################################################################

dates = ['January','February','March','April','May','June','July','August','September','October','November','December']

#NAO Data
NAO = np.zeros((66,13))
raw_nao = urllib2.urlopen("http://www.cpc.ncep.noaa.gov/products/precip/CWlink/pna/norm.nao.monthly.b5001.current.ascii.table")
unsplit_NAO = raw_nao.read().splitlines()

print 'Given a year and date, this program will tell you the NAO index for that time as well\
 as give you the plot of NAO vs. temperature standard deviation with time.\n'

#the following prompts the user to enter a month and a year 
user_competency = False
while user_competency == False:
    print 'Enter a year after 1950 you would like NAO data for:'
    year_input = raw_input()
    year_input = int(year_input)
    if year_input < 1950 or year_input>2015:
        print 'ERROR: years must be between 1950 and 2014'
        user_competency = False
    else:
        user_competency = True
        print 'Now enter the month: (enter 01 for Jan, 02 for Feb, etc.)' 
        month = raw_input()
        month = int(month)
        if month > 12 or month<1:
            user_competency = False
            while user_competency == False:
                print 'the month must be between 1 and 12:'
                month = raw_input()
                month = int(month)
                if 1< month <13:
                    user_competency = True
                else:
                    user_competency = False

title = 'Temperature'
#retrieving the NAO over 12 months for one year
data_series, NAOvalue = timeSeries(month,year_input,unsplit_NAO,title)
print 'The NAO for ',month,'/',year_input,' is',NAOvalue,'.'
total = 0
#computing the average NAO for a year
for i in range(len(data_series)):
    total = total + float(data_series[i])
print 'The average NAO value for this year is', round(total/len(data_series),3)

################################################################
#USING AVERAGE TEMPERATURE/PRECIPITATION WITH NAO DATA
###############################################################
##Temperature
fid = nc.Dataset('/home/jlc449/FinalNAO/airhighres.nc','r')
air = fid.variables['air'][:,:,:]   
lat = fid.variables['lat'][:]  
lon = fid.variables['lon'][:] 

fid.close()
##Precipitation
fid = nc.Dataset('/home/jlc449/FinalNAO/preciphighres.nc','r')
precip = fid.variables['precip'][565:,:,:] #565 is January 1948
fid.close()

#prompt asks the user to enter a latitude and a longitude they want to analyze
print ' *****************************************************************'
print 'To compare the NAO with temperature, we recommend a latitude bordering the West Atlantic.'
print 'Boston is 42.36 N and -71.06 ; NYC is 40.71 N and - 74.0059'
user_latitude = float(raw_input('Enter a latitude between -90 and 90:'))

print'For longitude, note that west of 0 degrees is negative.'
user_longitude = float(raw_input('Enter a longitude between -180(east) and 180(west):'))

if user_longitude < 0:
    user_longitude = 360 + user_longitude
else:
    user_longitude = 360 - user_longitude


#if the latitude/longitude is invalid, default to New York City
if (abs(user_latitude) > 90.5) or (abs(user_longitude)>360.5):
    laty = 98
    lony = 571

else:
    i=0
    difference = abs(user_latitude - lat[0])
    while difference > .5:
        i = i+1
        difference = abs(user_latitude - lat[i])

    laty = i        #index to use for latitude
  
    i=0
    difference = abs(user_longitude - lon[0])
    while difference > .5:
        i = i+1
        difference = abs(user_longitude - lon[i])
        
    lony = i      #index to use for longitude

print 'NOTICE: In all multi-lined plots, blue represents NAO values and red represents standard deviation'

theNAO = np.zeros((65,12))
#for use in StatisticalAnalysis function, needs 65 by 12 array of NAO
i=0
while i < 65:
    naolist, NAOvalue = timeSeries(0,i,unsplit_NAO,' ')
    theNAO[i,:] = naolist[:]
    i = i+1
#calculating the average deviation over 65 years, then the standard deviation
temp_average,standard_deviation = average_deviation(air,65)
temp_sd = sd_computation(temp_average,air,standard_deviation,65)

#plots temperature standard deviation (red) with the NAO(blue) for months from the user's given year
months = np.arange(1,13,1)
plt.plot(months,temp_sd[year_input-1950,:],'r--')
plt.show()

years = np.arange(1950,2015,1)
print 'Some months show a stronger NAO signal than others. This section will inform you\
 as to how the month you enter correlates with temperature standard deviation.\n'
print 'Enter a month to find out how temperature relates to the NAO for that month:'
month = int(raw_input())
negative_correlation,positive_correlation = statisticalAnalysis(month,years,temp_sd)
plt.title('NAO Index and Monthly Temperature Standard Deviation for %s' % dates[month-1])
plt.show()
print '============================================================'
print 'COMPUTING STATISTICAL ANALYSIS FOR EVERY MONTH (1950 - 2014)'
print '============================================================'

monthly_negative_correlation = [0]*12
monthly_positive_correlation = [0]*12

plotPanelCompilation(years,temp_sd,'Temp')

anomalyPlot(monthly_negative_correlation,monthly_positive_correlation, 'Temp')


print 'Continue analysis with precipitation? (y/n)'
answer = str(raw_input())
if answer == 'y':
    i=0
    title = 'Precipitation'
    years = np.arange(1950,2010,1)
    theNAO = np.zeros((60,12))
    if year_input >2010:
        year_input = int(raw_input('The precipitation dataset is only through the year 2010. Please enter a different year:'))
    
    while i < 60:
        naolist, NAOvalue = timeSeries(0,i,unsplit_NAO,' ')
        theNAO[i,:] = naolist[:]
        i = i+1
        
    plt.figure()
        
    plt.xlabel('Months')
    plt.ylabel('Index')
    plt.title('NAO Index and Monthly '+title+' Standard Deviation for %s' % year_input)
    plt.plot(months,theNAO[year_input-1950,:])
    plt.plot(months,months*0,'b--')   
    precip_average,standard_deviation = average_deviation(precip,60)

    
    precip_sd = sd_computation(precip_average,precip,standard_deviation,60)
    months = np.arange(1,13,1)
    plt.plot(months,precip_sd[year_input-1950,:],'r--')
    plt.show()
    
    print 'Enter a month to find out how precipitation correlates with the NAO through the years(1950-2010): '
    month = int(raw_input())
    
    monthly_negative_correlation,monthly_positive_correlation = statisticalAnalysis(month,years,precip_sd)
    plt.title('NAO Index and Monthly Precipitation Standard Deviation for %s' % dates[month-1])
    plt.show()
    #set these equal to zero because plotPanelCompilation refills them
    monthly_negative_correlation = [0]*12
    monthly_positive_correlation = [0]*12
    
    plotPanelCompilation(years,precip_sd,'Precip')   
    
    anomalyPlot(monthly_negative_correlation,monthly_positive_correlation, 'Precip')
else:
    print 'Thank you for letting me take you on a climatological journey. Long live Keith.'