#!/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_atmos = 96  # dimensions in the 2 spatial directions
    nlat_atmos = 72
    nc_atmos = 4  # number of grid cell vertices in the (i,j) plan
    il_nb_time_steps = 8  # number of time steps
    delta_t = 1800  # time step
    dp_pi = 3.14159265359
    dp_length = 1.2 * dp_pi

    ################# OASIS_INIT_COMP ##############################################
    comp = pyoasis.Component("atmos_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_atmos = 'atmos_python.out_'+str(w_unit)
    w_unit = open(comp_out_atmos, 'w')
    print(' -----------------------------------------------------------', file=w_unit)
    print(f' I am atmos 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_atmos, nlat_atmos, 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_atmos)
    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_atmos)
    else:
        raise Exception(f'Invalid CPPKEYDECOMP in {__file__}: {CPPKEYDECOMP}')

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  GRID DEFINITION
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #
    # Reading local grid arrays from input file atmos_mesh.nc
    grid_lon_atmos, grid_lat_atmos, grid_clo_atmos, grid_cla_atmos, grid_srf_atmos, grid_msk_atmos = read_grid(il_offsetx, il_offsety, il_extentx, il_extenty, nc_atmos, 'atmos_mesh.nc')

    ################# OASIS_WRITE_GRID  ##############################################
    grid = pyoasis.Grid('lmdz', nlon_atmos, nlat_atmos, grid_lon_atmos, grid_lat_atmos, partition)
    grid.set_corners(grid_clo_atmos, grid_cla_atmos)
    grid.set_mask(grid_msk_atmos)
    grid.write()
    print(f' grid_lat_atmos maximum and minimum', '%.5f' % np.max(grid_lat_atmos), '%.5f' % np.min(grid_lat_atmos), 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_ATM", partition, OASIS.IN)
    var_id[1] = pyoasis.Var("FIELD_SEND_ATM", partition, OASIS.OUT)
    print(f' var_id FRECVATM, var_id FSENDATM {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_atmos = pyoasis.asarray(np.full((il_extenty, il_extentx), -1.0))
        ######################## OASIS_GET ######################
        var_id[0].get(itap_sec, field_recv_atmos)
        print(' ', itap_sec, '%.5f' % np.min(field_recv_atmos), '%.5f' % np.max(field_recv_atmos), file=w_unit)

        # Definition of field produced by the component
        field_send_atmos = (ib+1)*(2.-np.cos(dp_pi*(np.arccos(np.cos(grid_lat_atmos*dp_pi/90.)*np.cos(grid_lon_atmos*dp_pi/90.))/dp_length)))

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

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

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

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

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