#!/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


def def_local_partition(nx, npes, mype):
    il_size = nx // npes
    il_offset = mype * il_size
    if (npes-mype) <= nx % npes:
        il_size += 1
        il_offset += (nx % npes) - (npes-mype)
    return il_size, il_offset


if __name__ == '__main__':
    
### Read in YAML file 
    current_path = os.getcwd()
    parser=argparse.ArgumentParser()
    #parser.add_argument("--run_dir",type=str)
    parser.add_argument("--conf",type=str)
    args = parser.parse_args()    
    yamel_file_path = os.path.join(current_path,args.conf)
    
    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
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    # for more than one processor is using a parallel partition and for single processor a serial 

    if npes > 1:
        il_size, il_offset = def_local_partition(simulation_nx, npes, mype)
        logging.debug('Local partition definition')
        logging.debug('il_size = {} il_offset = {} mype = {}'.format(il_size, il_offset, mype)) 
        partition = pyoasis.ApplePartition(il_offset, il_size)
        field_data_size = il_size
        logging.debug('Grid points in the current Apple - partition is: {},'.format(field_data_size))
    else:  
        partition = pyoasis.SerialPartition(simulation_nx)
        field_data_size = simulation_nx 
        logging.debug('Grid points in the current Serial - partition is: {},'.format(field_data_size))

       
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  DECLARATION OF THE COUPLING FIELDS
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    out_field_data  = pyoasis.asarray(np.full(field_data_size, -1.0))
    out_coupling_fields={}
    
    in_coupling_fields = {}  # Create a dictionary to store the variables
    in_field_data={}

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

    for name in coupling_in_vars:
        var = pyoasis.Var(name, partition, OASIS.IN)
        in_coupling_fields[name]  = var
        in_field_data[name] = pyoasis.asarray(np.full(field_data_size, -1.0))
        logging.debug('var_id {}, added as OASIS_IN var'.format(name))
        
            
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         TERMINATION OF DEFINITION PHASE
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    logging.debug('End of initialisation phase')

    comp.enddef()

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


    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, in_field_data[id])
            logging.debug('{} -> {:.3e} {:.3e} {:.3e}'.format( id, np.min(in_field_data[id]), np.mean(in_field_data[id]), np.max(in_field_data[id])))

        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":
                out_field_data[:]=1.4658e-06
                out_field_data[:]=100
            else:
                out_field_data[:]=sigma * np.random.randn(field_data_size) + mu
            field.put(itap_sec, out_field_data)
            logging.debug('{} put -> {:.3e} {:.3e} {:.3e}'.format( id, np.min(out_field_data), np.mean(out_field_data), np.max(out_field_data)))


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

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

    del comp



