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

from def_parallel_decomposition import def_local_partition, def_paral
from read_grid import read_grid

if __name__ == '__main__':
    # Python exclusive: set CPPKEYDECOMP
    CPPKEYDECOMP = 'DECOMP_APPLE'

    # CONSTANTS
    nlon_ocean = 182  # dimensions in the 2 spatial directions
    nlat_ocean = 149
    nc_ocean = 4  # number of grid cell vertices in the (i,j) plan
    il_nb_time_steps = 4  # number of time steps
    delta_t = 3600  # time step
    dp_pi = 3.14159265359
    dp_length = 1.2 * dp_pi

    ################# OASIS_INIT_COMP ##############################################
    comp = pyoasis.Component("ocean_component")

    ################# OASIS_GET_LOCALCOMM ######################################
    local_comm = comp.localcomm

    # Get rank in local communicator
    mype = comp.localcomm.rank
    npes = comp.localcomm.size

    # Unit for output messages : one file for each process
    w_unit = 100 + mype
    comp_out_ocean = 'ocean_python.out_'+str(w_unit)
    w_unit = open(comp_out_ocean, 'w')
    print(' -----------------------------------------------------------', file=w_unit)
    print(f' I am ocean 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_ocean, nlat_ocean, 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()

    ################# OASIS_DEF_PARTITION ######################################
    ig_paral = def_paral(CPPKEYDECOMP, il_offset, il_size, il_extentx, il_extenty, nlon_ocean)
    print(' ig_paral = ', ' '.join(map(str, ig_paral)), file=w_unit)
    if (CPPKEYDECOMP == 'DECOMP_APPLE'):
        partition = pyoasis.ApplePartition(il_offset, il_size)
    elif (CPPKEYDECOMP == 'DECOMP_BOX'):
        partition = pyoasis.BoxPartition(il_offset, il_extentx, il_extenty, nlon_ocean)
    else:
        raise Exception(f'Invalid CPPKEYDECOMP in {__file__}: {CPPKEYDECOMP}')

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  GRID DEFINITION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #
    # Reading local grid arrays from input file ocean_mesh.nc
    grid_lon_ocean, grid_lat_ocean, grid_clo_ocean, grid_cla_ocean, grid_srf_ocean, grid_msk_ocean = read_grid(il_offsetx, il_offsety, il_extentx, il_extenty, nc_ocean, 'ocean_mesh.nc')

    ################# OASIS_WRITE_GRID  ##############################################
    grid = pyoasis.Grid('torc', nlon_ocean, nlat_ocean, grid_lon_ocean, grid_lat_ocean, partition)
    grid.set_corners(grid_clo_ocean, grid_cla_ocean)
    grid.set_mask(grid_msk_ocean)
    grid.write()
    print(f' grid_lat_ocean maximum and minimum', '%.5f' % np.max(grid_lat_ocean), '%.5f' % np.min(grid_lat_ocean), file=w_unit)
    w_unit.flush()

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

    ################## OASIS_DEF_VAR #################################
    var_id = [None] * 2
    var_id[0] = pyoasis.Var("FIELD_RECV_OCN", partition, OASIS.IN)
    var_id[1] = pyoasis.Var("FIELD_SEND_OCN", partition, OASIS.OUT)
    print(f' var_id FRECVOCN, var_id FSENDOCN {var_id[0]._id} {var_id[1]._id}', file=w_unit)
    w_unit.flush()

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

    ################## OASIS_ENDDEF ##################################
    comp.enddef()

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

    for ib in range(il_nb_time_steps):
        itap_sec = delta_t * ib  # time in seconds

        field_recv_ocean = pyoasis.asarray(np.full((il_extenty, il_extentx), -1.0))
        ######################## OASIS_GET ######################
        var_id[0].get(itap_sec, field_recv_ocean)
        print(' ', itap_sec, '%.5f' % np.min(field_recv_ocean), '%.5f' % np.max(field_recv_ocean), file=w_unit)

        # Definition of field produced by the component
        field_send_ocean = (ib+1)*(2.-np.cos(dp_pi*(np.arccos(np.cos(grid_lat_ocean*dp_pi/180.)*np.cos(grid_lon_ocean*dp_pi/180.))/dp_length)))

        # print(' ', itap_sec, '%.5f' % np.min(field_send_ocean), '%.5f' % np.max(field_send_ocean), file=w_unit)
        w_unit.flush()

        ######################## OASIS_PUT ######################
        var_id[1].put(itap_sec, field_send_ocean)

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #         TERMINATION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

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

    ################## OASIS_TERMINATE ##################################
    del comp
