import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import numpy as np
import netCDF4 as nc

import sys
import os

from scene import test_scene


def main(scene_flag):
    scn = test_scene()
    lat = scn['lat']
    lon = scn['lon']

    plt.figure(figsize=[15, 10])
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([np.min(lon)-15, np.max(lon)+15,
                   np.min(lat)-15, np.max(lat)+15], crs=ccrs.PlateCarree())

    ax.coastlines(resolution='50m', color='black', linewidth=1)

    plt.pcolormesh(lon, lat, scn[scene_flag],
                   transform=ccrs.PlateCarree())

    plt.show()


def plot_tests(numpy_file='test_confidence.npz'):
    data = np.load(numpy_file)
    lat = data['lat']
    lon = data['lon']
    # confidence = data['confidence']
    conf_test = data['conf_test']
    confidence = data['confidence']
    cm = np.zeros(confidence.shape)
    cm[confidence > 0.66] = 1
    cm[confidence > 0.95] = 2
    cm[confidence > 0.99] = 3

    for i in range(9):
        print(f'Making plot {i}\n')
        plt.figure(figsize=[15, 10])
        ax = plt.axes(projection=ccrs.PlateCarree())
        ax.set_extent([np.min(lon)-5, np.max(lon)+5,
                       np.min(lat)-5, np.max(lat)+5], crs=ccrs.PlateCarree())

        ax.coastlines(resolution='50m', color='black', linewidth=1)

        plt.pcolormesh(lon, lat, conf_test[i, :, :], vmin=0, vmax=1,
                       transform=ccrs.PlateCarree())
        plt.title(f'Confidence {i}')
        plt.colorbar()

        plt.savefig(f'figures/confplot_{i}.png')

    plt.figure(figsize=[15, 10])
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([np.min(lon)-5, np.max(lon)+5,
                   np.min(lat)-5, np.max(lat)+5], crs=ccrs.PlateCarree())

    ax.coastlines(resolution='50m', color='black', linewidth=1)

    plt.pcolormesh(lon, lat, cm, vmin=0, vmax=3,
                   transform=ccrs.PlateCarree())
    plt.title('Cloud Mask')
    # plt.colorbar()
    cb = plt.colorbar(ticks=range(4))
    cb.set_ticklabels(['confident cloudy', 'probably cloudy', 'probably clear', 'confident clear'])

    plt.savefig('figures/total_cm.png')


def plot_confidence(numpy_file='test_confidence.npz'):
    data = np.load(numpy_file)
    lat = data['lat']
    lon = data['lon']
    # confidence = data['confidence']
    confidence = data['confidence']
    plt.figure(figsize=[15, 10])
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([np.min(lon)-5, np.max(lon)+5,
                   np.min(lat)-5, np.max(lat)+5], crs=ccrs.PlateCarree())

    ax.coastlines(resolution='50m', color='black', linewidth=1)

    plt.pcolormesh(lon, lat, confidence, vmin=0, vmax=1,
                   transform=ccrs.PlateCarree())
    plt.title('Confidence')
    plt.colorbar()

    plt.savefig('figures/total_confidence.png')


def plot_cloud_mask():
    f = nc.Dataset('/ships19/hercules/pveglio/mvcm_viirs_hires/CLDMSK_L2_VIIRS_SNPP.A2022173.1454.001.2022174035130.nc')
    lat = f['geolocation_data/latitude'][:]
    lon = f['geolocation_data/longitude'][:]
    # icm = f['geophysical_data/Integer_Cloud_Mask'][:]
    cm = f['geophysical_data/Cloud_Mask'][:]
    qa = f['geophysical_data/Quality_Assurance'][:]
    f.close()

    # cm_flag = np.array((np.array(cm[0, :, :], np.byte) &
    #                   (pow(2, 1) + pow(2, 2)))/pow(2, 1))
    cm_flag = read_bits(qa[:, :, 2], 2) & ~read_bits(cm[2, :, :], 2)

    plt.figure(figsize=[15, 10])
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([np.min(lon)-5, np.max(lon)+5,
                   np.min(lat)-5, np.max(lat)+5], crs=ccrs.PlateCarree())

    ax.coastlines(resolution='50m', color='black', linewidth=1)

    plt.pcolormesh(lon, lat, cm_flag, vmin=0, vmax=1,
                   transform=ccrs.PlateCarree())
    # cb = plt.colorbar(ticks=range(4))
    # cb.set_ticklabels(['confident cloudy', 'probably cloudy', 'probably clear', 'confident clear'])
    cb = plt.colorbar(ticks=range(2))
    cb.set_ticklabels(['No', 'Yes'])
    # plt.savefig('figures/ref_cloudmask.png')
    plt.savefig('figures/high_cloud_flag_11_12.png')


def read_bits(byte, bits):
    # orig_shape = byte.shape
    # if byte.ndim > 1:
    #     byte = byte.reshape((byte.shape[0]*byte.shape[1],))

    if type(bits) is int:
        flag = np.array((np.array(byte, np.byte) & pow(2, bits))/pow(2, bits),
                        dtype='bool')
    else:
        flag = (np.array(byte, np.byte) & pow(2, bits[0]) +
                pow(2, bits[1]))/pow(2, bits[0])

        flag = np.array(np.trunc(1.0-flag/2.0), dtype='bool')
    # flag = flag.reshape((orig_shape))
    return flag


if __name__ == "__main__":
    # main(sys.argv[1])
    plot_tests(sys.argv[1])
    # plot_confidence(sys.argv[1])
    # plot_cloud_mask()