#!/usr/bin/env python3
# Copyright 2022 - Barcelona Supercomputing Center
# Authors: Rodrigo Martín Posada, Etienne Tourigny
# MIT License
import pyoasis
from pyoasis import OASIS  # pylint: disable=no-name-in-module
import numpy as np
import logging
import math

from amip_forcing_mod import AMIPForcing
from amip_utils import read_namelist

# Constants
LOG_FILE_NAME = 'amip.log'
CPLNG_COMP_NAME = 'AMIPFORC'

#Options
write_grid = True
write_grid_mask = True
write_grid_corners = True
write_grid_areas = True

if __name__ == '__main__':
    # Read RunLengthSec,TimeStepSec from namelist
    RunLengthSec, TimeStepSec, StartYear, StartMonth, \
        StartDay, FixYear, LDebug, vars_config = read_namelist()
    if LDebug:
        logging.basicConfig(filename=LOG_FILE_NAME, format='%(message)s', level=logging.DEBUG)

    #  time axis unit = days
    deltaT = TimeStepSec/86400.

    data_yref_min = max([conf['yref_min'] for conf in vars_config])
    date_yref_max = min([conf['yref_max'] for conf in vars_config])

    # valid FixYear to get a seasonal cycle, or set one if none and running outside AMIP window
    if (FixYear > 0):
        FixYear = min(date_yref_max, max(data_yref_min, FixYear))
    else:
        if ((StartYear < data_yref_min) or (StartYear > date_yref_max)):
            # set pseudo fixyear,
            FixYear = min(date_yref_max, max(data_yref_min, StartYear))
            # but do not overshoot
            if ((StartYear < data_yref_min) and (RunLengthSec > (data_yref_min-StartYear+1)*3600*24*365)):
                raise Exception("RunLengthSec is too long for StartYear less than data_yref_min")

    # Setup
    amip = AMIPForcing(deltaT, FixYear, StartYear, vars_config)
    amip_vars = [v.id for v in amip.vars]
    logging.debug('amip vars: {}'.format(str(amip_vars)))

    # OASIS init
    comp = pyoasis.Component(CPLNG_COMP_NAME)

    # Write grids
    oasis_vars = {}
    grids_names = np.unique([v.grid_name for v in amip.vars])
    logging.debug('grids_names: {}'.format(str(grids_names)))
    for grid_name in grids_names:
        # Get OASIS vars within this grid_name
        grid_vars = list(filter(lambda x: x.grid_name == grid_name, amip.vars))
        v = grid_vars[0]
        grid_nx = v.grid_nx
        grid_ny = v.grid_ny
        # Create partition
        partition = pyoasis.SerialPartition(grid_nx*grid_ny)
        # Define the grid and write to files
        if write_grid:
            grid = pyoasis.Grid(grid_name, grid_nx, grid_ny, v.grid_lon, v.grid_lat, partition)
            # add grid mask
            if write_grid_mask:
                grid.set_mask(np.zeros(v.grid_lon.shape))
            # add grid corners
            if write_grid_corners:
                ncrn = 4
                clo = pyoasis.asarray(np.zeros((grid_nx, grid_ny, ncrn), dtype=np.float64))
                clo[:, :, 0] = v.grid_lon[:, :] - v.dx/2.0
                clo[:, :, 1] = v.grid_lon[:, :] + v.dx/2.0
                clo[:, :, 2] = clo[:, :, 1]
                clo[:, :, 3] = clo[:, :, 0]
                cla = pyoasis.asarray(np.zeros((grid_nx, grid_ny, ncrn), dtype=np.float64))
                cla[:, :, 0] = v.grid_lat[:, :] - v.dy/2.0
                cla[:, :, 1] = cla[:, :, 0]
                cla[:, :, 2] = v.grid_lat[:, :] + v.dy/2.0
                cla[:, :, 3] = cla[:, :, 2]
                grid.set_corners(clo, cla)
            # add grid areas
            if write_grid_areas:
                dp_conv = math.pi/180.
                area = np.zeros((grid_nx, grid_ny), dtype=np.float64)
                area[:, :] = dp_conv * \
                             np.abs(np.sin(cla[:, :, 2] * dp_conv) -
                                    np.sin(cla[:, :, 0] * dp_conv)) * \
                             np.abs(clo[:, :, 1] - clo[:, :, 0])
                grid.set_area(area)
            # write grid to files
            grid.write()

        # Set vars
        for v in grid_vars:
            if not v.oasis_name in oasis_vars:
                logging.debug('creating oasis var {} with grid {}'.format(v.oasis_name, grid_name))
                oasis_vars[v.oasis_name] = (v, pyoasis.Var(v.oasis_name, partition, OASIS.OUT))
            else:
                logging.debug('not creating oasis var {} with grid {} as it hs already been defined'.format(v.oasis_name, grid_name))
        logging.debug('oasis_vars: {}'.format(str(oasis_vars.keys())))

    comp.enddef()

    amip.setup_amip_forcing(StartYear, StartMonth, StartDay)

    # Time loop
    # Note that according to OASIS convention, the first time step is at time=0
    # and the last one is at runlength-timestep (see OASIS documentation)
    for time in range(0, RunLengthSec, TimeStepSec):
        amip.update_vars()

        # OASIS put for each var
        for var_name, (amip_v, oa_v) in oasis_vars.items():
            logging.debug('sending var {} at time {}'.format(str(var_name),time))
            #oa_v.put(time, amip_v.send_field)
            send_field = 0
            for v in amip.vars:
                if v.oasis_name == var_name:
                    send_field += v.send_field
                    logging.debug('accumulating {} from amip var {} for oasis var {} total= {}'.format(str(v.send_field), v.id, var_name,str(send_field)))
            oa_v.put(time, send_field)

    amip.finalise_amip_forcing()

    del comp
