
from landlab import Component
from ...utils.decorators import use_file_name_or_kwds
import numpy as np


_VALID_METHODS = set(['Grid'])


def _assert_method_is_valid(method):
    if method not in _VALID_METHODS:
        raise ValueError('%s: Invalid method name' % method)


class Radiation(Component):

    """Compute 1D and 2D total incident shortwave radiation.

    Landlab component that computes 1D and 2D total incident shortwave
    radiation. This code also computes relative incidence shortwave radiation
    compared to a flat surface. Ref: Bras, Rafael L. Hydrology: an
    introduction to hydrologic science. Addison Wesley Publishing Company,
    1990.

    .. codeauthor:: Sai Nudurupati & Erkan Istanbulluoglu

    Examples
    --------
    >>> from landlab import RasterModelGrid
    >>> from landlab.components import Radiation
    >>> import numpy as np

    >>> grid = RasterModelGrid((5, 4), spacing=(0.2, 0.2))
    >>> rad = Radiation(grid)
    >>> rad.name
    'Radiation'
    >>> rad.input_var_names
    ('topographic__elevation',)
    >>> sorted(rad.output_var_names) # doctest: +NORMALIZE_WHITESPACE
    ['radiation__incoming_shortwave_flux',
     'radiation__net_shortwave_flux',
     'radiation__ratio_to_flat_surface']
    >>> sorted(rad.units) # doctest: +NORMALIZE_WHITESPACE
    [('radiation__incoming_shortwave_flux', 'W/m^2'),
     ('radiation__net_shortwave_flux', 'W/m^2'),
     ('radiation__ratio_to_flat_surface', 'None'),
     ('topographic__elevation', 'm')]

    >>> rad.grid.number_of_node_rows
    5
    >>> rad.grid.number_of_node_columns
    4
    >>> rad.grid is grid
    True
    >>> np.all(grid.at_cell['radiation__ratio_to_flat_surface'] == 0.)
    True
    >>> np.all(grid.at_node['topographic__elevation'] == 0.)
    True

    >>> grid['node']['topographic__elevation'] = np.array([
    ...       0., 0., 0., 0.,
    ...       1., 1., 1., 1.,
    ...       2., 2., 2., 2.,
    ...       3., 4., 4., 3.,
    ...       4., 4., 4., 4.])
    >>> current_time = 0.5
    >>> rad.update(current_time)
    >>> np.all(grid.at_cell['radiation__ratio_to_flat_surface'] == 0.)
    False
    """

    _name = 'Radiation'

    _input_var_names = (
        'topographic__elevation',
    )

    _output_var_names = (
        'radiation__incoming_shortwave_flux',
        'radiation__ratio_to_flat_surface',
        'radiation__net_shortwave_flux',
    )

    _var_units = {
        'topographic__elevation': 'm',
        'radiation__incoming_shortwave_flux': 'W/m^2',
        'radiation__ratio_to_flat_surface': 'None',
        'radiation__net_shortwave_flux': 'W/m^2',
    }

    _var_mapping = {
        'topographic__elevation': 'node',
        'radiation__incoming_shortwave_flux': 'cell',
        'radiation__ratio_to_flat_surface': 'cell',
        'radiation__net_shortwave_flux': 'cell',
    }

    _var_doc = {
        'topographic__elevation':
            'elevation of the ground surface relative to some datum',
        'radiation__incoming_shortwave_flux':
            'total incident shortwave radiation over the time step',
        'radiation__ratio_to_flat_surface':
            'ratio of total incident shortwave radiation on sloped surface \
             to flat surface',
        'radiation__net_shortwave_flux':
            'net incident shortwave radiation over the time step',
    }

    @use_file_name_or_kwds
    def __init__(self, grid, method='Grid', cloudiness=0.2,
                 latitude=34., albedo=0.2, solar_constant=1366.67,
                 clearsky_turbidity=2., opt_airmass=0., **kwds):
        """
        Parameters
        ----------
        grid: RasterModelGrid
            A grid.
        method: {'Grid'}, optional
            Currently, only default is available.
        cloudiness: float, optional
            Cloudiness.
        latitude: float, optional
            Latitude (radians).
        albedo: float, optional
            Albedo.
        solar_constant: float, optional
            Solar Constant (W/m^2).
        clearsky_turbidity: float, optional
            Clear sky turbidity.
        opt_airmass: float, optional
            Optical air mass.
        """

        self._method = method
        self._N = cloudiness
        self._latitude = latitude
        self._A = albedo
        self._Io = solar_constant
        self._n = clearsky_turbidity
        self._m = opt_airmass

        _assert_method_is_valid(self._method)

        super(Radiation, self).__init__(grid, **kwds)

        for name in self._input_var_names:
            if name not in self.grid.at_node:
                self.grid.add_zeros(name, at='node',
                                    units=self._var_units[name])

        for name in self._output_var_names:
            if name not in self.grid.at_cell:
                self.grid.add_zeros(name, at='cell',
                                    units=self._var_units[name])

        if 'Slope' not in self.grid.at_cell:
            self.grid.add_zeros('Slope', at='cell', units='radians')

        if 'Aspect' not in self.grid.at_cell:
            self.grid.add_zeros('Aspect', at='cell', units='radians')

        self._nodal_values = self.grid['node']
        self._cell_values = self.grid['cell']
        self._slope, self._aspect = \
            grid.calculate_slope_aspect_at_nodes_burrough(
                vals='topographic__elevation')
#        self._slope = grid.calc_slope_of_node( \
#                                elevs = 'topographic__elevation')
#        self._aspect =
        self._cell_values['Slope'] = self._slope
        self._cell_values['Aspect'] = self._aspect

    def update(self, current_time, hour=12., **kwds):
        """Update fields with current loading conditions.

        Parameters
        ----------
        current_time: float
              Current time (years).
        hour: float, optional
              Hour of the day.
        """
        self._t = hour
        self._radf = self._cell_values['radiation__ratio_to_flat_surface']
        self._Rs = self._cell_values['radiation__incoming_shortwave_flux']
        self._Rnet = self._cell_values['radiation__net_shortwave_flux']

        self._julian = np.floor((current_time - np.floor(current_time)) *
                                365.25)    # Julian day

        self._phi = np.radians(self._latitude)    # Latitude in Radians

        self._delta = 23.45 * np.radians(
            np.cos(2*np.pi / 365 * (172 - self._julian)))   # Declination angle

        self._tau = (self._t + 12.0) * np.pi / 12.0     # Hour angle

        self._alpha = np.arcsin(np.sin(self._delta) * np.sin(self._phi) +
                                np.cos(self._delta) * np.cos(self._phi) *
                                np.cos(self._tau))     # Solar Altitude

        if self._alpha <= 0.25 * np.pi / 180.0:    # If altitude is -ve,
            self._alpha = 0.25 * np.pi / 180.0     # sun is beyond the horizon

        # Counting for Albedo, Cloudiness and Atmospheric turbidity
        self._Rgl = (self._Io * np.exp((-1) * self._n * (
            0.128 - 0.054 * np.log10(1. / np.sin(self._alpha)))*(
                1. / np.sin(self._alpha))))

        self._phisun = (np.arctan(- np.sin(self._tau) / (np.tan(self._delta) *
                        np.cos(self._phi) - np.sin(self._phi) *
                        np.cos(self._tau))))      # Sun's Azhimuth
        if (self._phisun >= 0 and - np.sin(self._tau) <= 0):
            self._phisun = self._phisun + np.pi

        elif (self._phisun <= 0 and - np.sin(self._tau) >= 0):
            self._phisun = self._phisun + np.pi

        self._flat = (np.cos(np.arctan(0)) * np.sin(self._alpha) +
                      np.sin(np.arctan(0)) * np.cos(self._alpha) *
                      np.cos(self._phisun - 0))   # flat surface reference

        # flat surface total incoming shortwave radiation
        self._Rsflat = self._Rgl * self._flat

        # flat surface Net incoming shortwave radiation
        self._Rnetflat = ((1 - self._A) * (1 - 0.65 * (self._N ** 2)) *
                          self._Rsflat)

        self._sloped = (np.cos(self._slope) * np.sin(self._alpha) +
                        np.sin(self._slope) * np.cos(self._alpha) *
                        np.cos(self._phisun - self._aspect))
        # Ratio of cosine of solar incidence angle of sloped surface to that
        # of a flat surface
        self._radf = self._sloped / self._flat

        self._radf[self._radf <= 0.] = 0.
        self._radf[self._radf > 6.] = 6.

        # Sloped surface Toatl Incoming Shortwave Radn
        self._Rs = self._Rsflat * self._radf
        self._Rnet = self._Rnetflat * self._radf

        self._cell_values['radiation__ratio_to_flat_surface'] = self._radf
        self._cell_values['radiation__incoming_shortwave_flux'] = self._Rs
        self._cell_values['radiation__net_shortwave_flux'] = self._Rnet
