import pandas as pd
import numpy as np
import statsmodels.api as sm
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error


from scipy import stats

import plotly.graph_objs as go
import cufflinks

cufflinks.go_offline()


def make_hist(df, x, category=None):
    """
    Make an interactive histogram, optionally segmented by `category`

    :param df: dataframe of data
    :param x: string of column to use for plotting
    :param category: string representing column to segment by

    :return figure: a plotly histogram to show with iplot or plot
    """
    if category is not None:
        data = []
        for name, group in df.groupby(category):
            data.append(go.Histogram(dict(x=group[x], name=name)))
    else:
        data = [go.Histogram(dict(x=df[x]))]

    layout = go.Layout(
        yaxis=dict(title="Count"),
        xaxis=dict(title=x.replace("_", " ").title()),
        title=f"{x.replace('_', ' ').title()} Distribution by {category.replace('_', ' ').title()}"
        if category
        else f"{x.replace('_', ' ').title()} Distribution",
    )

    figure = go.Figure(data=data, layout=layout)
    return figure


def make_cum_plot(df, y, category=None, ranges=False):
    """
    Make an interactive cumulative plot, optionally segmented by `category`

    :param df: dataframe of data, must have a `published_date` column
    :param y: string of column to use for plotting or list of two strings for double y axis
    :param category: string representing column to segment by
    :param ranges: boolean for whether to add range slider and range selector

    :return figure: a plotly plot to show with iplot or plot
    """
    if category is not None:
        data = []
        for i, (name, group) in enumerate(df.groupby(category)):
            group.sort_values("published_date", inplace=True)
            data.append(
                go.Scatter(
                    x=group["published_date"],
                    y=group[y].cumsum(),
                    mode="lines+markers",
                    text=group["title"],
                    name=name,
                    marker=dict(size=10, opacity=0.8, symbol=i + 2),
                )
            )
    else:
        df.sort_values("published_date", inplace=True)
        if len(y) == 2:
            data = [
                go.Scatter(
                    x=df["published_date"],
                    y=df[y[0]].cumsum(),
                    name=y[0].title(),
                    mode="lines+markers",
                    text=df["title"],
                    marker=dict(
                        size=10,
                        color="blue",
                        opacity=0.6,
                        line=dict(color="black"),
                    ),
                ),
                go.Scatter(
                    x=df["published_date"],
                    y=df[y[1]].cumsum(),
                    yaxis="y2",
                    name=y[1].title(),
                    mode="lines+markers",
                    text=df["title"],
                    marker=dict(
                        size=10,
                        color="red",
                        opacity=0.6,
                        line=dict(color="black"),
                    ),
                ),
            ]
        else:
            data = [
                go.Scatter(
                    x=df["published_date"],
                    y=df[y].cumsum(),
                    mode="lines+markers",
                    text=df["title"],
                    marker=dict(
                        size=12,
                        color="blue",
                        opacity=0.6,
                        line=dict(color="black"),
                    ),
                )
            ]
    if len(y) == 2:
        layout = go.Layout(
            xaxis=dict(title="Published Date", type="date"),
            yaxis=dict(title=y[0].replace("_", " ").title(), color="blue"),
            yaxis2=dict(
                title=y[1].replace("_", " ").title(),
                color="red",
                overlaying="y",
                side="right",
            ),
            font=dict(size=14),
            title=f"Cumulative {y[0].title()} and {y[1].title()}",
        )
    else:
        layout = go.Layout(
            xaxis=dict(title="Published Date", type="date"),
            yaxis=dict(title=y.replace("_", " ").title()),
            font=dict(size=14),
            title=f"Cumulative {y.replace('_', ' ').title()} by {category.replace('_', ' ').title()}"
            if category is not None
            else f"Cumulative {y.replace('_', ' ').title()}",
        )

    # Add a rangeselector and rangeslider for a data xaxis
    if ranges:
        rangeselector = dict(
            buttons=list(
                [
                    dict(count=1, label="1m", step="month", stepmode="backward"),
                    dict(count=6, label="6m", step="month", stepmode="backward"),
                    dict(count=1, label="1y", step="year", stepmode="backward"),
                    dict(step="all"),
                ]
            )
        )
        rangeslider = dict(visible=True)
        layout["xaxis"]["rangeselector"] = rangeselector
        layout["xaxis"]["rangeslider"] = rangeslider
        layout["width"] = 1000
        layout["height"] = 600

    figure = go.Figure(data=data, layout=layout)
    return figure


def make_scatter_plot(
    df,
    x,
    y,
    fits=None,
    xlog=False,
    ylog=False,
    category=None,
    scale=None,
    sizeref=2,
    annotations=None,
    ranges=False,
    title_override=None,
):
    """
    Make an interactive scatterplot, optionally segmented by `category`

    :param df: dataframe of data
    :param x: string of column to use for xaxis
    :param y: string of column to use for yaxis
    :param fits: list of strings of fits
    :param xlog: boolean for making a log xaxis
    :param ylog boolean for making a log yaxis
    :param category: string representing categorical column to segment by, this must be a categorical
    :param scale: string representing numerical column to size and color markers by, this must be numerical data
    :param sizeref: float or integer for setting the size of markers according to the scale, only used if scale is set
    :param annotations: text to display on the plot (dictionary)
    :param ranges: boolean for whether to add a range slider and selector
    :param title_override: String to override the title

    :return figure: a plotly plot to show with iplot or plot
    """
    if category is not None:
        title = f"{y.replace('_', ' ').title()} vs {x.replace('_', ' ').title()} by {category.replace('_', ' ').title()}"
        data = []
        for i, (name, group) in enumerate(df.groupby(category)):
            data.append(
                go.Scatter(
                    x=group[x],
                    y=group[y],
                    mode="markers",
                    text=group["title"],
                    name=name,
                    marker=dict(size=8, symbol=i + 2),
                )
            )

    else:
        if scale is not None:
            title = f"{y.replace('_', ' ').title()} vs {x.replace('_', ' ').title()} Scaled by {scale.title()}"
            data = [
                go.Scatter(
                    x=df[x],
                    y=df[y],
                    mode="markers",
                    text=df["title"],
                    marker=dict(
                        size=df[scale],
                        line=dict(color="black", width=0.5),
                        sizemode="area",
                        sizeref=sizeref,
                        opacity=0.8,
                        colorscale="Viridis",
                        color=df[scale],
                        showscale=True,
                        sizemin=2,
                    ),
                )
            ]
        else:

            df.sort_values(x, inplace=True)
            title = f"{y.replace('_', ' ').title()} vs {x.replace('_', ' ').title()}"
            data = [
                go.Scatter(
                    x=df[x],
                    y=df[y],
                    mode="markers",
                    text=df["title"],
                    marker=dict(
                        size=12, color="blue", opacity=0.8, line=dict(color="black")
                    ),
                    name="observations",
                )
            ]
            if fits is not None:
                for fit in fits:
                    data.append(
                        go.Scatter(
                            x=df[x],
                            y=df[fit],
                            text=df["title"],
                            mode="lines+markers",
                            marker=dict(size=8, opacity=0.6),
                            line=dict(dash="dash"),
                            name=fit,
                        )
                    )

                title += " with Fit"
    layout = go.Layout(
        annotations=annotations,
        xaxis=dict(
            title=x.replace("_", " ").title() + (" (log scale)" if xlog else ""),
            type="log" if xlog else None,
        ),
        yaxis=dict(
            title=y.replace("_", " ").title() + (" (log scale)" if ylog else ""),
            type="log" if ylog else None,
        ),
        font=dict(size=14),
        title=title if title_override is None else title_override,
    )

    # Add a rangeselector and rangeslider for a data xaxis
    if ranges:
        rangeselector = dict(
            buttons=list(
                [
                    dict(count=1, label="1m", step="month", stepmode="backward"),
                    dict(count=6, label="6m", step="month", stepmode="backward"),
                    dict(count=1, label="1y", step="year", stepmode="backward"),
                    dict(step="all"),
                ]
            )
        )
        rangeslider = dict(visible=True)
        layout["xaxis"]["rangeselector"] = rangeselector
        layout["xaxis"]["rangeslider"] = rangeslider
        layout["width"] = 1000
        layout["height"] = 600

    figure = go.Figure(data=data, layout=layout)
    return figure


def make_linear_regression(df, x, y, intercept_0):
    """
    Create a linear regression, either with the intercept set to 0 or
    the intercept allowed to be fitted

    :param df: dataframe with data
    :param x: string or list of stringsfor the name of the column with x data
    :param y: string for the name of the column with y data
    :param intercept_0: boolean indicating whether to set the intercept to 0
    """
    if isinstance(x, list):
        lin_model = LinearRegression()
        lin_model.fit(df[x], df[y])

        slopes, intercept, = (
            lin_model.coef_,
            lin_model.intercept_,
        )
        df["predicted"] = lin_model.predict(df[x])
        r2 = lin_model.score(df[x], df[y])
        rmse = np.sqrt(mean_squared_error(y_true=df[y], y_pred=df["predicted"]))
        equation = f'{y.replace("_", " ")} ='

        names = ["r2", "rmse", "intercept"]
        values = [r2, rmse, intercept]
        for i, (p, s) in enumerate(zip(x, slopes)):
            if (i + 1) % 3 == 0:
                equation += f'<br>{s:.2f} * {p.replace("_", " ")} +'
            else:
                equation += f' {s:.2f} * {p.replace("_", " ")} +'
            names.append(p)
            values.append(s)

        equation += f" {intercept:.2f}"
        annotations = [
            dict(
                x=0.4 * df.index.max(),
                y=0.9 * df[y].max(),
                showarrow=False,
                text=equation,
                font=dict(size=10),
            )
        ]

        df["index"] = list(df.index)
        figure = make_scatter_plot(
            df, x="index", y=y, fits=["predicted"], annotations=annotations
        )
        summary = pd.DataFrame({"name": names, "value": values})
    else:
        if intercept_0:
            lin_reg = sm.OLS(df[y], df[x]).fit()
            df["fit_values"] = lin_reg.fittedvalues
            summary = lin_reg.summary()
            slope = float(lin_reg.params)
            equation = f"${y.replace('_', ' ')} = {slope:.2f} * {x.replace('_', ' ')}$"

        else:
            lin_reg = stats.linregress(df[x], df[y])
            intercept, slope = lin_reg.intercept, lin_reg.slope
            params = ["pvalue", "rvalue", "slope", "intercept"]
            values = []
            for p in params:
                values.append(getattr(lin_reg, p))
            summary = pd.DataFrame({"param": params, "value": values})
            df["fit_values"] = df[x] * slope + intercept
            equation = f"${y.replace('_', ' ')} = {slope:.2f} * {x.replace('_', ' ')} + {intercept:.2f}$"

        annotations = [
            dict(
                x=0.75 * df[x].max(),
                y=0.9 * df[y].max(),
                showarrow=False,
                text=equation,
                font=dict(size=32),
            )
        ]
        figure = make_scatter_plot(
            df, x=x, y=y, fits=["fit_values"], annotations=annotations
        )
    return figure, summary


def make_poly_fits(df, x, y, degree=6):
    """
    Generate fits and make interactive plot with fits

    :param df: dataframe with data
    :param x: string representing x data column
    :param y: string representing y data column
    :param degree: integer degree of fits to go up to

    :return fit_stats: dataframe with information about fits
    :return figure: interactive plotly figure that can be shown with iplot or plot
    """

    # Don't want to alter original data frame
    df = df.copy()
    fit_list = []
    rmse = []
    fit_params = []

    # Make each fit
    for i in range(1, degree + 1):
        fit_name = f"fit degree = {i}"
        fit_list.append(fit_name)
        z, res, *rest = np.polyfit(df[x], df[y], i, full=True)
        fit_params.append(z)
        df.loc[:, fit_name] = np.poly1d(z)(df[x])
        rmse.append(np.sqrt(res[0]))

    fit_stats = pd.DataFrame({"fit": fit_list, "rmse": rmse, "params": fit_params})
    figure = make_scatter_plot(df, x=x, y=y, fits=fit_list)
    return figure, fit_stats


def make_extrapolation(df, y, years, degree=4):
    """
    Extrapolate `y` into the future `years` with `degree`  polynomial fit

    :param df: dataframe of data
    :param y: string of column to extrapolate
    :param years: number of years to extrapolate into the future
    :param degree: integer degree of polynomial fit

    :return figure: plotly figure for display using iplot or plot
    :return future_df: extrapolated numbers into the future
    """

    df = df.copy()
    x = "days_since_start"
    df["days_since_start"] = (
        (df["published_date"] - df["published_date"].min()).dt.total_seconds()
        / (3600 * 24)
    ).astype(int)

    cumy = f"cum_{y}"
    df[cumy] = df.sort_values(x)[y].cumsum()

    figure, summary = make_poly_fits(df, x, cumy, degree=degree)

    min_date = df["published_date"].min()
    max_date = df["published_date"].max()

    date_range = pd.date_range(
        start=min_date, end=max_date + pd.Timedelta(days=int(years * 365))
    )

    future_df = pd.DataFrame({"date": date_range})

    future_df[x] = (
        (future_df["date"] - future_df["date"].min()).dt.total_seconds() / (3600 * 24)
    ).astype(int)

    newcumy = f"cumulative_{y}"

    future_df = future_df.merge(df[[x, cumy]], on=x, how="left").rename(
        columns={cumy: newcumy}
    )

    z = np.poly1d(summary.iloc[-1]["params"])
    pred_name = f"predicted_{y}"
    future_df[pred_name] = z(future_df[x])
    future_df["title"] = ""

    last_date = future_df.loc[future_df["date"].idxmax()]
    prediction_text = f"On {last_date['date'].date()} the {y} will be {float(last_date[pred_name]):,.0f}."
    annotations = [
        dict(
            x=future_df["date"].quantile(0.4),
            y=0.8 * future_df[pred_name].max(),
            text=prediction_text,
            showarrow=False,
            font=dict(size=16),
        )
    ]

    title_override = f'{y.replace("_", " ").title()} with Extrapolation {years} Years into the Future'

    figure = make_scatter_plot(
        future_df,
        "date",
        newcumy,
        fits=[pred_name],
        annotations=annotations,
        ranges=True,
        title_override=title_override,
    )
    return figure, future_df