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

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

# Constants
NX_ATMO = 88838  # number of LAND grid cells at T255 resolution
COUPLING_INTERVAL = 86400
RUN_LENGTH_SEC = 31536000
NUMBER_TIME_STEPS = int(RUN_LENGTH_SEC / COUPLING_INTERVAL)


def def_local_partition(nlon, nlat, npes, mype):
    il_extentx = nlon
    il_extenty = int(nlat/npes)
    if (mype == npes-1):
        il_extenty = nlat - int(nlat/npes) * mype
    il_size = il_extentx * il_extenty
    il_offsetx = 0
    il_offsety = int(nlat/npes) * mype
    il_offset = nlon * il_offsety
    return il_extentx, il_extenty, il_size, il_offsetx, il_offsety, il_offset


if __name__ == '__main__':
    nlon = int(NX_ATMO / 86)
    nlat = int(NX_ATMO / nlon)
    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(f'I am IFS-mock process with rank : {mype}', file=w_unit)
    print(f'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_extentx, il_extenty, il_size, il_offsetx, il_offsety, il_offset = def_local_partition(nlon, nlat, npes, mype)
    print('Local partition definition', file=w_unit)
    print('il_extentx, il_extenty, il_size, il_offsetx, il_offsety, il_offset = ',
          il_extentx, il_extenty, il_size, il_offsetx, il_offsety, 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(f'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 and max value', file=w_unit)
    w_unit.flush()

    for ib in range(NUMBER_TIME_STEPS):
        itap_sec = COUPLING_INTERVAL * ib  # time in seconds

        field_recv_A_SST = pyoasis.asarray(np.full((il_extenty, il_extentx), -1.0))
        field_recv_A_Ice_frac = pyoasis.asarray(np.full((il_extenty, il_extentx), -1.0))

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

        print(f'ITAP_SEC: {itap_sec}', file=w_unit)
        print('A_SST ->', itap_sec, '%.5f' % np.min(field_recv_A_SST), '%.5f' % np.max(field_recv_A_SST), file=w_unit)
        print('A_Ice_frac ->', itap_sec, '%.5f' % np.min(field_recv_A_Ice_frac), '%.5f' %
              np.max(field_recv_A_Ice_frac), file=w_unit)
        print(f'--------------------------------------\n', file=w_unit)
        w_unit.flush()

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

    del comp