"""Sample aospy object library using the included example data."""
from datetime import datetime
import os

import aospy
from aospy import Model, Proj, Region, Run, Var
from aospy.data_loader import DictDataLoader
from aospy.internal_names import LAND_MASK_STR, LON_STR


rootdir = os.path.join(aospy.__path__[0], 'test', 'data', 'netcdf')


_file_map = {'monthly': os.path.join(rootdir,
                                     '000[4-6]0101.precip_monthly.nc')}
example_run = Run(
    name='example_run',
    description=(
        'Control simulation of the idealized moist model'
    ),
    default_start_date=datetime(4, 1, 1),
    default_end_date=datetime(6, 12, 31),
    data_loader=DictDataLoader(_file_map)
)


example_model = Model(
    name='example_model',
    grid_file_paths=(os.path.join(rootdir, '00040101.precip_monthly.nc'),
                     os.path.join(rootdir, 'im.landmask.nc')),
    runs=[example_run],
    grid_attrs={LAND_MASK_STR: 'custom_land_mask', LON_STR: 'custom_lon'}
)


def total_precip(precip_largescale, precip_convective):
    """Sum of convective and large-scale precipitation.

    Parameters
    ----------
    precip_largescale, precip_convective : xarray.DataArrays
        Precipitation from grid-scale condensation and from convective
        parameterization, respectively.

    Returns
    -------
    xarray.DataArray
    """
    return precip_largescale + precip_convective


def conv_precip_frac(precip_largescale, precip_convective):
    """Fraction of total precip that is from convection parameterization.

    Parameters
    ----------
    precip_largescale, precip_convective : xarray.DataArrays
        Precipitation from grid-scale condensation and from convective
        parameterization, respectively.

    Returns
    -------
    xarray.DataArray
    """
    total = total_precip(precip_largescale, precip_convective)
    # Mask using xarray's `where` method to prevent divide-by-zero.
    return precip_convective / total.where(total)


precip_largescale = Var(
    name='precip_largescale',
    alt_names=('condensation_rain',),
    def_time=True,
    description='Precipitation generated via grid-scale condensation',
)


precip_convective = Var(
    name='precip_convective',
    alt_names=('convection_rain',),
    def_time=True,
    description='Precipitation generated by convective parameterization',
)


precip_total = Var(
    name='precip_total',
    def_time=True,
    func=total_precip,
    variables=(precip_largescale, precip_convective),
)


precip_conv_frac = Var(
   name='precip_conv_frac',
   def_time=True,
   func=conv_precip_frac,
   variables=(precip_largescale, precip_convective),
)


globe = Region(
    name='globe',
    description='Entire globe',
    west_bound=0,
    east_bound=360,
    south_bound=-90,
    north_bound=90,
    do_land_mask=False
)


tropics = Region(
    name='tropics',
    description='Tropics, defined as 30S-30N',
    west_bound=0,
    east_bound=360,
    south_bound=-30,
    north_bound=30,
    do_land_mask=False
)


example_proj = Proj(
    'example_proj',
    direc_out='example-output',
    tar_direc_out='example-tar-output',
    models=[example_model],
    regions=(globe, tropics)
)


if __name__ == '__main__':
    pass
