#!/usr/bin/env python3
# Copyright 2023 - Barcelona Supercomputing Center
# Authors: Amirpasha Mozaffari
# MIT License

import yaml
from yaml.loader import SafeLoader
import os 
import argparse
import sys 



# class yaml_file:


def read_yaml_meta(yaml_file):
    '''
    Reads the metadata of the yaml file 
    it includes type  which could be:
    type 1 : DataCoupler  
    type 2 : NEMO / IFS / RUNOFF toy model 
    '''
    with open(yaml_file) as file:
        read_in = yaml.load(file, Loader=SafeLoader)
        file.close()

    # Model Information
    yaml_type      = read_in.get('type')
    return yaml_type 


def read_yaml_conf_type1(yaml_file):
    '''
    Reads the experiment yaml file 
    '''

    with open(yaml_file) as file:
        read_in = yaml.load(file, Loader=SafeLoader)
        file.close()
    
    # Model Information
    ModelNameSend       = read_in.get('ModelNameSend')
    ModelNameReceive      = read_in.get('ModelNameReceive')
    LogFileName     = read_in.get('LogFileName')
    PartitionType   = read_in.get('PartitionType')

    ## Run Information 

    RunLengthSec    = read_in.get('RunLengthSec')
    TimeStepSec     = read_in.get('TimeStepSec')
    StartYear       = read_in.get('StartYear')
    StartMonth      = read_in.get('StartMonth')
    StartDay        = read_in.get('StartDay')
    FixYear         = read_in.get('FixYear')
    GridInfo        = read_in.get('GridInfo')
    LDebug          = read_in.get('LDebug')

    # Create the log_file name and sanity-check 

    # Check the type of input in yaml file 
    if "FileInputVars" in read_in : 
        FileInputVars_readin = {}
        FileInputVars_readin   = read_in.get('FileInputVars')
        # Name of the Variable for this experiment 
        FileInputVarNames  = list(FileInputVars_readin.keys())
        # Variable configuration for this experiment 
        FileInputVarFields_conf = list(FileInputVars_readin.values())

    if "OasisOutputVars" in read_in :
        OasisOutputVars_readin = {} 
        OasisOutputVars_readin = read_in.get('OasisOutputVars')
        print("NOTE: OasisOutputVars_readin : {}".format(OasisOutputVars_readin), file=sys.stderr)
        # Name of the Variable for this experiment 
        OasisOutputVarsNames  = list(OasisOutputVars_readin.keys())

        # Variable configuration for this experiment 
        OasisOutputVars_conf = list(OasisOutputVars_readin.values())

    return ModelNameSend, ModelNameReceive, LogFileName, PartitionType, RunLengthSec, TimeStepSec, StartYear, StartMonth,\
             StartDay, FixYear, GridInfo, LDebug, FileInputVarFields_conf, OasisOutputVars_conf, OasisOutputVars_readin


def read_yaml_conf_type2(yaml_file): 

# Open the file and load the file 
    with open(yaml_file) as f:
        read_in_data = yaml.load(f, Loader=SafeLoader)
        f.close()

### Read in file 
    #read_yaml(read_in_data)
    model = read_in_data.get("model")
    model_name = model.get("ModelName")
    PartitionType = model.get("PartitionType")
    
    # Read in Simulation information
    simulation = read_in_data.get("simulation")
    simulation_coupling_interval = simulation.get("coupling_interval")
    simulation_run_length_sec = simulation.get("run_length_sec")
    simulation_nx = simulation.get("nx")

    # Read in Coupling infomration
    coupling = read_in_data.get("coupling")
    coupling_out_vars = coupling.get("out_vars")
    coupling_in_vars = coupling.get("in_vars")
    restart_file = coupling.get("restart_file")  # This feature is not used for now. TODO 

    return model, model_name, PartitionType,  simulation_coupling_interval, simulation_run_length_sec,\
              simulation_nx, coupling_out_vars, coupling_in_vars, restart_file



