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

NAMELIST_FILE_NAME = 'namelist.amip'
NX_ATMO = 88838  # number of LAND grid cells at T255 resolution

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


if __name__ == '__main__':
    # Constants from namelist
    nml = f90nml.read(NAMELIST_FILE_NAME)
    COUPLING_INTERVAL = nml['NAMAMIP']['TimeStepSec']
    RUN_LENGTH_SEC = nml['NAMAMIP']['RunLengthSec']

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

    # Unit for output messages
    out_ifs = 'IFS_mock.out_'+str(100+mype)
    w_unit = open(out_ifs, 'w')
    print('-----------------------------------------------------------', file=w_unit)
    print('I am IFS-mock process with rank :', mype, file=w_unit)
    print('in my local communicator gathering', npes, 'processes', file=w_unit)
    print('----------------------------------------------------------', file=w_unit)
    w_unit.flush()

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  PARTITION DEFINITION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    # Definition of the local partition
    il_size, il_offset = def_local_partition(npes, mype)
    print('Local partition definition', file=w_unit)
    print('il_size, il_offset = ', il_size, il_offset, file=w_unit)
    w_unit.flush()
    partition = pyoasis.ApplePartition(il_offset, il_size)

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

    A_SST = pyoasis.Var("A_SST", partition, OASIS.IN)
    A_Ice_frac = pyoasis.Var("A_Ice_frac", partition, OASIS.IN)
    print('var_id A_SST, var_id A_Ice_frac', A_SST._id, A_Ice_frac._id, file=w_unit)
    w_unit.flush()

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         TERMINATION OF DEFINITION PHASE
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    print('End of initialisation phase', file=w_unit)
    w_unit.flush()

    comp.enddef()

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  TIME STEP LOOP
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    print('Timestep, field min, mean and max value', file=w_unit)
    w_unit.flush()

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

        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))

        A_SST.get(itap_sec, field_recv_A_SST)
        A_Ice_frac.get(itap_sec, field_recv_A_Ice_frac)

        print('ITAP_SEC:', itap_sec, file=w_unit)
        print('A_SST ->', itap_sec,
              "{:.4e}".format(np.min(field_recv_A_SST)),
              "{:.4e}".format(np.mean(field_recv_A_SST)),
              "{:.4e}".format(np.max(field_recv_A_SST)),
              file=w_unit)
        print('A_Ice_frac ->', itap_sec,
              "{:.4e}".format(np.min(field_recv_A_Ice_frac)),
              "{:.4e}".format(np.mean(field_recv_A_Ice_frac)),
              "{:.4e}".format(np.max(field_recv_A_Ice_frac)),
              file=w_unit)
        print('--------------------------------------\n', file=w_unit)
        w_unit.flush()

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         TERMINATION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    print('End of the program', file=w_unit)
    w_unit.flush()

    del comp