#!/home/poker/miniconda3/envs/goes16_201710/bin/python

# This reads in a GRB netcdf4 file of radiances, converts to Brightness
# Temperature and plots using matplotlib and cartopy
# This particular script was used for ABI channel 13 (clean IR window)

# import necessary modules

import netCDF4
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import os,errno
import shutil
import sys
from PIL import Image

# Read in AOS logo to put in the corner of the image
aoslogo = Image.open('/home/poker/uw-aoslogo.png')
aoslogoheight = aoslogo.size[1]
aoslogowidth = aoslogo.size[0]

# We need a float array between 0-1, rather than
# a uint8 array between 0-255
aoslogo = np.array(aoslogo).astype(np.float) / 255

# netcdf file is first argument to this script, open as netcdf4 object

dt = sys.argv[1]
f = netCDF4.Dataset(dt)

# Optionally print out a bunch of stuff to interrogate the file
#print(f)
#print("f.variables[Rad] ",f.variables["Rad"])
#print("valid_pixel_count ",f.variables['valid_pixel_count'][0])
#print("undersaturated_pixel_count ",f.variables['undersaturated_pixel_count'][0])
#print("saturated_pixel_count ",f.variables['saturated_pixel_count'][0])
#print("min_radiance_value_of_valid_pixels ",f.variables['min_radiance_value_of_valid_pixels'][0])
#print("max_radiance_value_of_valid_pixels ",f.variables['max_radiance_value_of_valid_pixels'][0])
#print("mean_radiance_value_of_valid_pixels ",f.variables['mean_radiance_value_of_valid_pixels'][0])
#print("std_dev_radiance_value_of_valid_pixels ",f.variables['std_dev_radiance_value_of_valid_pixels'][0])
#print("nominal_satellite_height ",f.variables['nominal_satellite_height'][0])
#print("geospatial_lat_lon_extent ",f.variables['geospatial_lat_lon_extent'][0])
#print("band_id ",f.variables['band_id'][0])
#print("band_wavelength ",f.variables['band_wavelength'][0])
#print("esun ",f.variables['esun'][0])
#print("kappa0 ",f.variables['kappa0'][0])
#print("earth_sun_distance_anomaly_in_AU ",f.variables['earth_sun_distance_anomaly_in_AU'][0])
#print(" ")



# Convert radiance to Brightness Temp 
fk1 = f.variables['planck_fk1'][0]
fk2 = f.variables['planck_fk2'][0]
bc1 = f.variables['planck_bc1'][0]
bc2 = f.variables['planck_bc2'][0]

print("fk1 = ",fk1)
print("fk2 = ",fk2)
print("bc1 = ",bc1)
print("bc2 = ",bc2)
print(" ")


# Here's the scale factor and offset that are referred to - 
# I am not using them at all

scalefactor = f.variables['Rad'].scale_factor
add_offset = f.variables['Rad'].add_offset
print("scale factor is ",f.variables['Rad'].scale_factor)
print("add_offset   is ",f.variables['Rad'].add_offset)

# Reaed in Radiance values into data_var

data_var = f.variables['Rad'][:]

#Compute brightness temp from radiances
data_var = (fk2 / ( np.log((fk1 / data_var) + 1 )) - bc1) / bc2


print("BT values from y,x = 0-10, 0-10 (far northwest corner) ", data_var[0:10,0:10])
print("max of data_var ",np.max(data_var))
print("min of data_var ",np.min(data_var))

# Set a = the brightness temps to plot

a = data_var
print("max of data_var ",np.max(data_var))
print("min of data_var ",np.min(data_var))

# read in x and y values - these are angles in steradians

xa = f.variables['x'][:]
ya = f.variables['y'][:]

#    scale_factor: -1.4e-05
#    add_offset: 0.151865

# print("xa",xa)
# print("ya",ya)


Convert to meters by multiplying by satellite height
#print("satellite height", f.variables['nominal_satellite_height'][0])
xa = xa*35785831
ya = ya*35785831

print("xa",xa)
print("ya",ya)

import cartopy.crs as ccrs

globe = ccrs.Globe(semimajor_axis=6378137.,semiminor_axis=6356752.)

proj = ccrs.Geostationary(central_longitude=f.variables['nominal_satellite_subpoint_lon'][0],
                          satellite_height=f.variables['nominal_satellite_height'][0] * 1000,
                          globe=globe,sweep_axis='x')

##########################################################
# The rest of this is me defining some subregions to plot
# I crop off a North America region in addition to plotting
# the whole disk

image_rows=f.dimensions["y"].size
image_columns=f.dimensions["x"].size
namer_image_crop_top=0
namer_image_crop_bottom=-2400
namer_image_crop_left=400
namer_image_crop_right=-1400
#
namer_image_size_y=(image_rows+namer_image_crop_bottom-namer_image_crop_top)
namer_image_size_x=(image_columns+namer_image_crop_right-namer_image_crop_left)
#
print("namer image size")
print(namer_image_size_x, namer_image_size_y)
#
namer_image_size_x=15.1
namer_image_size_y=12.6

# Create a new figure with size specified. The 14.98,14.983 was determined
# with some playing around, trying to get even valued pixel amounts in x
# and y direction. Very difficult to make that happen
fig = plt.figure(figsize=(namer_image_size_x,namer_image_size_y),dpi=80.)
fig2 = plt.figure(figsize=(14.98,14.983),dpi=80.)
fig9 = plt.figure(figsize=(image_columns/78.,image_rows/78.))

# Put a single axes on this figure; set the projection for the axes to be our
# Fixed disk projection
ax = fig.add_subplot(1, 1, 1, projection=proj)
ax.outline_patch.set_edgecolor('none')
ax2 = fig2.add_subplot(1, 1, 1, projection=proj)
ax2.outline_patch.set_edgecolor('none')
ax9 = fig9.add_subplot(1, 1, 1, projection=proj)
ax9.outline_patch.set_edgecolor('none')



# Set up a color table to get same enhancement as McIDAS plots
import matplotlib as mpl

cdict = {'red': ((0.0, 0.0, 0.0),
                 (.001, 1.00, 1.00),
                 (.107, 1.00, 1.00),
                 (.113, 0.498, 0.498),
                 (.173, 1.00, 1.00),
                 (.179, 0.902, 0.902),
                 (.227, 0.102, 0.102),
                 (.233, 0.00, 0.00),
                 (.287, 0.902, 0.902),
                 (.293, 1.00, 1.00),
                 (.346, 1.00, 1.00),
                 (.352, 1.00, 1.00),
                 (.406, 0.101, 0.101),
                (.412, 0.00, 0.00),
                 (.481, 0.00, 0.00),
                 (.484, 0.00, 0.00),
                 (.543, 0.00, 0.00),
                 (.546, 0.773, 0.773),
                 (.994, 0.012, 0.012),
                 (.997, 0.004, 0.004),
                 (1.0, 0.0, 0.0)),
         'green': ((0.0, 0.0, 0.0),
                 (.001, 1.00, 1.00),
                 (.107, 1.00, 1.00),
                 (.113, 0.00, 0.00),
                 (.173, 0.498, 0.498),
                 (.179, 0.902, 0.902),
                 (.227, 0.102, 0.102),
                 (.233, 0.00, 0.00),
                 (.287, 0.00, 0.00),
                 (.293, 0.00, 0.00),
                 (.346, 0.902, 0.902),
                 (.352, 1.00, 1.00),
                 (.406, 1.00, 1.00),
                 (.412, 1.00, 1.00),
                 (.481, 0.00, 0.00),
                 (.484, 0.00, 0.00),
                 (.543, 1.00, 1.00),
                 (.546, 0.773, 0.773),
                 (.994, 0.012, 0.012),
                 (.997, 0.004, 0.004),
                   (1.0, 0.0, 0.0)),
         'blue': ((0.0, 0.00, 0.00),
                 (.001, 1.00, 1.00),
                 (.107, 0.00, 0.00),
                 (.113, 0.498, 0.498),
                 (.173, 0.786, 0.786),
                 (.179, 0.902, 0.902),
                 (.227, 0.102, 0.102),
                 (.233, 0.00, 0.00),
                 (.287, 0.00, 0.00),
                 (.293, 0.00, 0.00),
                 (.346, 0.00, 0.00),
                 (.352, 0.00, 0.00),
                 (.406, 0.00, 0.00),
                 (.412, 0.00, 0.00),
                 (.481, 0.451, 0.451),
                 (.484, 0.451, 0.451),
                 (.543, 1.00, 1.00),
                 (.546, 0.773, 0.773),
                 (.994, 0.012, 0.012),
                 (.997, 0.004, 0.004),
                  (1.0, 0.0, 0.0))}


my_cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,2048)

# Plot the data using my color map
# set the colormap to extend over a range of values from 162 to 330.
# Note, we save the image returned by imshow for later...

im = ax.imshow(a[namer_image_crop_top:namer_image_crop_bottom,namer_image_crop_left:namer_image_crop_right], extent=(xa[namer_image_crop_left],xa[namer_image_crop_right],ya[namer_image_crop_bottom],ya[namer_image_crop_top]), origin='upper',cmap=my_cmap, vmin=162., vmax=330.)

im2 = ax2.imshow(a[:], extent=(xa[0],xa[-1],ya[-1],ya[0]), origin='upper', cmap=my_cmap, vmin=162., vmax=330.)

im9 = ax9.imshow(a[:], extent=(xa[0],xa[-1],ya[-1],ya[0]), origin='upper', cmap=my_cmap, vmin=162., vmax=330.)

# put on coastlines

ax.coastlines(resolution='50m', color='green')
ax2.coastlines(resolution='50m', color='green')
ax9.coastlines(resolution='50m', color='green')

import cartopy.feature as cfeat

# Add country borders with a thick line.
ax.add_feature(cfeat.BORDERS, linewidth='1', edgecolor='green')
ax2.add_feature(cfeat.BORDERS, linewidth='1', edgecolor='green')
ax9.add_feature(cfeat.BORDERS, linewidth='1', edgecolor='green')

# Set up a feature for the state/province lines. Tell cartopy not to fill in the polygons
state_boundaries = cfeat.NaturalEarthFeature(category='cultural',
                                             name='admin_1_states_provinces_lakes',
                                             scale='50m', facecolor='none', edgecolor='red')

# Add the feature with dotted lines, denoted by ':'
ax.add_feature(state_boundaries, linestyle=':')
ax2.add_feature(state_boundaries, linestyle=':')
ax9.add_feature(state_boundaries, linestyle=':')

# axes for North America crop
cbaxes1 = fig.add_axes([0.135,0.12,0.755,0.02])
cbar1 = fig.colorbar(im, cax=cbaxes1, orientation='horizontal')
font_size = 14
cbar1.ax.tick_params(labelsize=font_size, labelcolor='yellow')
cbar1.ax.xaxis.set_ticks_position('top')
cbar1.ax.xaxis.set_label_position('top')

# axes for smaller full disk image
#cbaxes2 = fig2.add_axes([0.135,0.12,0.755,0.02])
cbaxes2 = fig2.add_axes([0.135,0.12,0.255,0.005])
cbar2 = fig2.colorbar(im2, cax=cbaxes2, orientation='horizontal')
font_size = 10
cbar2.ax.tick_params(labelsize=font_size, labelcolor='black')
cbar2.ax.xaxis.set_ticks_position('top')
cbar2.ax.xaxis.set_label_position('top')

# axes for full resolution
cbaxes9 = fig9.add_axes([0.135,0.15,0.755,0.02])
cbar9 = fig9.colorbar(im9, cax=cbaxes9, orientation='horizontal')
font_size = 18
#cbar9.set_label('Brightness Temperature (K)',size=18)
cbar9.ax.tick_params(labelsize=font_size, labelcolor='yellow')
cbar9.ax.xaxis.set_ticks_position('top')
cbar9.ax.xaxis.set_label_position('top')


# Grab info about the image time for the label and file name
import datetime
import calendar

time_var = f.time_coverage_start

iyear = time_var[0:4]
print("iyear ",iyear)
imonth = time_var[5:7]
print("imonth ",imonth)
cmonth = calendar.month_abbr[int(imonth)]
print("cmonth ",cmonth)
iday = time_var[8:10]
print("iday ",iday)
itime = time_var[11:19]
itimehr = time_var[11:13]
itimemn = time_var[14:16]

ctime_string = iyear +' '+cmonth+' '+iday+'  '+itime+' GMT'
ctime_file_string = iyear + imonth + iday + itimehr + itimemn

time_string = 'GOES-16 Band 13\nClean LW IR Window\n%s '%ctime_string
print(time_string)

# Put shadow effect around title so you can see in both light and dark

from matplotlib import patheffects
outline_effect = [patheffects.withStroke(linewidth=2, foreground='black')]

text = ax.text(0.005, 0.90, time_string,
    horizontalalignment='left', transform = ax.transAxes,
    color='yellow', fontsize='small', weight='bold')
text.set_path_effects(outline_effect)

text2 = ax2.text(0.005, 0.92, time_string,
    horizontalalignment='left', transform = ax2.transAxes,
    color='yellow', fontsize='small', weight='bold')
text2.set_path_effects(outline_effect)

text9 = ax9.text(0.50, 0.97, time_string,
    horizontalalignment='center', transform = ax9.transAxes,
    color='yellow', fontsize='large', weight='bold')
text9.set_path_effects(outline_effect)


# Filename to save to

filename1=ctime_file_string+"_namer.jpg"
filename2=ctime_file_string+"_fulldisk.jpg"
filename9=ctime_file_string+"_fulldisk_full.jpg"

# Put AOS logo in upper left hand corner. This is tricky to get right for
# some reason. Easier to put in lower left by setting (aoslogo, 0, 0, zorder=10)

fig.figimage(aoslogo,  0, fig.bbox.ymax - aoslogoheight - 28  , zorder=10)
fig2.figimage(aoslogo,  0, fig.bbox.ymax - aoslogoheight + 156  , zorder=10)
fig9.figimage(aoslogo,  0, fig.bbox.ymax - aoslogoheight - 300  , zorder=10)

# Save figure, no white border around the edge

fig.savefig(filename1, bbox_inches='tight', pad_inches=0)
fig2.savefig(filename2, bbox_inches='tight', pad_inches=0)
fig9.savefig(filename9, bbox_inches='tight', pad_inches=0)

quit()