################################################################################
# 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.stats


def hourly_index(daily_index, fill_gaps=False):
    index = pd.date_range(start=daily_index.min(), end=daily_index.max().replace(hour=23), freq='H')

    # remove days that are not in the daily index:
    if not fill_gaps and len(index) > len(daily_index) * 24:
        df = pd.DataFrame(index=index.date, data=dict(hour=index.hour, keep=True))
        df.index = df.index.to_datetime()
        dropdates = list(set(index.date) - set(daily_index.date))
        df.loc[dropdates, 'keep'] = False
        df = df[df.keep]
        index = pd.DatetimeIndex(
            [
                pd.datetime(y, m, d, h)
                for y, m, d, h in zip(df.index.year, df.index.month, df.index.day, df.hour)
            ]
        )

    return index


def distribute_equally(daily_data, divide=False):
    """Obtains hourly values by equally distributing the daily values.

    Args:
        daily_data: daily values
        divide: if True, divide resulting values by the number of hours in
            order to preserve the daily sum (required e.g. for precipitation).

    Returns:
        Equally distributed hourly values.
    """

    index = hourly_index(daily_data.index)
    hourly_data = daily_data.reindex(index)
    hourly_data = hourly_data.groupby(hourly_data.index.day).transform(
        lambda x: x.fillna(method='ffill', limit=23)
    )

    if divide:
        hourly_data /= 24

    return hourly_data


def vapor_pressure(temp, hum):
    """
    Calculates vapor pressure from temperature and humidity after Sonntag (1990).

    Args:
        temp: temperature values
        hum: humidity value(s). Can be scalar (e.g. for calculating saturation vapor pressure).

    Returns:
        Vapor pressure in hPa.
    """

    if np.isscalar(hum):
        hum = np.zeros(temp.shape) + hum

    assert temp.shape == hum.shape

    positives = np.array(temp >= 273.15)
    vap_press = np.zeros(temp.shape) * np.nan
    vap_press[positives] = (
        6.112
        * np.exp((17.62 * (temp[positives] - 273.15)) / (243.12 + (temp[positives] - 273.15)))
        * hum[positives]
        / 100.0
    )
    vap_press[~positives] = (
        6.112
        * np.exp((22.46 * (temp[~positives] - 273.15)) / (272.62 + (temp[~positives] - 273.15)))
        * hum[~positives]
        / 100.0
    )

    return vap_press


def dewpoint_temperature(temp, hum):
    """computes the dewpoint temperature

    Parameters
    ----
    temp :      temperature [K]
    hum :       relative humidity


    Returns
        dewpoint temperature in K
    """
    assert temp.shape == hum.shape

    vap_press = vapor_pressure(temp, hum)

    positives = np.array(temp >= 273.15)
    dewpoint_temp = temp.copy() * np.nan
    dewpoint_temp[positives] = (
        243.12
        * np.log(vap_press[positives] / 6.112)
        / (17.62 - np.log(vap_press[positives] / 6.112))
    )
    dewpoint_temp[~positives] = (
        272.62
        * np.log(vap_press[~positives] / 6.112)
        / (22.46 - np.log(vap_press[~positives] / 6.112))
    )

    return dewpoint_temp + 273.15


def linregress(x, y, return_stats=False):
    """linear regression calculation

    Parameters
    ----
    x :         independent variable (series)
    y :         dependent variable (series)
    return_stats : returns statistical values as well if required (bool)


    Returns
    ----
    list of parameters (and statistics)
    """
    a1, a0, r_value, p_value, stderr = scipy.stats.linregress(x, y)

    retval = a1, a0
    if return_stats:
        retval += r_value, p_value, stderr

    return retval


def get_sun_times(dates, lon, lat, time_zone):
    """Computes the times of sunrise, solar noon, and sunset for each day.

    Parameters
    ----
    dates:      datetime
    lat :       latitude in DecDeg
    lon :       longitude in DecDeg
    time_zone : timezone


    Returns
    ----
    DataFrame:  [sunrise, sunnoon, sunset, day length] in dec hours
    """

    df = pd.DataFrame(index=dates, columns=['sunrise', 'sunnoon', 'sunset', 'daylength'])

    doy = np.array([(d - d.replace(day=1, month=1)).days + 1 for d in df.index])  # day of year

    # Day angle and declination after Bourges (1985):
    day_angle_b = np.deg2rad((360.0 / 365.25) * (doy - 79.346))

    declination = np.deg2rad(
        0.3723
        + 23.2567 * np.sin(day_angle_b)
        - 0.7580 * np.cos(day_angle_b)
        + 0.1149 * np.sin(2 * day_angle_b)
        + 0.3656 * np.cos(2 * day_angle_b)
        - 0.1712 * np.sin(3 * day_angle_b)
        + 0.0201 * np.cos(3 * day_angle_b)
    )

    # Equation of time with day angle after Spencer (1971):
    day_angle_s = 2 * np.pi * (doy - 1) / 365.0
    eq_time = (
        12.0
        / np.pi
        * (
            0.000075
            + 0.001868 * np.cos(day_angle_s)
            - 0.032077 * np.sin(day_angle_s)
            - 0.014615 * np.cos(2 * day_angle_s)
            - 0.040849 * np.sin(2 * day_angle_s)
        )
    )

    #
    standard_meridian = time_zone * 15.0
    delta_lat_time = (lon - standard_meridian) * 24.0 / 360.0

    omega_nul_arg = -np.tan(np.deg2rad(lat)) * np.tan(declination)
    omega_nul = np.arccos(omega_nul_arg)
    sunrise = 12.0 * (1.0 - (omega_nul) / np.pi) - delta_lat_time - eq_time
    sunset = 12.0 * (1.0 + (omega_nul) / np.pi) - delta_lat_time - eq_time

    # as an approximation, solar noon is independent of the below mentioned
    # cases:
    sunnoon = 12.0 * (1.0) - delta_lat_time - eq_time

    # $kf 2015-11-13: special case midnight sun and polar night
    # CASE 1: MIDNIGHT SUN
    # set sunrise and sunset to values that would yield the maximum day
    # length even though this a crude assumption
    pos = omega_nul_arg < -1
    sunrise[pos] = sunnoon[pos] - 12
    sunset[pos] = sunnoon[pos] + 12

    # CASE 2: POLAR NIGHT
    # set sunrise and sunset to values that would yield the minmum day
    # length even though this a crude assumption
    pos = omega_nul_arg > 1
    sunrise[pos] = sunnoon[pos]
    sunset[pos] = sunnoon[pos]

    daylength = sunset - sunrise

    # adjust if required
    sunrise[sunrise < 0] += 24
    sunset[sunset > 24] -= 24

    df.sunrise = sunrise
    df.sunnoon = sunnoon
    df.sunset = sunset
    df.daylength = daylength

    return df


def detect_gaps(dataframe, timestep, print_all=False, print_max=5, verbose=True):
    """checks if a given dataframe contains gaps and returns the number of gaps

    This function checks if a dataframe contains any gaps for a given temporal
    resolution that needs to be specified in seconds. The number of gaps
    detected in the dataframe is returned.

    Args:
        dataframe: A pandas dataframe object with index defined as datetime
        timestep (int): The temporal resolution of the time series in seconds
            (e.g., 86400 for daily values)
        print_all (bool, opt): Lists every gap on the screen
        print_mx (int, opt): The maximum number of gaps listed on the screen in
            order to avoid a decrease in performance if numerous gaps occur
        verbose (bool, opt): Enables/disables output to the screen

    Returns:
        The number of gaps as integer. Negative values indicate errors.
    """
    gcount = 0
    msg_counter = 0
    warning_printed = False
    try:
        n = len(dataframe.index)
    except:
        print('Error: Invalid dataframe.')
        return -1
    for i in range(0, n):
        if i > 0:
            time_diff = dataframe.index[i] - dataframe.index[i - 1]
            if time_diff.delta / 1e9 != timestep:
                gcount += 1
                if print_all or (msg_counter <= print_max - 1):
                    if verbose:
                        print(
                            'Warning: Gap in time series found between %s and %s'
                            % (dataframe.index[i - 1], dataframe.index[i])
                        )
                    msg_counter += 1
                if msg_counter == print_max and verbose and not warning_printed:
                    print(
                        'Warning: Only the first %i gaps have been listed. Try to increase '
                        'print_max parameter to show more details.'
                        % msg_counter
                    )
                    warning_printed = True
    if verbose:
        print('%i gaps found in total.' % (gcount))
    return gcount


def drop_incomplete_days(dataframe, shift=0):
    """truncates a given dataframe to full days only

    This function truncates a given pandas dataframe (time series) to full days
    only, thus dropping leading and tailing hours of incomplete days. Please
    note that this methodology only applies to hourly time series.

    Args:
        dataframe: A pandas dataframe object with index defined as datetime
        shift (unsigned int, opt): First hour of daily recordings. For daily
            recordings of precipitation gages, 8 would be the first hour of
            the subsequent day of recordings since daily totals are
            usually recorded at 7. Omit defining this parameter if you intend
            to pertain recordings to 0-23h.

    Returns:
        A dataframe with full days only.
    """
    dropped = 0
    if shift > 23 or shift < 0:
        print("Invalid shift parameter setting! Using defaults.")
        shift = 0
    first = shift
    last = first - 1
    if last < 0:
        last += 24
    try:
        # todo: move this checks to a separate function
        n = len(dataframe.index)
    except:
        print('Error: Invalid dataframe.')
        return dataframe

    delete = list()

    # drop heading lines if required
    for i in range(0, n):
        if dataframe.index.hour[i] == first and dataframe.index.minute[i] == 0:
            break
        else:
            delete.append(i)
            dropped += 1

    # drop tailing lines if required
    for i in range(n - 1, 0, -1):
        if dataframe.index.hour[i] == last and dataframe.index.minute[i] == 0:
            break
        else:
            delete.append(i)
            dropped += 1
    # print("The following rows have been dropped (%i in total):" % dropped)
    # print(delete)
    return dataframe.drop(index=dataframe.index[delete])


def prepare_interpolation_data(data_daily, column_hours):
    start_date = data_daily.index[0]
    end_date = data_daily.index[-1]

    data = pd.Series()

    for column, hour in column_hours.items():
        index = pd.date_range(
            start=start_date.replace(hour=hour), end=end_date.replace(hour=hour), freq='D'
        )
        s = pd.Series(index=index, data=data_daily[column].values)
        data = data.append(s)

    data = data.sort_index()
    data = data.reindex(hourly_index(data_daily.index))

    return data


def daily_from_hourly(df):
    """Aggregates data (hourly to daily values) according to the characteristics
    of each variable (e.g., average for temperature, sum for precipitation)

    Args:
        df: dataframe including time series with one hour time steps

    Returns:
        dataframe (daily)
    """

    df_daily = pd.DataFrame()

    if 'temp' in df:
        df_daily['temp'] = df.temp.resample('D').mean()
        df_daily['tmin'] = df.temp.groupby(df.temp.index.date).min()
        df_daily['tmax'] = df.temp.groupby(df.temp.index.date).max()

    if 'precip' in df:
        df_daily['precip'] = df.precip.resample('D').sum()

    if 'glob' in df:
        df_daily['glob'] = df.glob.resample('D').mean()

    if 'hum' in df:
        df_daily['hum'] = df.hum.resample('D').mean()

    if 'hum' in df:
        df_daily['hum_min'] = df.hum.groupby(df.hum.index.date).min()

    if 'hum' in df:
        df_daily['hum_max'] = df.hum.groupby(df.hum.index.date).max()

    if 'wind' in df:
        df_daily['wind'] = df.wind.resample('D').mean()

    if 'ssd' in df:
        df_daily['ssd'] = df.ssd.resample('D').sum() / 60  # minutes to hours

    df_daily.index.name = None
    return df_daily


def calculate_mean_daily_course_by_month(data_hourly, normalize=False):
    data_hourly = data_hourly.copy()
    data_hourly.index.name = None

    df = data_hourly.groupby([data_hourly.index.month, data_hourly.index.hour]).mean()
    df = df.reset_index().pivot(columns='level_0', index='level_1')
    df.columns = df.columns.droplevel()  # remove MultiIndex
    df.columns.name = None
    df.index.name = None

    if normalize:
        df = (df - df.min()) / (df.max() - df.min())  # normalize values to 0-1 range

    return df
