import os
from datetime import datetime as dt
from datetime import timedelta as delta
import logging
import pandas as pd
from netCDF4 import MFDataset, MFTime
import numpy as np
import matplotlib.pyplot as plt
import math

CHOICES = ['air_temp', 'dewpoint',
   'rh', 'pressure', 'wind_speed', 'wind_dir', 'accum_precip', 'solar_flux']

TITLES = {'air_temp': 'Temperature and Dewpoint(°C)',
         'rh': 'Relative Humidity(%)', 'pressure': 'Pressure(hpa)', 
         'wind_speed': 'Wind Speed(m/s)', 'wind_dir': 'Wind Direction(°)',
         'accum_precip': 'Accumulated Precipitation Since 00Z (mm)', 
         'solar_flux': 'Solar Flux(W/m^2)'}

IND_TITLES = {'air_temp': 'Temperature and Dewpoint',
         'rh': 'Relative Humidity', 'pressure': 'Pressure', 
         'wind_speed': 'Wind Speed', 'wind_dir': 'Wind Direction',
         'accum_precip': 'Accumulated Precipitation Since 00Z', 
         'solar_flux': 'Solar Flux'}             

# The purpose of this method is to take a string in the format
# YYYY-MM-DDTHH:MM:SS and convert that to a datetime object
# or a string in the format YY-MM-DD and convert that to a datetime object
# used in coordination with argparse -s and -e params
# @param datetime string
# @return datetime object

def _dt_convert(datetime_str):
    #parse datetime string, return datetime object
    try:
        return dt.strptime(datetime_str, '%Y-%m-%dT%H:%M:%S')
   
    except:
        return dt.strptime(datetime_str, '%Y-%m-%d')

# The purpose of this method is to take a list of level_b1 filepaths
# and convert those into a pandas frame with only good data
# @param input_files - list of level_b1 filepaths
# @return frame - pandas dataframe with only good data

def get_data(input_files):
    files = MFDataset(input_files)

    #dictionary with var_name: good data pairs
    dataDict = {}

    #get the data from the files
    for name in CHOICES:
        data = files.variables[name][:]

        qc_data = files.variables['qc_' + name][:]
        
        dataDict[name] = data
        dataDict['qc_' + name] = qc_data

    #for some reason, time has floating point precision issues
    #so I'm using base_time and offsets instead
    #might need to fix later
    base_time = files.variables['base_time'][:]

    #convert base_time epoch format into date_time object
    base_time_obj = dt(1970,1,1) + delta(seconds=int(base_time))
    
    #get offsets from files
    off_sets = files.variables['time_offset']

    #convert that into offsets from the first file's base_time
    off_sets = MFTime(off_sets)[:]  

    #for each offset, convert that into a datetime object
    stamps = [base_time_obj + delta(seconds=int(s)) for s in off_sets]

    #append all stamps into the data frame
    dataDict['stamps'] = stamps

    return pd.DataFrame(dataDict).set_index(['stamps'])

# The purpose of this method is to use the qc mask
# to only get the good data
# @param qc_data - np array with nan for good values and various set bits for others
# @param stamps - all stamps corresponding to data
# @param data - data
# @return good data and their corresponding stamps   

def get_good_data(qc_data, stamps, data):
    #all good data and stamps associated with them
    good_data = []
    good_stamps = []

    #get good indices
    good_indices = np.where(np.isnan(qc_data))[0]

    for idx in good_indices:
        #get good_data
        good_data.append(data[idx])
        good_stamps.append(stamps[idx])

    return [good_data, good_stamps]

# The purpose of this method is to determines all the 12:00:00 days
# within a start and end date
# @param cur_dt - start time
# @param dates - end time
# @return dates in lists

def get_dates_in_range(cur_dt, end):
    curr = cur_dt
    while curr <= end:
        yield curr
        curr += delta(days=1)

# The purpose of this method is to get the min and max of 
# an array of stamps and determines all the 12:00:00 days
# in that array
# @param dewpoint stamps - time stamps for valid dewpoint data. Only specified if
# air temp and dew point need to be ploted in the same plot
# @param dates - time stamps for data
# @return datesInRange - all days with 12:00:00
def find_half_days(dewpoint_stamps, dates):
    date_max = np.amax(dates)
    date_min = np.amin(dates)

    if(dewpoint_stamps):
        dew_max = np.amax(dewpoint_stamps)
        dew_min = np.amin(dewpoint_stamps)

        date_max = max(date_max, dew_max)
        date_min = min(date_min, dew_min)    

    cur_dt = date_min.replace(hour=12, minute=0, second=0)
    
    if cur_dt < date_min:
        cur_dt += delta(days=1)

    datesInRange = list(get_dates_in_range(cur_dt, date_max))    

    return datesInRange           

# The purpose of this method is to create a thumbnail plot
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param dewpoint_stamps - timestamps for valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
# @param dewpoint_data - valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
#@param output - output filename pattern
def thumbnail_plot(dates, data, ymin, ymax,
    dewpoint_stamps, dewpoint_data, output, wind_dir):

    if(dewpoint_stamps):
        plt.plot(dates, data, 'r', dewpoint_stamps, dewpoint_data, 'b')
        plt.axis([min(dates), max(dates), ymin, ymax])
        #plt.plot(dewpoint_stamps, dewpoint_data, 'b')

    elif(wind_dir):
        plt.plot(dates, data, 'ko', markersize=0.1)
        plt.axis([min(dates), max(dates), 0, 360])

    else:
        plt.plot(dates, data, 'k')
        plt.axis([min(dates), max(dates), ymin, ymax])

    #scale
    fig = plt.gcf()
    dpi = fig.get_dpi()
    fig.set_size_inches(80/float(dpi), 80/float(dpi))

    #create axes
    axes = plt.gca()
    axes.axes.get_xaxis().set_ticklabels([])
    axes.axes.get_xaxis().set_tick_params(length=7)
    axes.axes.get_yaxis().set_ticklabels([])
    axes.axes.get_yaxis().set_visible(False)
    axes.spines['right'].set_visible(False)
    axes.spines['bottom'].set_visible(False)

    #got the basic shape of the plot
    #still need to draw a line in the middle
    datesInRange = find_half_days(dewpoint_stamps, dates)
    
    for date in datesInRange:
        plt.axvline(x=date, color='k')       

# The purpose of this method is to create a thumbnail
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param create_air_dew_plot - boolean specifying if air temp and dewpoint
# @param output - output filename pattern 
# should be in same plot
def create_thumbnail(frame, ymin, ymax, create_air_dew_plot, output):
    #see if we're going to need subplots
    #subplots needed when len is greater than 2 for one variable
    #or greater than 4 for two variables

    need_subplots = len(list(frame.columns.values)) > 2 and not create_air_dew_plot
    need_subplots = need_subplots or len(list(frame.columns.values)) > 4 and create_air_dew_plot

    if need_subplots:
        plots_created = 1

        if create_air_dew_plot:
            numberPlots = (len(list(frame.columns.values)) - 2) / 2
            plotNumber = numberPlots * 100 + 10

        else:
            numberPlots = len(list(frame.columns.values)) / 2
            plotNumber = numberPlots * 100 + 10    

    stamps = list(frame.index)

    #get dewpoint data
    if create_air_dew_plot:
        all_data = frame['dewpoint']
        qc_data = frame['qc_dewpoint']

        good_list = get_good_data(qc_data, stamps, all_data)
        dewpoint_data = good_list[0]
        dewpoint_stamps= good_list[1]    

    for name in CHOICES:
        if name == 'dewpoint' and create_air_dew_plot:
            continue

        if name not in list(frame.columns.values):
            continue        

        elif need_subplots:
            plt.subplot(plotNumber + plots_created)
            plots_created += 1


        all_data = frame[name]
        qc_data = frame['qc_' + name]

        good_list = get_good_data(qc_data, stamps, all_data)
        good_data = good_list[0]
        good_stamps= good_list[1]

        #if we don't need two lines in same plot, plot
        if not (create_air_dew_plot and name == 'air_temp'):
            if name != 'wind_dir':    
                thumbnail_plot(good_stamps, good_data, ymin[name], ymax[name], 
                    None, None, output, False)

            if name == 'wind_dir':
               thumbnail_plot(good_stamps, good_data, ymin[name], ymax[name], 
                    None, None, output, True)      

        else:        
            thumbnail_plot(good_stamps, good_data, ymin[name], ymax[name], 
                dewpoint_stamps, dewpoint_data, output, False)

    plt.savefig(output + '.thumbnail.png')

def full_plot_stamps(stamps):
    first_stamp = np.sort(stamps)[0]

    first_stamp = first_stamp.replace(hour=0, minute=0, second=0)

    #print([((s - first_stamp).total_seconds()/3600) for s  in stamps])

    #print([(int((s - first_stamp).total_seconds()) / 3600) for s  in stamps])

    return [((s - first_stamp).total_seconds()/3600) for s  in stamps]

# The purpose of this method is to change how the y ticks are shown
# @param min - minimum of the y range right now
# @param max - maximum of the y range right now
# @return new tick labels

def get_new_labels(mini, maxi):
    delta = math.ceil((maxi - mini) / 6)

    return np.arange(mini, (mini + delta*6), delta)                 
        
# The purpose of this method is to create a full_size plot
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param dewpoint_stamps - timestamps for valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
# @param dewpoint_data - valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
# @param output - output filename pattern
# @param wind_direction - says different yaxis labels
# @param accum_precip - different yaxis handling
# @param see if we need subplots
def full_plot(dates, data, ymin, ymax,
    dewpoint_stamps, dewpoint_data, output, wind_dir,
    accum_precip, need_subplots):

    dates = full_plot_stamps(dates)

    if(dewpoint_stamps):
        dewpoint_stamps = full_plot_stamps(dewpoint_stamps)

        plt.plot(dates, data, 'r', dewpoint_stamps, dewpoint_data, 'b')
        plt.axis([min(dates), max(dates), ymin, ymax])
        

        #labels = [math.ceil(float(item.get_text())) for item in ticks]
        #plt.plot(dewpoint_stamps, dewpoint_data, 'b')


    elif(wind_dir):
        plt.plot(dates, data, 'ko', markersize=1)
        plt.axis([min(dates), max(dates), 0, 360])

        axes = plt.gca()
        axes.get_yaxis().set_ticks([0,90,180,270])


    else:
        plt.plot(dates, data, 'k')
        plt.axis([min(dates), max(dates), ymin, ymax])

    if not need_subplots:
        return

    #get axes
    axes = plt.gca()

    #if not wind_direction, reset yaxis ticks
    if not wind_dir:
        if not accum_precip:
            mini, maxi = axes.get_ylim()
            new_labels = get_new_labels(mini, maxi)

        else:
            mini, maxi = axes.get_ylim()
            delta = ((maxi - mini)/6)
            new_labels = np.arange(mini, (mini + delta*6), delta)
            
            if len(new_labels) == 0:
                new_labels = [0, 0.05, 0.1]   

        plt.yticks(new_labels)
        axes.get_yaxis().get_major_ticks()[-1].set_visible(False)

    #hid last tick because for subplots above the last one, it's peeping out...
    # will show for the last subplot or when subplots are not needed
    # in other words, this code should only affect plots above the last plot
    # only when sublots are needed
    axes.get_xaxis().get_major_ticks()[-1].set_visible(False)
    axes.get_xaxis().get_major_ticks()[0].set_visible(False)

# The purpose of this method is to get the subtitle's y coordinate
# based upon how many subplots there are
# @param numSubPlots - number of subplots
# @return y-coordinate of subtitle
def get_subtitle_location(numSubPlots):
    return 1 - 0.05*numSubPlots

# The purpose of this method is to create a full size quicklook
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param create_air_dew_plot - boolean specifying if air temp and dewpoint 
# should be in same plot
def create_full_plot(frame, ymin, ymax, create_air_dew_plot, output):
    #see if we're going to need subplots
    #subplots needed when len is greater than 2 for one variable
    #or greater than 4 for two variables

    need_subplots = len(list(frame.columns.values)) > 2 and not create_air_dew_plot
    need_subplots = need_subplots or len(list(frame.columns.values)) > 4 and create_air_dew_plot

    #get need supplot vars
    if need_subplots:
        plots_created = 1

        if create_air_dew_plot:
            numberPlots = (len(list(frame.columns.values)) - 2) / 2
            plotNumber = numberPlots * 100 + 10

        else:
            numberPlots = len(list(frame.columns.values)) / 2
            plotNumber = numberPlots * 100 + 10    

    stamps = list(frame.index)

    #get dewpoint data
    if create_air_dew_plot:
        all_data = frame['dewpoint']
        qc_data = frame['qc_dewpoint']

        good_list = get_good_data(qc_data, stamps, all_data)
        dewpoint_data = good_list[0]
        dewpoint_stamps= good_list[1]       

    for name in CHOICES:
        #if dewpoint with temp or a name not in frame
        # continue
        if name == 'dewpoint' and create_air_dew_plot:
            continue

        if name not in list(frame.columns.values):
            continue 

        all_data = frame[name]
        qc_data = frame['qc_' + name]

        good_list = get_good_data(qc_data, stamps, all_data)
        good_data = good_list[0]
        good_stamps= good_list[1]

        dateString = good_stamps[0].strftime('%m/%d/%Y')

        # create plots
        if need_subplots:
            if plots_created == 1:
                plt.figure().suptitle('AO&SS Building Tower Meteorogram ' + dateString, fontsize=13)
                ax1 = plt.subplot(plotNumber + plots_created)
                ax1.set_title(TITLES[name], x=0.5, y=get_subtitle_location(numberPlots), fontsize=8)


            else:
                curr_plot = plt.subplot(plotNumber + plots_created, sharex=ax1)
                curr_plot.set_title(TITLES[name], x=0.5, y=get_subtitle_location(numberPlots), fontsize=8)

            plots_created  += 1

        else:
            plt.figure().suptitle('AO&SS Building Tower ' + IND_TITLES[name] + ' ' + dateString, fontsize=13)
            plt.ylabel(TITLES[name])

        #if we don't need two lines in same plot, plot
        if not (create_air_dew_plot and name == 'air_temp'):
            full_plot(good_stamps, good_data, ymin[name], ymax[name], 
                        None, None, output, name == 'wind_dir',
                        name == 'accum_precip', need_subplots)      

        else:        
            full_plot(good_stamps, good_data, ymin[name], ymax[name], 
                dewpoint_stamps, dewpoint_data, output, False, False,
                need_subplots)                  

    if need_subplots:
        plt.subplots_adjust(hspace=0, bottom=0.125)

    plt.xlabel('Time (UTC)')
 
    #reverse code done in create_full_plot
    axes = plt.gca()
    axes.get_xaxis().get_major_ticks()[-1].set_visible(True)            
    axes.get_xaxis().get_major_ticks()[0].set_visible(True)            
    
    plt.savefig(output + '.png')

# the purpose of this method is to use TW's algo to convert
# tempC and relhum to dewpoint
# @param tempC - temperature value in deg C
# @param relhum - relative humidity value in %

def calcDewpoint(tempC, relhum):
    """
    Algorithm from Tom Whittaker tempC is the temperature in degrees Celsius,
    relhum is the relative humidity as a percentage.

    :param tempC: temperature in celsius
    :param relhum: relative humidity as a percentage
    """
    if tempC is None or relhum is None:
        return np.NaN

    gasconst = 461.5
    latheat = 2500800.0

    dp = 1.0 / (1.0 / (273.15 + tempC) - gasconst * np.log((0.0 + relhum) / 100) /
                (latheat - tempC * 2397.5))

    return min(dp - 273.15, tempC)    

# The purpose of this method is to check to see if we have dewpoint data
# if we do not, fill in dewpoint data
# @param frame - pandas frame of nc data
# @return frame with dewpoint data and qc_dewpoint with only good data flagged

def check_dewpoint(frame):
    dewpoint = frame['dewpoint']
    qc_dewpoint = frame['qc_dewpoint']

    good_list = get_good_data(qc_dewpoint.tolist(), list(frame.index), dewpoint.tolist())

    # no dewpoint data that is viable
    if len(good_list[0]) == 0:
        temp = frame['air_temp'].tolist()
        rel_hum = frame['rh'].tolist()
        qc_temp = frame['qc_air_temp'].tolist()
        qc_rel_hum = frame['qc_rh'].tolist()
        stamps = list(frame.index)

        newDewpoint = {}

        qc_dew = {}

        for idx, rh_quality in enumerate(qc_rel_hum):
            rh = rel_hum[idx]
            temp_quality = qc_temp[idx]
            air_temp = temp[idx]
            stamp = stamps[idx]

            if not math.isnan(temp_quality):
                air_temp = None

            elif not math.isnan(rh_quality):
                rh = None

            dew = calcDewpoint(air_temp, rh)    

            newDewpoint[stamp] = dew

            if math.isnan(dew):
                qc_dew[stamp] = 1

            else:
                qc_dew[stamp] = np.NaN    
        
        newDewpoint = pd.Series(newDewpoint)
        qc_dew = pd.Series(qc_dew)

        frame['dewpoint'] = newDewpoint
        frame['qc_dewpoint'] = qc_dew
    
    return frame                 

def main():
    import argparse

    #argparse description
    parser = argparse.ArgumentParser(description="Use data from level_b1 netCDF files to create netCDF files")

    #argparse verbosity info
    parser.add_argument('-v', '--verbose', action='count', 
                       default=int(os.environ.get("VERBOSITY", 2)), 
                       dest='verbosity',
                       help=('each occurence increases verbosity 1 level through'
                             + ' ERROR-WARNING-INFO-DEBUG (default INFO)'))

    #argparse start and end times
    parser.add_argument('-s', '--start-time', type=_dt_convert,
    help="Start time of plot. If only -s is given, a plot of " +
        "only that day is created. Formats allowed: \'YYYY-MM-DDTHH:MM:SS\', \'YYYY-MM-DD\'")

    parser.add_argument('-e', '--end-time', type=_dt_convert,
    help="End time of plot. If only -e is given, a plot of only that day is " +
          "created. Formats allowed: \'YYYY-MM-DDTHH:MM:SS\', \'YYYY-MM-DD\'")

    #netcdf file paths
    parser.add_argument("input_files", nargs="+", help="aoss_tower_level_b1 paths")

    #output filename pattern
    parser.add_argument('-o', '--output', help="filename pattern")

    #thumbnail or full
    parser.add_argument('-t', '--thumb-nail', action='store_true', help="if specified, script creates a thumbnail")

    #plot names
    parser.add_argument('--varnames', nargs="+", choices=CHOICES,
    required=True,
    help="the variable names for the desired plot. Valid choices: air_temp, " +
    "dewpoint, rh, pressure, wind_speed, wind_dir, accum_precip, solar_flux")

    #--daily flag
    parser.add_argument('-d', '--daily', action='store_true', 
    help="creates a plot for every day. Usually used to create plots " +
    "that will line up for aoss tower quicklooks page")

    args = parser.parse_args()

    levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]

    level = levels[min(3, args.verbosity)]
   
    logging.basicConfig(level=level)

    var_names = args.varnames

    create_air_dew_plot = False

    #see if we need to plot air_temp and dewpoint on the same plot
    if 'air_temp' in var_names and 'dewpoint' in var_names:
        create_air_dew_plot = True
        var_names.remove('dewpoint')

    #get data
    frame = get_data(args.input_files)

    frame = check_dewpoint(frame)
    
    #only have the data we need
    for name in CHOICES:
        if name == 'dewpoint' and create_air_dew_plot:
            continue

        if name not in var_names:
           del frame[name]
           del frame['qc_' + name]

    #frame only contains data from start-end times
    if(args.start_time and args.end_time):
        start_filter = args.start_time.strftime('%Y-%m-%d %H:%M:%S')
        end_filter = args.end_time.strftime('%Y-%m-%d %H:%M:%S')
        frame = frame[start_filter: end_filter]

    #frame only contains data from start-end of that day
    elif(args.start_time):
        end_time = args.start_time.replace(hour=23, minute=59, second=59)
        start_filter = args.start_time.strftime('%Y-%m-%d %H:%M:%S')
        end_filter = end_time.strftime('%Y-%m-%d %H:%M:%S')
        frame = frame[start_filter: end_filter]

    ymin = {}
    ymax = {}    

    if not args.daily:
        for name in CHOICES:
            if name == 'dewpoint' and create_air_dew_plot:
                continue

            if name not in frame:
                continue

            ymin[name] = None
            ymax[name] = None    

        if(args.thumb_nail):
            create_thumbnail(frame, ymin, ymax, create_air_dew_plot, args.output)

        else:
            create_full_plot(frame, ymin, ymax, create_air_dew_plot, args.output)

    else:
        for name in CHOICES:
            if name == 'dewpoint' and create_air_dew_plot:
                continue

            if name not in frame:
                continue

            if name == 'accum_precip':    
                ymin[name] = frame[name].min()
                ymax[name] = frame[name].max()

            elif name == 'air_temp' and create_air_dew_plot:
                ymin[name] = math.floor(frame[name].min())
                ymax[name] = math.ceil(frame[name].max())

                dew_min = math.floor(frame['dewpoint'].min())
                dew_max = math.ceil(frame['dewpoint'].max())

                ymin[name] = min(dew_min, ymin[name])
                ymax[name] = max(dew_max, ymax[name])


            else:
                ymin[name] = math.floor(frame[name].min())
                ymax[name] = math.ceil(frame[name].max())    

        frameList = [(group[1]) for group in frame.groupby(frame.index.day)] 

        counter = 0

        for frame in frameList:
            if args.thumb_nail:
                create_thumbnail(frame, ymin, ymax, create_air_dew_plot, args.output.format(str(counter)))
            else:
                create_full_plot(frame, ymin, ymax, create_air_dew_plot, args.output.format(str(counter)))
            counter += 1
            plt.gcf().clear()          

if __name__ == "__main__":
    main()