import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import matplotlib.dates as mdates

#creates line plot from an array
def plot_line(timestamp, arr: np.ndarray, ylabel: str):
    plt.plot(timestamp, arr)
    plt.ylabel(ylabel)
    plt.show()


#formats and creates ice mass balance plot
def plot_simb(dtc: np.ndarray, surface_distance, bottom_distance, timestamp, initial_snow_depth, distance_between_rangefinders):

    #calculate height of surface rangefinder above the ice surface
    height_above_ice = surface_distance[0] + initial_snow_depth

    #create arrays that store values for each timestamp
    snow_height = height_above_ice - surface_distance
    ice_thickness = distance_between_rangefinders - height_above_ice - bottom_distance
    ice_thickness = - ice_thickness #make negative for plotting because it's below waterline

    #define temperature string position (height above ice)
    temp_string_length = 3.84 #standard SIMB3 temperature string has 2 cm spacing and 192 thermistors
    temp_string_offset = 0 #the axial distance between the surface rangefinder and the first thermister on the DTC
    temp_string_top = height_above_ice - temp_string_offset #distance from the top of the ice to the first thermistor in the temperature string
    temp_string_bottom = temp_string_length - temp_string_top
    temp_string_bottom = temp_string_bottom * (-1) #make negative for plotting
    

    #shift excel serial date (1900 epoch) to the Python 1970 epoch by adding the number of days between 01/01/1900 and 01/01/1970 
    #then subtract 1 to account for the famous 1900 excel leap year timestamp bug
    timestamp = timestamp - 25568 - 1 

    #set plot bounds
    upper_bound = 1
    lower_bound = -3 
    
    #create mass balance plot
    fig, ax = plt.subplots()
    im = ax.imshow(dtc, cmap='jet', vmin=-30, vmax=0, extent=[timestamp[0], timestamp[-1], temp_string_bottom, temp_string_top],
                   aspect='auto', zorder=0)     #format ice temperature heatmap from dtc (temperature string) data
    plt.colorbar(im, ax=ax)

    #color in between the lines
    ax.fill_between(timestamp, snow_height, 0, color="grey", zorder=1)                  #plot snow height on top of ice
    ax.fill_between(timestamp, temp_string_top, snow_height, color='white', zorder=2)       #create whitespace above snow
    ax.fill_between(timestamp, ice_thickness, lower_bound, color='b', zorder=3)         #plot lower boundary of ice

    #bound y axis
    ax.set_ylim(lower_bound, upper_bound)       

    #format x-axis as date
    formatter = mdates.DateFormatter("%Y-%m-%d")
    ax.xaxis.set_major_formatter(formatter)

    #only label each month
    locator = mdates.MonthLocator()
    ax.xaxis.set_major_locator(locator)

    plt.show()


#function that reads csv file and calls plot functions
def plot_from_csv(path_to_data: str, initial_snow_depth: float, plot_range):

    #try to read csv at input location
    try:
        simb_df = pd.read_csv(path_to_data,skiprows=plot_range[0], nrows=plot_range[1])      #create pandas dataframe from csv
        timestamp = simb_df.time_stamp.to_numpy()       #extracts time stamp from simb_df dataframe, converts to numpy array
    except FileNotFoundError:
        print('Invalid Path, please check file location.')
        return

    # Create plot from input

    # Distance between sounders in meters, used for bounding the Ice Mass Balance Plot
    sounder_dist = 4.05

    # Extract relevant data from dataframe, convert to numpy array
    snow_dist = simb_df.surface_distance.to_numpy()                     # distance from upper sounder to snow
    water_depth = simb_df.bottom_distance.to_numpy()                    # distance from lower sounder to bottom of ice
    dtc_values = simb_df.filter(like='dtc_values_').to_numpy().T        # temperature string data at 2 cm intervals

    # Call plot function for ice mass balance plot
    plot_simb(dtc_values, snow_dist, water_depth, timestamp, initial_snow_depth, sounder_dist)


# plot_from_csv(location (str), initial snow depth (meters), plot type (str))
# Plot type can be 'simb' (default), 'air_temp', 'air_pressure', or 'water_temp'
plot_from_csv(r'SIMB3sampleDataSheet.csv', .2, [0, 3378])