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

L080_NX = 35718  # number of LAND grid cells at T159 resolution
L128_NX = 88838  # number of LAND grid cells at T255 resolution
L256_NX = 348528  # number of LAND grid cells at T511 resolution


if __name__ == '__main__':

    COUPLING_INTERVAL = 2700
    RUN_LENGTH_SEC = 86400

    # Init OASIS
    comp = pyoasis.Component("IFS_TOY")
    local_comm = comp.localcomm
    # Get rank in local communicator
    mype = comp.localcomm.rank
    npes = comp.localcomm.size

    # Unit for output messages
    out_ifs = 'IFS_toy.out_'+str(100+mype)
    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
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

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

    out_coupling_fields={}
    out_coupling_vars=["A_TauX_oce","A_TauY_oce","A_TauX_ice","A_TauY_ice","A_Qns_mix","A_Qs_mix","A_Evap_total","A_Precip_liquid","A_Precip_solid","A_Qs_ice","A_Qns_ice","A_dQns_dT","A_Evap_ice","A_Runoff","A_Calving"]
    for v in out_coupling_vars:
        out_coupling_fields[v] = pyoasis.Var(v, partition, OASIS.OUT)
        logging.debug('var_id {}, added as OASIS_OUT var'.format(v))

    in_coupling_fields={}
    in_coupling_vars=["A_SST","A_Ice_temp","A_Ice_albedo","A_Ice_frac","A_Ice_thickness","A_Snow_thickness"]
    for v in in_coupling_vars:
        in_coupling_fields[v] = pyoasis.Var(v, partition, OASIS.IN)
        logging.debug('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(nx, -1.0))

    for itap_sec in range(0, RUN_LENGTH_SEC, 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(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')

    del comp
