#!/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(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__':
    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_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(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, mean 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_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(f'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(f'--------------------------------------\n', file=w_unit)
        w_unit.flush()

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

    del comp