#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 16 16:18:01 2019

@author: cacquist
"""

# coding: utf-8

# In[1]:

# ------------------------------------------------------------------------
# date      : 12.04.2018
# author    : Claudia Acquistapace
# goal      : routine to read 1D meteogram for a given date and site ( Joyce ) and extract data for the site and also
# level2 variables for the site Store them in a ncdf file to be copied on ostro for comparison 1to1 with observations from the ground
# DAYS WITH BOUNDARY LAYER CLOUDS OF INTEREST:
# - 20130502 (folder 20130502-default )
# - 20130505 (folder 20130505-default-redone_v1)
# - 20130511 (folder 20130511-default )
# - 20160603 (folder 20160603-default-redone_v2 )

# ------------------------------------------------------------------------

# In[1]:

# ---- importing libraries
import numpy as np
import matplotlib
import scipy
import numpy.ma as ma
import pandas as pd
import netCDF4 as nc4
import glob
from netCDF4 import Dataset
import matplotlib.dates as mdates
import xarray as xr
from myFunctions import f_closest
import matplotlib.pyplot as plt
from myFunctions import f_calcPblHeightRN
from myFunctions import f_calcWvariance
from myFunctions import f_runningMeanSkewnessVarianceStd_W
from myFunctions import f_PBLClass
from myFunctions import f_calcCloudBaseTopPBLcloudsV2
from myFunctions import f_calcCloudBaseTopPBLclouds
from myFunctions import f_calcPblHeightTW
from myFunctions import f_cloudmask
from myFunctions import f_calcWindSpeed_Dir
from myFunctions import f_calculateCloudBaseTopThickness

def f_processModelOutput(path_icon, \
                         iconFilename, \
                         modelInputParameters, \
                         date, \
                         humanInfo, \
                         debuggingFlag, \
                         verboseFlag, \
                         pathDebugFig, \
                         pathOut, \
                         domSel):

    print('processing meteograms for the '+date)

    # ---- reading datafile selected
    data          = Dataset(path_icon+iconFilename, mode='r')
    time          = data.variables['time'][:].copy()
    datetime_ICON = nc4.num2date(data.variables['time'][:],data.variables['time'].units)
    Qi            = data.variables['QI'][:].copy()
    Qc            = data.variables['QC'][:].copy()
    T             = data.variables['T'][:].copy() # in [K]
    zonalWind     = data.variables['U'][:].copy()
    merWind       = data.variables['V'][:].copy()
    vertWind      = data.variables['W'][:].copy()
    LWP           = data.variables['TQC'][:].copy()
    IWV           = data.variables['TQV'][:].copy()
    thetaV        = data.variables['THETAV'][:].copy()
    height        = data.variables['height'][:].copy()
    P             = data.variables['P'][:].copy()        # [Pa]
    RH            = data.variables['REL_HUM'][:].copy()
    q             = data.variables['QV_DIA'][:].copy()   # [kg/kg]
    Hsurf         = float(data.station.split('_hsurf=')[-1].split('\n')[0])
    height2       = data.variables['height_2'][:].copy()
    rho           = data.variables['RHO'][:].copy()
    SWSurfFlux    = data.variables['SOBS'][:].copy() # shortwave net flux at surface
    LWSurfFlux    = data.variables['THBS'][:].copy() # longwave net flux at surface
    LHFL          = data.variables['LHFL'][:].copy() # latent heat flux (surface)
    SHFL          = data.variables['SHFL'][:].copy() # sensible heat flux (surface)
    TempSurf      = data.variables['T_S'][:]

    #print(Hsurf)
    #print(height2[-1])
    #print(height2[0])
    #print(len(height2))
    # subtracting from model height arrays the height of the ground level at JOYCE
    # and make it comparable with the observations
    height2 = height2 - np.repeat(Hsurf, len(height2))
    height  = height -np.repeat(Hsurf, len(height))

    # --- reading dimension of height and time arrays
    dimTime       = len(datetime_ICON)
    dimHeight     = len(height2)

    if verboseFlag == 1:
        print('variable extracted from the data')
        print('data loaded for '+date)
        print('dimension for height_2 :', dimHeight)
        print('dimension for time :', dimTime)

    # ------------------------------------------------------------------
    # plot meridional and zonal wind for checking fields
    # ------------------------------------------------------------------
    if debuggingFlag == 1:
        if verboseFlag == 1:
            print('no plots')

# =============================================================================
#         fig, ax = plt.subplots(figsize=(14,6))
#         ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M"))
#         ax.xaxis.set_minor_formatter(mdates.DateFormatter("%H:%M"))
#         ax.spines["top"].set_visible(False)
#         ax.spines["right"].set_visible(False)
#         ax.get_xaxis().tick_bottom()
#         ax.get_yaxis().tick_left()
#         ax.xaxis_date()
#         ax.set_xlabel("time [hh:mm]", fontsize=16)
#         ax.set_ylabel("Fluxes at the surface [W/m2]", fontsize=16)
#         plt.plot(datetime_ICON, SWSurfFlux, label='Shortwave net flux')
#         plt.plot(datetime_ICON, LWSurfFlux, label='Longwave net flux')
#         plt.legend()
#         plt.savefig(pathDebugFig+'surface_LWSW_surfFlux_iconlem_'+date+'.png', format='png')
#
# =============================================================================

# =============================================================================
#         fig, ax = plt.subplots(figsize=(14,6))
#         ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M"))
#         ax.xaxis.set_minor_formatter(mdates.DateFormatter("%H:%M"))
#         ax.spines["top"].set_visible(False)
#         ax.spines["right"].set_visible(False)
#         ax.get_xaxis().tick_bottom()
#         ax.get_yaxis().tick_left()
#         ax.xaxis_date()
#         ax.set_xlabel("time [hh:mm]", fontsize=16)
#         ax.set_ylabel("Latent/sensible heat fluxes at the surface [W/m2]", fontsize=16)
#         plt.plot(datetime_ICON, LHFL, label='Latent heat flux')
#         plt.plot(datetime_ICON, SHFL, label='Sensible heat flux')
#         plt.legend()
#         plt.savefig(pathDebugFig+'surface_LatentSensible_heatFlux_iconlem_'+date+'.png', format='png')
# =============================================================================
        if verboseFlag == 1:
            print('end of plotting graphs for debugging in debugging mode')

    # ------------------------------------------------------------------
    # defining constants needed for calculations
    # ------------------------------------------------------------------
    Rw    = 462.
    Rl    = 287.
    g     = 9.81
    P_0   = 100*1000.
    const = 0.286 # R/Cp
    P_0   = 100*1000.
    const = 0.286 # R/Cp
    Lv    = 2260 # J / kg
    Cp    = 1005.7 # /K Kg

    # ------------------------------------------------------------------
    # derivation of water vapor mixing ratio
    # ------------------------------------------------------------------
    r     = np.zeros((dimTime, dimHeight))
    for itempo in range(dimTime):
        for ih in range(dimHeight):
            r[itempo,ih] = q[itempo,ih]/(1. - q[itempo,ih] )
    if verboseFlag == 1:
        print('water vapor mixing ratio calculated')

    # ------------------------------------------------------------------
    # --- calculating cloud mask for ice and liquid clouds using thresholds on Qi, Qc
    # ------------------------------------------------------------------
    QcThreshold            = modelInputParameters['QcThresholdVar']
    QiThreshold            = modelInputParameters['QiThresholdVar']
    cloudMask              = f_cloudmask(time,height2,Qc,Qi,QiThreshold,QcThreshold)

# =============================================================================
#
#     for indT in range(len(datetime_ICON)):#
#         if (~np.isnan(CBMatrix_ICON[indT,0]) == True) and (~np.isnan(CTMatrix_ICON[indT,0])== True):
#
#             indCB          = f_closest(height, CBMatrix_ICON[indT,0])
#             indCT          = f_closest(height, CTMatrix_ICON[indT,0])
#
#             if (indCB == 0) or (indCT == 0):
#                 CT_array_ICON[indT] = np.nan
#                 CB_array_ICON[indT] = np.nan
#             else:
#                 CT_array_ICON[indT] = height[indCT]                                 # saving cloud top height
#                 CB_array_ICON[indT] = height[indCB]                                 # saving cloud base height
#
# =============================================================================
    # calculating cloud base , cloud top and cloud thickness for all clouds and for pbl clouds
    clouds, PBLclouds = f_calculateCloudBaseTopThickness(cloudMask, datetime_ICON, height2, humanInfo)
    # deriving lowest cloud base and corresponding cloud top for PBL clouds
    CBarr = np.zeros(dimTime)
    CBarr.fill(np.nan)
    CTarr = np.zeros(dimTime)
    CTarr.fill(np.nan)
    iPBL = 0
    for itime in range(dimTime):
        if iPBL < len(PBLclouds.time.values):
            if clouds.time.values[itime] == PBLclouds.time.values[iPBL]:
                CBarray = PBLclouds.cloudBase.values[iPBL, :]
                if CBarray.size - np.count_nonzero(np.isnan(CBarray)) != 0:
                    minCB = np.nanmin(PBLclouds.cloudBase.values[iPBL, :])
                    CBarr[itime] = minCB
                    indexLevelMin = np.nanargmin(PBLclouds.cloudBase.values[iPBL, :])
                    CTarr[itime] = PBLclouds.cloudTop[iPBL, indexLevelMin]
                iPBL = iPBL + 1

    print('cloud base and cloud top for ICON-LEM calculated ')

    # ------------------------------------------------------------------
    # ---- calculating potential temperature and equivalent potential temperature
    # ------------------------------------------------------------------
    theta   = np.zeros((dimTime, dimHeight))
    theta_e = np.zeros((dimTime, dimHeight))
    theta.fill(np.nan)
    theta_e.fill(np.nan)

    for iTime in range(dimTime):
        for iHeight in range(dimHeight):
            if height[iHeight] < Hsurf:
                theta[iTime, iHeight] = 0.
            else:
                theta[iTime, iHeight] = T[iTime, iHeight] * (float(P_0)/float(P[iTime, iHeight]))**(const)
    if verboseFlag == 1:
        print('potential temperature calculated')

    for iTime in range(dimTime):
        for iHeight in range(dimHeight):
            lv = (2500.-2.42*(T[iTime, iHeight]-273.15))*1000. # latent heat of vaporization in J/kg
            theta_e[iTime, iHeight] = theta[iTime, iHeight]+(lv*r[iTime, iHeight]/Cp)*         (np.power(100000./P[iTime, iHeight], Rl/Cp)) # equivalent potential temperature in K
    if verboseFlag == 1:
        print('equivalent potential temperature calculated')

    # ------------------------------------------------------------------
    # --- Calculating Boundary layer height using the richardson number derivation according to Seidel Et al, 2010
    # ------------------------------------------------------------------
    device = 'mod'
    PBLHeightArrRN    = f_calcPblHeightRN(thetaV,zonalWind,merWind,height2,time, device)
    if verboseFlag == 1:
        print('height of the PBL (RN) calculated')

    # ------------------------------------------------------------------
    # --- calculation of the variance, std, skewness of the vertical velocity using a running mean window
    # ------------------------------------------------------------------
    timeWindowSk    = modelInputParameters['timeWindowSkVar']
    runningWindow   = modelInputParameters['runningWindowVar']
    resultDyn       = f_runningMeanSkewnessVarianceStd_W(time, timeWindowSk, runningWindow, height2, vertWind)
    # output of the function : varianceW, stdWmatrix, SKmatrix
    varianceWmatrix = resultDyn[0]
    stdWmatrix      = resultDyn[1]
    SKmatrix        = resultDyn[2]

    if verboseFlag == 1:
        print('variance, std and skewness of w calculated')

    print('std max = '+str(np.nanmax(stdWmatrix)))
    # ------------------------------------------------------------------
    # --- Calculating Boundary layer height using the threshold on variance of w ()
    # ------------------------------------------------------------------
    device         = 'mod'
    sigmaW         = stdWmatrix
    sigmaThreshold = modelInputParameters['SigmaWThresStd'] #  m/s, threshold for std of w from Schween et al, 2014.AMT
    PBLHeightArrTW = f_calcPblHeightTW(sigmaW,sigmaThreshold,height2,time, device)
    if verboseFlag == 1:
        print('height of the PBL (TW) calculated')

    # ------------------------------------------------------------------
    # --- Calculating variance over the timewindow using running mean
    # ------------------------------------------------------------------
    #timewindow      = modelInputParameters['timewindowVar']
    #varianceWmatrix = f_calcWvariance(vertWind,time,height2,timewindow)
    #if verboseFlag == 1:
    #    print('variance of vertical velocity calculated')

    # ------------------------------------------------------------------
    # --- calculation of the connection of the turbulence to the surface.
    # ------------------------------------------------------------------
    #Turbulence is connected to the surface if checks if variance at 200 m of height is greater than 0.03 for turbulence
    # calculating the time serie of difference of the sigmaW and the threshold value at 200 m height
    deltaSigma         = np.subtract(varianceWmatrix, 0.03)[:,f_closest(height,200.)]
    connection2Surface = []                    # array indicating connection of the turbulence to the surface
    # calculating connection to the surface. =0 ( not connected, if sigmaW(200)-sigmaGround)<0,
    # =1 (connected thus turbulent, if sigmaW(200)-sigmaGround)>0)
    for itime in range(dimTime):
        if deltaSigma[itime] < 0.:
            connection2Surface.append(0)
        else:
            connection2Surface.append(1)
    if verboseFlag == 1:
        print('connection of turbulence with the surface calculated')

    # ------------------------------------------------------------------
    #  ---- calculation of the stability array
    # ------------------------------------------------------------------
    stabilityArr = []
    # difference of temperature between 150m and closest level to surface
    deltaT       = np.subtract(T, T[f_closest(height,Hsurf),:])[:,f_closest(height,150.)]
    for itime in range(dimTime):

        #print(Tarray[indRef]-Tarray[indGround])
        if deltaT[itime] < 0.3:
            stabilityArr.append(1)
        else:
            stabilityArr.append(0)
    if verboseFlag == 1:
        print('stability at the surface calculated')

    # ------------------------------------------------------------------
    # --- Calculation of wind shear as done for PBL  ( running mean over 30 min of sqrt(Delta U^2 + delta V^2))/delta H
    # where variations are calculated over 5 range gates
    # ------------------------------------------------------------------
    windData      = f_calcWindSpeed_Dir(datetime_ICON, height2, zonalWind, merWind)
    windSpeed     = windData['windSpeed']
    windDirection = windData['windDirection']

# =============================================================================
# --- calculating shear of horizontal wind
    u_rm = np.zeros((len(datetime_ICON), len(height2)))
    v_rm = np.zeros((len(datetime_ICON), len(height2)))
# --- defining running mean values of zonal and meridional wind
    for indH in range(0,len(height2)):
         zonal = pd.Series(zonalWind[:,indH])
         mer   = pd.Series(merWind[:,indH])
         #u_rm[:,indH] = pd.rolling_mean(zonalWind[:,indH], window=200)
         #v_rm[:,indH] = pd.rolling_mean(merWind[:,indH], window=200)
         u_rm[:,indH] = zonal.rolling(200).mean()
         v_rm[:,indH] = mer.rolling(200).mean()
#
# calculating wind shear and horizontal wind
    shear_ICON = np.zeros((len(datetime_ICON), len(height2)))
    for indT in range(0,len(datetime_ICON)):
        for indH in range(0,len(height2)):

             if (indH < 2.) or (indH > len(height2)-3):
                 shear_ICON[indT, indH] = 0.
             else:
                 deltaV = (np.absolute(v_rm[indT, indH+2] - v_rm[indT, indH-2]))**2
                 deltaU = (np.absolute(u_rm[indT, indH+2] - u_rm[indT, indH-2]))**2
                 deltaH = np.absolute(height[indH+2] - height[indH-2])
                 shear_ICON[indT, indH] = (np.sqrt(deltaU + deltaV))/deltaH
# =============================================================================
    if verboseFlag == 1:
        print('horizontal wind speed, direction and shear calculated')
    # ------------------------------------------------------------------
    # ----calculating boundary layer classification (version from submitted paper)
    # ------------------------------------------------------------------
    ylim        = np.repeat(3000, dimTime)     # defining array of heights up to which PBL classification is calculated
    gradWindThr = 0.01
    SigmaWThres = 0.2

    outputClass = f_PBLClass(datetime_ICON, \
                             height2, \
                             gradWindThr, \
                             SigmaWThres, \
                             ylim, \
                             cloudMask, \
                             varianceWmatrix, \
                             SKmatrix, \
                             stabilityArr, \
                             connection2Surface, \
                             shear_ICON, \
                             CBarr)
    PBLclass     = outputClass[0]

    if verboseFlag == 1:
        print('PBL classification calculated')

    if debuggingFlag == 1:
        print('dimensions of PBL class')
        print(np.shape(PBLclass))

# =============================================================================
#         # plotting classification
#         fig, ax = plt.subplots(figsize=(10,4))
#         ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M"))
#         ax.xaxis.set_minor_formatter(mdates.DateFormatter("%H:%M"))
#         ax.spines["top"].set_visible(False)
#         ax.spines["right"].set_visible(False)
#         ax.get_xaxis().tick_bottom()
#         ax.get_yaxis().tick_left()
#         ax.xaxis_date()
#         cax = ax.pcolormesh(datetime_ICON, height2, PBLclass.transpose(), vmin=0., vmax=6., cmap=plt.cm.get_cmap("jet", 7))
#         ax.set_ylim(Hsurf,3000.)                                               # limits of the y-axes
#         #ax.set_xlim(0,24)                                 # limits of the x-axes
#         ax.set_title("PBL classification", fontsize=14)
#         ax.set_xlabel("time [UTC]", fontsize=12)
#         ax.set_ylabel("height [m]", fontsize=12)
#         cbar = fig.colorbar(cax, ticks=[0, 1, 2, 3, 4, 5, 6], orientation='vertical')
#         cbar.ticks=([0,1,2,3,4,5,6])
#         cbar.ax.set_yticklabels(['no class','in cloud','non turb','cloud driven','convective', 'intermittent','wind shear'])
#         cbar.set_label(label="PBL type",size=12)
#         cbar.ax.tick_params(labelsize=12)
#         cbar.aspect=20
#         plt.savefig(pathDebugFig+'PBLclassification_iconlem_'+date+'.png', format='png')
# =============================================================================

    # ------------------------------------------------------------------
    # --- calculation of the LCL
    # ------------------------------------------------------------------
    # determining P, T and RH at the surface
    Psurf    =  data.variables['P_SFC'][:].copy()
    Tsurf    = data.variables['T2M'][:].copy()
    RHsurf   = RH[:,149]
    LCLarray = []
    from myFunctions import lcl
    for iTime in range(dimTime):
        LCLarray.append(lcl(Psurf[iTime],Tsurf[iTime],RHsurf[iTime]/100.))
    if verboseFlag == 1:
        print('LCL calculated')

    # ------------------------------------------------------------------
    # calculate LTS index for lower tropospheric stability (Wood and Bretherton, 2006)
    # ------------------------------------------------------------------
    LTS  = np.zeros(dimTime)
    H700 = np.zeros(dimTime)
    Pthr = 700 * 100. # Pressure level of 700 Hpa used as a reference
    # calculating height of the surface

    indSurf = 146# f_closest(height,Hsurf)
    for iTime in range(dimTime):
        indP700 = f_closest(P[iTime,:],Pthr)
        LTS[iTime] = theta[iTime, indP700] - theta[iTime, indSurf]
        H700[iTime] = height[indP700]
    if verboseFlag == 1:
        print('LTS calculated')
    #print(theta[4500, indP700])
    #print(theta[4500, indSurf])
    #print(theta[4500, :])

    # ------------------------------------------------------------------
    # ---- calculating liquid potential temperature
    # ------------------------------------------------------------------
    theta_liquid = np.zeros((dimTime, dimHeight))
    theta_liquid.fill(np.nan)

    for iTime in range(dimTime):
        for iHeight in range(dimHeight):
            if height[iHeight] < Hsurf:
                theta_liquid[iTime, iHeight] = 0.
            else:
                theta_liquid[iTime, iHeight] = theta[iTime, iHeight] - (Lv/Cp)* Qc[iTime, iHeight]

    if verboseFlag == 1:
        print('liquid potential temperature calculated')

    # ------------------------------------------------------------------
    # ------- saving mean outputs as ncdf file
    # ------------------------------------------------------------------

    f = nc4.Dataset(pathOut+'icon_lem_derivedproperties'+date+'.nc','w', format='NETCDF4')      # creates a netCDF file for writing
    tempgrp = f.createGroup('Temp_data')                                                    # creates a group: A netCDF group is basically a directory or folder within the netCDF dataset

    # specify dimensions of the data ( each dimension of multidimensiona array needs to be given a name and a length)
    tempgrp.createDimension('dimH', len(height2))                                 # dimension for height
    tempgrp.createDimension('dimHlong', len(height))                                 # dimension for height

    tempgrp.createDimension('dimHsurf', 1)                                         # dimension for scalar values
    tempgrp.createDimension('dimT', len(datetime_ICON))                                   # dimension for time
    tempgrp.createDimension('NclassesPBL', 8)                # dimension for the number of cloud layers found
    tempgrp.createDimension('dimHlarger', len(height))                                 # dimension for height
    tempgrp.createDimension('nchar', 5)

    # preallocating netCDF variables for data storage
    varHeight2             = tempgrp.createVariable('height2', 'f4', 'dimH')
    varHeight             = tempgrp.createVariable('height', 'f4', 'dimHlong')
    vardomain             = tempgrp.createVariable('domain', 'S1', 'nchar')
    vartime               = tempgrp.createVariable('datetime_ICON', 'f4', 'dimT')
    varLTS                = tempgrp.createVariable('LTS', 'f4', 'dimT')
    varPBLheight          = tempgrp.createVariable('PBLHeightArrRN', 'f4', 'dimT')
    varPBLheight2         = tempgrp.createVariable('PBLHeightArrTW', 'f4', 'dimT')
    varCloudLayers        = tempgrp.createVariable('NcloudLayers', 'f4', 'dimT')
    varHsurf              = tempgrp.createVariable('HeightSurface', 'f4', 'dimHsurf')
    varLCL                = tempgrp.createVariable('LCLarray', 'f4', 'dimT')
    varLWP                = tempgrp.createVariable('LWP', 'f4', 'dimT')
    varIWV                = tempgrp.createVariable('IWV', 'f4', 'dimT')
    varLHFL               = tempgrp.createVariable('LHFL', 'f4', 'dimT')
    varSHFL               = tempgrp.createVariable('SHFL', 'f4', 'dimT')
    varLWSF               = tempgrp.createVariable('LWSurfFlux', 'f4', 'dimT')
    varSWSF               = tempgrp.createVariable('SWSurfFlux', 'f4', 'dimT')
    # PBL class and connected flags, LTS clouds, SW clouds, PBL height, CB height
    varPBL_class          = tempgrp.createVariable('PBLclass', 'f4', ('dimT','dimH'))
    varflagCloud          = tempgrp.createVariable('flagCloud', 'f4', ('dimT','dimH'))
    varQc                 = tempgrp.createVariable('Qc', 'f4', ('dimT','dimH'))
    varQi                 = tempgrp.createVariable('Qi', 'f4', ('dimT','dimH'))
    varflagTurb           = tempgrp.createVariable('flagTurb', 'f4', ('dimT','dimH'))
    varflagcloudDriven    = tempgrp.createVariable('flagcloudDriven', 'f4', ('dimT','dimH'))
    varflagInstability    = tempgrp.createVariable('flagInstability', 'f4',('dimT','dimH'))
    varflagWindShear      = tempgrp.createVariable('flagWindShear', 'f4', ('dimT','dimH'))
    varflagSurfDriven     = tempgrp.createVariable('flagSurfaceDriven', 'f4', ('dimT','dimH'))
    varvarianceW          = tempgrp.createVariable('varianceW', 'f4', ('dimT','dimH'))
    varHwind              = tempgrp.createVariable('windSpeed', 'f4', ('dimT','dimH'))
    varWindDirection      = tempgrp.createVariable('windDirection', 'f4', ('dimT','dimH'))
    varShearHwind         = tempgrp.createVariable('shearHwind', 'f4', ('dimT','dimH'))
    varcloudMask          = tempgrp.createVariable('cloudMask', 'f4', ('dimT','dimH'))
    varthetaPot           = tempgrp.createVariable('theta', 'f4', ('dimT','dimH'))
    varskewnessW          = tempgrp.createVariable('skewnessW', 'f4', ('dimT','dimH'))
    varstdWmatrix         = tempgrp.createVariable('stdWmatrix', 'f4', ('dimT','dimH'))
    varMixingRatio        = tempgrp.createVariable('r', 'f4', ('dimT','dimH'))
    varthetaL             = tempgrp.createVariable('theta_liquid', 'f4', ('dimT','dimH'))
    varthetaPot_e         = tempgrp.createVariable('theta_e', 'f4', ('dimT','dimH'))
    varw                  = tempgrp.createVariable('vertWind', 'f4', ('dimT','dimHlarger'))
    varP                  = tempgrp.createVariable('P', 'f4', ('dimT','dimH'))
    varRH                 = tempgrp.createVariable('RH', 'f4', ('dimT','dimH'))
    varQ                  = tempgrp.createVariable('q', 'f4', ('dimT','dimH'))
    varT                  = tempgrp.createVariable('T', 'f4', ('dimT','dimH'))
    varMerWind            = tempgrp.createVariable('merWind', 'f4', ('dimT','dimH'))
    varZonWind            = tempgrp.createVariable('zonalWind', 'f4', ('dimT','dimH'))
    varRho                = tempgrp.createVariable('rho', 'f4', ('dimT','dimH'))
    varT_surf             = tempgrp.createVariable('TempSurf', 'f4', 'dimT')

    # passing data into the variables
    varHeight2[:]         = height2
    varHeight[:]          = height
    vardomain             = domSel
    vartime[:]            = time
    varLTS[:]             = LTS
    varPBLheight[:]       = PBLHeightArrRN
    varPBLheight2[:]      = PBLHeightArrTW
    varHsurf              = Hsurf
    varLCL[:]             = LCLarray
    varLWP[:]             = LWP
    varIWV[:]             = IWV
    varLHFL[:]            = LHFL
    varSHFL[:]            = SHFL
    varLWSF[:]            = LWSurfFlux
    varSWSF[:]            = SWSurfFlux
    varPBL_class[:,:]     = PBLclass
    varflagCloud[:]       = outputClass[1]
    varflagTurb[:]        = outputClass[2]
    varflagcloudDriven[:] = outputClass[3]
    varflagInstability[:] = outputClass[4]
    varflagWindShear[:]   = outputClass[5]
    varflagSurfDriven[:]  = outputClass[6]
    varvarianceW[:,:]     = varianceWmatrix
    varHwind[:,:]         = windSpeed
    varWindDirection[:,:] = windDirection
    varShearHwind[:,:]    = shear_ICON
    varcloudMask[:,:]     = cloudMask
    varthetaPot[:,:]      = theta
    varskewnessW[:,:]     = SKmatrix
    varstdWmatrix[:,:]    = stdWmatrix
    varMixingRatio[:,:]   = r
    varthetaL[:,:]        = theta_liquid
    varthetaPot_e[:,:]    = theta_e
    varw[:,:]             = vertWind
    varP[:,:]             = P
    varRH[:,:]            = RH
    varQ[:,:]             = q
    varT[:,:]             = T
    varMerWind[:,:]       = merWind
    varZonWind[:,:]       = zonalWind
    varRho[:,:]           = rho
    varQc[:,:]            = Qc
    varQi[:,:]            = Qi
    varT_surf[:]          = TempSurf

    #Add global attributes
    f.description = "icon lem model derived physical quantities and PBL classification"
    f.history = "Created by Claudia Acquistapace cacquist@meteo.uni-koeln.de - University of Cologne"

    #Add local attributes to variable instances
    varPBL_class.units = '1=in cloud, 2=non turb, 3=cloud driven, 4=convective, 5=intermittent, 6=wind shear'
    vartime.units = 'seconds since '+date[0:4]+'-'+date[4:6]+'-'+date[6:8]+' 00:00:00'

    # closing ncdf file
    f.close()

    print('File Saved ')


"""Change column distance_bin to distance_cat

Revision ID: 2524785502b4
Revises: c137e7385dd7
Create Date: 2020-03-20 16:47:15.648707

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '2524785502b4'
down_revision = 'c137e7385dd7'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('race', sa.Column('distance_cat', sa.String(), nullable=True))
    op.create_foreign_key(None, 'race', 'runner_contact', ['runner_contact_id'], ['id'])
    op.drop_column('race', 'distance_bin')
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('race', sa.Column('distance_bin', sa.VARCHAR(), autoincrement=False, nullable=True))
    op.drop_constraint(None, 'race', type_='foreignkey')
    op.drop_column('race', 'distance_cat')
    # ### end Alembic commands ###


import re
import logging
import ROOT
import plottingConfig as cfg

class Config(cfg.PlottingConfig):

    def __init__ (self, options):
        self.options = options
        super(Config, self).__init__()

        sigma = 1 # at mu=1 (arbitrary for AZh)
        sigma_units = 'fb'
        # self.force_mu = (True, 0.16) # 700 GeV
        self.force_mu = (True, 10) # 600 GeV

        # for child classes to use
        # self.loggingLvl = logging.INFO
        self.loggingLvl = logging.DEBUG
        self.verbose = False
        self.formats = [ 'eps', 'pdf', 'png', 'root', 'C' ]
        self.blind = True
        self.thresh_drop_legend = 0.01
        self.restrict_to = []
        self.excludes = []
        self.additionalPlots = []
        self.add_sig_to_ratio_plot = True
        self.use_exp_sig = True
        # self.transferResults_fitName = "HiggsNorm"
        # self.get_binning_hist_removal = ["_meas2l2q2v2q"]
        self.bkg_substr_name = "Diboson"
        self.bkg_substr_list = ["diboson", "Diboson", "WZ", "ZZ", "VZ"]
        self.file_tags = ["Y", "L", "J", "T", "TType", "Flv", "Sgn", "isMVA", "dist", "Spc", "D", "nAddTag", "BMax", "BMin", "Fat", "incFat", "incJet", "incAddTag"]
        self.weight_tags = ["Higgsweighted", "Dibosonweighted"]
        self.sig_names = ["VH"]
        self.signal = ["A#rightarrow Zh (best fit)", self._STACK, ROOT.kRed + 1, 1] # last = mult factor
        self.expected_signal = ["VHbb", self._STACK, ROOT.kRed +1, self.force_mu[1]] # last = expected mu
        #self.expected_signal = ["A#rightarrow Zh (#sigma={0} {1})".format(int(sigma*self.force_mu[1]), sigma_units), self._STACK, ROOT.kRed +1, self.force_mu[1]] # last = expected mu
        # self.additional_signal = ["A#rightarrow Zh", self._OVERPRINT, ROOT.kRed +1, 1.]
        self.bkg_tuple = {'ttbar': ("t#bar{t}", 42, ROOT.kOrange, []),
                'stopt': ("t, s+t chan", 41, ROOT.kOrange - 1, ["stops"]),
                'stops': ("t, s+t chan", 41, ROOT.kOrange - 1, ["stopt"]),
                'stopWt': ("Wt", 40, ROOT.kYellow - 7, []),
                'stop': ("Single top", 40, ROOT.kOrange - 1, []),
                'Zbb': ("Z+bb", 25, ROOT.kAzure + 3, []),
                'Zbc': ("Z+bc", 24, ROOT.kAzure + 2, []),
                'Zclbl': ("Z+(bl,cl)", 23, ROOT.kAzure + 1, []),
                'Zbl': ("Z+bl", 23, ROOT.kAzure + 1, []),
                'Zcl': ("Z+cl", 21, ROOT.kAzure - 8, []),
                'Zcc': ("Z+cc", 22, ROOT.kAzure - 4, []),
                'Zhf': ("Z+(bb,bc,cc)", 22, ROOT.kAzure + 2, []),
                'Zl': ("Z+l", 20, ROOT.kAzure - 9, []),
                'Wbl': ("W+bl", 33, ROOT.kGreen + 2, []),
                'Wbb': ("W+bb", 35, ROOT.kGreen + 4, []),
                'Wbc': ("W+bc", 34, ROOT.kGreen + 3, []),
                'Wcc': ("W+cc", 32, ROOT.kGreen + 1, []),
                'Whf': ("W+(bb,bc,cc,bl)", 32, ROOT.kGreen + 3, []),
                'Wcl': ("W+cl", 31, ROOT.kGreen - 6, []),
                'Wl': ("W+l", 30, ROOT.kGreen - 9, []),
                'WZ': ("WZ", 53, ROOT.kGray + 1, ["ZZ"]),
                'ZZ': ("ZZ", 52, ROOT.kGray + 1, ["WZ"]),
                'VZ': ("VZ", 51, ROOT.kGray + 1, []),
                'diboson': ("Diboson", 51, ROOT.kGray + 1, []),
                'WW': ("WW", 50, ROOT.kGray + 3, []),
                'Diboson': ("Diboson", 50, ROOT.kGray + 1, []),
                #'VH125': ("Vh", 49, ROOT.kRed - 6, []),
                'multijet': ("Multijet", 39, ROOT.kViolet-9, ["multijetMu", "multijetEl"]),
                'multijetEl': ("Multijet", 39, ROOT.kViolet-9, ["multijetMu", "multijet"]),
                'multijetMu': ("Multijet", 39, ROOT.kViolet-9, ["multijetEl", "multijet"])}
        # self.ATLAS_suffix = "Internal"
        # self.ATLAS_suffix = "Simulation"
        self.ATLAS_suffix = "Preliminary"
        # self.ATLAS_suffix = ""
        # for yields
        self.make_slides = False
        self.window = None
        self.priorities = {
        "data" : 80,
        "S/sqrt(S+B)" : 73,
        "S/B" : 72,
        "Bkg" : 60,
        "MC" : 75,
        "SignalExpected" : 71,
        "Signal" : 70,
        "VH125" : 57,
        "ZvvH125" : 67,
        "ggZvvH125" : 67,
        "qqZvvH125" : 67,
        "WlvH125" : 68,
        "ZllH125" : 69,
        "ggZllH125" : 69,
        "qqZllH125" : 69,
        "ZvvH150" : 67,
        "ggZvvH150" : 67,
        "qqZvvH150" : 67,
        "WlvH150" : 68,
        "ZllH150" : 69,
        "AZhllbb1200" : 70,
        "AZhvvbb1200" : 70,
        "AZhllbb1000" : 70,
        "AZhvvbb1000" : 70,
        "AZhllbb400" : 70,
        "AZhvvbb400" : 70,
        "AZhllbb300" : 70,
        "AZhvvbb300" : 70,
        "AZhllbb600" : 70,
        "AZhvvbb600" : 70,
        "bbAZhllbb600" : 70,
        "bbAZhvvbb600" : 70,
        "ggZllH150" : 69,
        "qqZllH150" : 69,
        "ttbar" : 45,
        "stops" : 43,
        "stopt" : 42,
        "stopst" : 41,
        "stopWt" : 40,
        "stop" : 40,
        "Zhf" : 27,
        "Zb" : 24,
        "Zbl" : 25,
        "Zbb" : 27,
        "Zbc" : 26,
        "Zc" : 21,
        "Zcl" : 100,
        "Zclbl" : 22,
        "Zcc" : 23,
        "Zl" : 20,
        "Whf" : 37,
        "Wb" : 34,
        "Wbl" : 35,
        "Wbb" : 37,
        "Wbc" : 36,
        "Wcc" : 33,
        "Wc" : 31,
        "Wcl" : 32,
        "Wl" : 30,
        "WZ" : 53,
        "ZZ" : 52,
        "VZ" : 51,
        "WW" : 50,
        "Diboson" : 50,
        "diboson" : 50,
        "multijet" : 45,
        "multijetEl" : 45,
        "multijetMu" : 45,
        "MJ0lep" : 45,
        "MJ1lep" : 45,
        "MJ2lep" : 45,
        "MJ2lepEl" : 45,
        "MJ1lepEl" : 45,
        "MJ1lepMu" : 45,
        }
        # for reduced diag plots only
        self.exclude_str = 'HiggsNorm'
        self.cov_classification = {
        "BTag": [False, ["SysFT_EFF_Eigen", "SysFT_EFF_extrapolation"], []],
        "Top": [False, ["SysWt", "SysTop", "SysTtbar", "SysMVH"], []],
        "ModelBoson": [False, ["SysVV", "SysWM","SysZM","SysWD","SysZD","SysWP","SysZP","SysVj"], []],
        "Norm": [False, ["Norm","Ratio"], []],
        "norm": [False, ["norm"], []],
        "Lepton": [False, ["SysMUON","SysEL","SysEG"], []],
        "Jet": [False, ["SysJET","FATJET"], []],
        "MET": [False, ["SysMET"], []],
        "LUMI": [False, ["LUMI"], []],
        "Shifted": [True, [], ["blablabla"]]
        }
        self.cov_special = {
        "noMCStat": [[], ["gamma"]],
        "JES": [["SigX", "norm_", "Jet"], []],
        "BTag": [["SigX", "norm_", "BTag"], []],
        "Mbb": [["SigX", "norm_", "Mbb"], []],
        "Modelling": [["SigX", "norm_", "Norm", "Ratio", "PtBi"], []],
        "SF": [["SigX", "norm_"], []],
        "Norm": [["3JNorm", "norm_", "Norm", "Ratio"], []]
        }
        self.syst_to_study = ["JetEResol", "Mbb_Whf", "V_Whf", "METScale", "TChanP",
        "ttbarHigh", "BJetReso", "ZblZbb", "BTagB1", "norm_Wbb", "WblWbbRatio"]
        self.suspicious_syst = ["norm_"]
        # for yield ratios only
        self.category_condenser = {
        # "_HistSyst": ["_Exp", False],
        # "_dist(mva|mjj)": ["_dist", False],
        # "_distMV1cBTag": ["_dist", False],
        "_distmV": ["_dist", False],
        # "_isMVA[01]": ["_isMVA", False],
        # "_B[0-5]_": ["_B9_", False],
        "_B(Max500_BMin0|BMin500)_": ["_Bresolvedmerged_", False],
        # "_TType(ll|mm|tt|xx)": ["_TType", False],
        "_T[012]": ["_Tx", False],
        "_(incJet1_J|incFat1_Fat|J)[1235]": ["_Jx", False],
        # "_Spc[0-9a-z]*top[a-z]*cr": ["_TType", False],
        # "(multijet)(.*_L)([0123])(.*)": [r'MJ\3lep\2\3\4', False],
        "_L[012]": ["_Lx", False],
        "_D(SR|topemucr)": ["_DallRegions", False],
        # "_W(bb|bl|bc|cc)_": ["_Whf_", True],
        # "_Z(bb|bl|bc|cc)_": ["_Zhf_", True]
        }

        logging.basicConfig(format='%(levelname)s in %(module)s: %(message)s', level=self.loggingLvl)

    def do_rebinning (self, prop):
        # NOTE: JWH - ED board requests
        if prop["dist"] == "mVH":
            if "mBBcr" in prop["D"] or "topemucr" in prop["D"]:
                if prop["L"] == "2" or prop["L"] == "0":
                    if prop.get("incFat", "-1") == "1" or prop.get("incJet", "-1") == "1":
                        return False
            if "SR" in prop["D"]:
                if prop["L"] == "2" or prop["L"] == "0":
                    if prop.get("incFat", "-1") == "1":
                        return False
                    if prop["L"] == "0":
                        return False
        return True

    def is_signal(self, compname):
        """ Check if a component is Higgs. If yes, return mass """
        # Spyros: Add ggA to list of signal names - has to be first in list otherwise we get problems
        signames = self.sig_names
        has_mass = False
        mass = ""
        # Spyros: if sg in compname matches also mVH so doesn't work for resonance analyses
        # remove mVH from compname
        compname = re.sub('mVH', '', compname)
        for sg in signames:
            if sg in compname:
                has_mass = True
                pos = compname.find(sg) + len(sg)
                mass = int(re.sub("[^0-9]", "", compname[pos:pos + compname[pos:].find('_')]))
                break
        return has_mass, mass

    def blind_data (self, setup):
        def _do_blinding (title):
            #return False, []
            return "T2" in title, [110, 140]

        do_blinding, blind_range = _do_blinding(setup.title)
        if do_blinding:
            # blind entire range
            if blind_range[0] == 0 and blind_range[1] == 0:
                blind_range[0] = setup.data.h.GetXaxis().GetXmin()
                blind_range[1] = setup.data.h.GetXaxis().GetXmax()
            setup.data.blind(blind_range[0], blind_range[1])
        #else:
        # # Add general blinding at 2% S/B
        # for i in range(1, setup.hsum.GetNbinsX()+1):
        #     if setup.hsum.GetBinContent(i) > 0:
        #         sob = setup.exp_sig.h.GetBinContent(i) / ( setup.hsum.GetBinContent(i) )
        #         if sob > 0.02:
        #             setup.data.blind(setup.hsum.GetBinLowEdge(i), setup.hsum.GetBinLowEdge(i+1))
        #     elif setup.exp_sig.h.GetBinContent(i) > 0:
        #         setup.data.blind(setup.hsum.GetBinLowEdge(i), setup.hsum.GetBinLowEdge(i+1))

    def preprocess_main_content_histogram (self, hist, setupMaker):
        return hist
        # def change_MeV_GeV(hist):
        #     if isinstance(hist, ROOT.TH1):
        #         new_hist = hist.Clone()
        #         bins = new_hist.GetXaxis().GetXbins()
        #         for i in range(bins.GetSize()):
        #             bins[i] /= 1000.
        #         new_hist.SetBins(bins.GetSize()-1, bins.GetArray())
        #         for i in range(new_hist.GetNbinsX()+2):
        #             new_hist.SetBinContent(i, hist.GetBinContent(i))
        #             new_hist.SetBinError(i, hist.GetBinError(i))
        #     elif isinstance(hist, ROOT.TGraph):
        #         new_hist = hist
        #         xbins = new_hist.GetX()
        #         for i in range(new_hist.GetN()):
        #             xbins[i] /= 1000.
        #         if isinstance(hist, ROOT.TGraphAsymmErrors):
        #             xbinsup = new_hist.GetEXhigh()
        #             xbinsdo = new_hist.GetEXlow()
        #             for i in range(new_hist.GetN()):
        #                 xbinsup[i] /= 1000.
        #                 xbinsdo[i] /= 1000.
        #     return new_hist
        #
        # new_hist = hist
        # props = sm.setup.properties
        # if props:
        #     # Changes for MeV/GeV
        #     affected_dists = ["MEff", "MEff3", "MET", "mLL", "mTW", "pTB1", "pTB2", "pTJ3", "pTV", "mBB", "mBBJ"]
        #     if props["L"] == "1" and props["dist"] in affected_dists:
        #         new_hist = change_MeV_GeV(hist)
        #
        # return new_hist

    def make_sum_plots (self, func):
        #add MET for 0 lepton merged+resolved signal region
        #add mBB for 0 mbbcr+SR
        for tag_i in ["1", "2"] :
                func("Region_BMax500_BMin0_incJet1_J2_T"+tag_i+"_L2_Y2015_distmBB_Dtopemucr",
                     rt=["_L2", "_T"+tag_i, "_distmBB", "_Dtopemucr"], ea=[])
                func("Region_BMax500_BMin0_incJet1_J2_T"+tag_i+"_L2_Y2015_distmBB",
                     rt=["_L2", "_T"+tag_i, "_distmBB"], ea=["_Dtopemucr"])
                func("Region_BMax500_BMin150_incJet1_J2_T"+tag_i+"_L0_Y2015_distmBB",
                     rt=["_L0", "_T"+tag_i, "_distmBB"], ea=[])
                func("Region_BMin150_T"+tag_i+"_L0_Y2015_distMET_DSR",
                     rt=["_L0","_T"+tag_i, "_distMET","_DSR"], ea=["_L2","_DmBBcr","_Dtopemucr"])
                func("Region_BMin0_T"+tag_i+"_L2_Y2015_distpTV_DSR",
                     rt=["_L2","_T"+tag_i, "_distpTV","_DSR"], ea=["_DmBBcr","_Dtopemucr"])

    def get_run_info (self):
        lumi = {}
        if self._year == "4023":
            lumi["2011"] = ["4.7", 7]
            lumi["2012"] = ["20.3", 8]
        if self._year == "2011":
            lumi["2011"] = ["4.7", 7]
        if self._year == "2012":
            lumi["2012"] = ["20.3", 8]
        if self._year == "2015":
            lumi["2015"] = ["3.2", 13]
        return lumi

    def get_title_height (self):
        return 3.5 if self._year == "4023" else 2

    def draw_category_ids (self, props, l, pos, nf):

        merged = False
        plural_jets = False
        nf += 0.25*nf # a bit more vertical spacing

        nleps = props.get("L", "-100")
        if nleps == '3':
            nleps = "0+1+2"
        njets = props.get("J", "-1")
        nincjets = props.get("incJet", "-1")
        if njets == "23":
            plural_jets = True
            njets = "2+3"
        elif nincjets == '1':
            plural_jets = True
            # njets += '+'
            njets = '#geq {}'.format(njets)
        elif int(njets) > 1:
            plural_jets = True
        nfatjets = props.get("Fat", "-1")
        nincfatjets = props.get("incFat", "-1")
        if int(nfatjets) > 0 and nincfatjets == '1':
            plural_jets = True
            merged = True
            # nfatjets += '+'
            nfatjets = '#geq {}'.format(nfatjets)
            # nfatjets += ' #leq'
        elif int(nfatjets) > 1:
            plural_jets = True
        ntags = props.get("T", "-100")

        region = ""
        if not nleps == '-100':
            if len(region) > 0:
                region += ', '
            region += "{} lep.".format(nleps)
        if not njets == '-1' or not nfatjets == '-1':
            if len(region) > 0:
                region += ', '
            region += "{} {}jet{}".format(nfatjets if merged else njets,
                                             "large-R " if merged else "",
                                             "s" if plural_jets else "")
        if not ntags == '-100':
            if len(region) > 0:
                region += ', '
            region += "{} tag{}".format(ntags,
                                        "s" if not int(ntags) == 1 else "")

        pTVBin = ""
        pTVmin = props.get("BMin", "-999")
        pTVmax = props.get("BMax", "-999")
        if not pTVmin == "-999" and pTVmax == "-999" and not pTVmin == "0":
            pTVBin = "{0} GeV #leq p_{{T}}^{{V}}".format(pTVmin)
        elif (pTVmin == "0" or pTVmin == "-999") and not pTVmax == "-999":
            pTVBin = "p_{{T}}^{{V}} < {0} GeV".format(pTVmax)
        elif not pTVmin == "-999" and not pTVmax == "-999":
            pTVBin = "{0} GeV #leq p_{{T}}^{{V}} < {1} GeV".format(pTVmin, pTVmax)

        signalControl = props.get("D", "")
        if not signalControl == "":
            def add_strings (base, addition):
                if base == "":
                    return addition
                else:
                    return base + ", " + addition
            temp = signalControl
            signalControl = ""
            reduce_SR_CR_mBB = props["dist"] == "pTV" or props["dist"] == "MET"
            if temp.find('SR') == 0:
                if reduce_SR_CR_mBB: signalControl = "m_{b#bar{b}} SR"
                elif merged: signalControl = add_strings(signalControl, "75 GeV #leq m_{b#bar{b}} < 145 GeV")
                else: signalControl = add_strings(signalControl, "110 GeV #leq m_{b#bar{b}} < 140 GeV")
                temp = temp[2:]
            if "highmBBcr" in temp:
                if reduce_SR_CR_mBB: signalControl = "m_{b#bar{b}} upper CR"
                elif merged: signalControl = add_strings(signalControl, "145 GeV #leq m_{b#bar{b}}")
                else: signalControl = add_strings(signalControl, "140 GeV #leq m_{b#bar{b}}")
                temp = temp.replace("highmBBcr", "")
            if "lowmBBcr" in temp:
                if reduce_SR_CR_mBB: signalControl = "m_{b#bar{b}} lower CR"
                elif merged: signalControl = add_strings(signalControl, "m_{b#bar{b}} < 75 GeV")
                else: signalControl = add_strings(signalControl, "m_{b#bar{b}} < 110 GeV")
                temp = temp.replace("lowmBBcr", "")
            if "mBBcr" in temp:
                if reduce_SR_CR_mBB: signalControl = "m_{b#bar{b}} CR"
                elif merged: signalControl = add_strings(signalControl, "m_{b#bar{b}} #leq 75 GeV, 145 GeV < m_{b#bar{b}}")
                else: signalControl = add_strings(signalControl, "m_{b#bar{b}} #leq 110 GeV, 140 GeV < m_{b#bar{b}}")
                temp = temp.replace("mBBcr", "")
            if "topemucr" in temp:
                signalControl = add_strings(signalControl, "e#mu")
                temp = temp.replace("topemucr", "")
            if "topaddbjetcr" in temp:
                signalControl = add_strings(signalControl, "+1 b-jet")
                temp = temp.replace("topaddbjetcr", "")

        pos_next = pos[1] - 0.1*nf # a bit more spacing

        l.DrawLatex(pos[0], pos_next, region)
        if not pTVBin == "":
            pos_next -= nf
            l.DrawLatex(pos[0], pos_next, pTVBin)
        if not signalControl == "":
            pos_next -= nf
            l.DrawLatex(pos[0], pos_next, signalControl)

        pos_next -= nf
        return (pos[0], pos_next)

    def force_mu_value (self):
        return self.force_mu

    def get_year_str (self):
        return self._year if int(self._year) < 2015 else ""

    def get_xbound_from_properties (self, prop):
        return (40, 400) if prop["dist"] == "pTB1" else None

    def get_legend_pos_from_properties (self, prop):
        result = None
        if prop["L"] == '0' and prop["dist"] == "VpT":
            result = [0.155, 0.13, 0.375, 0.65]
        if prop["dist"] == "dPhiVBB":
            result = [0.16, 0.16, 0.38, 0.68]
        return result

    def get_yscale_factor_from_properties (self, prop, logy):
        # if prop["dist"] == "MV1cB1" or prop["dist"] == "MV1cB2" or prop["dist"] == "MV1cBTag":
        #     if not logy: return 1.5
        # if prop["dist"] == "dPhiVBB" :
        #     if logy: return 5
        #     else : return 0.7
        # if prop["dist"] == "dPhiLBmin" :
        #     if not logy: return 1.3
        # if prop["dist"] == "mjj" :
        #     if not logy: return 1.1
        # if prop["dist"] == "dRBB" :
        #     if logy: return 500
        # if prop["dist"] == "MV1cBTag" :
        #     if not logy: return 0.75
        # if prop["L"] == "0" :
        #     if prop["dist"] == "MV1cB1" or prop["dist"] == "MV1cB2" or prop["dist"] == "mjj" :
        #         if not logy: return 1.1
        #     if prop["dist"] == "MET" :
        #         if not logy: return 1.0/1.15
        return 1.0

    def postprocess_main_content_histogram (self, prop, hist):
        # draw line denoting the transition of merged and resolved
        if prop["dist"] == "MET" or prop["dist"] == "pTV":
            max_value = hist.GetMaximum()
            min_value = 0#hist.GetYaxis().GetXmin()
            x_value = hist.GetXaxis().GetBinLowEdge(hist.GetXaxis().FindBin(500))
            l = ROOT.TLine(x_value, min_value, x_value, max_value)
            l.SetLineStyle(2)
            l.SetLineWidth(4)
            l.SetNDC(False)
            l.DrawLine(x_value, min_value, x_value, max_value)
            logging.debug("drawing line with endpoint coordinates ({},{}) and ({},{})".format(x_value, min_value, x_value, max_value))
        return hist

    def get_xTitle (self, prop, data_hist):
        """ get title of X-axis from properties """
        if not prop:
            return ""
        varname = prop["dist"]
        result = varname
        labels = {
            # new
            "MV1cB1": "MV1c(b_{1}) OP",
            "MV1cB2": "MV1c(b_{2}) OP",
            "MV1cBTag": "MV1c(b) OP",
            "dEtaBB": "#Delta#eta(b_{1},b_{2})",
            "dEtaVBB": "#Delta#eta(V,bb)",
            "dPhiLBmin": "#Delta#phi(lep,b)_{min}",
            "dPhiVBB": "#Delta#phi(V,bb)",
            "dRBB": "#DeltaR(b_{1},b_{2})",
            #"MEff": "M_{eff} [GeV]",
            #"MEff3": "M_{eff3} [GeV]",
            "MEff": "H_{T} [GeV]",
            "MEff3": "H_{T} [GeV]",
            "MET": "E_{T}^{miss} [GeV]",
            "mLL": "M_{ll} [GeV]",
            "mTW": "m_{T}(W) [GeV]",
            "mva": "BDT_{VH}",
            "mvaVZ": "BDT_{VZ}",
            "pTB1": "p_{T}(b_{1}) [GeV]",
            "pTB2": "p_{T}(b_{2}) [GeV]",
            "pTJ3": "p_{T}(j_{3}) [GeV]",
            "pTV": "p_{T}^{V} [GeV]",
            "VpT": "p_{T}^{V} [GeV]",
            "mVH": "m_{T}(Vh) [GeV]"
            }
        if "mjj" in varname:
            # nominal
            tmp_extra = ""
            tmp_extra2 = " [GeV]"
            # hack for mjj trafo D
            #tmp_extra = "Transformed "
            #tmp_extra2 = ""
            #
            if prop["T"] == "2":
                result = tmp_extra+"m_{bb}"+tmp_extra2
            elif prop["T"] == "1":
                result = tmp_extra+"m_{bj}"+tmp_extra2
            else:
                result = tmp_extra+"m_{jj}"+tmp_extra2
        elif "mBBJ" in varname:
            if prop["T"] == "2":
                result = "m_{bbj} [GeV]"
            elif prop["T"] == "1":
                result = "m_{bjj} [GeV]"
            else:
                result = "m_{jjj} [GeV]"
        elif "mBB" in varname:
            if prop["T"] == "2":
                result = "m_{bb} [GeV]"
            elif prop["T"] == "1":
                result = "m_{bj} [GeV]"
            else:
                result = "m_{jj} [GeV]"
        elif "mVH" in varname:
            if prop["L"] == "1" or prop["L"] == "0":
                result = "m_{T}(Vh) [GeV]"
            else:
                result = "m(Vh) [GeV]"
        elif varname in labels:
            result = labels[varname]
            #for k in labels:
                #if k in varname:
                    #return labels[k]
        return result

    def get_yTitle_tag (self, prop, data_hist):
        extra_unit = ""
        if prop["dist"] == "MEff" :  extra_unit = " GeV"
        if prop["dist"] == "MEff3" : extra_unit = " GeV"
        if prop["dist"] == "MET" :   extra_unit = " GeV"
        if prop["dist"] == "mLL" :   extra_unit = " GeV"
        if prop["dist"] == "mTW" :   extra_unit = " GeV"
        if prop["dist"] == "pTB1" :  extra_unit = " GeV"
        if prop["dist"] == "pTB2" :  extra_unit = " GeV"
        if prop["dist"] == "pTJ3" :  extra_unit = " GeV"
        if prop["dist"] == "pTV" :   extra_unit = " GeV"
        #if prop["dist"] == "VpT" :   extra_unit = " GeV" # new
        if prop["dist"] == "mjj" :   extra_unit = " GeV" # hack -> comment when trafoD
        if prop["dist"] == "mBB" :   extra_unit = " GeV"
        if prop["dist"] == "mBBJ" :  extra_unit = " GeV"
        if prop["dist"] == "mVH" :  extra_unit = " GeV"

        # NOTE: JWH - ED board requests
        if not self.do_rebinning(prop):
        # if not (prop["dist"] == "mVH" and prop.get("incFat", "-1") == "-1" and
        #     prop.get("D", "") == "SR" and prop.get("L", "0") == "2") :
            extra_number = str(data_hist.GetBinWidth(1))
            if not extra_number.find('.') == -1: extra_number = extra_number[:extra_number.find('.')]
            extra_unit = " " + extra_number + extra_unit

        y_ratio = round(data_hist.GetBinWidth(1), 2)
        if (y_ratio*10) % 10 == 0 and (y_ratio*100) % 100 == 0: y_ratio = int(y_ratio)
        if prop["dist"] == "VpT": extra_str = " / bin" # new
        elif prop["dist"] == "mVH":  extra_str = " /" + extra_unit
        else: extra_str = " / " + str(y_ratio) + extra_unit # new

        if prop["dist"] == "MV1cB1":   extra_str = ""
        if prop["dist"] == "MV1cB2":   extra_str = ""
        if prop["dist"] == "MV1cBTag": extra_str = ""

        return extra_str

    def set_y_range (self, hist, nlegend_items, miny, maxy, log_scale, prop):
        # if log_scale and prop["dist"] == "mVH":
        #     hist.SetMaximum(maxy * 100)
        #     hist.SetMinimum(0.001)
        #     return

        bottom_padding = 1.0/16.0
        content_faction = 4.0/7.0 if nlegend_items <= 8 else 3.0/7.0
        if prop["dist"] == "mVH":
            # figures 2)a-d in conf note
            if (prop["L"] == "0" or prop["L"] == "2") and log_scale:
                if prop["T"] == "1" or prop["T"] == "2":
                    if prop["D"] == "mBBcr":
                        if prop.get("BMax", "-999") == "500":
                            content_faction *= 1.25
            # figures 3)a,b in conf note
            if prop["D"] == "topemucr" and log_scale:
                if prop["T"] == "1":
                    content_faction *= 1.15
                if prop["T"] == "2":
                    content_faction *= 1.25
            if "SR" in prop["D"]:
                # figures 6)a-d in conf note
                if prop.get("BMax", "-999") == "500" and log_scale:
                    if prop["L"] == "0":
                        if prop["T"] == "1":
                            content_faction *= 1.15
                        if prop["T"] == "2":
                            content_faction *= 1.25
                    if prop["L"] == "2":
                        content_faction *= 1.25
                # figures 7)a,c,d in conf note
                if prop.get("BMin", "-999") == "500" and not log_scale:
                    if prop["L"] == "0":
                        if prop["T"] == "1":
                            content_faction *= 1.5
                    if prop["L"] == "2":
                        if prop["T"] == "1":
                            content_faction *= 2.15
                        if prop["T"] == "2":
                            content_faction *= 1.15
        # figures 4)a-d in conf note
        if prop["dist"] == "mBB" and not log_scale:
            if prop.get("BMax", "-999") == "500" and not (prop.get("D", "") == "topemucr"):
            # if prop["L"] == "0":
            #     if prop["T"] == "1":
                content_faction *= 1.5
            if prop.get("BMax", "-999") == "500" and prop.get("D", "") == "topemucr":
                content_faction *= 1.15
        # figures 10)a-d in conf note
        if (prop["dist"] == "MET" or prop["dist"] == "pTV") and log_scale:
            content_faction *= 1.25
        if not log_scale:
            if miny < 1e-6: miny = 0
            plot_scale = (maxy - miny)
            bottom = miny - bottom_padding*plot_scale
            top = bottom + plot_scale/content_faction

            # hist.SetMinimum(bottom)
            # hist.SetMaximum(top)

            hist.GetYaxis().SetLimits(bottom, top)

            # hist.GetHistogram().GetYaxis().SetRangeUser(bottom, top)
            logging.debug("setting plot y-range to ({0}, {1})".format(hist.GetHistogram().GetYaxis().GetXmin(), hist.GetHistogram().GetYaxis().GetXmax()))
            return
        else:
            log_miny = ROOT.TMath.Log10(miny)
            log_maxy = ROOT.TMath.Log10(maxy)
            plot_scale = (log_maxy - log_miny)
            # 0.25 is just fine tuning
            # bottom = log_miny - 0.25*bottom_padding*plot_scale
            bottom = log_miny
            top = bottom + plot_scale/content_faction

            # hist.SetMinimum(ROOT.TMath.Power(10, bottom))
            # hist.SetMaximum(ROOT.TMath.Power(10, top))

            hist.GetYaxis().SetLimits(ROOT.TMath.Power(10, bottom), ROOT.TMath.Power(10, top))

            # hist.GetHistogram().GetYaxis().SetRangeUser(ROOT.TMath.Power(10, bottom), ROOT.TMath.Power(10, top))
            logging.debug("setting log scale plot y-range to ({0}, {1})".format(hist.GetHistogram().GetYaxis().GetXmin(), hist.GetHistogram().GetYaxis().GetXmax()))
            return

        # if not log_scale and miny > 0:
        #     miny = 0
        # if log_scale and miny <= 1:
        #     miny = 0.25
        # mini = miny
        #
        # if mini < 0:
        #     hist.SetMinimum(mini*1.25)
        # else:
        #     mini = 0
        #     # fix 0 cut in the Y axis
        #     #hist.SetMinimum(0.01)
        # if log_scale:
        #     hist.SetMaximum(maxy * 100)
        #     hist.SetMinimum(miny / 2.5)
        # else:
        #     hist.SetMaximum(mini + (maxy - mini) * 1.5)

    def auto_compute_ratio_yscale_from_properties (self, prop):
        return (prop["dist"] == "mva" or prop["dist"] == "mvaVZ")

    def scale_all_yvals(self, prop):
        return prop["dist"] == "mva", 0.05

    def postprocess_dataMC_ratio_histogram (self, prop, hist):
        return hist

    def determine_year_from_title (self, title):
        if "2015" in title:
            return "2015"
        elif "2012" in title:
            return "2012"
        elif "2011" in title:
            return "2011"
        elif "both" in title:
            return "4023"

    def add_additional_signal_info_to_legend (self, legend, signal):
        if signal.mode == self._STACK:
            legend.AddEntry(ROOT.NULL, "m_{H}=" + str(signal.mass) + " GeV", "")
        else:
            legend.AddEntry(ROOT.NULL, "m_{H}=" + str(signal.mass) + " GeV", "")


import os, argparse, traceback, glob, random, itertools, time, torch, threading, queue
import numpy as np
import torch.optim as optim
from models.tacotron import post_CBHG
from torch.nn import L1Loss
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from util.hparams import *

data_dir = './data'
mel_list = sorted(glob.glob(os.path.join(data_dir + '/mel', '*.npy')))
spec_list = sorted(glob.glob(os.path.join(data_dir + '/spec', '*.npy')))
mel_len = np.load(os.path.join(data_dir + '/mel_len.npy'))

def DataGenerator():
    while True:
        idx_list = np.random.choice(len(mel_list), batch_group, replace=False)
        idx_list = sorted(idx_list)
        idx_list = [idx_list[i : i + batch_size] for i in range(0, len(idx_list), batch_size)]
        random.shuffle(idx_list)

        for idx in idx_list:
            random.shuffle(idx)

            mel = [torch.from_numpy(np.load(mel_list[mel_len[i][1]])) for i in idx]
            spec = [torch.from_numpy(np.load(spec_list[mel_len[i][1]])) for i in idx]

            mel = pad_sequence(mel, batch_first=True)
            spec = pad_sequence(spec, batch_first=True)

            yield [mel, spec]

class Generator(threading.Thread):
    def __init__(self, generator):
        threading.Thread.__init__(self)
        self.queue = queue.Queue(8)
        self.generator = generator
        self.start()

    def run(self):
        for item in self.generator:
            self.queue.put(item)
        self.queue.put(None)

    def next(self):
        next_item = self.queue.get()
        if next_item is None:
             raise StopIteration
        return next_item

def train(args):
    train_loader = Generator(DataGenerator())

    model = post_CBHG(K=8, conv_dim=[256, mel_dim]).cuda()

    optimizer = optim.Adam(model.parameters())

    step, epochs = 0, 0
    if args.checkpoint is not None:
        ckpt = torch.load(args.checkpoint)
        model.load_state_dict(ckpt['model'])
        optimizer.load_state_dict(ckpt['optimizer'])
        step = ckpt['step'],
        step = step[0]
        epoch = ckpt['epoch']
        print('Load Status: Epoch %d, Step %d' % (epoch, step))

    torch.backends.cudnn.benchmark = True

    try:
        for epoch in itertools.count(epochs):
            for _ in range(batch_group):
                start = time.time()
                mel, target = train_loader.next()
                mel = mel.float().cuda()
                target = target.float().cuda()

                pred = model(mel)
                loss = L1Loss()(pred, target)

                model.zero_grad()
                loss.backward()
                optimizer.step()

                step += 1
                print('step: {}, loss: {:.5f}, {:.3f} sec/step'.format(step, loss, time.time() - start))

                if step % checkpoint_step == 0:
                    save_dir = './ckpt/' + args.name + '/2'
                    torch.save({
                        'model': model.state_dict(),
                        'optimizer': optimizer.state_dict(),
                        'step': step,
                        'epoch': epoch
                    }, os.path.join(save_dir, 'ckpt-{}.pt'.format(step)))

    except Exception as e:
        traceback.print_exc()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', '-c', default=None)
    parser.add_argument('--name', '-n', required=True)
    args = parser.parse_args()
    save_dir = os.path.join('./ckpt/' + args.name, '2')
    os.makedirs(save_dir, exist_ok=True)
    train(args)


import math

import torch
import torch.nn as nn
import torch.nn.functional as F

class Bottleneck(nn.Module):
    def __init__(self, nChannels, growthRate, dropout_rate):
        super(Bottleneck, self).__init__()
        self.dropout_rate = dropout_rate
        interChannels = 4 * growthRate
        self.bn1 = nn.BatchNorm2d(nChannels)
        self.conv1 = nn.Conv2d(nChannels, interChannels, kernel_size=1,
                               bias=False)
        self.bn2 = nn.BatchNorm2d(interChannels)
        self.conv2 = nn.Conv2d(interChannels, growthRate, kernel_size=3,
                               padding=1, bias=False)

    def forward(self, x):
        out = self.conv1(F.relu(self.bn1(x)))
        out = self.conv2(F.relu(self.bn2(out)))
        if self.dropout_rate > 0:
            out = F.dropout(out, p=self.dropout_rate, training=self.training)
        out = torch.cat((x, out), 1)
        return out

class SingleLayer(nn.Module):
    def __init__(self, nChannels, growthRate, dropout_rate):
        super(SingleLayer, self).__init__()
        self.dropout_rate = dropout_rate
        self.bn1 = nn.BatchNorm2d(nChannels)
        self.conv1 = nn.Conv2d(nChannels, growthRate, kernel_size=3,
                               padding=1, bias=False)

    def forward(self, x):
        out = self.conv1(F.relu(self.bn1(x)))
        if self.dropout_rate > 0:
            out = F.dropout(out, p=self.dropout_rate, training=self.training)
        out = torch.cat((x, out), 1)
        return out

class Transition(nn.Module):
    def __init__(self, nChannels, nOutChannels):
        super(Transition, self).__init__()
        self.bn1 = nn.BatchNorm2d(nChannels)
        self.conv1 = nn.Conv2d(nChannels, nOutChannels, kernel_size=1,
                               bias=False)

    def forward(self, x):
        out = self.conv1(F.relu(self.bn1(x)))
        out = F.avg_pool2d(out, 2)
        return out

class Model(nn.Module):
    def __init__(self, nClasses=256, growthRate=12, depth=40, bottleneck=False,
                 dropout_rate=0.0, n_layers=3, **kwargs):
        super().__init__()

        self.n_layers = n_layers
        print("n_layers", n_layers)

        # dense blocks per layer
        nDenseBlocks = (depth - 4) // n_layers
        if bottleneck:
            nDenseBlocks //= 2

        if bottleneck:
            reduction = 0.5
        else:
            reduction = 1.0

        # initial convolution
        nChannels = 2 * growthRate
        self.conv1 = nn.Conv2d(3, nChannels, kernel_size=3, padding=1,
                               bias=False)

        for layer_n in range(1, n_layers + 1):
            dense_layer = self._make_dense(
                nChannels, growthRate, nDenseBlocks, bottleneck, dropout_rate)
            setattr(self, f'dense{layer_n}', dense_layer)
            nChannels += nDenseBlocks * growthRate
            if layer_n < n_layers:
                nOutChannels = int(math.floor(nChannels * reduction))
                trainsition_layer = Transition(nChannels, nOutChannels)
                setattr(self, f'trans{layer_n}', trainsition_layer)
                nChannels = nOutChannels

        self.bn1 = nn.BatchNorm2d(nChannels)
        self.fc = nn.Linear(nChannels, nClasses)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.bias.data.zero_()

    def _make_dense(self, nChannels, growthRate, nDenseBlocks, bottleneck, dropout_rate):
        layers = []
        for i in range(int(nDenseBlocks)):
            if bottleneck:
                layers.append(Bottleneck(nChannels, growthRate, dropout_rate))
            else:
                layers.append(SingleLayer(nChannels, growthRate, dropout_rate))
            nChannels += growthRate
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv1(x)
        for i in range(1, self.n_layers):
            dense_layer = getattr(self, f'dense{i}')
            trans_layer = getattr(self, f'trans{i}')
            out = trans_layer(dense_layer(out))
        last_dense_layer = getattr(self, f'dense{self.n_layers}')
        out = last_dense_layer(out)

        out = F.avg_pool2d(F.relu(self.bn1(out)), out.size()[-1])
        out = torch.squeeze(torch.squeeze(out, 2), 2)
        out = F.log_softmax(self.fc(out))
        return out

    def save(self, path):
        torch.save(self.state_dict(), path)

    def load(self, path):
        state_dict = torch.load(path)
        self.load_state_dict(state_dict)


import os
from zipfile import ZipFile, ZIP_DEFLATED

class Zip():
    """Zip up all the contents of a directory into the output file."""
    def __init__(self, input_directory, output_file):
      self.input_directory = input_directory
      self.output_file = output_file

    def zip(self):
      try:
        zip_file = ZipFile(self.output_file, 'w', ZIP_DEFLATED)
        for root, dirs, files in os.walk(self.input_directory, topdown=True):
          dir_prefix = root[len(os.path.commonprefix([self.input_directory, root]))+1:]
          if len(dirs) is 0 and len(files) is 0:
            zip_file.write(root, dir_prefix[:-1])
          else:
            for name in files:
              zip_file.write(os.path.join(root, name), os.path.join(dir_prefix, name))

      finally:
        zip_file.close()


#!/usr/bin/env python
"""A script to normalized interview transcripts.
It outputs a single text file with cleaned lines
one sentence per line"""

import argparse
import re
import string
import spacy

fillers = [
    "eh",
    "m",
    "mm",
    "mmm",
    "ah",
    "ahm",
    "ehm",
    "yy",
    "y",
    "aha",
    "a-ha",
    "aa",
    "e",
    "ee",
    "łyy",
    "ym",
    "yym",
    "ymm",
    "yyym",
    "oh",
    "am",
    "oo",
    "hm",
    "em",
    "emm",
    "eem",
    "yyo",
    "ii",
    "nnn",
    "nn",
    "no",
    "mhm",
    "am",
    "amm",
    "aam",
    "eey",
    "eeyy",
    "mmyy",
    "yhm",
    "ymhm",
    "mmy",
    "yynn",
    "li",
    "cc",
]

nlp = spacy.load("pl_core_news_lg")
punctuation = string.punctuation + '…' +'–' + '’' + "‘"

def get_data(path: str) -> list:
    """reads .txt file into a list of strings"""
    list_of_lines = []
    with open(path, "r") as source:
        for line in source:
            if line == False:
                continue
            else:
                line = line.lstrip().rstrip()
                list_of_lines.append(line)
    return list_of_lines

def write_data(data: list, path: str):
    """writes data line by line into a file"""
    with open(path, "w") as sink:
        for sentence in data:
            print(sentence, file=sink)

def remove_fillers(line: str) -> str:
    """removes filler expresisons"""
    tokens = line.split()
    for word in tokens:
        if word in fillers:
            tokens.remove(word)
    no_fillers = " ".join(word for word in tokens)
    return no_fillers

def pre_tokenization(data: list) -> list:
    """data normalization to be performed before sentence tokenization"""
    cleaned = []
    for line in data:
        # replace ';' and '%' with 'ł'
        add_l = re.sub(r"[;%]", "ł", line)
        # replace the elipses with whitespaces to account for certain types of stutters
        no_elipses = re.sub(r"[…]", " ", add_l)
        # replace two period elipses with whitespace
        two_periods = re.sub(r"\.\.", " ", no_elipses)
        # remove hyphenated stutters
        no_stutters = re.sub(r"\b[a-zA-ZżźćńółęąśŻŹĆĄŚĘŁÓŃ]+-+\W", "", two_periods)
        # remove digits and numbers
        no_numbers = re.sub(r"(?:[+-]|\()?\$?\d+(?:,\d+)*(?:\.\d+)?\)?", "", no_stutters)
        # remove bracketed content
        no_brackets = re.sub(r"\[.*?\]", "", no_numbers)
        # remove content in parentheses
        no_parens = re.sub(r"\(.*?\)", "", no_brackets)
        # remove all duplicate words
        # retain only the first word
        no_duplicates = re.sub(r"\b(\w+)(?:\W+\1\b)+", r"\1", no_parens)
        # append only non-empty strings
        if no_duplicates:
            cleaned.append(no_duplicates)
    return cleaned

def make_continuous(data: list) -> string:
    """joins a list of strings into one long string"""
    one_line = " ".join(data)
    return one_line

def post_tokenization(data: list) -> list:
    """data normaization to be performed after sentence tokenization"""
    cleaned = []
    for sentence in data:
        # casefold and strip
        casefolded = sentence.lower().lstrip().rstrip()
        # standardize quotation marks
        standard_quotes = re.sub(r"[„”“]", '"', casefolded)
        # remove punctuation
        no_polish_punctuation = standard_quotes.translate(str.maketrans("", "", punctuation))
        # remove the hyphens
        no_hyphens = re.sub(r"-", " ", no_polish_punctuation)
        # remove the fillers
        no_fillers = remove_fillers(no_hyphens)
        # remove duplicates left over after the fillers were removed
        # leave only the first word
        no_duplicates = re.sub(r"\b(\w+)(?:\W+\1\b)+", r"\1", no_fillers)
         # remove multiple white spaces
        single_spaces = re.sub(" +", ' ', no_duplicates)
        if single_spaces:
            cleaned.append(single_spaces)
    return cleaned

def main(args: argparse.Namespace) -> None:
    lines = get_data(args.input)
    cleaned = pre_tokenization(lines)
    one_line = make_continuous(cleaned)
    to_tokenize = nlp(one_line)
    sent_tokenized = [sentence.text for sentence in to_tokenize.sents]
    cleaned_again = post_tokenization(sent_tokenized)
    write_data(cleaned_again, args.output)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", help="path to source file")
    parser.add_argument("--output", help="path to output")
    main(parser.parse_args())


"""
@创建日期 ：2022/4/25
@修改日期 ：2022/4/26
@作者 ：jzj
@功能 ：模型库，输出统一以字典格式
       dqn 输出 value
       a2c 输出 policy value
       fixme: 可能会抽象为参数构建的模式，不确定
"""

from typing import List
import tensorflow as tf
import tensorflow.keras.layers as layers
import tensorflow.keras.models as models

def make_model(id, args):
    if id == "cartpole_dqn":
        return CartPoleDQN(**args)
    elif id == "cartpole_a2c":
        return CartPoleA2C(**args)
    elif id == "flappybirdsimple_dqn":
        return FlappyBirdSimpleDqn(**args)
    elif id == "flappybirdsimple_a2c":
        return FlappyBirdSimpleA2C(**args)
    else:
        raise NotImplementedError

class ModelWrapper(models.Model):
    """fixme: dev ing"""
    def __init__(self, model):
        super(ModelWrapper, self).__init__()
        self.model = model

    def call(self, inputs):
        return self.model(inputs)

    def inference(self, x):
        if hasattr(self.model, "inference"):
            return self.model.inference(x)

        x = tf.expand_dims(x, 0)
        outputs = self.call(x)
        return outputs

# CartPole DQN
class CartPoleDQN(models.Model):
    def __init__(self, action_dim, hidden_dims: List):
        super(CartPoleDQN, self).__init__()
        self.input_layers = layers.InputLayer(input_shape=(4,))
        self.hidden_layers = []
        for hidden_dim in hidden_dims:
            self.hidden_layers.append(layers.Dense(hidden_dim, activation="tanh"))
        self.output_layer = layers.Dense(action_dim)

    def call(self, inputs):
        x = self.input_layers(inputs)
        for layer in self.hidden_layers:
            x = layer(x)
        x = self.output_layer(x)
        return {"value": x}

# CartPole A2C
class CartPoleA2C(tf.keras.Model):
    def __init__(self, num_action=2, num_hidden_units=128):
        super(CartPoleA2C, self).__init__()
        self.common = layers.Dense(num_hidden_units, activation=None)
        self.activation = layers.ReLU()
        self.actor = layers.Dense(num_action)
        self.critic = layers.Dense(1)

    def call(self, inputs: tf.Tensor):
        x = self.common(inputs)
        x = self.activation(x)
        return {"policy": self.actor(x), "value": self.critic(x)}

# FlappyBirdRGB A2C
class ConvBlock(layers.Layer):
    def __init__(self, filter, kernel_size, stride=1):
        super(ConvBlock, self).__init__()
        self.conv = layers.Conv2D(filter, kernel_size, stride, padding="same")
        self.bn = layers.BatchNormalization()
        self.activation = layers.ReLU()

    def call(self, inputs):
        return self.activation(self.bn(self.conv(inputs)))

class ResidualBlock(layers.Layer):
    def __init__(self, filter, kernel_size, stride, squeeze_factor, se=False):
        """fixme: 添加Se支持"""
        super(ResidualBlock, self).__init__()
        self.conv_block1 = ConvBlock(filter//squeeze_factor, kernel_size, stride)
        self.conv_block2 = ConvBlock(filter, kernel_size, stride)
        self.short_cut = ConvBlock(filter, 1)
        self.output_bn = layers.BatchNormalization()
        self.output_ac = layers.ReLU()

    def call(self, inputs):
        x = self.conv_block1(inputs)
        x = self.conv_block2(x)
        x = x + self.short_cut(inputs)
        x = self.output_ac(self.output_bn(x))
        return x

class PolicyHead(layers.Layer):
    def __init__(self, policy_dim):
        super(PolicyHead, self).__init__()
        self.conv = layers.Conv2D(1, kernel_size=3, strides=1, padding="same")
        self.bn = layers.BatchNormalization()
        self.dense = layers.Dense(policy_dim)

    def call(self, inputs):
        b, h, w, c = inputs.shape
        x = self.bn(self.conv(inputs))
        x = tf.reshape(x, (-1, h*w))
        x = self.dense(x)
        return x

class ValueHead(layers.Layer):
    def __init__(self, value_dim):
        super(ValueHead, self).__init__()
        self.conv = layers.Conv2D(1, kernel_size=3, strides=1, padding="same")
        self.bn = layers.BatchNormalization()
        self.dense = layers.Dense(value_dim)

    def call(self, inputs):
        b, h, w, c = inputs.shape
        x = self.bn(self.conv(inputs))
        x = tf.reshape(x, (-1, h*w))
        x = self.dense(x)
        return x

class FlappyBirdA2C(models.Model):
    """
    简单模型，注意policy输出的是logit值为非概率
    """
    def __init__(self, filters=[32, 64, 128], blocks=[2, 2, 4]):
        super(FlappyBirdA2C, self).__init__()
        self.conv1 = layers.Conv2D(32, 5, 2, padding="same")
        self.bn1 = layers.BatchNormalization()
        self.ac1 = layers.ReLU()
        self.pool1 = layers.MaxPooling2D(pool_size=3, strides=2, padding="same")

        self.middle_layers = []

        for filter, block in zip(filters, blocks):
            for n in range(block):
                self.middle_layers.append(ResidualBlock(filter, 3, 1, 4))
            self.middle_layers.append(layers.MaxPooling2D(pool_size=3, strides=2, padding="same"))

        self.policy_head = PolicyHead(policy_dim=2)
        self.value_head = ValueHead(value_dim=1)

    def call(self, inputs):
        x = self.pool1(self.ac1(self.bn1(self.conv1(inputs))))
        for layer in self.middle_layers[:-1]:
            x = layer(x)
        policy = self.policy_head(x)
        value = self.value_head(x)
        return {"policy": policy, "value": value}

# FlappyBirdSimple A2C
class FlappyBirdSimpleA2C(models.Model):
    def __init__(self, policy_dim=2, value_dim=1, hidden_dims=[32, 64]):
        super(FlappyBirdSimpleA2C, self).__init__()
        self.input_layers = layers.InputLayer(input_shape=(2,))
        self.hidden_layers = []
        for hidden_dim in hidden_dims:
            self.hidden_layers.append(layers.Dense(hidden_dim, activation="tanh",))

        self.policy_head = layers.Dense(policy_dim)
        self.value_head = layers.Dense(value_dim)

    def call(self, inputs):
        x = self.input_layers(inputs)
        for layer in self.hidden_layers:
            x = layer(x)
        policy = self.policy_head(x)
        value = self.value_head(x)
        return {"policy": policy, "value": value}

# FlappyBirdSimple DQN
class FlappyBirdSimpleDqn(models.Model):
    def __init__(self, value_dim=2, hidden_dims=[256, 256]):
        super(FlappyBirdSimpleDqn, self).__init__()
        self.input_layers = layers.InputLayer(input_shape=(2,))
        self.hidden_layers = []
        for hidden_dim in hidden_dims:
            self.hidden_layers.append(layers.Dense(hidden_dim, activation="tanh"))

        self.value_head = layers.Dense(value_dim)

    def call(self, inputs):
        x = self.input_layers(inputs)
        for layer in self.hidden_layers:
            x = layer(x)
        value = self.value_head(x)
        return {"value": value}


"""
Filename:     plot_zonal_mean.py
Author:       Damien Irving, irving.damien@gmail.com
Description:

"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import matplotlib.pyplot as plt
from matplotlib import gridspec
import iris
import iris.plot as iplt
from iris.experimental.equalise_cubes import equalise_attributes
import seaborn

# Import my modules

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'ocean-analysis':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)

try:
    import general_io as gio
    import timeseries
    import grids
    import convenient_universal as uconv
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')

# Define functions

experiment_colors = {}
experiment_colors['historical'] = 'green'
experiment_colors['piControl'] = 'black'
experiment_colors['historicalAA'] = 'blue'
experiment_colors['historicalGHG'] = 'red'
experiment_colors['historicalnoAA'] = 'orange'

basins = {'atlantic': 2,
          'pacific': 3,
          'indian': 5}

def scale_data(cube, var, reverse_sign=False):
    """Scale data"""

    if var == 'precipitation_minus_evaporation_flux':
        cube.data = cube.data * 86400
        units = 'mm/day'
    else:
        units = cube.units

    if reverse_sign:
        cube.data = cube.data * -1

    return cube, units

def set_plot_grid(tas_trend=False):
    """Set the grid of plots.

    Args:
      tas_trend (bool): Include a panel for the tas trend?

    """

    if tas_trend:
        nrows = 4
        heights = [2, 1, 1, 1]
    else:
        nrows = 3
        heights = [2, 1, 1]

    gs = gridspec.GridSpec(nrows, 1, height_ratios=heights)

    return gs

def calculate_climatology(cube, time_bounds, experiment):
    """Calculate annual mean climatology"""

    if not experiment == 'piControl':
        time_constraint = gio.get_time_constraint(time_bounds)
        cube = cube.extract(time_constraint)

    cube = cube.collapsed('time', iris.analysis.MEAN)
    cube.remove_coord('time')

    return cube

def calc_linear_trend(data, xaxis):
    """Calculate the linear trend.
    polyfit returns [a, b] corresponding to y = a + bx
    """

    if data.mask[0]:
        return data.fill_value
    else:
        return numpy.polynomial.polynomial.polyfit(xaxis, data, 1)[-1]

def get_trend_cube(cube, xaxis='time'):
    """Get the trend data.
    Args:
      cube (iris.cube.Cube)
      xaxis (iris.cube.Cube)
    """

    coord_names = [coord.name() for coord in cube.dim_coords]
    assert coord_names[0] == 'time'

    if xaxis == 'time':
        trend_data = timeseries.calc_trend(cube, per_yr=True)
        trend_unit = ' yr-1'
    else:
        trend_data = numpy.ma.apply_along_axis(calc_linear_trend, 0, cube.data, xaxis.data)
        trend_data = numpy.ma.masked_values(trend_data, cube.data.fill_value)
        trend_unit = ' '+str(xaxis.units)+'-1'

    trend_cube = cube[0, ::].copy()
    trend_cube.data = trend_data
    trend_cube.remove_coord('time')
    trend_cube.units = str(cube.units) + trend_unit

    return trend_cube

def get_scale_factor(tas_cube):
    """Calculate scale factor (linear warming).

    Multiplies the linear trend (K / yr) by the number of years

    """

    linear_trend = get_trend_cube(tas_cube)
    scale_factor = linear_trend.data * tas_cube.shape[0]

    return scale_factor

def plot_climatology(climatology_dict, var, units, legloc, aggregation='Zonal mean'):
    """Plot the zonal mean climatology"""

    for experiment in ['historical', 'historicalGHG', 'historicalAA', 'historicalnoAA', 'piControl']:
        if climatology_dict[experiment]:
            color = experiment_colors[experiment]
            iplt.plot(climatology_dict[experiment], color=color, alpha=0.8, label=experiment)

    plt.legend(loc=legloc)
    plt.ylabel('%s %s (%s)' %(aggregation, var.replace('_', ' '), units) )

def check_lats(climatology_dict, experiment):
    """Sometimes the latitude axes are not exactly equal after regridding."""

    experiment_lats = climatology_dict[experiment].coord('latitude')
    control_lats = climatology_dict['piControl'].coord('latitude')
    if not control_lats == experiment_lats:
        diffs = experiment_lats.points - control_lats.points
        assert numpy.abs(diffs).max() < 0.0001, "%s and control have very different latitude axes" %(experiment)
        climatology_dict[experiment].coord('latitude').points = control_lats.points
        climatology_dict[experiment].coord('latitude').bounds = control_lats.bounds
        assert climatology_dict[experiment].coord('latitude') == climatology_dict['piControl'].coord('latitude'), \
        "Problem with %s latitude axis" %(experiment)

    return climatology_dict[experiment]

def plot_difference(climatology_dict):
    """Plot the difference between experiment and control climatology"""

    assert climatology_dict['piControl'], 'must have control data for difference plot'

    for experiment in ['historical', 'historicalGHG', 'historicalAA', 'historicalnoAA']:
        if climatology_dict[experiment]:
            climatology_dict[experiment] = check_lats(climatology_dict, experiment)
            diff_cube = climatology_dict[experiment] - climatology_dict['piControl']
            iplt.plot(diff_cube, color=experiment_colors[experiment], alpha=0.8)

    plt.ylabel('Experiment - piControl')

def plot_trend(trend_dict, units, scaled=False):
    """Plot the trend"""

    for experiment in ['historical', 'historicalGHG', 'historicalAA', 'historicalnoAA', 'piControl']:
        if trend_dict[experiment]:
            iplt.plot(trend_dict[experiment], color=experiment_colors[experiment], alpha=0.8)

    if not scaled:
        plt.ylabel('Trend ($%s \enspace yr^{-1}$)' %(units) )
    else:
        plt.ylabel('Trend ($%s \enspace yr^{-1}$) scaled by warming' %(units) )

def read_data(inargs):
    """Read input data into appropriate dictionaries."""

    file_dict = {'historical': inargs.historical_files,
                 'historicalGHG': inargs.historicalghg_files,
                 'historicalAA': inargs.historicalaa_files,
                 'historicalnoAA': inargs.historicalnoaa_files,
                 'piControl': inargs.picontrol_files}

    tas_dict = {'historical': inargs.historical_tas_file,
                'historicalGHG': inargs.historicalghg_tas_file,
                'historicalAA': inargs.historicalaa_tas_file,
                'historicalnoAA': inargs.historicalnoaa_tas_file,
                'piControl': None}

    area_dict = {'historical': inargs.historical_area_file,
                 'historicalGHG': inargs.historicalghg_area_file,
                 'historicalAA': inargs.historicalaa_area_file,
                 'historicalnoAA': inargs.historicalnoaa_area_file,
                 'piControl': inargs.picontrol_area_file}

    basin_dict = {'historical': inargs.historical_basin_file,
                  'historicalGHG': inargs.historicalghg_basin_file,
                  'historicalAA': inargs.historicalaa_basin_file,
                  'historicalnoAA': inargs.historicalnoaa_basin_file,
                  'piControl': inargs.picontrol_basin_file}

    return file_dict, tas_dict, area_dict, basin_dict

def get_areacello_data(cube):
    """Generate an area data array."""

    dim_coord_names = [coord.name() for coord in cube.dim_coords]
    assert 'latitude' in dim_coord_names
    assert 'longitude' in dim_coord_names

    if not cube.coord('latitude').has_bounds():
        cube.coord('latitude').guess_bounds()

    if not cube.coord('longitude').has_bounds():
        cube.coord('longitude').guess_bounds()

    area_data = iris.analysis.cartography.area_weights(cube)
    area_data = numpy.ma.masked_where(numpy.ma.getmask(cube.data), area_data)

    return area_data

def area_ajustment(data_cube, area_file, metadata_dict):
    """Multipy a data cube by its cell area."""

    if area_file:
        area_cube = iris.load_cube(area_file[0])
        area_data = uconv.broadcast_array(area_cube.data, [1, 2], data_cube.shape)
        metadata_dict[area_file[0]] = area_cube.attributes['history']
    else:
        area_data = get_areacello_data(data_cube)

    data_cube.data = data_cube.data * area_data
    if 'm-2' in str(data_cube.units):
        units = str(data_cube.units).replace('m-2', "")
    else:
        units = str(data_cube.units) + ' m2'

    return data_cube, units, metadata_dict

def main(inargs):
    """Run the program."""

    file_dict, tas_dict, area_dict, basin_dict = read_data(inargs)

    metadata_dict = {}
    climatology_dict = {}
    time_trend_dict = {}
    tas_scaled_trend_dict = {}
    branch_dict = {}
    for experiment in ['historical', 'historicalGHG', 'historicalAA', 'historicalnoAA', 'piControl']:
        filenames = file_dict[experiment]
        if not filenames:
            climatology_dict[experiment] = None
            time_trend_dict[experiment] = None
            tas_scaled_trend_dict[experiment] = None
        else:
            print(experiment)
            try:
                time_constraint = gio.get_time_constraint(inargs.total_time)
            except (AttributeError, TypeError):
                time_constraint = iris.Constraint()

            with iris.FUTURE.context(cell_datetime_objects=True):
                cube = iris.load(filenames, gio.check_iris_var(inargs.var))

                # Merge cubes
                metadata_dict[filenames[0]] = cube[0].attributes['history']
                equalise_attributes(cube)
                iris.util.unify_time_units(cube)
                cube = cube.concatenate_cube()
                cube = gio.check_time_units(cube)

                # Time extraction and branch time info
                coord_names = [coord.name() for coord in cube.dim_coords]
                assert coord_names[0] == 'time'

                if 'historical' in experiment:
                    original_time_length = cube.shape[0]
                    cube = cube.extract(time_constraint)
                    new_time_length = cube.shape[0]
                    branch_time_index_offset = original_time_length - new_time_length

                    branch_time = cube.attributes['branch_time']
                    time_length = cube.shape[0]
                    branch_dict[experiment] = (branch_time, time_length, branch_time_index_offset)

                elif experiment == 'piControl':
                    branch_time, time_length, branch_time_index_offset = branch_dict['historical']
                    start_index, error = uconv.find_nearest(cube.coord('time').points, float(branch_time) + 15.5, index=True)
                    if abs(error) > 15:
                        print("WARNING: Large error of %f in locating branch time"  %(error))
                        start_index = 0
                    start_index = start_index + branch_time_index_offset
                    cube = cube[start_index:start_index+time_length, ::]

                # Temporal smoothing
                cube = timeseries.convert_to_annual(cube, full_months=True)

                # Mask marginal seas
                if basin_dict[experiment]:
                    basin_cube = iris.load_cube(basin_dict[experiment])
                    cube = uconv.mask_marginal_seas(cube, basin_cube)

                # Regrid and select basin
                cube, coord_names, regrid_status = grids.curvilinear_to_rectilinear(cube)
                if not inargs.basin == 'globe':
                    if basin_dict[experiment] and not regrid_status:
                        ndim = cube.ndim
                        basin_array = uconv.broadcast_array(basin_cube.data, [ndim - 2, ndim - 1], cube.shape)
                    else:
                        basin_array = uconv.create_basin_array(cube)
                    cube.data.mask = numpy.where((cube.data.mask == False) & (basin_array == basins[inargs.basin]), False, True)

                # Scale
                cube, units = scale_data(cube, inargs.var, reverse_sign=inargs.reverse_sign)

                # Zonal statistic
                if inargs.area_adjust:
                    if regrid_status:
                        area_dict[experiment] = None
                    cube, units, metadata_dict = area_ajustment(cube, area_dict[experiment], metadata_dict)
                    zonal_cube = cube.collapsed('longitude', iris.analysis.SUM)
                    aggregation = 'Zonally integrated'
                else:
                    zonal_cube = cube.collapsed('longitude', iris.analysis.MEAN)
                    aggregation = 'Zonal mean'
                zonal_cube.remove_coord('longitude')

                # Climatology and trends
                climatology_dict[experiment] = calculate_climatology(zonal_cube, inargs.climatology_time, experiment)
                time_trend_dict[experiment] = get_trend_cube(zonal_cube)
                if tas_dict[experiment]:
                    tas_cube = iris.load_cube(tas_dict[experiment], 'air_temperature' & time_constraint)
                    scale_factor = get_scale_factor(tas_cube)
                    print(experiment, 'warming:', scale_factor)
                    tas_scaled_trend_dict[experiment] = time_trend_dict[experiment] * (1. / abs(scale_factor))
                    metadata_dict[tas_dict[experiment][0]] = tas_cube.attributes['history']
                else:
                    tas_scaled_trend_dict[experiment] = None

    # Create the plots

    tas_scaled_trend_flag = tas_scaled_trend_dict['historicalGHG'] and tas_scaled_trend_dict['historicalAA']

    fig = plt.figure(figsize=[15, 20])
    gs = set_plot_grid(tas_trend=tas_scaled_trend_flag)

    ax_main = plt.subplot(gs[0])
    plt.sca(ax_main)
    plot_climatology(climatology_dict, inargs.var, units, inargs.legloc, aggregation)
    plt.title('%s (%s), %s' %(inargs.model, inargs.run, inargs.basin) )

    ax_diff = plt.subplot(gs[1])
    plt.sca(ax_diff)
    plot_difference(climatology_dict)

    ax_time_trend = plt.subplot(gs[2])
    plt.sca(ax_time_trend)
    plot_trend(time_trend_dict, units)

    if tas_scaled_trend_flag:
        ax_tas_trend = plt.subplot(gs[3])
        plt.sca(ax_tas_trend)
        plot_trend(tas_scaled_trend_dict, units, scaled=True)

    plt.xlabel('latitude')

    plt.savefig(inargs.outfile, bbox_inches='tight')
    gio.write_metadata(inargs.outfile, file_info=metadata_dict)

if __name__ == '__main__':

    extra_info ="""
author:
  Damien Irving, irving.damien@gmail.com

note:

"""

    description=''
    parser = argparse.ArgumentParser(description=description,
                                     epilog=extra_info,
                                     argument_default=argparse.SUPPRESS,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("outfile", type=str, help="Output file name")
    parser.add_argument("var", type=str, help="Variable standard_name")
    parser.add_argument("model", type=str, help="Model name")
    parser.add_argument("run", type=str, help="Run (e.g. r1)")
    parser.add_argument("basin", type=str, choices=('atlantic', 'pacific', 'indian', 'globe'), help="Ocean basin")

    parser.add_argument("--historical_files", type=str, default=None, nargs='*',
                        help="Input files for the historical experiment")
    parser.add_argument("--historicalghg_files", type=str, default=None, nargs='*',
                        help="Input files for the historicalGHG experiment")
    parser.add_argument("--historicalaa_files", type=str, default=None, nargs='*',
                        help="Input files for the historicalAA experiment")
    parser.add_argument("--historicalnoaa_files", type=str, default=None, nargs='*',
                        help="Input files for the historicalnoAA experiment")
    parser.add_argument("--picontrol_files", type=str, default=None, nargs='*',
                        help="Input files for the piControl experiment")

    parser.add_argument("--historical_tas_file", type=str, default=None, nargs='*',
                        help="Global mean surface temperature file for historical experiment")
    parser.add_argument("--historicalghg_tas_file", type=str, default=None, nargs='*',
                        help="Global mean surface temperature file for historicalGHG experiment")
    parser.add_argument("--historicalaa_tas_file", type=str, default=None, nargs='*',
                        help="Global mean surface temperature file for historicalAA experiment")
    parser.add_argument("--historicalnoaa_tas_file", type=str, default=None, nargs='*',
                        help="Global mean surface temperature file for historicalnoAA experiment")

    parser.add_argument("--historical_area_file", type=str, default=None, nargs='*',
                        help="Cell area file for historical experiment")
    parser.add_argument("--historicalghg_area_file", type=str, default=None, nargs='*',
                        help="Cell area file for historicalGHG experiment")
    parser.add_argument("--historicalaa_area_file", type=str, default=None, nargs='*',
                        help="Cell area file for historicalAA experiment")
    parser.add_argument("--historicalnoaa_area_file", type=str, default=None, nargs='*',
                        help="Cell area file for historicalnoAA experiment")
    parser.add_argument("--picontrol_area_file", type=str, default=None, nargs='*',
                        help="Cell area file for piControl experiment")

    parser.add_argument("--historical_basin_file", type=str, default=None, nargs='*',
                        help="Cell basin file for historical experiment")
    parser.add_argument("--historicalghg_basin_file", type=str, default=None, nargs='*',
                        help="Cell basin file for historicalGHG experiment")
    parser.add_argument("--historicalaa_basin_file", type=str, default=None, nargs='*',
                        help="Cell basin file for historicalAA experiment")
    parser.add_argument("--historicalnoaa_basin_file", type=str, default=None, nargs='*',
                        help="Cell basin file for historicalnoAA experiment")
    parser.add_argument("--picontrol_basin_file", type=str, default=None, nargs='*',
                        help="Cell basin file for piControl experiment")

    parser.add_argument("--area_adjust", action="store_true", default=False,
                        help="Adjust plots for area [default=False]")
    parser.add_argument("--reverse_sign", action="store_true", default=False,
                        help="Multiple the data by -1 (CCSM4 has wrong sign for wind stress) [default=False]")

    parser.add_argument("--climatology_time", type=str, nargs=2, metavar=('START_DATE', 'END_DATE'),
                        default=('1986-01-01', '2005-12-31'), help="Time period for climatology [default = entire]")
    parser.add_argument("--total_time", type=str, nargs=2, metavar=('START_DATE', 'END_DATE'),
                        default=None, help="Time period for entire analysis. Must go right to end of experiment for control overlap period to be calculated correctly. [default = entire]")

    parser.add_argument("--legloc", type=int, default=8,
                        help="Legend location")

    args = parser.parse_args()
    main(args)