import numpy as np
import matplotlib.pyplot as plt


# Example
# do_plot(['5', '10', '15', '20', '30', '40', '60', '80', '100', '120', '140', '160'], [values, values_bt, values_bt_frefl], ['OPD', 'BT+OPD', 'BT+REFL+OPD'],
# ['green', 'blue', 'black'], title='OPD SRES Estimation Model Comparison', xlabel='OPD', ylabel='MAE %', barWidth=0.3)

def do_plot(group_names, group_values, value_labels, bar_colors, title=None, xlabel=None, ylabel=None, ylim=None, barWidth=0.1):
    num_groups = len(group_names)
    num_bars = len(group_values)

    bars_x = []
    for i in range(num_bars):
        x = [x + i*barWidth for x in np.arange(num_groups)]
        bars_x.append(x)

    fig, ax = plt.subplots()
    if ylim is not None:
        ax.set_ylim(ylim)

    for k in range(num_bars):
        plt.bar(bars_x[k], group_values[k], color=bar_colors[k], width=barWidth, edgecolor='white', label=value_labels[k])

    # Add xticks on the middle of the group bars
    if xlabel is not None:
        plt.xlabel(xlabel, fontweight='bold')
    plt.xticks([r + barWidth for r in range(num_groups)], group_names)

    if ylabel is not None:
        plt.ylabel(ylabel, fontweight='bold')

    if title is not None:
        plt.title(title, fontweight='bold')

    # Create legend & Show graphic
    plt.legend()
    plt.show()


def do_plot_simple(x_values, y_values, color, title=None, xlabel=None, ylabel=None, ylim=None, barWidth=0.1):
    fig, ax = plt.subplots()
    if ylim is not None:
        ax.set_ylim(ylim)

    plt.bar(np.arange(len(x_values)), y_values, color=color, width=barWidth, edgecolor='black')

    # Add xticks on the middle of the group bars
    if xlabel is not None:
        plt.xlabel(xlabel, fontweight='bold')
    plt.xticks([r for r in range(len(x_values))], labels=[int(xlbl) for xlbl in x_values])

    if ylabel is not None:
        plt.ylabel(ylabel, fontweight='bold')

    if title is not None:
        plt.title(title, fontweight='bold')

    # Create legend & Show graphic
    plt.legend()
    plt.show()


def run_plot_cld_layers(cld_layer_bot, cld_layer_top, hgt_ref, x_values, start, end):
    plot_cld_layers(cld_layer_bot[start:end], cld_layer_top[start:end], hgt_ref[start:end], x_values[start:end])


def plot_cld_layers(cld_layer_bot, cld_layer_top, hgt_ref, x_values):
    x_len = len(x_values)

    for k in range(x_len):
        if hgt_ref[k] is not None:
            plt.plot(x_values[k], hgt_ref[k], marker='o', color='red')

        if cld_layer_bot[k] is not None:
            bots = cld_layer_bot[k]
            tops = cld_layer_top[k]
            nlyrs = len(bots)
            for i in range(nlyrs):
                plt.vlines(x_values[k], bots[i], tops[i], color='black')
                plt.plot(x_values[k], bots[i], marker='^', color='black')
                plt.plot(x_values[k], tops[i], marker='v', color='black')

    plt.show()

# import cartopy.crs as ccrs
# xrds = xr.open_dataset('/Users/rink/data/OR_ABI-L1b-RadC-M3C14_G16_s20190900022158_e20190900024531_c20190900024578.nc')
# geos = ccrs.Geostationary(central_longitude=-75.0, sweep_axis='x')
# img = xrds.Rad[:,:]
# x = xrds.x[:]
# y = xrds.y[:]

# keep: lats > 30 and lats < 50
# get_breaks(times[keep], 3600)
# 4628:4646


def img_plot(img, proj, x, y):
    fig = plt.figure(figsize=(7, 5))
    ax = fig.add_subplot(1, 1, 1, projection=proj)
    ax.imshow(img, origin='upper', extent=(x.min(), x.max(), y.min(), y.max()))


import pickle


def get_cld_layers():
    f = open('/Users/rink/layer_bot.pkl', 'rb')
    layer_bot = pickle.load(f)
    f.close()

    f = open('/Users/rink/layer_top.pkl', 'rb')
    layer_top = pickle.load(f)
    f.close()

    f = open('/Users/rink/grd_hgt.pkl', 'rb')
    cld_hgt = pickle.load(f)
    f.close()

    f = open('/Users/rink/lat.pkl', 'rb')
    lats = pickle.load(f)
    f.close()

    f = open('/Users/rink/lon.pkl', 'rb')
    lons = pickle.load(f)
    f.close()

    f = open('/Users/rink/times.pkl', 'rb')
    times = pickle.load(f)
    f.close()

    return layer_bot, layer_top, cld_hgt, lats, lons, times