# -*- coding: utf-8 -*-
  
#EDITOR : Thomas PONTHIEU
#VERSION : V1 (2024)
#LAST MODIFICATION : 22/10/2024

import netCDF4 as nc
import matplotlib.pyplot as plt
import argparse
import sys
import os

print("")
print("###################################################################################################################")
print("####################################### AIRCORE NETCDF READING PROGRAM  ###########################################")
print("###################################################################################################################")
    
###################
## CONFIGURATION ##
###################
OUTPUT_FIGURES_DIRECTORY = "./"

###############
## ARGUMENTS ##
###############
flag_help = 0
if len(sys.argv) != 3:
    if len(sys.argv) == 1:
        flag_help = 1
    else:
        flag_help = 1
        print("ERROR: wrong number of arguments")
        print("")
else:
    parser = argparse.ArgumentParser()
    parser.add_argument('INPUTFILE', type=str)
    parser.add_argument('MODE', type=int)
    args = parser.parse_args()
    inputfile = args.INPUTFILE
    display_mode = args.MODE

    if display_mode != 0 and display_mode != 1:
        flag_help = 1
        print("ERROR: MODE argument must be equal to 0 or 1")
        print("")

if flag_help == 1:
    print("###################################################################################################################")
    print("#                                                                                                                 #")
    print("#  Help: python aircore_database_read.py                                                                          #")
    print("#                                                                                                                 #")
    print("#  Syntax: python aircore_database_read.py [PATH/]FILE.nc MODE                                                    #")
    print("#    where:                                                                                                       #")
    print("#     MODE = 0 : displaying netCDF architecture (global attributes + groups Flight/AirCore/Meteosonde/Trajectory) #")
    print("#      Example: python aircore_database_read.py AC_TRN_202312151008_012_v2024.nc 0                                #")
    print("#     MODE = 1 : plotting if available CO2, CH4, CO, CO_G53, N2O, T, RH vs Pressure                               #")
    print("#      Example: python aircore_database_read.py AC_TRN_202312151008_012_v2024.nc 1                                #")
    print("#       Output: figures in your working directory                                                                 #")
    print("#                                                                                                                 #")
    print("###################################################################################################################")

else:

    print("")
    if not os.path.isfile(inputfile):
        print(f"ERROR: inputfile {inputfile} doesn't exist")
        sys.exit()
    print("-------------------------")
    print("---- INPUT ARGUMENTS ----")
    print("-------------------------")
    print("INPUTFILE : ", inputfile)
    print("")
    if display_mode == 0:
        print("MODE : displaying netCDF architecture (global attributes + groups Flight/AirCore/Meteosonde/Trajectory)")
    elif display_mode == 1:
        print("MODE : plotting if available CO2, CH4, CO, CO_G53, N2O, T, RH vs Pressure")
    print("")

    ###############
    ## FUNCTIONS ##
    ###############
    def Avail_data(dataset, var):
        avail_var = dataset.groups["Availability"].variables[f"{var}"].value
        return avail_var

    def Avail_AircoreVar(dataset, var):
        avail_var = dataset.groups["Availability"].groups["AirCore_fields"].variables[f"{var}"].value
        return avail_var

    def Subgroup_avail_data(dataset, subgroup):
        avail_fields = dataset.groups["Availability"].groups[f"{subgroup}"]
        avail_fields_names = avail_fields.variables.keys()
        var_tab = []
        for var_name in avail_fields_names:
            variable_values = avail_fields.variables[var_name].value
            lines = [var_name, variable_values]
            var_tab.append(lines)
        return var_tab

    def Read_netCDF_File(dataset, var):
        pressure = dataset.groups["AirCore"].variables["P"][:]
        atm_var = dataset.groups["AirCore"].variables[f"{var}"][:]
        return pressure, atm_var

    ##########
    ## MAIN ##
    ##########
    
    dataset = nc.Dataset(inputfile, "r")

    #GLOBAL ATTRIBUTES
    if display_mode == 0:
        print("---------------------------")
        print("---- GLOBAL ATTRIBUTES ----")
        print("---------------------------")
        global_attributes = dataset.ncattrs()
        for attr_list in global_attributes:
            attr_value = getattr(dataset, attr_list)
            print(f"   {attr_list} = {attr_value}")
         
    #FLIGHT
    if display_mode == 0:
        print("----------------------")
        print("---- GROUP FLIGHT ----")
        print("----------------------")
        group_flight = dataset.groups["Flight"].variables.keys()
        for flight_list in group_flight:
            if flight_list == "gas_calibration_scale":
                flight_comment = dataset.groups["Flight"].variables[flight_list].value
                print("  ", flight_list ," = ", flight_comment)
            else:
                flight_comment = dataset.groups["Flight"].variables[flight_list].value
                print("  ", flight_list ," = ", flight_comment)
         
    #VARIABLES AVAILABILITY
    if display_mode == 0:
        print("----------------------------")
        print("---- GROUPS + VARIABLES ----")
        print("----    AVAILABILITY    ----")
        print("----------------------------")
        var =["AirCore","Meteosonde_ascent","Meteosonde_descent","Trajectory"]
        for var_name in var:
            Avail_value = Avail_data(dataset, var_name)
            print(f"{var_name} = {Avail_value}")
            
            if Avail_value == "yes":
                
                if var_name == "AirCore":
                    subgroup_var_tab = Subgroup_avail_data(dataset, "AirCore_fields")
                    size_name = max(len(item[0]) for item in subgroup_var_tab)
                    size_value = max(len(item[1]) for item in subgroup_var_tab)
                    for i in range(len(subgroup_var_tab)):
                        line_format = "{:<{size_name}} = {:<{size_value}}"
                        print("  ",line_format.format(subgroup_var_tab[i][0], subgroup_var_tab[i][1], size_name=size_name, size_value=size_value))
                
                if var_name == "Meteosonde_descent" or var_name == "Meteosonde_ascent":
                    subgroup_var_tab = Subgroup_avail_data(dataset, "Meteosonde_fields")
                    size_name = max(len(item[0]) for item in subgroup_var_tab)
                    size_value = max(len(item[1]) for item in subgroup_var_tab)
                    j = 11
                    for j in range(len(subgroup_var_tab)):
                        line_format = "{:<{size_name}} = {:<{size_value}}"
                        print("  ",line_format.format(subgroup_var_tab[j][0], subgroup_var_tab[j][1], size_name=size_name, size_value=size_value))
                
                if var_name == "Trajectory":
                    subgroup_var_tab = Subgroup_avail_data(dataset, "Trajectory_fields")
                    size_name = max(len(item[0]) for item in subgroup_var_tab)
                    size_value = max(len(item[1]) for item in subgroup_var_tab)
                    k = 20
                    for k in range(len(subgroup_var_tab)):
                        line_format = "{:<{size_name}} = {:<{size_value}}"
                        print("  ",line_format.format(subgroup_var_tab[k][0], subgroup_var_tab[k][1], size_name=size_name, size_value=size_value))
                
                print("")

    if display_mode == 1:
        var_list = ['CO2', 'CH4', 'CO', 'CO_G53', 'N2O', 'T', 'RH']
        for var in var_list:
            avail_value = Avail_AircoreVar(dataset, var)
            if avail_value == "yes":
                print("-------------------------------")
                print(f"---- PLOTTING {var} PROFILES ----")
                print("-------------------------------")
                 
                figure = plt.figure(figsize=(8, 6))
                linewidth = 0.75
                 
                pressure, pltvar = Read_netCDF_File(dataset, var)
                plt.plot(pltvar, pressure, lw=linewidth, color='black')
                 
                if var == 'CO2':
                    plt.xlabel(f"{var} concentration (ppm)")
                elif var == 'T':
                    plt.xlabel(f"{var} (kelvin)")
                elif var == 'RH':
                    plt.xlabel(f"{var} (%)")
                else:
                    plt.xlabel(f"{var} concentration (ppb)")
                
                plt.ylabel("Pressure (hPa)")
                plt.ylim(20, 1020)
                plt.title(f"{var} vs Pressure")
                plt.grid()
                plt.grid(which='both', axis='y', color='gray', linestyle='--', linewidth=0.5)
                plt.gca().invert_yaxis()
                plt.yscale("log")

                site = inputfile.split('/')[-1].split('_')[1]
                yyyymmdd = inputfile.split('/')[-1].split('_')[2][:8]
                flight_nb = inputfile.split('/')[-1].split('_')[3]
                output_figure = OUTPUT_FIGURES_DIRECTORY + '/' + site + '_' + yyyymmdd + '_' + flight_nb + '_' + var + 'vsP.jpg'
                plt.savefig(output_figure)
                print("")
                print(f"Figure : {output_figure}")
                print("")

        plt.show()

