#!/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 DataCouplerUtils import read_yaml_conf , read_yaml_meta
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 


L080_NX = 35718  # number of grid cells at T159 resolution
L128_NX = 88838  # number of grid cells at T255 resolution
L256_NX = 348528  # number of grid cells at T511 resolution
CO2_CMODE = False # is CO2 concentration-mode active?
CO2_EMODE = False # is CO2 emission-mode active?


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

## will be replaced with YAML version in step 5 
def read_name_couple(namcouple): 
    with open(namcouple, 'r') as f:
        contents = ""
        CO2_CMODE = None 
        CO2_EMODE = None 

        for line in f:
            li=line.strip()
            if not li.startswith("#"):
                contents += line
        # TODOD find a better way to detect CO2_CMODE
        if 'GLOBAL GLOBAL' in contents:
            CO2_CMODE = True
        if 'A_CO2_emis ' in contents:
            CO2_EMODE = True
        if 'L080' in contents:
            nx = L080_NX
        elif 'L128' in contents:
            nx = L128_NX
        elif 'L256' in contents:
            nx = L256_NX
        else:
            raise Exception('Invalid IFS grid in namcouple')
    return nx


if __name__ == '__main__':

    try:
        parser=argparse.ArgumentParser()
        parser.add_argument("--conf",type=str)
        args = parser.parse_args()
        yaml_conf_file = args.conf 
    except:
        print("NOTE:TOY  Could note read in the input yaml file : {}".format(yaml_conf_file), file=sys.stderr)
    else:
        print("NOTE:TOY input yaml file is : {}".format(yaml_conf_file), file=sys.stderr)


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

    #

    if yaml_type == 1 : 
        ModelName, OasisCPLNG, LogFileName, RunLengthSec, TimeStepSec, StartYear, StartMonth, \
            StartDay, FixYear, GridInfo, LDebug, fileIn_vars_config, oasisOut_vars_config , OasisOutputVars_readin = read_yaml_conf_type1(yaml_conf_file)

        COUPLING_INTERVAL = TimeStepSec
        RUN_LENGTH_SEC = RunLengthSec
        #nx = GridInfo 
 

    elif yaml_type == 2: 
        model, model_name, 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 
    

    # 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
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    if yaml_type == 1 : 
        
        namcouple = 'namcouple'
        nx = read_name_couple(namcouple)

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  PARTITION DEFINITION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
        # Definition of the local partition
        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)
    
    else : 

        # Definition of the local partition for toy model 
        partition = pyoasis.SerialPartition(simulation_nx)


    # Extract the send_id, receive_id, send_grid_name, receive_grid_name from the OasisConf.
    OasisOutputVars = []
    Oasis_receive_id = [] 
    for v in OasisOutputVars_readin.keys() :
        OasisOutputVars.append([OasisOutputVars_readin[v]['send_id'], OasisOutputVars_readin[v]['receive_id'], OasisOutputVars_readin[v]['send_grid_name'], OasisOutputVars_readin[v]['receive_grid_name']])
        Oasis_receive_id.append(OasisOutputVars_readin[v]['receive_id'])

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  DECLARATION OF THE COUPLING FIELDS
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    variables = {}  # Create a dictionary to store the variables

    for name in Oasis_receive_id:
        var = pyoasis.Var(name, partition, OASIS.IN)
        variables[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')

    fields = {}

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

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


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

    del comp
