################################################################################
# This file is part of MELODIST - MEteoroLOgical observation time series       #
# DISaggregation Tool.                                                         #
#                                                                              #
# Copyright (C) 2016-2023 Florian Hanzer, Kristian Förster, Benjamin Winter,   #
# Thomas Marke                                                                 #
#                                                                              #
# MELODIST is free software: you can redistribute it and/or modify it under    #
# the terms of the GNU General Public License as published by the Free         #
# Software Foundation, either version 3 of the License, or (at your option)    #
# any later version.                                                           #
#                                                                              #
# MELODIST is distributed in the hope that it will be useful, but WITHOUT ANY  #
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS    #
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more        #
# details.                                                                     #
#                                                                              #
# You should have received a copy of the GNU General Public License along with #
# this program.  If not, see <http://www.gnu.org/licenses/>.                   #
################################################################################
import numpy as np
import pandas as pd
import scipy.optimize

import melodist.util

"""
This routine diaggregates daily values of global radiation data to hourly values
"""


def disaggregate_radiation(
    data_daily,
    sun_times=None,
    pot_rad=None,
    method='pot_rad',
    angstr_a=0.25,
    angstr_b=0.5,
    bristcamp_a=0.75,
    bristcamp_c=2.4,
    mean_course=None,
):
    """general function for radiation disaggregation

    Args:
        daily_data: daily values
        sun_times: daily dataframe including results of the util.sun_times function
        pot_rad: hourly dataframe including potential radiation
        method: keyword specifying the disaggregation method to be used
        angstr_a: parameter a of the Angstrom model (intercept)
        angstr_b: parameter b of the Angstrom model (slope)
        mean_course: monthly values of the mean hourly radiation course

    Returns:
        Disaggregated hourly values of shortwave radiation.
    """
    # check if disaggregation method has a valid value
    if method not in ('pot_rad', 'pot_rad_via_ssd', 'pot_rad_via_bc', 'mean_course'):
        raise ValueError('Invalid option')

    glob_disagg = pd.Series(index=melodist.util.hourly_index(data_daily.index), dtype=float)

    if method == 'mean_course':
        assert mean_course is not None

        pot_rad = pd.Series(index=glob_disagg.index, dtype=float)
        pot_rad[:] = (
            mean_course.unstack().loc[list(zip(pot_rad.index.month, pot_rad.index.hour))].values
        )
    else:
        assert pot_rad is not None

    pot_rad_daily = pot_rad.resample('D').mean()

    if method in ('pot_rad', 'mean_course'):
        globalrad = data_daily.glob
    elif method == 'pot_rad_via_ssd':
        # in this case use the Angstrom model
        globalrad = pd.Series(index=data_daily.index, data=0.0)
        dates = sun_times.index[sun_times.daylength > 0]  # account for polar nights
        globalrad[dates] = angstroem(
            data_daily.ssd[dates],
            sun_times.daylength[dates],
            pot_rad_daily[dates],
            angstr_a,
            angstr_b,
        )
    elif method == 'pot_rad_via_bc':
        # using data from Bristow-Campbell model
        globalrad = bristow_campbell(
            data_daily.tmin, data_daily.tmax, pot_rad_daily, bristcamp_a, bristcamp_c
        )

    globalrad_equal = globalrad.reindex(
        pot_rad.index, method='ffill'
    )  # hourly values (replicate daily mean value for each hour)
    pot_rad_daily_equal = pot_rad_daily.reindex(pot_rad.index, method='ffill')
    glob_disagg = pot_rad / pot_rad_daily_equal * globalrad_equal
    glob_disagg[glob_disagg < 1e-2] = 0.0

    return glob_disagg


def potential_radiation(
    dates,
    lon,
    lat,
    timezone,
    terrain_slope=0,
    terrain_slope_azimuth=0,
    cloud_fraction=0,
    split=False,
):
    """
    Calculate potential shortwave radiation for a specific location and time.

    This routine calculates global radiation as described in:
    Liston, G. E. and Elder, K. (2006): A Meteorological Distribution System for
    High-Resolution Terrestrial Modeling (MicroMet), J. Hydrometeorol., 7, 217–234.

    Corrections for eccentricity are carried out following:
    Paltridge, G.W., Platt, C.M.R., 1976. Radiative processes in Meteorology and Climatology.
    Elsevier Scientific Publishing Company, Amsterdam, Oxford, New York.

    Parameters
    ----------
    dates : DatetimeIndex or array-like
        The dates for which potential radiation shall be calculated
    lon : float
        Longitude (degrees)
    lat : float
        Latitude (degrees)
    timezone : float
        Time zone
    terrain_slope : float, default 0
        Terrain slope as defined in Liston & Elder (2006) (eq. 12)
    terrain_slope_azimuth : float, default 0
        Terrain slope azimuth as defined in Liston & Elder (2006) (eq. 13)
    cloud_fraction : float, default 0
        Cloud fraction between 0 and 1
    split : boolean, default False
        If True, return a DataFrame containing direct and diffuse radiation,
        otherwise return a Series containing total radiation
    """
    solar_constant = 1367.0
    days_per_year = 365.25
    tropic_of_cancer = np.deg2rad(23.43697)
    solstice = 173.0

    dates = pd.DatetimeIndex(dates)
    dates_hour = np.array(dates.hour)
    dates_minute = np.array(dates.minute)
    day_of_year = np.array(dates.dayofyear)

    # compute solar decline in rad
    solar_decline = tropic_of_cancer * np.cos(
        2.0 * np.pi * (day_of_year - solstice) / days_per_year
    )

    # compute the sun hour angle in rad
    standard_meridian = timezone * 15.0
    delta_lat_time = (lon - standard_meridian) * 24.0 / 360.0
    hour_angle = np.pi * (((dates_hour + dates_minute / 60.0 + delta_lat_time) / 12.0) - 1.0)

    # get solar zenith angle
    cos_solar_zenith = np.sin(solar_decline) * np.sin(np.deg2rad(lat)) + np.cos(
        solar_decline
    ) * np.cos(np.deg2rad(lat)) * np.cos(hour_angle)
    cos_solar_zenith = cos_solar_zenith.clip(min=0)
    solar_zenith_angle = np.arccos(cos_solar_zenith)

    # compute transmissivities for direct and diffus radiation using cloud fraction
    transmissivity_direct = (0.6 + 0.2 * cos_solar_zenith) * (1.0 - cloud_fraction)
    transmissivity_diffuse = (0.3 + 0.1 * cos_solar_zenith) * cloud_fraction

    # modify solar constant for eccentricity
    beta = 2.0 * np.pi * (day_of_year / days_per_year)
    radius_ratio = (
        1.00011
        + 0.034221 * np.cos(beta)
        + 0.00128 * np.sin(beta)
        + 0.000719 * np.cos(2.0 * beta)
        + 0.000077 * np.sin(2 * beta)
    )
    solar_constant_times_radius_ratio = solar_constant * radius_ratio

    mu = np.arcsin(np.cos(solar_decline) * np.sin(hour_angle) / np.sin(solar_zenith_angle))
    cosi = np.cos(terrain_slope) * cos_solar_zenith + np.sin(terrain_slope) * np.sin(
        solar_zenith_angle
    ) * np.cos(mu - terrain_slope_azimuth)

    # get total shortwave radiation
    direct_radiation = solar_constant_times_radius_ratio * transmissivity_direct * cosi
    diffuse_radiation = (
        solar_constant_times_radius_ratio * transmissivity_diffuse * cos_solar_zenith
    )
    direct_radiation = direct_radiation.clip(min=0)

    df = pd.DataFrame(index=dates, data=dict(direct=direct_radiation, diffuse=diffuse_radiation))

    if split:
        return df
    else:
        return df.direct + df.diffuse


def bristow_campbell(tmin, tmax, pot_rad_daily, A, C):
    """calculates potential shortwave radiation based on minimum and maximum temperature

    This routine calculates global radiation as described in:
    Bristow, Keith L., and Gaylon S. Campbell: On the relationship between
    incoming solar radiation and daily maximum and minimum temperature.
    Agricultural and forest meteorology 31.2 (1984): 159-166.

    Args:
        daily_data: time series (daily data) including at least minimum and maximum temeprature
        pot_rad_daily: mean potential daily radiation
        A: parameter A of the Bristow-Campbell model
        C: parameter C of the Bristow-Campbell model
    Returns:
        series of potential shortwave radiation
    """

    assert tmin.index.equals(tmax.index)

    temp = pd.DataFrame(data=dict(tmin=tmin, tmax=tmax))
    temp = temp.reindex(pd.date_range(start=temp.index[0], end=temp.index[-1], freq='D'))
    temp['tmin_nextday'] = temp.tmin
    temp.loc[temp.index[:-1], 'tmin_nextday'] = temp['tmin'].iloc[1:].values

    temp = temp.loc[tmin.index]
    pot_rad_daily = pot_rad_daily.loc[tmin.index]

    dT = temp.tmax - (temp.tmin + temp.tmin_nextday) / 2

    dT_m_avg = dT.groupby(dT.index.month).mean()
    B = 0.036 * np.exp(-0.154 * dT_m_avg[temp.index.month])
    B.index = temp.index

    if isinstance(A, pd.Series):
        months = temp.index.month
        A = A.loc[months].values
        C = C.loc[months].values

    transmissivity = A * (1 - np.exp(-B * dT**C))
    R0 = transmissivity * pot_rad_daily

    return R0


def fit_bristow_campbell_params(tmin, tmax, pot_rad_daily, obs_rad_daily):
    """
    Fit the A and C parameters for the Bristow & Campbell (1984) model using observed daily
    minimum and maximum temperature and mean daily (e.g. aggregated from hourly values) solar
    radiation.

    Parameters
    ----------
    tmin : Series
        Observed daily minimum temperature.

    tmax : Series
        Observed daily maximum temperature.

    pot_rad_daily : Series
        Mean potential daily solar radiation.

    obs_rad_daily : Series
        Mean observed daily solar radiation.
    """

    def bc_absbias(ac):
        return np.abs(np.mean(bristow_campbell(df.tmin, df.tmax, df.pot, ac[0], ac[1]) - df.obs))

    df = pd.DataFrame(data=dict(tmin=tmin, tmax=tmax, pot=pot_rad_daily, obs=obs_rad_daily)).dropna(
        how='any'
    )
    res = scipy.optimize.minimize(bc_absbias, [0.75, 2.4])  # i.e. we minimize the absolute bias

    return res.x


def angstroem(ssd, day_length, pot_rad_daily, a, b):
    """
    Calculate mean daily radiation from observed sunshine duration according to Angstroem (1924).

    Parameters
    ----------
    ssd : Series
        Observed daily sunshine duration.

    day_length : Series
        Day lengths as calculated by ``calc_sun_times``.

    pot_rad_daily : Series
        Mean potential daily solar radiation.

    a : float
        First parameter for the Angstroem model (originally 0.25).

    b : float
        Second parameter for the Angstroem model (originally 0.75).
    """
    if isinstance(a, pd.Series):
        months = ssd.index.month
        a = a.loc[months].values
        b = b.loc[months].values

    glob_day = (a + b * ssd / day_length) * pot_rad_daily

    return glob_day


def fit_angstroem_params(ssd, day_length, pot_rad_daily, obs_rad_daily):
    """
    Fit the a and b parameters for the Angstroem (1924) model using observed daily
    sunshine duration and mean daily (e.g. aggregated from hourly values) solar
    radiation.

    Parameters
    ----------
    ssd : Series
        Observed daily sunshine duration.

    day_length : Series
        Day lengths as calculated by ``calc_sun_times``.

    pot_rad_daily : Series
        Mean potential daily solar radiation.

    obs_rad_daily : Series
        Mean observed daily solar radiation.
    """
    df = pd.DataFrame(
        data=dict(ssd=ssd, day_length=day_length, pot=pot_rad_daily, obs=obs_rad_daily)
    ).dropna(how='any')

    def angstroem_opt(x, a, b):
        return angstroem(x[0], x[1], x[2], a, b)

    x = np.array([df.ssd, df.day_length, df.pot])
    popt, pcov = scipy.optimize.curve_fit(angstroem_opt, x, df.obs, p0=[0.25, 0.75])

    return popt
