#!/usr/bin/env python3
# Author: Alexander Mahura : 2023-06
#
#############################################################################
#                             LOADING LIBRARIES
#############################################################################
import sys
import numpy as np
import os
import copy
import time
from datetime import datetime as dt
import warnings
import matplotlib.pyplot as plt

#############################################################################
#                             FUNCTIONS
#############################################################################

#----------------------------------------------------------------------------
#  Function to read file with observations extracted from .odb-file
#  note: (YYYY MM DD HH MM SS) 
#----------------------------------------------------------------------------

def readobs(type_obs, path_tmp_obs, tmpfilename):
    """
    Read extracted observations from .odb-file for ground-based (SYNOP), 
         radiosounding (TEMP), and satellite (AMSU-A) data
    Input
    -------
    type_obs: synop, radsound, satellite
        Acronyms in names of files with temporarily extracted observations
    Output
    -------
    ds: dataset with extracted observations
    """
    
    if type_obs=='synop':
        #filename = path_tmp_obs + "OBSSYNOP/" + tmpfilename  # Puhti
        filename = path_tmp_obs + tmpfilename  # Lumi
        print('obs synop tmp-filename is:', filename)
        if os.path.isfile(filename)==True:
            # 2023-08-25 - no need to skip 1st row with header ! skiprows=1
            ds = np.genfromtxt(filename) 
            #, dtype={'names': ('year', 'month', 'day', 
            #            'hour','minute','second','longitude','latitude','variable','value','station'),
            #            'formats': ('i4', 'i2', 'i2', 'i2', 'i2', 'i2', 'f8', 'f8', 'i4', 'f8', 'S6')})
            #print('OBS at Synop stations\n',ds)
        else:
            print('OBS: The specified synop file does not exist !!!')
    elif type_obs=='radsound':
        #filename = path_tmp_obs + "OBSRADSOUND/" + tmpfilename # Puhti
        filename = path_tmp_obs + tmpfilename  # Lumi
        print('obs radsound tmp-filename is:', filename)
        if os.path.isfile(filename)==True:
            # 2023-08-25 - no need to skip 1st row with header ! skiprows=1
            ds = np.genfromtxt(filename)
            #ds = np.loadtxt(filename, dtype={'names': ('year', 'month', 'day', 
            #            'hour','minute','second','longitude','latitude','variable','value','station'),
            #            'formats': ('i4', 'i2', 'i2', 'i2', 'i2', 'i2', 'f8', 'f8', 'i4', 'f8', 'S6')})
            print('OBS from Radiosounding\n',ds)
        else:
            print('OBS: The specified radsound file does not exist !!!')
    else:
        #filename = path_tmp_obs + "OBSSATELLITE/" +  tmpfilename  # Puhti
        filename = path_tmp_obs + tmpfilename  # Lumi 
        print('obs satellite tmp-filename is:', filename)
        if os.path.isfile(filename)==True:
            # 2023-08-25 - no need to skip 1st row with header ! skiprows=1
            ds = np.genfromtxt(filename)
            #ds = np.loadtxt(filename, dtype={'names': ('year', 'month', 'day', 
            #            'hour','minute','second','longitude','latitude','variable','value','station'),
            #            'formats': ('i4', 'i2', 'i2', 'i2', 'i2', 'i2', 'f8', 'f8', 'i4', 'f8', 'S6')})
            print('OBS from Satellite\n',ds)
        else:
            print('OBS: The specified satellite file does not exist !!!')
    return ds

#----------------------------------------------------------------------------
#  Function to read file with modeled results extracted 
#  (& interpolated to points of observatio) from gsv-file
#  (YYYY MM DD HH MM SS)
#----------------------------------------------------------------------------

def readmod(type_mod, path_tmp_mod, tmpfilename):
    """
    Read extracted & interpolated to points of observatiobs the modelled results 
         from gsv-file for ground-based (SYNOP), radiosounding (TEMP), and 
         satellite (AMSU-A) data
    Input
    -------
    type_mod: synop, radsound, satellite
        Acronyms in names of files with temporarily extracted modelled results
    Output
    -------
    ds: dataset with extracted & interpolated modelled results
    """
    
    if type_mod=='synop':
        #filename = path_tmp_mod + "MODSYNOP/" + tmpfilename # Puhti
        filename = path_tmp_mod + tmpfilename # Lumi
        print('mod synop tmp-filename is:', filename)
        if os.path.isfile(filename)==True:
            # 2023-08-30 - no need to skip 1st row with header ! skiprows=1
            # 2020-01-20 00:00:00 24.4916 60.999199 273.143829345703  101150 39
            # 2020 01 20 00 00 00 24.4916 60.999199 273.143829345703  39 101150
            ds = np.genfromtxt(filename)
            #, dtype={'names': ('year-month-day', 
            #            'hour-minute-second','longitude','latitude','value','station','variable'),
            #            'formats': ('S10', 'S8', 'f8', 'f8', 'f8', 'S6', 'i4')})   # v#1
            #, dtype={'names': ('year', 'month', 'day', 'hour', 'minute', 'second',
            #            'longitude','latitude','value','variable','station'),
            #            'formats': ('S10', 'S8', 'f8', 'f8', 'f8', 'S6', 'i4')})   # v#2      
            #print('MOD for Synop stations\n',ds)
        else:
            print('MOD: The specified synop file does not exist !!!')
    elif type_obs=='radsound':
        #filename = path_tmp_mod + "MODRADSOUND/" + tmpfilename # Puhti
        filename = path_tmp_mod + tmpfilename # Lumi
        print('mod radsound tmp-filename is:', filename)
        if os.path.isfile(filename)==True:
            # 2023-08-30 - no need to skip 1st row with header ! skiprows=1
            ds = np.loadtxt(filename, dtype={'names': ('year', 'month', 'day', 
                        'hour','minute','second','longitude','latitude','value','variable','station'),
                        'formats': ('i4', 'i2', 'i2', 'i2', 'i2', 'i2', 'f8', 'f8', 'f8', 'S6', 'i4')})
            print('MOD for Radiosounding\n',ds)
        else:
            print('MOD: The specified radsound file does not exist !!!')
    else:
        #filename = path_tmp_mod + "MODSATELLITE/" + tmpfilename # Puhti
        filename = path_tmp_mod + tmpfilename # Lumi
        print('mod satellite tmp-filename is:', filename)
        if os.path.isfile(filename)==True:
            # 2023-08-30 - no need to skip 1st row with header ! skiprows=1
            ds = np.loadtxt(filename, dtype={'names': ('year', 'month', 'day', 
                        'hour','minute','second','longitude','latitude','value','variable','station'),
                        'formats': ('i4', 'i2', 'i2', 'i2', 'i2', 'i2', 'f8', 'f8', 'f8', 'S6', 'i4')})
            print('MOD for Satellite\n',ds)
        else:
            print('MOD: The specified satellite file does not exist !!!')
    return ds

#############################################################################
#                             MAIN PROGRAM
#############################################################################
# Passing parameters to program
try:
    fnamedate = sys.argv[1]; met_param = sys.argv[2]
    print (fnamedate)
    print (met_param)
    print ('Command line: py-script-name YYYYMMDDHH met_param-acronym')
except:
    print ("Usage:",sys.argv[0], "fnamedate met_aparm"); sys.exit(1)

##fnamedate = '2020012000'    
# list to list meteorological variables (add here):
##met_param="2t"

# Paths to temporarily storing of extracted OBS and MOD data ...
# ----- on Puhti
##path_tmp_obs='/scratch/project_2000640/wrk/ama/DEOBS/OBSDATA/'
##path_tmp_mod='/scratch/project_2000640/wrk/ama/DEOBS/MODDATA/'
# ----- on Lumi
#path_tmp_obs='/projappl/project_465000454/ama/DEOBS/SYNOP/DATAOBS/'
#path_tmp_mod='/projappl/project_465000454/ama/DEOBS/SYNOP/DATAMOD/'
path_tmp_obs='DATAOBS/'
path_tmp_mod='DATAMOD/'

print ('Path to dir with extacted OBS data: ', path_tmp_obs)
print ('Path to dir with extacted MOD data: ', path_tmp_mod)


# Constructing exisitng names of files with OBS & MOD data
# for SYNOP (fx: synop_2020012000 and mod_synop_2020012000_2t.dat)
tmpfname_obs_synop     = 'obs_synop_'+fnamedate+'.dat'
tmpfname_mod_synop     = 'mod_synop_'+fnamedate+'_'+met_param+'.dat'
print ('Name of file with synop OBS data: ', tmpfname_obs_synop)
print ('Name of file with synop MOD data: ', tmpfname_mod_synop)
#tmpfname_obs_synop='synop_2020012000.dat'
#tmpfname_mod_synop='mod_synop_2020012000_2t.dat'
#print (tmpfname_obs_synop, tmpfname_mod_synop)

# for RADSOUND
tmpfname_obs_radsound  = 'obs_radsound_'+fnamedate+'.dat'
tmpfname_mod_radsound  = 'mod_radsound_'+fnamedate+'_'+met_param+'.dat'
# for SATELLITE
tmpfname_obs_satellite  = 'obs_satellite_'+fnamedate+'.dat'
tmpfname_mod_satellite  = 'mod_satellite_'+fnamedate+'_'+met_param+'.dat'

##############################################################################
# Observations ...
for ii in range(0,1):  #(0,3) # testing on SYNOP obs data
    if ii == 0:
        type_obs='synop'
        tmpfilename=tmpfname_obs_synop
        ds_obs_synop=readobs(type_obs, path_tmp_obs, tmpfilename)
        rec_ds_obs_synop=len(ds_obs_synop)
        #t print('OBS from Synop stations: ',rec_ds_obs_synop,' records \n',ds_obs_synop)
    elif ii == 1:
        type_obs='radsound'
        tmpfilename=tmpfname_obs_radsound
        ds_obs_radsound=readobs(type_obs, path_tmp_obs, tmpfilename)
        rec_ds_obs_radsound=len(ds_obs_radsound)
        ##print('OBS from Radiosounding: ',rec_ds_obs_radsound,' records \n',ds_obs_radsound)
    else:
        type_obs='satellite'
        tmpfilename=tmpfname_obs_satellite
        ds_obs_satellite=readobs(type_obs, path_tmp_obs, tmpfilename)
        rec_ds_obs_satellite=len(ds_obs_satellite)
        ##print('OBS from Satellites: ',rec_ds_obs_satellite,' records \n',ds_obs_satellite)
        
##############################################################################
# Modelled results ...
for jj in range(0,1):  #(0,3) # testing on SYNOP mod data
    if jj == 0:
        type_mod='synop'
        tmpfilename=tmpfname_mod_synop
        ds_mod_synop=readmod(type_mod, path_tmp_mod, tmpfilename)
        rec_ds_mod_synop=len(ds_mod_synop)
        #t print('MOD for Synop stations: ',rec_ds_mod_synop,' records \n',ds_mod_synop)
    elif jj == 1:
        type_mod='radsound'
        tmpfilename=tmpfname_mod_radsound
        ds_mod_radsound=readobs(type_mod, path_tmp_mod, tmpfilename)
        rec_ds_mod_radsound=len(ds_mod_radsound)
        ##print('MOD for Radiosounding: ',rec_ds_mod_radsound,' records \n',ds_mod_radsound)
    else:
        type_mod='satellite'
        tmpfilename=tmpfname_mod_satellite
        ds_mod_satellite=readobs(type_mod, path_tmp_mod, tmpfilename)
        rec_ds_mod_satellite=len(ds_mod_satellite)
        ##print('MOD for Satellites: ',rec_ds_mod_satellite,' records \n',ds_mod_satellite)

##############################################################################
# Calculate mismatch/difference/bias between synop OBS and MOD

#OBS: ('year', 'month', 'day', 'hour','minute','second','longitude','latitude','value','variable','station')
#MOD: ('year-month-day', 'hour-minute-second','longitude','latitude','value','variable','station')

# Checking number of records (obs & mod) => must be the same 
nr_obs = ds_obs_synop.shape[0]
nr_mod = ds_mod_synop.shape[0]
#t print(nr_obs, nr_mod)

val_obs_synop=np.zeros(nr_obs)
val_mod_synop=np.zeros(nr_mod)

# Converting to modelled values to same scale as observed ... 
val_obs_synop = np.array(ds_obs_synop [0:nr_obs,  9])           # synop obs (2t in deg C)
#t print(val_obs_synop)

if met_param == '2t':
    val_mod_synop = np.array(ds_mod_synop [0:nr_mod,  8] - 273.15)  # synop mod (in deg K -> convert to T-273.15)
else:
    val_mod_synop = np.array(ds_mod_synop [0:nr_mod,  8]) # no need to convert modelled data

#t print(val_mod_synop)
    
# Calculating a mismatch/difference/bias ...
dif_synop = val_obs_synop - val_mod_synop

#t print(val_obs_synop, val_mod_synop, dif_synop)

# Calculating statistcis on synop obs & mod & difference data (self-control & a possibility of flagging)
# Maximum
obs_max = max(val_obs_synop); mod_max = max(val_mod_synop); dif_synop_max = max(dif_synop)
# Minimum
obs_min = min(val_obs_synop); mod_min = min(val_mod_synop); dif_synop_min = min(dif_synop)
# Mean
obs_mean = np.mean(val_obs_synop); mod_mean = np.mean(val_mod_synop); dif_synop_mean = np.mean(dif_synop)
# Median
obs_med = np.median(val_obs_synop); mod_med = np.median(val_mod_synop); dif_synop_med = np.median(dif_synop)
# Variance
obs_var = np.var(val_obs_synop); mod_var = np.var(val_mod_synop); dif_synop_var = np.var(dif_synop)

# Correl. coefficient - (!) not needed for climate simulations
#r_cor = np.corrcoef(val_obs_synop, val_mod_synop)
#cor_obs_mod = r_cor[1, 0]
#t print(cor_obs_mod)

#t print (obs_max, mod_max, obs_min, mod_min, obs_mean, mod_mean, obs_med, mod_med, obs_var, mod_var)
#t print (dif_synop_max, dif_synop_min, dif_synop_mean, dif_synop_median, dif_synop_var)

# Time slot for which stats is calculated (used in plotting fraph)
val_yy = str(int(np.array(ds_obs_synop [1,0])))
val_mm = str(int(np.array(ds_obs_synop [1,1])))
val_dd = str(int(np.array(ds_obs_synop [1,2])))
val_hh = str(int(np.array(ds_obs_synop [1,3])))
print ('Time slot:', val_yy, val_mm, val_dd, val_hh)

print ('          OBS      MOD      DIF')
print ('Max  {:8.3f} {:8.3f} {:8.3f}'.format(obs_max, mod_max, dif_synop_max))
print ('Min  {:8.3f} {:8.3f} {:8.3f}'.format(obs_min, mod_min, dif_synop_min))
print ('Mean {:8.3f} {:8.3f} {:8.3f}'.format(obs_mean, mod_mean, dif_synop_mean))
print ('Med  {:8.3f} {:8.3f} {:8.3f}'.format(obs_med, mod_med, dif_synop_med))
print ('Var  {:8.3f} {:8.3f} {:8.3f}'.format(obs_var, mod_var, dif_synop_var))

emsp='     ' # empty space for text arrangement on graphs
rec_max=emsp +str('{:8.3f}'.format(obs_max)) +str('{:8.3f}'.format(mod_max)) +str('{:8.3f}'.format(dif_synop_max)) +'   Max'
rec_min=emsp +str('{:8.3f}'.format(obs_min)) +str('{:8.3f}'.format(mod_min)) +str('{:8.3f}'.format(dif_synop_min)) +'   Min'
rec_mean=emsp+str('{:8.3f}'.format(obs_mean))+str('{:8.3f}'.format(mod_mean))+str('{:8.3f}'.format(dif_synop_mean))+'   Mean'
rec_med=emsp +str('{:8.3f}'.format(obs_med)) +str('{:8.3f}'.format(mod_med)) +str('{:8.3f}'.format(dif_synop_med)) +'   Median'
rec_var=emsp +str('{:8.3f}'.format(obs_var)) +str('{:8.3f}'.format(mod_var)) +str('{:8.3f}'.format(dif_synop_var)) +'   Variance'
rec_nr_obs=emsp+emsp+emsp+str(int('{:5d}'.format(nr_obs)))+'  -  # of stations'

#print ('{:8.3f} {:8.3f} {:8.3f} {:8.3f} {:8.3f} {:8.3f} {:8.3f} {:8.3f}'.
#       format(obs_max, mod_max, obs_min, mod_min, obs_mean, mod_mean, obs_med, mod_med))

# Plotting : scatter plot for observed vs. modelled values
# Selecting type of plot
if type_obs == 'synop':
    type_title = 'SYNOP'
elif type_obs == 'radsound':
    type_title = 'RADIOSOUNDING'
else:
    type_title = 'SATELLITE'

txt_time_slot='Time slot (YYYY, MM, DD, HH): '+val_yy+'-'+val_mm+'-'+val_dd+'-'+val_hh
txt_title=type_title+': Meteorological Variable: '+ met_param

fig=plt.figure()

plt.scatter(val_mod_synop, val_obs_synop)
plt.xlabel('Modelled Value (MOD)')
plt.ylabel('Observed Value (OBS)')
plt.suptitle(txt_title, fontsize=16, weight='bold')
plt.title(txt_time_slot, fontsize=12)
plt.text(mod_max,obs_max,'        OBS    MOD    DIF', weight='bold')
plt.text(mod_max,obs_max-0.5, rec_max)
plt.text(mod_max,obs_max-1.0, rec_min)
plt.text(mod_max,obs_max-1.5, rec_mean)
plt.text(mod_max,obs_max-2.0, rec_med)
plt.text(mod_max,obs_max-2.5, rec_var)
plt.text(mod_max,obs_max-3.0, rec_nr_obs)

##plt.show()

# Saving figure to png-format-file
fnamesave=type_obs+'_'+fnamedate+'_'+met_param+'.png'
plt.savefig(fnamesave, dpi = 300, bbox_inches="tight")

print('Output graphical file is: ', fnamesave)

#exit()
