"""Define an interface for the RRTMG radiation scheme (through CliMT). """
from copy import deepcopy
import numpy as np
import datetime
from sympl import DataArray
from typhon.physics import vmr2specific_humidity
import climt
import logging

from .radiation import Radiation
from konrad.cloud import ClearSky, CloudEnsemble

__all__ = [
    'RRTMG',
]


class RRTMG(Radiation):
    """RRTMG radiation scheme using the CliMT python wrapper."""

    def __init__(self, *args, solar_constant=510, mcica=False,
                 **kwargs):
        """
        Parameters:
            zenith_angle (float): angle of the Sun [degrees].
            In konrad, with no diurnal cycle, this should represent a diurnal
            mean zenith angle. With a diurnal cycle, this represents latitude.
            For single radiation calculations, this is the angle to the Sun at
            the time and location of interest.

            bias (dict-like): include bias corrections to the fluxes and/or
                heating rates

            solar_constant (int): [W m^-2]

            mcica (bool):

                * :code:`False`
                    use the nomcica version of RRTMG (clear-sky or overcast)

                * :code:`True`
                    use the mcica version of RRTMG (needed for partly cloudy
                    skies)

        """
        super().__init__(*args, **kwargs)
        self._state_lw = None
        self._state_sw = None

        self._rad_lw = None
        self._rad_sw = None

        self._is_mcica = mcica

        # These are set in the first call and depend on properties set in the
        # cloud class or instance.
        self._cloud_optical_properties = None
        self._cloud_ice_properties = None

        self.solar_constant = solar_constant

    def init_radiative_state(self, atmosphere, surface):

        climt.set_constants_from_dict({"stellar_irradiance": {
                "value": self.solar_constant, "units": 'W m^-2'}})
        if self._is_mcica:
            overlap = 'maximum_random'
        else:
            overlap = 'random'
        self._rad_lw = climt.RRTMGLongwave(
            cloud_optical_properties=self._cloud_optical_properties,
            cloud_ice_properties=self._cloud_ice_properties,
            cloud_liquid_water_properties='radius_dependent_absorption',
            cloud_overlap_method=overlap,
            mcica=self._is_mcica)
        self._rad_sw = climt.RRTMGShortwave(
            ignore_day_of_year=True,
            cloud_optical_properties=self._cloud_optical_properties,
            cloud_ice_properties=self._cloud_ice_properties,
            cloud_liquid_water_properties='radius_dependent_absorption',
            cloud_overlap_method=overlap,
            mcica=self._is_mcica)
        state_lw = {}
        state_sw = {}

        plev = atmosphere['plev']
        phlev = atmosphere['phlev']
        numlevels = len(plev)

        for state0 in state_lw, state_sw:
            state0['mid_levels'] = DataArray(
                    np.arange(0, numlevels),
                    dims=('mid_levels',),
                    attrs={'label': 'mid_levels'})

            state0['interface_levels'] = DataArray(
                    np.arange(0, numlevels+1),
                    dims=('interface_levels',),
                    attrs={'label': 'interface_levels'})

            state0['air_pressure'] = DataArray(
                    plev,
                    dims=('mid_levels',),
                    attrs={'units': 'Pa'})

            state0['air_pressure_on_interface_levels'] = DataArray(
                    phlev,
                    dims=('interface_levels',),
                    attrs={'units': 'Pa'})

            state0['time'] = datetime.datetime(2000, 1, 1)

        ### Aerosols ###
        # TODO: Should all the aerosol values be zero?!
        # Longwave specific
        num_lw_bands = self._rad_lw.num_longwave_bands

        state_lw['longwave_optical_thickness_due_to_aerosol'] = DataArray(
                np.zeros((num_lw_bands, numlevels)),
                dims=('num_longwave_bands', 'mid_levels'),
                attrs={'units': 'dimensionless'})
        state_lw['surface_longwave_emissivity'] = DataArray(
            surface.longwave_emissivity * np.ones((num_lw_bands,)),
            dims=('num_longwave_bands'),
            attrs={'units': 'dimensionless'})

        # Shortwave specific changes
        num_sw_bands = self._rad_sw.num_shortwave_bands

        num_aerosols = self._rad_sw.num_ecmwf_aerosols
        state_sw['aerosol_optical_depth_at_55_micron'] = DataArray(
                np.zeros((num_aerosols, numlevels)),
                dims=('num_ecmwf_aerosols', 'mid_levels'),
                attrs={'units': 'dimensionless'})
        state_sw['solar_cycle_fraction'] = DataArray(
            0, attrs={'units': 'dimensionless'})
        state_sw['flux_adjustment_for_earth_sun_distance'] = DataArray(
            1, attrs={'units': 'dimensionless'})

        for surface_albedo in ['surface_albedo_for_diffuse_near_infrared',
                               'surface_albedo_for_direct_near_infrared',
                               'surface_albedo_for_diffuse_shortwave',
                               'surface_albedo_for_direct_shortwave']:
            state_sw[surface_albedo] = DataArray(
                np.array(float(surface.albedo)),
                attrs={'units': 'dimensionless'})

        for quant in ['shortwave_optical_thickness_due_to_aerosol',
                      'single_scattering_albedo_due_to_aerosol',
                      'aerosol_asymmetry_parameter']:
            state_sw[quant] = DataArray(
                np.zeros((num_sw_bands, numlevels)),
                dims=('num_shortwave_bands', 'mid_levels'),
                attrs={'units': 'dimensionless'})

        return state_lw, state_sw

    def update_cloudy_radiative_state(self, cloud, state0, sw=True):

        # Take cloud quantities from cloud class.
        props = (
           'cloud_ice_particle_size',
           'mass_content_of_cloud_liquid_water_in_atmosphere_layer',
           'mass_content_of_cloud_ice_in_atmosphere_layer',
           'cloud_water_droplet_radius',
           'cloud_area_fraction_in_atmosphere_layer',
        )

        lw_props = (
            'longwave_optical_thickness_due_to_cloud',
        )

        sw_props = (
            'cloud_forward_scattering_fraction',
            'cloud_asymmetry_parameter',
            'shortwave_optical_thickness_due_to_cloud',
            'single_scattering_albedo_due_to_cloud',
        )

        for varname in props + sw_props if sw else props + lw_props:
            state0[varname] = cloud[varname]

    def update_radiative_state(self, atmosphere, surface, state0, sw=True):
        """ Update CliMT formatted atmospheric state using parameters from our
        model.

        Parameters:
            atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model.
            surface (konrad.surface): Surface model.
            state0 (dictionary): atmospheric state in the format for climt
            sw (bool): Toggle between shortwave and longwave calculations.

        Returns:
            dictionary: updated state
        """

        state0['air_temperature'] = DataArray(
            atmosphere['T'][0, :],
            dims=('mid_levels',),
            attrs={'units': 'degK'})

        vmr_h2o = atmosphere['H2O'][0, :]
        specific_humidity = vmr2specific_humidity(vmr_h2o)
        state0['specific_humidity'] = DataArray(
            specific_humidity,
            dims=('mid_levels',),
            attrs={'units': 'g/g'})

        # CliMT/konrad name mapping
        gas_name_mapping = [
            ('mole_fraction_of_methane_in_air', 'CH4'),
            ('mole_fraction_of_carbon_dioxide_in_air', 'CO2'),
            ('mole_fraction_of_nitrous_oxide_in_air', 'N2O'),
            ('mole_fraction_of_ozone_in_air', 'O3'),
            ('mole_fraction_of_cfc11_in_air', 'CFC11'),
            ('mole_fraction_of_cfc12_in_air', 'CFC12'),
            ('mole_fraction_of_cfc22_in_air', 'CFC22'),
            ('mole_fraction_of_carbon_tetrachloride_in_air', 'CCl4'),
            ('mole_fraction_of_oxygen_in_air', 'O2'),
        ]

        for climt_key, konrad_key in gas_name_mapping:
            vmr = atmosphere.get(konrad_key, default=0, keepdims=False)
            state0[climt_key] = DataArray(
                vmr,
                dims=('mid_levels',),
                attrs={'units': 'mole/mole'})

        # Surface quantities
        state0['surface_temperature'] = DataArray(
            np.array(surface['temperature'][-1]),
            attrs={'units': 'degK'})

        if sw:  # properties required only for shortwave
            state0['zenith_angle'] = DataArray(
                np.array(np.deg2rad(self.current_solar_angle)),
                attrs={'units': 'radians'})

        return state0

    def radiative_fluxes(self, atmosphere, surface, cloud):
        """Returns shortwave and longwave fluxes and heating rates.

        Parameters:
            atmosphere (konrad.atmosphere.Atmosphere): atmosphere model
            surface (konrad.surface): surface model
            cloud (konrad.cloud): cloud model

        Returns:
            tuple: containing two dictionaries, one of air temperature
            values and the other of fluxes and heating rates
        """
        if self._state_lw is None or self._state_sw is None:  # first time only
            self._cloud_optical_properties = cloud._rrtmg_cloud_optical_properties
            self._cloud_ice_properties = cloud._rrtmg_cloud_ice_properties
            self._state_lw, self._state_sw = self.init_radiative_state(
                    atmosphere, surface)
            self.update_cloudy_radiative_state(cloud, self._state_lw, sw=False)
            self.update_cloudy_radiative_state(cloud, self._state_sw, sw=True)

        # if there are clouds update the cloud properties for the radiation
        if not isinstance(cloud, ClearSky):
            self.update_cloudy_radiative_state(cloud, self._state_lw, sw=False)
            self.update_cloudy_radiative_state(cloud, self._state_sw, sw=True)

        self.update_radiative_state(atmosphere, surface, self._state_lw,
                                    sw=False)
        self.update_radiative_state(atmosphere, surface, self._state_sw,
                                    sw=True)

        lw_fluxes = self._rad_lw(self._state_lw)
        sw_fluxes = self._rad_sw(self._state_sw)

        return lw_fluxes, sw_fluxes

    def calc_cloudy_nomcica_radiation(self, atmosphere, surface, cloud):

        cloud_fraction = deepcopy(cloud['cloud_area_fraction_in_atmosphere_layer'][:])

        if self._state_sw is None:  # first time only
            cf_cloudy = cloud_fraction[cloud_fraction != 0]
            if cf_cloudy.size > 0:
                if not np.all(cf_cloudy == cf_cloudy[0]):
                    logging.warning(
                        'The konrad implementation of nomcica with partial '
                        'cloud fraction is only suitable for a single '
                        'rectangular cloud. Consider using mcica instead.')

        cf_max = np.max(cloud_fraction)

        # Calculate overcast and clear sky radiative fluxes.
        # Make all the cloudy layers overcast, so that the nomcica version of
        # RRTMG can be used in the shortwave - all wavelengths see cloud.
        # Use the same approach for the longwave for consistency.
        cloud['cloud_area_fraction_in_atmosphere_layer'][cloud_fraction != 0] = 1
        lw_overcast, sw_overcast = self.radiative_fluxes(
            atmosphere, surface, cloud
        )
        lw_fluxes, sw_fluxes = lw_overcast[1], sw_overcast[1]

        # Combine overcast and clear sky fluxes and heating rates to get the
        # all sky fluxes and heating rates.
        for fluxes in lw_fluxes, sw_fluxes:
            for key in fluxes.keys():
                if 'clear_sky' not in key:
                    clear_part = fluxes[key + '_assuming_clear_sky'][:]
                    fluxes[key][:] *= cf_max  # weighted by cloud area fraction
                    fluxes[key][:] += (1 - cf_max) * clear_part

        cloud['cloud_area_fraction_in_atmosphere_layer'][:] *= cloud_fraction

        return lw_fluxes, sw_fluxes

    def calc_radiation(self, atmosphere, surface, cloud):
        """Updates the shortwave, longwave and net heatingrates.
        Converts output from radiative_fluxes to be in the format required for
        our model.

        Parameters:
            atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model.
            surface (konrad.surface): Surface model.
            cloud (konrad.cloud): cloud model
        """
        if not self._is_mcica and not isinstance(cloud, ClearSky):
            if isinstance(cloud, CloudEnsemble):
                weights, cloud_combinations = cloud.get_combinations()

                lw_fluxes = {}
                sw_fluxes = {}
                for weight, cloud_case in zip(weights, cloud_combinations):
                    lw_temp, sw_temp = self.radiative_fluxes(
                        atmosphere, surface, cloud_case)

                    for key, value in lw_temp[1].items():
                        if key not in lw_fluxes.keys():
                            lw_fluxes[key] = weight * value.data
                        else:
                            lw_fluxes[key][:] += weight * value.data

                    for key, value in sw_temp[1].items():
                        if key not in sw_fluxes.keys():
                            sw_fluxes[key] = weight * value.data
                        else:
                            sw_fluxes[key][:] += weight * value.data

            else:
                lw_fluxes, sw_fluxes = self.calc_cloudy_nomcica_radiation(
                    atmosphere, surface, cloud)
        else:
            lw_dT_fluxes, sw_dT_fluxes = self.radiative_fluxes(
                atmosphere, surface, cloud)
            lw_fluxes = lw_dT_fluxes[1]
            sw_fluxes = sw_dT_fluxes[1]

        self['lw_htngrt'] = np.expand_dims(
                lw_fluxes['air_temperature_tendency_from_longwave'].data, 0)
        self['lw_htngrt_clr'] = np.expand_dims(
                lw_fluxes['air_temperature_tendency_from_longwave_assuming_clear_sky'].data, 0)
        self['lw_flxu'] = np.expand_dims(
                lw_fluxes['upwelling_longwave_flux_in_air'].data, 0)
        self['lw_flxd'] = np.expand_dims(
                lw_fluxes['downwelling_longwave_flux_in_air'].data, 0)
        self['lw_flxu_clr'] = np.expand_dims(
                lw_fluxes['upwelling_longwave_flux_in_air_assuming_clear_sky'].data, 0)
        self['lw_flxd_clr'] = np.expand_dims(
                lw_fluxes['downwelling_longwave_flux_in_air_assuming_clear_sky'].data, 0)
        self['sw_htngrt'] = np.expand_dims(
                sw_fluxes['air_temperature_tendency_from_shortwave'].data, 0)
        self['sw_htngrt_clr'] = np.expand_dims(
                sw_fluxes['air_temperature_tendency_from_shortwave_assuming_clear_sky'].data, 0)
        self['sw_flxu'] = np.expand_dims(
                sw_fluxes['upwelling_shortwave_flux_in_air'].data, 0)
        self['sw_flxd'] = np.expand_dims(
                sw_fluxes['downwelling_shortwave_flux_in_air'].data, 0)
        self['sw_flxu_clr'] = np.expand_dims(
                sw_fluxes['upwelling_shortwave_flux_in_air_assuming_clear_sky'].data, 0)
        self['sw_flxd_clr'] = np.expand_dims(
                sw_fluxes['downwelling_shortwave_flux_in_air_assuming_clear_sky'].data, 0)

        self.coords={
            'time': np.array([0]),
            'phlev': atmosphere['phlev'],
            'plev': atmosphere['plev'],
        }
