# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_plotting.ipynb.

# %% auto 0
__all__ = ['plot_distribution', 'plot_power', 'plot_betas', 'plot_binary_power', 'plot_continuous_power']

# %% ../nbs/03_plotting.ipynb 5
import numpy as np
from scipy.stats import binom
import pandas as pd

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.figure_factory as ff

from scipy.stats import beta

from .power import sample_size_binary, sample_size_continuous

# %% ../nbs/03_plotting.ipynb 7
def plot_distribution(df, show_rug=True):
    """
    Plots the distribution for both groups of generate_continuous_data
    """
    fig = ff.create_distplot(
        [
            df[df["group"] == 0]["target"].values,
            df[df["group"] == 1]["target"].values,
        ],
        ["Control", "Variant"],
        show_hist=False,
        show_rug=show_rug,
    )

    fig.update_layout(
        template="simple_white",
        xaxis_title="Target",
        yaxis_title="PDF",
        legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.95),
    )
    return fig

# %% ../nbs/03_plotting.ipynb 8
def plot_power(simulation, added_lines=[], is_effect=True):
    """
    Takes simulation dict and plots the power over sample sizes
    Added lines: plot horizontal lines for additional sample size estimations, takes a list of dicts with sample_size=sample_size,label=label
    is_effect: True: there is an effect in reality, hence we are interested in power. False: we are interested in false positives.
    """
    approaches = simulation["approaches"]
    if is_effect == True:
        ytitle = "Power"
        yrange = [0, 1]
    elif is_effect == False:
        ytitle = "False positives"
        yrange = [0, 0.1]
    else:
        raise Exception(
            "Non boolean for is effect encountered in power plot function."
        )

    fig = go.Figure()
    for approach in approaches:
        fig.add_trace(
            go.Scatter(
                x=simulation["sizes"] / 2.0,
                y=simulation[f"power_{approach}"],
                mode="lines+markers",
                name=approach,
            )
        )

    for approach in approaches:
        fig.add_trace(
            go.Scatter(
                x=2 * [simulation[f"sample_size_{approach}"] / 2.0],
                y=[0, 1],
                mode="lines",
                line_dash="dash",
                name=f"Sample size ({approach})",
            )
        )
    for added_line in added_lines:
        fig.add_trace(
            go.Scatter(
                x=2 * [added_line["sample_size"]],
                y=[0, 1],
                mode="lines",
                line_dash="dash",
                name=added_line["label"],
            )
        )
    if is_effect == False:
        fig.add_trace(
            go.Scatter(
                x=[0, simulation["sizes"].max() / 2.0],
                y=[0.05, 0.05],
                mode="lines",
                line_dash="dash",
                line_color="gray",
                line_width=1,
            )
        )

    fig.update_layout(
        template="simple_white",
        xaxis_title="Sample size per variation",
        yaxis_title=ytitle,
        yaxis_range=yrange,
        legend=dict(yanchor="top", y=1, xanchor="left", x=0.0),
    )
    fig.update_xaxes(showspikes=True, spikemode="across", spikethickness=1)
    fig.update_yaxes(showspikes=True, spikemode="across", spikethickness=1)
    fig.update_layout(spikedistance=1000, hoverdistance=100)

    # fig.update_layout(showlegend=True)
    return fig

# %% ../nbs/03_plotting.ipynb 10
def plot_betas(df_contingency, xmin=0.0, xmax=0.2, names=["A", "B"]):
    """
    Plots two beta distributions, one for control and for variant.
    Takes as input a contigency table as dataframe
    """
    betas = dict()
    for group in [0, 1]:
        betas[group] = beta(
            df_contingency.loc[group].converted + 1,
            df_contingency.loc[group].not_converted + 1,
        )

    x = np.linspace(xmin, xmax, 500)

    fig = make_subplots()

    for group in [0, 1]:
        fig.add_trace(
            go.Scatter(
                x=x, y=betas[group].pdf(x), mode="lines", name=names[group]
            )
        )

    for group in [0, 1]:
        fig.add_vline(
            x=df_contingency.loc[group].cvr, line_width=2, line_dash="dash"
        )

    fig.update_layout(
        title="",
        template="simple_white",
        xaxis_title="CVR",
        yaxis_title="PDF",
        legend=dict(yanchor="top", y=1, xanchor="left", x=0.0),
    )

    return fig

# %% ../nbs/03_plotting.ipynb 11
def plot_binary_power(
    cr0=0.01,
    cr1=0.012,
    alpha=0.05,
    one_sided=True,
    vline_power=0.8,
    powers=np.arange(0.1, 0.91, 0.01),
):
    """
    Generate a chart that shows the
    """
    sizes = []
    powers = powers[powers > alpha * 1.1]
    for power in powers:
        size = sample_size_binary(
            cr0=cr0, cr1=cr1, alpha=alpha, power=power, one_sided=one_sided
        )
        sizes.append(size)
        fig = make_subplots()

    fig.add_trace(
        go.Scatter(
            x=sizes,
            y=powers,
            mode="lines",
        )
    )

    fig.add_vline(
        x=sample_size_binary(
            cr0=cr0,
            cr1=cr1,
            alpha=alpha,
            power=vline_power,
            one_sided=one_sided,
        ),
        line_width=2,
        line_dash="dash",
        line_color="gray",
    )

    fig.update_layout(
        title="Statistical power vs sample size (binary)",
        template="simple_white",
        xaxis_title="Sample size per variant",
        yaxis_title="Power",
        legend=dict(yanchor="top", y=1, xanchor="left", x=0.0),
    )
    fig.update_xaxes(showspikes=True, spikemode="across", spikethickness=1)
    fig.update_yaxes(showspikes=True, spikemode="across", spikethickness=1)
    fig.update_layout(spikedistance=1000, hoverdistance=100)
    return fig

# %% ../nbs/03_plotting.ipynb 13
def plot_continuous_power(
    mu1=5.0, mu2=5.05, sigma=1, alpha=0.05, one_sided=True, vline_power=0.8
):
    powers = np.arange(0.1, 0.91, 0.025)
    sizes = []
    for power in powers:
        size = sample_size_continuous(
            mu1=mu1,
            mu2=mu2,
            sigma=sigma,
            alpha=alpha,
            power=power,
            one_sided=one_sided,
        )
        sizes.append(size)
        fig = make_subplots()

    fig.add_trace(
        go.Scatter(
            x=sizes,
            y=powers,
            mode="lines",
        )
    )

    fig.add_vline(
        x=sample_size_continuous(
            mu1=mu1,
            mu2=mu2,
            sigma=sigma,
            alpha=alpha,
            power=vline_power,
            one_sided=one_sided,
        ),
        line_width=2,
        line_dash="dash",
        line_color="gray",
    )

    fig.update_layout(
        title="Statistical power vs sample size (continuous)",
        template="simple_white",
        xaxis_title="Sample size per variant",
        yaxis_title="Power",
        hovermode="y unified",
        legend=dict(yanchor="top", y=1, xanchor="left", x=0.0),
    )
    fig.update_xaxes(showspikes=True, spikemode="across", spikethickness=1)
    fig.update_yaxes(showspikes=True, spikemode="across", spikethickness=1)
    fig.update_layout(spikedistance=1000, hoverdistance=100)
    return fig
