from datetime import timedelta
from typing import Callable

import numpy as np
from vsm_datascience_common import constants
from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
from vsm_datascience_common.cycles.sequences import get_timestamps, get_values

from utils.smoothing import get_curve_composition


def get_day_of_week_encoded(cycle: dict, shift: int = 0) -> dict:
    timestamps = get_timestamps(cycle)
    weekdays = np.array([x.weekday() for x in timestamps])
    sine_encoded = np.sin(weekdays * (2 * np.pi / 7))
    cosine_encoded = np.cos(weekdays * (2 * np.pi / 7))
    return {
        "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None],
        "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None]
    }


def get_hour_of_day_encoded(cycle: dict, shift: int = 0) -> dict:
    timestamps = get_timestamps(cycle)
    hours = np.array([x.hour for x in timestamps])
    sine_encoded = np.sin(hours * (2 * np.pi / 24))
    cosine_encoded = np.cos(hours * (2 * np.pi / 24))
    return {
        "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None],
        "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None]
    }


def get_month_of_year_encoded(cycle: dict, shift: int = 0) -> dict:
    timestamps = get_timestamps(cycle)
    months = np.array([x.month for x in timestamps])
    sine_encoded = np.sin(months * (2 * np.pi / 12))
    cosine_encoded = np.cos(months * (2 * np.pi / 12))
    return {
        "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None],
        "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None]
    }


def get_hours_from_start(cycle: dict, shift: int = 0) -> np.ndarray:
    timestamps = get_timestamps(cycle)
    return np.arange(len(timestamps))[shift if shift > 0 else 0:shift if shift < 0 else None]


def get_temperature(cycle: dict, shift: int = 0) -> np.ndarray:
    return get_values(cycle)[shift if shift > 0 else 0:shift if shift < 0 else None]


def get_fertility_curve(cycle_length: int,
                        ov_index: int,
                        fertility_base_curve: np.ndarray | list,
                        base_curve_offset: int) -> np.ndarray:
    fert_prob = np.zeros((cycle_length,))
    if ov_index + base_curve_offset < 0:
        length_in_cycle = len(fertility_base_curve) + ov_index + base_curve_offset
        fert_prob[:length_in_cycle] = fertility_base_curve[-length_in_cycle:]
    elif ov_index + base_curve_offset + len(fertility_base_curve) > cycle_length:
        length_in_cycle = cycle_length - ov_index - base_curve_offset
        fert_prob[ov_index + base_curve_offset:] = fertility_base_curve[:length_in_cycle]
    else:
        fert_prob[ov_index + base_curve_offset:
                  ov_index + base_curve_offset + len(fertility_base_curve)] = fertility_base_curve
    return fert_prob


def get_fertility_probability_base_curve() -> (np.ndarray, int):
    """
    This function returns the base curve for the fertility probability as well as the offset to the ovulation day in measurements
    :return: numpy array with the base curve resampled to the number of measurements per day
    """
    fertilization_chance_curve = [0, 0.07, 0.12, 0.25, 0.3, 0.18, 0]
    offset_to_ov = -6

    new_indices = np.arange(len(fertilization_chance_curve) * constants.MEASUREMENTS_PER_DAY)
    fertility_chance_curve_resampled = np.interp(new_indices,
                                                 np.linspace(0, len(new_indices), num=len(fertilization_chance_curve)),
                                                 fertilization_chance_curve)
    ov_offset_resampled = offset_to_ov * constants.MEASUREMENTS_PER_DAY
    return fertility_chance_curve_resampled, ov_offset_resampled


def get_fertility_probability(cycle: dict, shift: int = 0) -> np.ndarray:
    # check, if biphasic
    if "classification_results" in cycle and "results" in cycle["classification_results"][0]:
        if cycle["classification_results"][0]["results"]["predicted_class"] != "biphasic":
            return np.zeros((len(get_values(cycle)) - abs(shift, )))
    else:
        return np.zeros((len(get_values(cycle)) - abs(shift, )))

    if "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]:
        try:
            ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"]
            timestamps = get_timestamps(cycle)
            ov_index = np.where(timestamps >= ov_timestamp)[0][0]
            fertility_chance_curve_resampled, ov_offset_resampled = get_fertility_probability_base_curve()
            return get_fertility_curve(len(get_values(cycle)), ov_index, fertility_chance_curve_resampled,
                                       ov_offset_resampled)[shift if shift > 0 else 0:shift if shift < 0 else None]
        except:
            return np.zeros((len(get_values(cycle)) - abs(shift, )))
    else:
        return np.zeros((len(get_values(cycle)) - abs(shift, )))


def get_ov_over_probability(cycle: dict, shift: int = 0) -> np.ndarray:
    """
    Get the ovulation over probability for a cycle
    :param cycle: cycle data
    :param shift: shift to apply
    :return: ovulation over probability
    """
    if "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]:
        try:
            ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"]
            timestamps = get_timestamps(cycle)
            ov_index = np.where(timestamps > ov_timestamp)[0][0]
            ov_over_probability = np.zeros(len(get_values(cycle)))
            ov_over_probability[ov_index:] = 1
            return ov_over_probability[shift if shift > 0 else 0:shift if shift < 0 else None]
        except:
            return np.zeros(len(get_values(cycle)))[shift if shift > 0 else 0:shift if shift < 0 else None]
    else:
        return np.zeros(len(get_values(cycle)))[shift if shift > 0 else 0:shift if shift < 0 else None]


def get_days_relative_to_ov(cycle: dict) -> np.ndarray:
    """
    Get the time in days until the next OV event.
    """
    timestamps = get_timestamps(cycle)
    if ("ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]) and (
            "classification_results" in cycle and "results" in cycle["classification_results"][0] and
            cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"):
        try:
            ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"]
            ov_index = np.where(timestamps > ov_timestamp)[0][0]
            ov_timestamp = timestamps[ov_index]
            days_until_ov = np.empty(len(timestamps))
            for i in range(len(timestamps)):
                current_day = timestamps[i]
                day_diff = ov_timestamp - current_day
                if day_diff.days < -1:
                    pass
                # days_until_ov[i] = max(day_diff.days -1)
                # use - days, as days relative to event
                days_until_ov[i] = - day_diff.days

            return np.array(days_until_ov)
        except:
            pass

    return np.full((len(timestamps),), np.nan)


def get_ov_day(cycle: dict) -> np.ndarray:
    """
    Get the ovulation day for a cycle
    :param cycle: cycle data
    :return: ovulation day
    """
    values = get_values(cycle)
    seq_len = len(values)
    is_biphasic = "classification_results" in cycle and "results" in cycle["classification_results"][0] and \
                  cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"
    has_ov_results = "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]
    if is_biphasic and has_ov_results:
        try:
            ov_day = cycle["ov_detection_results"][0]["results"]["ovulation_day"]
            return np.full((seq_len,), ov_day)
        except:
            pass

    return np.full((seq_len,), np.nan)


def get_is_biphasic(cycle: dict) -> np.ndarray:
    """
    Get the is_biphasic flag for a cycle
    :param cycle: cycle data
    :return: is_biphasic flag
    """
    values = get_values(cycle)
    seq_len = len(values)
    is_biphasic = "classification_results" in cycle and "results" in cycle["classification_results"][0] and \
                  cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"
    if is_biphasic:
        return np.full((seq_len,), 1)

    return np.full((seq_len,), 0)


def get_curve_composition_as_features(cycle: dict, shift: int = 0) -> dict:
    composition = get_curve_composition(get_values(cycle))
    return {
        "trend": composition[0][shift if shift > 0 else 0:shift if shift < 0 else None],
        "seasonal": composition[1][shift if shift > 0 else 0:shift if shift < 0 else None],
        "residual": composition[2][shift if shift > 0 else 0:shift if shift < 0 else None],
        "smoothed": composition[3][shift if shift > 0 else 0:shift if shift < 0 else None]
    }


def get_rolling_average_with_padding(cycle: dict,
                                     shift: int = 0,
                                     rolling_average_window_length: int = constants.MEASUREMENTS_PER_DAY) -> np.ndarray:
    """
    Get the rolling average of the temperature values with padding, so that the rolling average has the same length as the
    original temperature values
    Args:
        cycle (dict): cycle data
        shift (int): shift to apply
        rolling_average_window_length (int): window length for the rolling average, default is one day (288 measurements)

    Returns:
        np.ndarray: rolling average of the temperature values with padding

    """

    values = get_values(cycle)
    values_padded = np.concatenate(
        [np.full((constants.MEASUREMENTS_PER_DAY,), np.mean(values[:constants.MEASUREMENTS_PER_DAY])), values])
    rolling_average = np.convolve(values_padded, np.ones(rolling_average_window_length) / rolling_average_window_length,
                                  mode="valid")[-len(values):]
    return rolling_average[shift if shift > 0 else 0:shift if shift < 0 else None]


def get_window_fn(cycle: dict,
                  window_size: int,
                  fn: Callable) -> np.ndarray:
    """
    Apply a window function to a time series.
    :param cycle: cycle data
    :param window_size: window size
    :return: array with minimum value in window for each value in input array
    """
    values = get_values(cycle)
    values_padded = np.concatenate([np.full(window_size - 1, values[0]), values])

    windows = np.lib.stride_tricks.sliding_window_view(values_padded, window_shape=window_size)
    return fn(windows)


def get_cycle_length_stats(cycle: dict) -> dict:
    user_id = cycle["user_id"]
    cycle_length = len(get_values(cycle))
    cycles_dates = list(get_cycles_collection().find({
        "user_id": user_id, "ends_at": {"$exists": True},
        "starts_at": {"$lt": cycle["starts_at"]},
    },
        {"starts_at": 1, "ends_at": 1}))
    if len(cycles_dates) == 0:
        # use the average cycle length as fallback
        average_cycle_length = 37.13
        cycle_length_std = 12.42
    else:
        cycle_lengths = [x["ends_at"] - x["starts_at"] for x in cycles_dates]
        average_cycle_length = np.mean(cycle_lengths).total_seconds() / (60 * 60 * 24)
        cycle_length_std = np.std([x.total_seconds() for x in cycle_lengths]) / (60 * 60 * 24)

    return {
        "average_cycle_length": np.full((cycle_length,), average_cycle_length),
        "cycle_length_std": np.full((cycle_length,), cycle_length_std),
    }


def get_average_ovulation_day(cycle: dict) -> float:
    user_id = cycle["user_id"]
    cycle_length = len(get_values(cycle))
    cycles_dates = list(
        get_cycles_collection().find({
            "user_id": user_id,
            "ov_detection_results.0.results": {"$exists": True},
            "starts_at": {"$lt": cycle["starts_at"]},
        },
            {"ov_detection_results": 1}))
    if len(cycles_dates) == 0:
        # use the average ovulation day as fallback
        average_ovulation_day = 18.9
    else:
        ovulation_days = [x["ov_detection_results"][0]["results"]["ovulation_day"] for x in cycles_dates]
        average_ovulation_day = np.mean(ovulation_days)
    return np.full((cycle_length,), average_ovulation_day)


def get_ovulation_std(cycle: dict) -> float:
    user_id = cycle["user_id"]
    cycle_length = len(get_values(cycle))
    cycles_dates = list(
        get_cycles_collection().find({
            "user_id": user_id,
            "ov_detection_results.0.results": {"$exists": True},
            "starts_at": {"$lt": cycle["starts_at"]},
        },
            {"ov_detection_results": 1}))
    if len(cycles_dates) == 0:
        # use the average ovulation std as fallback
        ovulation_std = 4.07
    else:
        ovulation_days = [x["ov_detection_results"][0]["results"]["ovulation_day"] for x in cycles_dates]
        ovulation_std = np.std(ovulation_days)
    return np.full((cycle_length,), ovulation_std)


def get_biphasic_fraction(cycle: dict) -> float:
    user_id = cycle["user_id"]
    cycle_length = len(get_values(cycle))
    cycle_dates = list(
        get_cycles_collection().find({
            "user_id": user_id,
            "classification_results.0.results": {"$exists": True},
            "starts_at": {"$lt": cycle["starts_at"]},
        },
            {"classification_results": 1}))
    if len(cycle_dates) == 0:
        # use the average cycle length as fallback
        biphasic_fraction = 0.9449
    else:
        biphasic = [x for x in cycle_dates if
                    x["classification_results"][0]["results"]["predicted_class"] == "biphasic"]
        biphasic_fraction = len(biphasic) / len(cycle_dates)
    return np.full((cycle_length,), biphasic_fraction)


def get_num_cycles(cycle: dict) -> int:
    user_id = cycle["user_id"]
    cycle_length = len(get_values(cycle))
    cycles_dates = list(get_cycles_collection().find({"user_id": user_id,
                                                      "ends_at": {"$exists": True},
                                                      "starts_at": {"$lt": cycle["starts_at"]}
                                                      },
                                                     {"starts_at": 1, "ends_at": 1}))
    num_cycles = len(cycles_dates)
    return np.full((cycle_length,), num_cycles)


def get_average_temperatures(cycle: dict) -> dict:
    user_id = cycle["user_id"]
    cycle_length = len(get_values(cycle))
    cycles_dates = list(
        get_cycles_collection().find({"user_id": user_id, "ov_detection_results.0.results": {"$exists": True}},
                                     {"starts_at": 1, "ov_detection_results": 1, "measurements": 1}))

    if len(cycles_dates) == 0:
        return {
            "pre_ov_temperatures": np.full((cycle_length,), 37.07),
            "post_ov_temperatures": np.full((cycle_length,), 37.37),
        }

    pre_ov_temperatures = []
    post_ov_temperatures = []
    for x in cycles_dates:
        ov_detection = x["ov_detection_results"][0]["results"]
        values = get_values(x)
        timestamps = get_timestamps(x)
        cutoff_day = x["starts_at"] + timedelta(days=ov_detection["ovulation_day"])
        cutoff_index = np.where(np.array(timestamps) >= cutoff_day)[0][0]
        pre_ov_temperatures.append(values[:cutoff_index])
        post_ov_temperatures.append(values[cutoff_index:])

    pre_ov_temperatures = np.concatenate(pre_ov_temperatures)
    post_ov_temperatures = np.concatenate(post_ov_temperatures)

    pre_ov_temperatures = np.mean(pre_ov_temperatures)
    post_ov_temperatures = np.mean(post_ov_temperatures)

    return {
        "pre_ov_temperatures": np.full((cycle_length,), pre_ov_temperatures),
        "post_ov_temperatures": np.full((cycle_length,), post_ov_temperatures),
    }