# -*- coding: utf-8 -*-

import xarray as xr
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import matplotlib.dates as mdates
import datetime
import numpy as np

def convert_times_to_numbers(times):
    epoch = np.datetime64("1970-01-01")
    second = np.timedelta64(1000000000, "ns")
    times_datetime = [datetime.datetime.utcfromtimestamp((i - epoch) / second)
                      for i in times.values]
    times_num = mdates.date2num(times_datetime)
    return times_num

def select_meaning(var, meaning):
    meanings = dict(zip(var.attrs["flag_meanings"].split(" "),
                        var.attrs["flag_values"]))
    return var == meanings[meaning]

def plot_specMACS_cloud_mask(image, masks, title, figsize=(40,20), aspect='auto', cmap='gray', vmin=None, vmax=None, gamma_corr=1, lw=1):
    colors = ["mediumvioletred", "darkorange"]
    angles = image.T.angle
    times = image.T.time
    times_num = convert_times_to_numbers(times)

    fig, ax = plt.subplots(len(masks.keys()), 1, figsize=figsize, sharex=True)
    for i, mask_title in enumerate(masks.keys()):
        ax[i].pcolormesh(times_num, angles, image.T**gamma_corr, cmap=cmap, vmin=vmin, vmax=vmax)
        ax[i].set_title("{} {}".format(mask_title, title), fontsize=20)
        
        #invert y_axis for comparison with specMACS videos (made from measurements of (polarization) 2D cameras)
        ax[i].invert_yaxis()
        ax[i].yaxis.set_visible(False)
        mask = masks[mask_title]
        for j, mask_label in enumerate(mask.keys()):
            ax[i].contour(times_num, angles, mask[mask_label].T, origin='lower', colors=colors[j], linewidths=lw)

    lines = [Line2D([0], [0], color=c, linewidth=3, linestyle='solid') for c in colors]
    labels = list(mask.keys())
    ax[0].legend(lines, labels, loc="upper right", prop=dict(size=18))

    ax[2].xaxis_date()
    date_format = mdates.DateFormatter('%H:%M:%S')
    ax[2].xaxis.set_major_formatter(date_format)
    ax[2].tick_params(axis='x', which='major', labelsize=20)
    plt.show()

def main():
    #open dataset with OPeNDAP and select sequence
    ds = xr.open_dataset("https://macsserver.physik.uni-muenchen.de/products/dap/eurec4a/cloudmask/EUREC4A_HALO_specMACS_cloud_mask_20200122T150000-20200122T215959_v1.1.nc")
    start_time = "2020-01-22T15:39:00"
    end_time = "2020-01-22T15:45:00"
    ds_slice = ds.sel(time = slice(start_time, end_time))
    print("Plot specMACS cloudmask ({} - {})".format(start_time, end_time))

    #select masks
    swir_image = ds_slice.swir_radiance.isel(radiation_wavelength = 0)
    combined_mask = ds_slice.cloud_mask
    brightness_mask = ds_slice.brightness_cloud_mask
    comb_brightness_watervapor_mask = ds_slice.combined_brightness_watervapor_mask

    #"or"-Operator: |
    masks = {
    "optimal_mask":{
        "probably cloudy and most likely cloudy": select_meaning(combined_mask, "probably_cloudy") | select_meaning(combined_mask, "most_likely_cloudy"),
        "most likely cloudy": select_meaning(combined_mask, "most_likely_cloudy")
        },
    "brtns_cloud_mask":{
        "probably cloudy and most likely cloudy": select_meaning(brightness_mask, "probably_cloudy") | select_meaning(brightness_mask, "most_likely_cloudy"),
        "most likely cloudy": select_meaning(brightness_mask, "most_likely_cloudy")
        },
    "combined_brtns_wv_cloud_mask":{
        "probably cloudy and most likely cloudy": select_meaning(comb_brightness_watervapor_mask, "probably_cloudy") | select_meaning(comb_brightness_watervapor_mask, "most_likely_cloudy"),
        "most likely cloudy": select_meaning(comb_brightness_watervapor_mask, "most_likely_cloudy")
        }
    }
    
    title = "{} - {}".format(start_time, end_time)
    plot_specMACS_cloud_mask(swir_image, masks, title, gamma_corr=0.4)


if __name__ == "__main__":
    main()