#!/usr/bin/env python3

import pyoasis  # pylint: disable=import-error
from pyoasis import OASIS  # pylint: disable=import-error,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__':
    nlon_atmos = 96
    nlat_atmos = 72
    nc_atmos = 4
    CPPKEYDECOMP = 'DECOMP_APPLE'
    ################# 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)
    atmos_file = open(comp_out_atmos, 'w')
    print(' -----------------------------------------------------------', file=atmos_file)
    print(f' I am atmos process with rank : {mype}', file=atmos_file)
    print(f' in my local communicator gathering {npes} processes', file=atmos_file)
    print(' ----------------------------------------------------------', file=atmos_file)

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  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=atmos_file)
    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=atmos_file)

    ################# 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=atmos_file)
    if (CPPKEYDECOMP == 'DECOMP_APPLE'):
        partition = pyoasis.ApplePartition(il_offset, il_size)
    elif (CPPKEYDECOMP == 'DECOMP_BOX'):
        partition = pyoasis.BoxPartition(il_offset, il_extenty, il_extentx, nlon_atmos)
    else:
        raise Exception(f'CPPKEYDECOMP not set 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)
    print(f' grid_lat_atmos maximum and minimum {np.max(grid_lat_atmos)} {np.min(grid_lat_atmos)}', file=atmos_file)

    #+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    #  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=atmos_file)

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


    
    exit()
    partition = pyoasis.SerialPartition(n_points)

    var_1 = pyoasis.Var("FRECVATM_1", partition, OASIS.IN)
    print("Recv_one ", var_1)

    var_2 = pyoasis.Var("FRECVATM_2", partition, OASIS.IN)
    print("Recv_one ", var_2)

    comp.enddef()

    intra_one = comp.get_intracomm("sender-serial")
    print("Recv_one intra_one: rank = {} of {}".format(intra_one.rank,
                                                       intra_one.size))

    field = pyoasis.asarray(np.zeros(n_points, dtype=np.float64))

    print("Recv_one: coupling frequencies for {} = ".format(var_1.name),
          var_1.cpl_freqs)
    print("Recv_one: coupling frequencies for {} = ".format(var_2.name),
          var_2.cpl_freqs)

    for date in range(43200):
        if any([date % freq == 0 for freq in var_1.cpl_freqs]):
            var_1.get(date, field)
            if abs((field - date).sum()) < 1.e-8:
                print("Recv_one: field 1 received successfully at time {}".format(date))
            else:
                print("Warning: Recv_one at time {} got {} instead of {}".format(
                    date, field[0], date))

        if any([date % freq == 0 for freq in var_2.cpl_freqs]):
            var_2.get(date, field)
            if np.abs(field + date).sum() < 1.e-8:
                print("Recv_one: field 2 received successfully at time {}".format(date))
            else:
                print("Warning: Recv_one at time {} got {} instead of {}".format(date,
                                                                                 field[0], date))

    del comp
