"""
InsuLearner: Estimating Insulin Pump Settings via Machine Learning

Code underlying this blog post:
https://www.cameronsummers.com/how_I_calculate_my_sons_insulin_pump_settings_with_machine_learning
"""

import argparse
import datetime as dt
import warnings

import matplotlib.cbook
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scipy
import seaborn as sns
from sklearn.linear_model import LinearRegression

from InsuLearner.carbohydrate_sensitivity_factor import estimate_csf
from InsuLearner.nightscout.nightscout_api import NightscoutAPI
from InsuLearner.nightscout.nightscout_user_model import NightscoutUser
from InsuLearner.tidepool.tidepool_api import TidepoolAPI
from InsuLearner.tidepool.tidepool_user_model import TidepoolUser
from InsuLearner.util import get_logger

warnings.filterwarnings("ignore", category=matplotlib.MatplotlibDeprecationWarning)

__author__ = "Cameron Summers"
logger = get_logger(__name__)


def compute_aace_pump_settings(weight_kg, prepump_tdd):
    """
    Get pump settings using the American Association of Clinical Endocrinologists/American College of
    Endocrinology.

    Other references:
    https://diabetesed.net/wp-content/uploads/2019/09/Insulin-Pump-Calculations-Sept-2019-slides.pdf

    Review of insulin dosing formulas
    https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4960276/

    Args:
        weight_kg (float): weight in kg
        prepump_tdd (float): units of insulin per day, (CS: presuming this is a life-time average)

    Returns:
        (float, float, float): basal rate, carb insulin ratio, insulin sensitivity factor
    """
    tdd_method1 = weight_kg * 0.5
    tdd_method2 = prepump_tdd * 0.75
    starting_pump_tdd = (tdd_method1 + tdd_method2) / 2

    basal_rate = starting_pump_tdd * 0.5 / 24
    cir = 450.0 / starting_pump_tdd
    isf = 1700.0 / starting_pump_tdd

    return basal_rate, cir, isf


def pd_1d_series_to_X(series):
    """helper function"""
    return np.array(series.values.tolist())[:, np.newaxis]


def estimate_therapy_settings_from_window_stats_lr(aggregated_df,
                                                   K,
                                                   period_window_size_hours,
                                                   target_bg=110,
                                                   x="total_carbs",
                                                   y="total_insulin",
                                                   do_plots=True,
                                                   trained_model=None,
                                                   weight_scheme=None):
    """
    Fit the linear model and estimate settings
    """
    X_carbs = pd_1d_series_to_X(aggregated_df[x])
    y_insulin = aggregated_df[y]

    if weight_scheme is None:
        sample_weights = None
    else:
        if weight_scheme == "CGM Weighted":
            sample_weights = 1.0 / np.array(np.maximum(1.0, abs(target_bg - aggregated_df["cgm_geo_mean"])) / aggregated_df["cgm_percent_tir"]) * aggregated_df["cgm_percent_available"]
        elif weight_scheme == "Carb Uncertainty":
            sample_weights = scipy.stats.norm.pdf(aggregated_df["total_carbs"], aggregated_df["total_carbs"].mean(), aggregated_df["total_carbs"].std())
        elif weight_scheme == "Recency":
            sample_weights = 1 / ((aggregated_df["end_date"].max() - aggregated_df["end_date"]).dt.days + 5)
            # print(sample_weights)
        else:
            raise Exception("Unknown weight scheme {}.".format(weight_scheme))

        nan_mask = np.isnan(sample_weights)
        X_carbs = X_carbs[~nan_mask]
        y_insulin = y_insulin[~nan_mask]
        sample_weights = sample_weights[~nan_mask]

        if do_plots:
            plt.figure()
            plt.title("Sample Weight Distribution")
            plt.hist(sample_weights)

    if trained_model is None:
        lm_carb_to_insulin = LinearRegression()
        lm_carb_to_insulin.fit(X_carbs, y_insulin, sample_weight=sample_weights)
    else:
        lm_carb_to_insulin = trained_model

    basal_insulin_estimate = lm_carb_to_insulin.intercept_

    r2_fit = lm_carb_to_insulin.score(X_carbs, y_insulin)
    logger.info(f"Linear Model: Fit R^2 {np.round(r2_fit, 2)}. Intercept {np.round(lm_carb_to_insulin.intercept_, 2)}. Slope g/U {np.round(1.0 / lm_carb_to_insulin.coef_, 2)}")

    cir_estimate_slope = 1 / lm_carb_to_insulin.coef_[0]
    isf_estimate_slope = cir_estimate_slope * K

    basal_rate_estimate_hourly = basal_insulin_estimate / period_window_size_hours

    logger.info("Total Period Basal={:.2f}U. (Mean %Daily Total: {:.2f}%)".format(basal_insulin_estimate, np.nanmean(basal_insulin_estimate / aggregated_df[y]) * 100))

    logger.info("\n\n\tSettings Estimates:\n")

    logger.info(f"\tEstimated CIR={round(cir_estimate_slope, 2)} g/U.")
    logger.info(f"\tEstimated Hourly Basal={round(basal_rate_estimate_hourly, 3)} U/hr")
    logger.info(f"\tCSF={round(K, 2)} mg/dL / g")
    logger.info(f"\tEstimated ISF={round(isf_estimate_slope, 2)} mg/dL/ U")

    # TMP
    # y_insulin = pd_1d_series_to_X(aggregated_df[y])
    # X_carbs = aggregated_df[x]
    # lm_carb_to_insulin_inverted = LinearRegression()
    # lm_carb_to_insulin_inverted.fit(y_insulin, X_carbs, sample_weight=sample_weights)
    # r2_fit_inverted = lm_carb_to_insulin_inverted.score(y_insulin, X_carbs)
    # cir_estimate_slope_inverted = lm_carb_to_insulin_inverted.coef_[0]
    # isf_estimate_slope_inverted = cir_estimate_slope_inverted * K
    # basal_rate_estimate_inverted = -lm_carb_to_insulin_inverted.intercept_ / (cir_estimate_slope_inverted)
    # logger.info(f"\tSettings Inverted: R2={r2_fit_inverted} CIR={cir_estimate_slope_inverted} ISF={isf_estimate_slope_inverted} Basal Hourly={basal_rate_estimate_inverted/24}.")
    # TMP

    settings = (cir_estimate_slope, isf_estimate_slope, basal_insulin_estimate, lm_carb_to_insulin, r2_fit, K)
    if do_plots:
        plot_aggregated_scatter(aggregated_df,
                                period_window_size_hours,
                                lr_model=lm_carb_to_insulin,
                                settings=settings,
                                plot_aace=True,
                                weight_scheme=weight_scheme)

    return settings


def plot_aggregated_scatter(aggregated_df, period_window_size_hours, lr_model=None, settings=None, plot_aace=True, weight_scheme=None):
    """
    Plot the linear model on the data.
    """
    period_window_size_days = period_window_size_hours / 24

    fig, ax = plt.subplots(figsize=(8, 8))
    win_start_dates_str = aggregated_df["start_date"].dt.strftime("%Y-%m-%d")
    win_end_dates_str = aggregated_df["end_date"].dt.strftime("%Y-%m-%d")
    plt.title(f"Insulin Prediction Modeling, Aggr Period={period_window_size_days} Days, {win_start_dates_str.values[0]} to {win_end_dates_str.values[-1]}, {len(aggregated_df)} data points")

    hue_col = "cgm_geo_mean"
    # hue_col = None

    vars_to_plot = ["total_carbs", "total_insulin", "cgm_geo_mean"]
    scatter_df = aggregated_df[vars_to_plot]

    sns.scatterplot(data=scatter_df, x="total_carbs", y="total_insulin", hue=hue_col, ax=ax)
    plt.figure()

    # y_max_insulin = aggregated_df["total_insulin"].max() * 1.1
    y_max_insulin = 40
    ax.set_ylim(0, y_max_insulin)

    # y_max_carbs = aggregated_df["total_carbs"].max() * 1.1
    y_max_carbs = 250
    ax.set_xlim(0, y_max_carbs)

    if settings:

        cir_estimate_slope, isf_estimate_slope, basal_insulin_estimate, lm_carb_to_insulin, r2_fit, K = settings
        basal_glucose_lr = -basal_insulin_estimate / lr_model.coef_[0]

        x1, y1 = basal_glucose_lr, 0
        x2, y2 = aggregated_df["total_carbs"].max(), lr_model.predict([[aggregated_df["total_carbs"].max()]])
        ax.plot([x1, x2], [y1, y2[0]], label="Insulin Prediction LR Model (Weights: {})".format(weight_scheme))

        ax.set_xlabel("Total Exogenous Glucose in Period T")
        ax.set_ylabel("Total Insulin in Period T")

        # Equations and Settings
        tot_U_100g_day = 100/cir_estimate_slope + basal_insulin_estimate
        ax.text(0.02, 0.95, "y={:.4f}*x + {:.2f}, (R^2={:.2f})".format(lr_model.coef_[0], lr_model.intercept_, r2_fit), ha="left", va="top", transform=ax.transAxes)
        ax.text(0.02, 0.9, "CIR={:.2f} g/U \nBasal Rate={:.2f}U/hr \nISF={:.2f} mg/dL/U (K={:.2f})\nU for 100g Day: {:.2f}".format(cir_estimate_slope,
                                                                                                   basal_insulin_estimate / period_window_size_hours,
                                                                                                   isf_estimate_slope,
                                                                                                   K,
                                                                                                    tot_U_100g_day),
                                                                                                   ha="left", va="top", transform=ax.transAxes)

        # Stars
        ax.plot(0, basal_insulin_estimate, marker="*", markersize=12, color="green", label="Basal Insulin LR Estimate")
        ax.plot(basal_glucose_lr, 0, marker="*", markersize=12, color="orange", label="Endogenous Glucose LR Estimate")

        mean_insulin = aggregated_df["total_insulin"].mean()
        mean_carbs = aggregated_df["total_carbs"].mean()
        ax.plot(mean_carbs, mean_insulin, marker="*", markersize=12, color="red", label="Mean Insulin/Carbs")

        # Shades
        ax.fill_between([0, x2], [basal_insulin_estimate, basal_insulin_estimate], color="blue", alpha=0.2, label="Endogenous")
        ax.fill_between([0, x2], [basal_insulin_estimate, basal_insulin_estimate], [basal_insulin_estimate, y2[0]],
                        color="orange", alpha=0.2, label="Exogenous")

        # AACE line if using 1-day windows
        if plot_aace and period_window_size_hours == 1:
            tdd_mean = aggregated_df["total_insulin"].mean()
            aace_basal_insulin_estimate = tdd_mean / 2
            cir_aace = 450 / tdd_mean
            aace_basal_glucose_estimate = -aace_basal_insulin_estimate / (1/cir_aace)
            x2 = aggregated_df["total_carbs"].max()

            x1, y1 = (aace_basal_glucose_estimate, 0)
            y2 = 1.0 / (cir_aace) * x2 + aace_basal_insulin_estimate
            star_description = "AACE Basal Estimate (mean(TDD) / 2)"
            line_description = "AACE (*mean(TDD) only)"

            ax.plot([x1, x2], [y1, y2], label=line_description, color="gray", linestyle="--")
            ax.plot(0, aace_basal_insulin_estimate, marker="*", markersize=12, color="gray", label=star_description)

        if 1:
            current_basal_total = (0.6 * 24)
            current_carb_ratio = 9.9
            current_basal_glucose_estimate = -current_basal_total / (1/current_carb_ratio)
            x1_current, y1_current = (current_basal_glucose_estimate, 0)
            y2 = 1.0 / current_carb_ratio * x2 + current_basal_total
            ax.plot([x1_current, x2], [y1_current, y2], label="Current Settings", color="yellow", linestyle="-.")

        ax.legend(loc="lower right")

    plt.show()


def analyze_settings_lr(user, data_start_date, data_end_date,
                        K=12.5,
                        do_plots=True,
                        use_circadian_hour_estimate=True,
                        agg_period_window_size_hours=24,
                        agg_period_hop_size_hours=24,
                        weight_scheme=None):
    """
    Aggregate the data and fit a linear model
    """
    window_stats = user.compute_window_stats(data_start_date, data_end_date,
                                             use_circadian=use_circadian_hour_estimate,
                                             window_size_hours=agg_period_window_size_hours,
                                             hop_size_hours=agg_period_hop_size_hours,
                                             plot_hop_raw=False)
    window_df = pd.DataFrame(window_stats)

    # drop_indices = np.random.choice(window_df.index, 1, replace=False)
    # window_df = window_df.drop(drop_indices)

    # TMP
    # q = window_df["cgm_mean"].quantile(0.50)
    # window_df = window_df[window_df["cgm_mean"] < q]

    # window_df["total_carbs"] += 15
    # /TMP

    logger.debug(f'Mean of CGM Mean, {np.round(window_df["cgm_mean"].mean(), 2)}')
    logger.debug(f'Mean of CGM Geo Mean, {np.round(window_df["cgm_geo_mean"].mean(), 2)}')
    logger.debug(f'Total Period Insulin Mean: {np.round(window_df["total_insulin"].mean(), 2)}')
    logger.debug(f"{len(window_df)} Data Rows")

    settings = estimate_therapy_settings_from_window_stats_lr(window_df,
                                                              K,
                                                              period_window_size_hours=agg_period_window_size_hours,
                                                              target_bg=110,
                                                              x="total_carbs",
                                                              y="total_insulin",
                                                              do_plots=do_plots,
                                                              weight_scheme=weight_scheme)

    cir_estimate, isf_estimate, basal_insulin_estimate, lr_model, lr_score, K = settings

    return cir_estimate, basal_insulin_estimate, isf_estimate, lr_score


def load_user_data(username, password, data_start_date, data_end_date, estimation_window_size_days):
    """
    Load the Tidepool user_obj data for the given user_obj and parameters.

    Args:
        username: tidepool username
        password: tidepool password
        data_start_date: start date of analysis
        data_end_date: end date of analysis
        estimation_window_size_days: size of estimation window in days

    Returns:
        TidepoolUser object
    """

    try:
        tp_api_obj = TidepoolAPI(username, password)
        user = TidepoolUser()
        user.load_from_api(tp_api_obj, data_start_date, data_end_date,
                           user_id=None,  # For a user_obj that is sharing their Tidepool account with this one
                           save_data=False)
    except Exception as e:
        raise RuntimeError(
            "Failed to load Tidepool data. Check username/password, network connectivity, and date range."
        ) from e

    total_basal_days = user.get_num_days_span(data_type="basal")
    total_bolus_days = user.get_num_days_span(data_type="bolus")
    total_cgm_days = user.get_num_days_span(data_type="cgm")
    total_food_days = user.get_num_days_span(data_type="food")

    if np.sum([total_basal_days, total_bolus_days, total_cgm_days, total_food_days]) != (estimation_window_size_days*4):
        logger.warning(f"*** Warning *** : Num data days span not the size of estimation window size of {estimation_window_size_days} days")
        logger.warning(f"Basal Days Span {total_basal_days}. Bolus Days Span {total_bolus_days}. CGM Days Span {total_cgm_days}. Food Days Span {total_food_days}")

    user.analyze_duplicates(time_diff_thresh_sec=60 * 60)

    return user


def load_nightscout_user_data(base_url, token, data_start_date, data_end_date, estimation_window_size_days, api_secret=None):
    """
    Load Nightscout user data for the given date range.

    Returns:
        NightscoutUser object
    """
    try:
        ns_api_obj = NightscoutAPI(base_url=base_url, token=token, api_secret=api_secret)
        user = NightscoutUser()
        user.load_from_api(ns_api_obj, data_start_date, data_end_date, save_data=False)
    except Exception as e:
        raise RuntimeError(
            "Failed to load Nightscout data. Check URL/token/api-secret, network connectivity, and date range."
        ) from e

    total_basal_days = user.get_num_days_span(data_type="basal")
    total_bolus_days = user.get_num_days_span(data_type="bolus")
    total_cgm_days = user.get_num_days_span(data_type="cgm")
    total_food_days = user.get_num_days_span(data_type="food")

    if np.sum([total_basal_days, total_bolus_days, total_cgm_days, total_food_days]) != (estimation_window_size_days * 4):
        logger.warning(f"*** Warning *** : Num data days span not the size of estimation window size of {estimation_window_size_days} days")
        logger.warning(f"Basal Days Span {total_basal_days}. Bolus Days Span {total_bolus_days}. CGM Days Span {total_cgm_days}. Food Days Span {total_food_days}")

    user.analyze_duplicates(time_diff_thresh_sec=60 * 60)

    return user


def main():

    parser = argparse.ArgumentParser("InsuLearner: Estimate Insulin Pump Settings with Linear Regression")

    parser.add_argument("--source", choices=["tidepool", "nightscout"], default="tidepool",
                        help="Data source to use. Defaults to tidepool.")
    parser.add_argument("tp_username", nargs="?", default=None, help="Email username for Tidepool Account")
    parser.add_argument("tp_password", nargs="?", default=None, help="Password for Tidepool Account")
    parser.add_argument("--nightscout_url", type=str, default=None, help="Nightscout base URL (for --source nightscout)")
    parser.add_argument("--nightscout_token", type=str, default=None, help="Nightscout token (optional)")
    parser.add_argument("--nightscout_api_secret", type=str, default=None, help="Nightscout api-secret (optional)")
    parser.add_argument("-ht", "--height_inches", type=float, help="Your height in inches")
    parser.add_argument("-wt", "--weight_lbs", type=float, help="Your weight in pounds")
    parser.add_argument("-g", "--gender", choices=["male", "female"])
    parser.add_argument("--num_days", type=int, help="Number of days in the past to analyze data", default=60)
    parser.add_argument("--CSF", type=float, help="If entered, will use this CSF instead of estimating it from height and weight.")
    parser.add_argument("-eb", "--estimate_agg_boundaries", "-eb", action="store_true", default=True,
                        help="Use an autocorrelation-like algorithm to estimate aggregation boundaries to denoise the fit.")

    parser.add_argument("-aw", "--agg_period_window_size_hours", type=int,
                        help="The size in hours of each period to aggregate for fitting the model.", default=24)
    parser.add_argument("-ah", "--agg_period_hop_size_hours", "-ah", type=int,
                        help="The size in hours to hop each period for aggregation.", default=24)

    args = parser.parse_args()

    source = args.source
    tp_username = args.tp_username
    tp_password = args.tp_password
    nightscout_url = args.nightscout_url
    nightscout_token = args.nightscout_token
    nightscout_api_secret = args.nightscout_api_secret
    estimation_window_size_days = args.num_days
    height_inches = args.height_inches
    weight_lbs = args.weight_lbs
    gender = args.gender
    CSF = args.CSF
    estimate_agg_boundaries = args.estimate_agg_boundaries
    agg_period_window_size_hours = args.agg_period_window_size_hours
    agg_period_hop_size_hours = args.agg_period_hop_size_hours

    logger.debug("Args:")
    logger.debug(f"source: {source}")
    logger.debug(f"estimation_window_size_days: {estimation_window_size_days}")
    logger.debug(f"height_inches: {height_inches}")
    logger.debug(f"weight_lbs: {weight_lbs}")
    logger.debug(f"gender: {gender}")
    logger.debug(f"CSF: {CSF}")
    logger.debug(f"estimate_agg_boundaries: {estimate_agg_boundaries}")
    logger.debug(f"agg_period_window_size_hours: {agg_period_window_size_hours}")
    logger.debug(f"agg_period_hop_size_hours: {agg_period_hop_size_hours}")

    if source == "tidepool":
        if not tp_username or not tp_password:
            raise ValueError("For --source tidepool, provide tp_username and tp_password positional arguments.")
    elif source == "nightscout":
        if not nightscout_url:
            raise ValueError("For --source nightscout, provide --nightscout_url.")

    K = CSF
    if CSF is None:
        # Estimate CSF from blood volume based on height and weight
        K = estimate_csf(height_inches, weight_lbs, gender, metabolism_efficiency_percentage=0.23)
        logger.info(f"CSF estimated to be {np.round(K, 2)} for height_inches {height_inches}, weight_lbs {weight_lbs}, and gender {gender}")
    else:
        logger.info(f"Provided CSF={K}")

    # Get date info
    # Run settings analysis
    cirs = []
    basals = []
    isfs = []
    lr_scores = []
    tot_ins_100g = []

    window_sizes = [estimation_window_size_days]

    # window_start =
    for estimation_window_size_days in window_sizes:
        # todays = [dt.datetime.now() - dt.timedelta(days=i) for i in range(450, -15, -15)]
        todays = [dt.datetime.now()]
        for today in todays:
            # today = dt.datetime.now()
            data_start_date = today - dt.timedelta(days=estimation_window_size_days + 1)
            data_end_date = today - dt.timedelta(days=1)

            # Uncomment if specific dates desired
            # data_start_date = dt.datetime(year=2022, month=12, day=1)
            # data_end_date = dt.datetime(year=2023, month=1, day=26)

            logger.info(f"Running for dates {data_start_date} to {data_end_date}")

            # Load user_obj data into an object
            if source == "tidepool":
                user_obj = load_user_data(tp_username, tp_password, data_start_date, data_end_date, estimation_window_size_days)
            else:
                user_obj = load_nightscout_user_data(
                    base_url=nightscout_url,
                    token=nightscout_token,
                    api_secret=nightscout_api_secret,
                    data_start_date=data_start_date,
                    data_end_date=data_end_date,
                    estimation_window_size_days=estimation_window_size_days,
                )

            for i in range(1):

                cir_estimate, basal_insulin_estimate, isf_estimate, lr_score = analyze_settings_lr(user_obj,
                                    data_start_date=data_start_date,
                                    data_end_date=data_end_date,
                                    K=K,
                                    do_plots=True,
                                    use_circadian_hour_estimate=estimate_agg_boundaries,
                                    agg_period_window_size_hours=agg_period_window_size_hours,
                                    agg_period_hop_size_hours=agg_period_hop_size_hours,
                                    # weight_scheme="Recency"
                                    weight_scheme = None
                )

                cirs.append(cir_estimate)
                basals.append(basal_insulin_estimate/agg_period_window_size_hours)
                isfs.append(isf_estimate)
                lr_scores.append(lr_score)
                tot_ins_100g.append(100/cir_estimate + basal_insulin_estimate)

    # plt.plot(cirs)
    # plt.title("CIR")
    # plt.figure()
    # plt.plot(basals)
    # plt.title("Basal U/hr")
    # plt.figure()
    # plt.plot(tot_ins_100g)
    # plt.title("Total U for 100g in a Day")
    # plt.show()


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error(str(e))
        raise SystemExit(1)
