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

from amip_forcing_mod import AMIPForcing
from amip_utils import read_namelist

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

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)

    # OASIS init
    comp = pyoasis.Component(CPLNG_COMP_NAME)
    # Write grids
    oasis_vars = {}
    grids_names = np.unique([v.grid_name for v in amip.vars])
    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_ny, grid_nx = v.grid_lat.shape
        # Create partition
        partition = pyoasis.SerialPartition(grid_nx*grid_ny)
        # Write the grid
        grid = pyoasis.Grid(grid_name, grid_nx, grid_ny, v.grid_lon, v.grid_lat, partition)
        grid.set_mask(np.zeros(v.grid_lon.shape))
        grid.write()
        # Set vars
        for v in grid_vars:
            oasis_vars[v] = pyoasis.Var(v.oasis_name, partition, OASIS.OUT)
    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 amip_v, oa_v in oasis_vars.items():
            oa_v.put(time, amip_v.send_field)

    amip.finalise_amip_forcing()

    del comp
