#!/usr/bin/env python3
# Copyright 2023 - Barcelona Supercomputing Center
# Author: Rodrigo Martín Posada (2021), updated by : Amirpasha Mozaffari (2023) 
# TBD

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

import pyoasis
from   pyoasis import OASIS  # pylint: disable=no-name-in-module
import numpy as np

if __name__ == '__main__':

### Read in YAML file 
    current_path = os.getcwd()
    parser=argparse.ArgumentParser()
    parser.add_argument("--run_dir",type=str)
    parser.add_argument("--yaml_file_path",type=str)
    args = parser.parse_args()    
    yamel_file_path = os.path.join(current_path,args.yaml_file_path)
    if os.path.isfile(yamel_file_path) == False: 
        print("YAML file does not exist. Aborted")
        raise NameError("YAML file does not exist.")
        
    # Open the file and load the file 
    with open(yamel_file_path) 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")
    
    # 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 

### Pass the values to original code 
    
    # Init OASIS
    comp = pyoasis.Component(model_name)
    local_comm = comp.localcomm
    # Get rank in local communicator
    mype = comp.localcomm.rank
    npes = comp.localcomm.size  
    
    # BUG : shouldnt it be the npes and not mpye (?)
    mype_value = str(100+mype)
    out_model = '{model_name}.out_{mype_value}'.format(model_name = model_name, mype_value = mype_value)

    logging.basicConfig(filename=out_model, format='%(message)s', level=logging.DEBUG)
    logging.debug('-----------------------------------------------------------')
    logging.info('I am {model_name} process with rank : {mype}'.format(model_name = model_name,
                                                                         mype = mype))
    #logging.debug('I am IFS-toy process with rank : {}'.format(mype))
    logging.debug('in my local communicator gathering {} processes'.format(npes))
    logging.debug('----------------------------------------------------------')

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  PARTITION DEFINITION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    # Definition of the local partition
    partition = pyoasis.SerialPartition(simulation_nx)
    
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  DECLARATION OF THE COUPLING FIELDS
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    out_coupling_fields={}

    for v in coupling_out_vars:
        out_coupling_fields[v] = pyoasis.Var(v, partition, OASIS.OUT)
        logging.warning('var_id {}, added as OASIS_OUT var'.format(v))


    in_coupling_fields={}

    for v in coupling_in_vars:
        in_coupling_fields[v] = pyoasis.Var(v, partition, OASIS.IN)
        logging.info('var_id {}, added as OASIS_IN var'.format(v))


    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         TERMINATION OF DEFINITION PHASE
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    logging.debug('End of initialisation phase')

    comp.enddef()

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  TIME STEP LOOP
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    logging.debug('Timestep, field min, mean and max value')

    field_data = pyoasis.asarray(np.full(simulation_nx, -1.0))

    for itap_sec in range(0, simulation_run_length_sec, simulation_coupling_interval):

        logging.debug('ITAP_SEC: {} receiving fields'.format(itap_sec))
        logging.debug(str(in_coupling_fields.keys()))
        for id, field in in_coupling_fields.items():
            field.get(itap_sec, field_data)      
            logging.debug('{} -> {:.3e} {:.3e} {:.3e}'.format( id, np.min(field_data), np.mean(field_data), np.max(field_data)))
        logging.debug('--------------------------------------\n')

        logging.debug('ITAP_SEC: {} sending fields'.format(itap_sec))
        logging.debug(str(out_coupling_fields.keys()))
        for id, field in out_coupling_fields.items():
            sigma=1
            mu=1
            # send dummy data (100 for runoff, random for other variables)
            #TODO send dummy data read from oasis restart file
            if id == "A_Runoff":
                field_data[:]=1.4658e-06
                field_data[:]=100
            else:
                field_data[:]=sigma * np.random.randn(simulation_nx) + mu
            field.put(itap_sec, field_data)      
            logging.debug('{} put -> {:.3e} {:.3e} {:.3e}'.format( id, np.min(field_data), np.mean(field_data), np.max(field_data)))
        logging.debug('--------------------------------------\n')

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         TERMINATION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    logging.debug('End of the program')

    print('Model {} is finished successfully.'.format(model_name))

    del comp



