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


# 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

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)

    OasisModel, OasisCPLNG, LogFileName, RunLengthSec, TimeStepSec, StartYear, StartMonth, \
        StartDay, FixYear, LDebug, fileIn_vars_config, oasisOut_vars_config , OasisOutputVars_readin = read_yaml_conf(yaml_conf_file)


    # Constants from namelist

    COUPLING_INTERVAL = TimeStepSec 
    RUN_LENGTH_SEC = RunLengthSec 

    # Init OASIS
    comp = pyoasis.Component(OasisModel)

    # Get rank in local communicator
    mype = comp.localcomm.rank
    npes = comp.localcomm.size
    value = str(100+mype)  # ? Not sure what this means , same was in toy_model 
    # Unit for output messages
    out_ifs = ("{OasisModel}.out_{value}").format(OasisModel=OasisModel,value=value)
    logging.basicConfig(filename=out_ifs, format='%(message)s', level=logging.DEBUG)
    logging.debug('-----------------------------------------------------------')
    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
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    with open('namcouple', 'r') as f:
        #contents = f.read()
        contents = ""
        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')
    # 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)
    if CO2_CMODE or CO2_EMODE:
        #nx_co2 = 2 # for 2 hemisphere box model
        nx_co2 = 1
        if mype == 0:
            logging.debug('Local serial partition')
            logging.debug('nx_co2 = {} mype = {}'.format(nx_co2, mype))
            partition2 = pyoasis.SerialPartition(nx_co2)
    else:
        nx_co2 = 0
        partition2 = None


        # Prepare the list that is necessary for 
        OasisOutputVars = [] 
        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']])
            variable = OasisOutputVars_readin[v]['receive_id']
            #variable = pyoasis.Var(OasisOutputVars_readin[v]['receive_id'], partition, OASIS.IN)
            # print("NOTE: {variable_name} : {pyoasis_output}".format(variable_name =v, pyoasis_output=variable), file=sys.stderr)
            print("NOTE:OASIS VARIABLE NAME : {pyoasis_output}".format( pyoasis_output= variable), file=sys.stderr)


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

    A_SST = pyoasis.Var("A_SST", partition, OASIS.IN)
    print("NOTE:OASIS ASST : {pyoasis_output}".format( pyoasis_output=A_SST), file=sys.stderr)
    A_Ice_frac = pyoasis.Var("A_Ice_frac", partition, OASIS.IN)
    print("NOTE:OASIS A_Ice_frac : {pyoasis_output}".format( pyoasis_output=A_Ice_frac), file=sys.stderr)
    logging.debug('var_id A_SST, var_id A_Ice_frac {} {}'.format(A_SST._id, A_Ice_frac._id))
    
    if CO2_CMODE:
        if mype == 0:
            A_CO2_cconc = pyoasis.Var("A_CO2_cconc", partition2, OASIS.IN)
            logging.debug('var_id A_CO2_cconc {}'.format(A_CO2_cconc._id))
            A_CO2_econc = pyoasis.Var("A_CO2_econc", partition2, OASIS.IN)
            logging.debug('var_id A_CO2_econc {}'.format(A_CO2_cconc._id))
    if CO2_EMODE:
        A_CO2_emis = pyoasis.Var("A_CO2_emis", partition, OASIS.IN)
        logging.debug('var_id A_CO2_emis {}'.format(A_CO2_emis._id))
        A_CO2_land = pyoasis.Var("A_CO2_land", partition, OASIS.IN)
        logging.debug('var_id A_CO2_land {}'.format(A_CO2_land._id))
        A_CO2_ocean = pyoasis.Var("A_CO2_ocean", partition, OASIS.IN)
        logging.debug('var_id A_CO2_ocean {}'.format(A_CO2_ocean._id))

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         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_recv_A_SST = pyoasis.asarray(np.full(il_size, -1.0))
    field_recv_A_Ice_frac = pyoasis.asarray(np.full(il_size, -1.0))
    if CO2_CMODE and mype == 0:
        field_recv_A_CO2_cconc = pyoasis.asarray(np.full(nx_co2, -1.0))
        field_recv_A_CO2_econc = pyoasis.asarray(np.full(nx_co2, -1.0))
    if CO2_EMODE:
        field_recv_A_CO2_emis = pyoasis.asarray(np.full(il_size, -1.0))
        field_recv_A_CO2_land = pyoasis.asarray(np.full(il_size, -1.0))
        field_recv_A_CO2_ocean = pyoasis.asarray(np.full(il_size, -1.0))

    for itap_sec in range(0, RUN_LENGTH_SEC, COUPLING_INTERVAL):
        A_SST.get(itap_sec, field_recv_A_SST)
        A_Ice_frac.get(itap_sec, field_recv_A_Ice_frac)

        logging.debug('ITAP_SEC: {}'.format(itap_sec))
        logging.debug('A_SST -> {:.3e} {:.3e} {:.3e}'.format(
            np.min(field_recv_A_SST),
            np.mean(field_recv_A_SST),
            np.max(field_recv_A_SST)))
        logging.debug('A_Ice_frac -> {:.3e} {:.3e} {:.3e}'.format(
            np.min(field_recv_A_Ice_frac),
            np.mean(field_recv_A_Ice_frac),
            np.max(field_recv_A_Ice_frac)))
        if CO2_CMODE:
            if mype == 0:
                A_CO2_cconc.get(itap_sec, field_recv_A_CO2_cconc)
                logging.debug('A_CO2_cconc -> {:.3e} {:.3e} {:.3e}'.format( np.min(field_recv_A_CO2_cconc),  np.mean(field_recv_A_CO2_cconc), np.max(field_recv_A_CO2_cconc)))
                A_CO2_econc.get(itap_sec, field_recv_A_CO2_econc)
                logging.debug('A_CO2_econc -> {:.3e} {:.3e} {:.3e}'.format( np.min(field_recv_A_CO2_econc),  np.mean(field_recv_A_CO2_econc), np.max(field_recv_A_CO2_econc)))
        if CO2_EMODE:
            A_CO2_emis.get(itap_sec, field_recv_A_CO2_emis)
            logging.debug('A_CO2_emis -> {:.3e} {:.3e} {:.3e}'.format( np.min(field_recv_A_CO2_emis), np.mean(field_recv_A_CO2_emis), np.max(field_recv_A_CO2_emis)))
            A_CO2_land.get(itap_sec, field_recv_A_CO2_land)
            logging.debug('A_CO2_land -> {:.3e} {:.3e} {:.3e}'.format( np.min(field_recv_A_CO2_land), np.mean(field_recv_A_CO2_land), np.max(field_recv_A_CO2_land)))
            A_CO2_ocean.get(itap_sec, field_recv_A_CO2_ocean)
            logging.debug('A_CO2_ocean -> {:.3e} {:.3e} {:.3e}'.format( np.min(field_recv_A_CO2_ocean), np.mean(field_recv_A_CO2_ocean), np.max(field_recv_A_CO2_ocean)))
        logging.debug('--------------------------------------\n')

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

    del comp
