#!/usr/bin/env python

"""
marginal_distributions_temp.py
"""
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
import sys
import _pickle as cPickle
from matplotlib import rcParams
rcParams['legend.fontsize']='small'

cpath_errorstats = '/Volumes/Backup Plus/python/fcst_stats/'

clead = sys.argv[1]

# Reference std
stdrefs = {"Jan-Feb-Mar":1.0,   "Apr-May-Jun":1.0,\
    "Jul-Aug-Sep":1.0, "Oct-Nov-Dec":1.0}

models = ['raw', 'dav', 'qmap', 'mos', 'mvmos']

# --- read land-surface mask, lons, lats

infilename = 'lsmask_0p5.cPick'
print (infilename)
inf = open(infilename, 'rb')
lsmask = cPickle.load(inf)
lats = cPickle.load(inf)
lons = cPickle.load(inf)
inf.close()

# ---- populate the histogram of correlations, std devs

xstd = np.arange(0.5,1.501,0.01)
xcorr = np.arange(0.0,1.001,0.01)
nxstd = len(xstd)
nxcorr = len(xcorr)
corr_winter_hist = ma.zeros((nxcorr,5), dtype=np.float64)
corr_spring_hist = ma.zeros((nxcorr,5), dtype=np.float64)
corr_summer_hist = ma.zeros((nxcorr,5), dtype=np.float64)
corr_fall_hist   = ma.zeros((nxcorr,5), dtype=np.float64)
std_winter_hist = ma.zeros((nxstd,5), dtype=np.float64)
std_spring_hist = ma.zeros((nxstd,5), dtype=np.float64)
std_summer_hist = ma.zeros((nxstd,5), dtype=np.float64)
std_fall_hist   = ma.zeros((nxstd,5), dtype=np.float64)

for imodel, model in enumerate(models):

    # ---- read in the standardized anomaly and correlation from files
    #      prepared for Taylor diagram generation

    if model == 'raw':
        ftype = 'raw_forecast_errorstats_2019_'
        title = 'Raw\nmodel output\n+'+clead+' h'
        color = 'Red'
    elif model ==  'dav':
        ftype = 'decayavg_forecast_errorstats_2019_'
        title = 'Decaying-average\nbias correction\n+'+clead+' h'
        color = 'Peru'
    elif model == 'qmap':
        ftype = 'quantile_mapping_errorstats_2019_'
        title = 'Quantile\nmapping\n+'+clead+' h'
        color = 'Blue'
    elif model == 'mos':
        ftype = 'MOS_GEFSv12_2019_'
        title = 'Univariate\nMOS\n+'+clead+' h'
        color = 'Green'
    elif model == 'mvmos':
        ftype = 'MOS_mvr_GEFSv12_2019_'
        title = 'Multi-variate\nMOS\n+'+clead+' h'
        color = 'Gray'

    statsfile = cpath_errorstats + ftype + clead+'h_taylor.cPick'
    inf = open(statsfile, 'rb')
    corr_winter = cPickle.load(inf)
    corr_spring = cPickle.load(inf)
    corr_summer = cPickle.load(inf)
    corr_fall = cPickle.load(inf)
    std_normalized_winter = cPickle.load(inf)
    std_normalized_spring = cPickle.load(inf)
    std_normalized_summer = cPickle.load(inf)
    std_normalized_fall = cPickle.load(inf)
    std_obs_winter = cPickle.load(inf)
    std_obs_spring = cPickle.load(inf)
    std_obs_summer = cPickle.load(inf)
    std_obs_fall = cPickle.load(inf)
    inf.close()
    nlats, nlons = np.shape(corr_fall)

    for jy in range(nlats):
        for ix in range(nlons):
            if lsmask[jy,ix] == 1:
                corr_idx = np.int(corr_winter[jy,ix]*100.0)
                corr_winter_hist[corr_idx,imodel] = corr_winter_hist[corr_idx,imodel] + 1.0
                corr_idx = np.int(corr_spring[jy,ix]*100.0)
                corr_spring_hist[corr_idx,imodel] = corr_spring_hist[corr_idx,imodel] + 1.0
                corr_idx = np.int(corr_summer[jy,ix]*100.0)
                corr_summer_hist[corr_idx,imodel] = corr_summer_hist[corr_idx,imodel] + 1.0
                corr_idx = np.int(corr_winter[jy,ix]*100.0)
                corr_fall_hist[corr_idx,imodel] = corr_fall_hist[corr_idx,imodel] + 1.0

                std_idx = np.amin([nxstd-1, np.int(std_normalized_winter[jy,ix]*100.0 - 50.0)])
                std_winter_hist[std_idx,imodel] = std_winter_hist[std_idx,imodel] + 1.0
                std_idx = np.amin([nxstd-1, np.int(std_normalized_spring[jy,ix]*100.0 - 50.0)])
                std_spring_hist[std_idx,imodel] = std_spring_hist[std_idx,imodel] + 1.0
                std_idx = np.amin([nxstd-1, np.int(std_normalized_summer[jy,ix]*100.0 - 50.0)])
                std_summer_hist[std_idx,imodel] = std_summer_hist[std_idx,imodel] + 1.0
                std_idx = np.amin([nxstd-1, np.int(std_normalized_fall[jy,ix]*100.0 - 50.0)])
                std_fall_hist[std_idx,imodel] = std_fall_hist[std_idx,imodel] + 1.0

fig1 = plt.figure(figsize=(8.5,8))
fig1.suptitle('Histograms of forecast vs. analyzed correlation, +'+clead+' h', size='xx-large')

# ---- plot the histograms in fig1

for iseason, season in enumerate(['Jan-Feb-Mar','Apr-May-Jun','Jul-Aug-Sep','Oct-Nov-Dec']):
    for imodel, model in enumerate(models):

        if season == 'Jan-Feb-Mar':
            corrcoef = corr_winter_hist[:,imodel]
            cletter = '(a) '
            axlocs = [0.08,0.54,0.4,0.36]
        elif season == 'Apr-May-Jun':
            corrcoef = corr_spring_hist[:,imodel]
            cletter = '(b) '
            axlocs = [0.57,0.54,0.4,0.36]
        elif season == 'Jul-Aug-Sep':
            corrcoef = corr_summer_hist[:,imodel]
            cletter = '(c) '
            axlocs = [0.08,0.06,0.4,0.36]
        elif season == 'Oct-Nov-Dec':
            corrcoef = corr_fall_hist[:,imodel]
            cletter = '(d) '
            axlocs = [0.57,0.06,0.4,0.36]
            print ('iseason, imodel, corr = ', iseason, imodel, corrcoef)

        if model == 'raw':
            ftype = 'raw_forecast_errorstats_2019_'
            label = 'Raw model output'
            color = 'Red'
            linestyle = '--'
        elif model ==  'dav':
            ftype = 'decayavg_forecast_errorstats_2019_'
            label = 'Decaying-average'
            color = 'Peru'
            linestyle = '-'
        elif model == 'qmap':
            ftype = 'quantile_mapping_errorstats_2019_'
            label = 'Quantile mapping'
            color = 'Blue'
            linestyle = '-.'
        elif model == 'mos':
            ftype = 'MOS_GEFSv12_2019_'
            label = 'Univariate MOS'
            color = 'Green'
            linestyle = '-'
        elif model == 'mvmos':
            ftype = 'MOS_mvr_GEFSv12_2019_'
            label = 'Multi-variate MOS'
            color = 'Gray'
            linestyle = ':'

        # --- add to correlation plots

        corrcoef = corrcoef / np.sum(corrcoef)
        corrcoef = ma.masked_where(corrcoef == 0.0, corrcoef)
        if imodel == 0: ax1 = fig1.add_axes(axlocs)
        if imodel == 0: ax1.set_title(cletter+season, fontsize=15)
        if season == 'Jul-Aug-Sep':
            ax1.plot(xcorr, corrcoef, color=color,linestyle=linestyle,lw=1.5,label=label)
        else:
            ax1.plot(xcorr, corrcoef, color=color,linestyle=linestyle,lw=1.5)
        if imodel == 0: ax1.set_xlim(0.7,1.0)
        if imodel == 0: ax1.set_xticks([0.7,0.75,0.8,0.85,0.9,0.95,1.0])
        #if imodel == 0: ax1.set_ylim(0.0,0.2)
        if imodel == 0: ax1.set_xlabel('Correlation',fontsize=12)
        if imodel == 0: ax1.set_ylabel('Relative frequency',fontsize=12)
        #if imodel == 0: ax1.set_yticks([0.001,0.005, 0.01,0.05, 0.1, 0.5])
        #if imodel == 0: ax1.set_yticks([0.0,0.04,0.08,0.12,0.16,0.2])
        if imodel == 0: ax1.grid(True, lw=0.25)

        if season == 'Jul-Aug-Sep' and model == 'mvmos':
            ax1.legend(loc=0)

# ---- set plot titles and write to files

outfile1 = 'corr_histogram_4panel_'+clead+'h.pdf'
print ('saving to ', outfile1)
plt.savefig(outfile1)

fig2 = plt.figure(figsize=(8.5,8))
fig2.suptitle('Histograms of forecast vs. analyzed normalized standard deviation, +'+clead+' h', size='x-large')

# ---- plot the histograms in fig1

for iseason, season in enumerate(['Jan-Feb-Mar','Apr-May-Jun','Jul-Aug-Sep','Oct-Nov-Dec']):

    for imodel, model in enumerate(models):

        if season == 'Jan-Feb-Mar':
            stddev = std_winter_hist[:,imodel]
            cletter = '(a) '
            axlocs = [0.08,0.54,0.4,0.37]
        elif season == 'Apr-May-Jun':
            stddev = std_spring_hist[:,imodel]
            cletter = '(b) '
            axlocs = [0.58,0.54,0.4,0.37]
        elif season == 'Jul-Aug-Sep':
            stddev = std_summer_hist[:,imodel]
            cletter = '(c) '
            axlocs = [0.08,0.06,0.4,0.37]
        elif season == 'Oct-Nov-Dec':
            stddev = std_fall_hist[:,imodel]
            cletter = '(d) '
            axlocs = [0.58,0.06,0.4,0.37]

        if model == 'raw':
            ftype = 'raw_forecast_errorstats_2019_'
            label = 'Raw model output +'+clead+' h'
            color = 'Red'
            linestyle = '--'
        elif model ==  'dav':
            ftype = 'decayavg_forecast_errorstats_2019_'
            label = 'Decaying-average'
            color = 'Peru'
            linestyle = '-'
        elif model == 'qmap':
            ftype = 'quantile_mapping_errorstats_2019_'
            label = 'Quantile mapping '+clead+' h'
            color = 'Blue'
            linestyle = '-.'
        elif model == 'mos':
            ftype = 'MOS_GEFSv12_2019_'
            label = 'Univariate MOS '+clead+' h'
            color = 'Green'
            linestyle = '-'
        elif model == 'mvmos':
            ftype = 'MOS_mvr_GEFSv12_2019_'
            label = 'Multi-variate MOS '+clead+' h'
            color = 'Gray'
            linestyle = ':'

        # --- add to stddev plots

        stddev = stddev / np.sum(stddev)
        stddev = ma.masked_where(stddev == 0.0, stddev)
        if imodel == 0: ax2 = fig2.add_axes(axlocs)
        if imodel == 0: ax2.set_title(cletter+season, fontsize=14)
        if season == 'Jan-Feb-Mar':
            ax2.semilogy(xstd, stddev, color=color,linestyle=linestyle,lw=1.5,label=label)
        else:
            ax2.semilogy(xstd, stddev, color=color,linestyle=linestyle,lw=1.5)
        if imodel == 0: ax2.set_xlim(0.5,1.5)
        if imodel == 0: ax2.set_ylim(0.001,0.1)
        if imodel == 0: ax2.set_xlabel(r'$\sigma_f$ / $\sigma_a$',fontsize=12)
        if imodel == 0: ax2.set_ylabel('Relative frequency',fontsize=12)
        if imodel == 0: ax2.set_xticks([0.5,0.75,1.0,1.25,1.5])
        if imodel == 0: ax2.set_yticks([0.001,0.005, 0.01,0.05, 0.1])
        if imodel == 0: ax2.grid(True, lw=0.25)

        if season == 'Jan-Feb-Mar' and model == 'mvmos':
            ax2.legend(loc=3)

outfile2 = 'stddev_histogram_4panel_'+clead+'h.pdf'
print ('saving to ', outfile2)
plt.savefig(outfile2)


"""
File: image_editor_thingy_oct28.py
Author: Xander Hunt
Purpose: Make a rudimentary green screen.
"""

# This line imports or includes the library we will need
from PIL import Image

# This line opens the image and loads it into a variable called "image_original"
image_original = Image.open("snowscape.jpg")
image_other = Image.open("cat_small.jpg")

# This line attempts to open a new window to display the image.
# image_original.show()

# Record the original image's dimensions as variables.
(width1, height1) = image_original.size
print("Background dimensions: ", width1, height1)
(width2, height2) = image_other.size
print("Other dimensions: ", width2, height2)

# Create a pixels variable, check it too.
pixels1 = image_original.load()
pixels2 = image_other.load()
print(pixels2[0, 0])

# Change the green background.
for column in range(width2):
    for row in range(height2):
        (r, g, b) = pixels2[column, row]
        if (g > 180) and (r < 100) and (b < 100):
            pixels2[column, row] = pixels1[column, row]

# Return the color value of the pixel at (400, 300).
#pixels = image_original.load()
#print(pixels[400, 300])

# Draw some neat diagonal lines.
"""y_val = 0
for column in range(width1):
    pixels[column, y_val] = (0, 0, 0)
    y_val += 1
    if y_val >= 600:
        y_val = 0

y_val = 200
for column in range(width1):
    pixels[column, y_val] = (0, 0, 0)
    y_val += 1
    if y_val >= 600:
        y_val = 0

y_val = 599
for column in range(width1):
    pixels[column, y_val] = (0, 0, 0)
    y_val -= 1
    if y_val <= 0:
        y_val = 599

y_val = 399
for column in range(width1):
    pixels[column, y_val] = (0, 0, 0)
    y_val -= 1
    if y_val <= 0:
        y_val = 599"""

# Stick a red line down the middle.
"""for row in range(height1):
    pixels[400, row] = (255, 0, 0)"""

# Put on a dark filter
"""for row in range(height1):
    for column in range(width1):
        (r,g,b) = pixels[column, row]
        if r + g + b < 400:
            pixels[column, row] = (0, 0, 0)"""

#image_original.show()
image_other.show()