#!/usr/bin/env python3
# Copyright 2021 - Barcelona Supercomputing Center
# Author: Rodrigo Martín Posada
# MIT License

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

import yaml
from yaml.loader import SafeLoader
import logging
import argparse
import os
import sys 
from YAMLlUtils import read_yaml_meta, read_yaml_conf_type1, read_yaml_conf_type2

# TODO 
# It will read in the YAML file instead of the namelist. It will receive the name of the YAMl file as argparse value 
# it should turn to class in next step round of refactoring 


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


def variable_exists(var_name):
    return var_name in globals()


if __name__ == '__main__':

    
    parser=argparse.ArgumentParser()
    parser.add_argument("--conf",type=str)
    # parser.add_argument("--namcouple",type=str)
    args = parser.parse_args()
    yaml_conf_file = args.conf 
    if os.path.isfile(yaml_conf_file) == False: 
        print("YAML file does not exist. Aborted")
        raise NameError("YAML file does not exist.")

    # Determine the type of yaml file 
    yaml_type = read_yaml_meta(yaml_conf_file)

    # Read in DataCoupler YAML file type 

    if yaml_type == 1 :
         
        ModelNameSend, ModelNameReceive, LogFileName, PartitionType, RunLengthSec, TimeStepSec, StartYear, StartMonth,\
             StartDay, FixYear, GridInfo, LDebug, FileInputVarFields_conf, OasisOutputVars_conf, OasisOutputVars_readin = read_yaml_conf_type1(yaml_conf_file)

        ModelName = ModelNameReceive
        COUPLING_INTERVAL = TimeStepSec
        RUN_LENGTH_SEC = RunLengthSec
        nx = GridInfo 
        
        # Extract the send_id, receive_id, send_grid_name, receive_grid_name from the OasisConf.
        coupling_in_vars = [] 
        for v in OasisOutputVars_readin.keys() :
            coupling_in_vars.append(OasisOutputVars_readin[v]['receive_id'])
 
    # Read in ToModel Yaml file type 
    elif yaml_type == 2: 
        model, model_name, PartitionType, simulation_coupling_interval, simulation_run_length_sec, \
            simulation_nx, coupling_out_vars, coupling_in_vars, restart_file = read_yaml_conf_type2(yaml_conf_file)
        
        ModelName         =    model_name
        COUPLING_INTERVAL = simulation_coupling_interval
        RUN_LENGTH_SEC =  simulation_run_length_sec
        nx = simulation_nx 
    

    # Init OASIS
    comp = pyoasis.Component(ModelName)

    # 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 = '{ModelName}.out_{mype_value}'.format(ModelName = ModelName, mype_value = mype_value)

    logging.basicConfig(filename=out_model, format='%(message)s', level=logging.DEBUG)
    logging.debug('-----------------------------------------------------------')
    logging.info('I am {ModelName} process with rank : {mype}'.format(ModelName = ModelName,
                                                                         mype = mype_value))
    #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
        
    if yaml_type == 1 : #and PartitionType == "ApplePartition":
        il_size, il_offset = def_local_partition(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)

    elif yaml_type == 2 : #and PartitionType == "SerialPartition":  
        partition = pyoasis.SerialPartition(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 = {}  # Create a dictionary to store the variables

    for name in coupling_in_vars:
        var = pyoasis.Var(name, partition, OASIS.IN)
        in_coupling_fields[name] = var    
        logging.debug('var_id {}: {}'.format(name, var._id))


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

    comp.enddef()

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

    if yaml_type == 1 : 
        fields = {}

        for name in coupling_in_vars:
            fields[name] = pyoasis.asarray(np.full(il_size, -1.0))

        for itap_sec in range(0, RUN_LENGTH_SEC, COUPLING_INTERVAL):
            for name in coupling_in_vars:
                in_coupling_fields[name].get(itap_sec, fields[name])

    elif yaml_type == 2 :

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

        for itap_sec in range(0, RUN_LENGTH_SEC, COUPLING_INTERVAL):

            for id, field in in_coupling_fields.items():
                field.get(itap_sec, field_data)      

            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(nx) + mu
                field.put(itap_sec, field_data)      


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

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

    del comp
