# -*- coding: utf-8 -*-
"""
Created on Tue Sep 27 12:31:28 2022

@author: hgarcia
"""

from typing import Dict, Optional, Callable, Tuple, Union, List
from numpy import exp
import numpy
from numpy.core.fromnumeric import repeat, shape
import pandas
import plotly.express as px
import plotly.io as pio
import plotly.graph_objects as graph_objects


# Set the default theme
template =  graph_objects.layout.Template()
template.layout = graph_objects.Layout(
                                    title_x=0.5,
                                    # border width and size
                                    margin=dict(l=2, r=2, b=2, t=30),
                                    height=400,
                                    # Interaction
                                    hovermode="closest",
                                    # axes
                                    xaxis_showline=True,
                                    xaxis_linewidth=2,
                                    yaxis_showline=True,
                                    yaxis_linewidth=2,
                                    # Pick a slightly different P.O.V from default
                                    # this avoids the extremities of the y and x axes
                                    # being cropped off
                                    scene_camera=dict(eye=dict(x=1.5, y=1.5, z=0.1))
                                    )

template.data.scatter = [graph_objects.Scatter(marker=dict(opacity=0.8))]
template.data.scatter3d = [graph_objects.Scatter3d(marker=dict(opacity=0.8))]
template.data.surface = [graph_objects.Surface()]
template.data.histogram = [graph_objects.Histogram(marker=dict(line=dict(width=1)))]
template.data.box = [graph_objects.Box(boxpoints='outliers', notched=False)]


pio.templates["custom_template"] = template
pio.templates.default = "plotly_white+custom_template"

# Trendline colors
# Take note that the text for this course often refers to colours explicitly
# such as "looking at the red line". Changing the variable below may result 
# in this text being inconsistent
colours_trendline = px.colors.qualitative.Set1  

def _to_human_readable(text:str):
    '''
    Converts a label into a human readable form
    '''
    return text.replace("_", " ")


def _prepare_labels(df:pandas.DataFrame, labels:List[Optional[str]], replace_nones:bool=True):
    '''
    Ensures labels are human readable.
    Automatically picks data if labels not provided explicitly
    '''

    human_readable = {}

    if isinstance(replace_nones, bool):
        replace_nones = [replace_nones] * len(labels)

    for i in range(len(labels)):
        lab = labels[i]
        if replace_nones[i] and (lab is None):
            lab = df.columns[i]
            labels[i] = lab

        # make human-readable
        if lab is not None:
            human_readable[lab] = _to_human_readable(lab)

    return labels, human_readable


def box_and_whisker(df:pandas.DataFrame,
                label_x:Optional[str]=None,
                label_y:Optional[str]=None,
                label_x2:Optional[str]=None,
                title=None,
                show:bool=False):
    '''
    Creates a box and whisker plot and optionally shows it. Returns the figure for that plot.

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    df: The data
    label_x: What to group by. Defaults to None
    label_y: What to plot on the y axis. Defaults to count of df.columns[0]
    label_x2: If provided, splits boxplots into 2+ per x value, each with its own colour
    title: Plot title
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured

    '''

    # Automatically pick columns if not specified
    selected_columns, axis_labels = _prepare_labels(df, [label_x, label_y, label_x2], replace_nones=[False, True, False])

    fig = px.box(df,
                    x=selected_columns[0],
                    y=selected_columns[1],
                    color=label_x2,
                    labels=axis_labels,
                    title=title)

    # Show the plot, if requested
    if show:
        fig.show()

    # return the figure
    return fig


def histogram(df:pandas.DataFrame,
                label_x:Optional[str]=None,
                label_y:Optional[str]=None,
                label_colour:Optional[str]=None,
                nbins:Optional[int]=None,
                title=None,
                include_boxplot=False,
                histfunc:Optional[str]=None,
                show:bool=False):
    '''
    Creates a 2D histogram and optionally shows it. Returns the figure for that histogram.

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    df: The data
    label_x: What to bin by. Defaults to df.columns[0]
    label_y: If provided, the sum of these numbers becomes the y axis. Defaults to count of label_x
    label_colour: If provided, creates a stacked histogram, splitting each bar by this column
    title: Plot title
    nbins: the number of bins to show. None for automatic
    histfunc: How to calculate y. See plotly for options
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured

    '''

    # Automatically pick columns if not specified
    selected_columns, axis_labels = _prepare_labels(df, [label_x, label_y, label_colour], replace_nones=[True, False, False])


    fig = px.histogram(df,
                        x=selected_columns[0],
                        y=selected_columns[1],
                        nbins=nbins,
                        color=label_colour,
                        labels=axis_labels,
                        title=title,
                        marginal="box" if include_boxplot else None,
                        histfunc=histfunc
                        )

    # Set the boxplot notches to False by default to deal with plotting bug
    # But only call this line if the user wants to include a boxplot
    if include_boxplot:
        fig.data[1].notched = False

    # Show the plot, if requested
    if show:
        fig.show()

    # return the figure
    return fig


def multiple_histogram(df:pandas.DataFrame,
                label_x:str,
                label_group:str,
                label_y:Optional[str]=None,
                histfunc:str='count',
                nbins:Optional[int]=None,
                title=None,
                show:bool=False):
    '''
    Creates a 2D histogram and optionally shows it. Returns the figure for that histogram.

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    df: The data
    label_x: What to bin by. Defaults to df.columns[0]
    label_y: If provided, the sum of these numbers becomes the y axis. Defaults to count of label_x
    title: Plot title
    nbins: the number of bins to show. None for automatic
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured

    '''

    assert (histfunc != 'count') or (label_y == None), "Set histfunc to a value such as sum or avg if using label_y"

    # Automatically pick columns if not specified
    selected_columns, axis_labels = _prepare_labels(df,  [label_x, label_y, label_group], replace_nones=[True, False, False])

    fig = graph_objects.Figure(layout=dict(
                                    title=title,
                                    xaxis_title_text=axis_labels[label_x],
                                    yaxis_title_text=histfunc if label_y is None else (histfunc + " of " + axis_labels[label_y]))
                                )

    group_values = sorted(set(df[label_group]))

    for group_value in group_values:
        dat = df[df[label_group] == group_value]
        x = dat[selected_columns[0]]

        if label_y is None:
            y = None
        else:
            y = dat[selected_columns[1]]

        fig.add_trace(graph_objects.Histogram(
            x=x,
            y=y,
            histfunc=histfunc,
            name=group_value, # name used in legend and hover labels
            nbinsx=nbins))

    #Place legend title
    fig.update_layout(legend_title_text=label_group)

    # Show the plot, if requested
    if show:
        fig.show()

    # return the figure
    return fig


def line_2D(
                trendline:Union[Tuple[str,Callable],List[Tuple[str,Callable]], Dict[str,List[float]]],
                x_range:List[float]=[0,1],
                label_x:str='x',
                label_y:str='y',
                legend_title:str="Line",
                title=None,
                show:bool=False):
    '''
    Creates a 2D line plot *using functions* and optionally shows it. Returns the figure for that plot.
    If you simply want a line plot using data, call scatter_2D then write fig.update_traces(mode='lines')

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    trendline:  (name, function) tuples. The functions accept X (a numpy array) and return Y (an iterable). Alternatively a dict of pre-calculated values
    x_range:    Sets the x-axis range. If this has more than three values, it is interpeted as each x-value to be graphed
    label_x:    The title for the x-axis
    label_y:    The title for the y-axis
    legend_title: The title for the legend
    title:      The plot title. If None and a single function is provided, the title is automatically set. Use "" to avoid
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured 

    '''

    if isinstance(trendline, tuple):
        trendline = [trendline]

    x = numpy.array([])
    y = numpy.array([])

    if len(x_range) == 2:
        x_vals = numpy.linspace(x_range[0], x_range[1], num=200)
    else:
        # X-range is interpreted as x_vals
        x_vals = numpy.array(x_range)
        x_vals.sort()

        # Rewrite x_range to actually be an x-axis range
        x_range = [x_vals[0], x_vals[-1]]

    names = []

    if isinstance(trendline, dict):
        for cur in trendline.items():
            name = cur[0]
            x = numpy.concatenate([x, x_vals])
            names = names + ([name] * len(x_vals))
            y = numpy.concatenate([y, cur[1]])
    else:
        for cur in trendline:
            name = cur[0]
            x = numpy.concatenate([x, x_vals])
            names = names + ([name] * len(x_vals))
            y = numpy.concatenate([y, cur[1](x=x_vals)])
    
    data = dict()
    data[label_x] = x
    data[label_y] = y
    data[legend_title] = names

    df = pandas.DataFrame(data)

    # Pick a title if none provided and we only have one function
    if (title is None) and (len(trendline) == 1):
        title = trendline[0][0]

    # Create as a 2d scatter but with lines
    fig = scatter_2D(df, label_colour=legend_title, title=title, show=False, x_range=x_range)
    fig.update_traces(mode='lines')

    # Don't show a legend if we only have one function plotted
    if len(trendline) == 1:
        fig.update_layout(showlegend=False)

    if show:
        fig.show()

    return fig


def scatter_2D(df:pandas.DataFrame,
                label_x:Optional[str]=None,
                label_y:Optional[str]=None,
                label_colour:Optional[str]=None,
                label_size:Optional[str]=None,
                size_multiplier:float=1,
                title=None,
                show:bool=False,
                x_range:Optional[List[float]]=None,
                trendline:Union[Callable,List[Callable],None]=None):
    '''
    Creates a 2D scatter plot and optionally shows it. Returns the figure for that scatter.

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    df: The data
    label_x: The label to extract from df to plot on the x axis. Defaults to df.columns[0]
    label_y: The label to extract from df to plot on the y axis. Defaults to df.columns[1]
    label_colour: The label to extract from df to colour points by
    title: Plot title
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured 
    x_range:    Overrides the x-axis range
    trendline:  A function that accepts X (a numpy array) and returns Y (an iterable)

    '''

    # Automatically pick columns if not specified
    selected_columns, axis_labels = _prepare_labels(df, [label_x, label_y, label_colour], [True, True, False])


    # Create the figure and plot
    fig = px.scatter(df,
                x=selected_columns[0],
                y=selected_columns[1],
                color=selected_columns[2],
                labels=axis_labels,
                hover_data=[label_size],
                title=title
                )

    if label_size is None:
        # User a marker size inversely proportional to the number of points
        size = int((round(22.0 - 19/(1+exp(-(df.shape[0]/100-2)))) * size_multiplier))
    else:
        # Set the size based on a label
        size = df[label_size]*size_multiplier

    fig.update_traces(marker={'size': size})

    if x_range is not None:
        fig.update_xaxes(range=[x_range[0], x_range[1]])

    # Create trendlines
    if trendline is not None:
        if isinstance(trendline, Callable):
            trendline = [trendline]
        x_min = min(df[selected_columns[0]]) if x_range is None else x_range[0]
        x_max = max(df[selected_columns[0]]) if x_range is None else x_range[1]
        evaluate_for = numpy.linspace(x_min, x_max, num=200)
        shapes = []
        for t,colour in zip(trendline,colours_trendline):
            y_vals = t(evaluate_for)
            path = "M" + " L ".join([str(c[0]) + " " + str(c[1]) for c in zip(evaluate_for,y_vals)])
            shapes.append(dict(
                                type="path",
                                path=path,
                                line_color=colour,
                            )
                        )

        fig.update_layout(shapes=shapes)

    # Show the plot, if requested
    if show:
        fig.show()

    # return the figure
    return fig


def scatter_3D(df:pandas.DataFrame,
                label_x:Optional[str]=None,
                label_y:Optional[str]=None,
                label_z:Optional[str]=None,
                label_colour:Optional[str]=None,
                title=None,
                show:bool=False):
    '''
    Creates a 3D scatter plot and optionally shows it. Returns the figure for that scatter.

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    df: The data
    label_x: The label to extract from df to plot on the x axis. Defaults to df.columns[0]
    label_y: The label to extract from df to plot on the y axis. Defaults to df.columns[1]
    label_z: The label to extract from df to plot on the z axis. Defaults to df.columns[2]
    label_colour: The label to extract from df to colour points by. Defaults to label_x
    title: Plot title
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured
    '''

    # Automatically pick columns if not specified
    selected_columns, axis_labels = _prepare_labels(df, [label_x, label_y, label_z])

    if label_colour is None:
        # Colour by the Z dimension
        label_colour = selected_columns[2]
    else:
        axis_labels[label_colour] = _to_human_readable(label_colour)

    # Create the figure and plot
    fig = px.scatter_3d(df,
                x=selected_columns[0],
                y=selected_columns[1],
                z=selected_columns[2],
                color=label_colour,
                labels=axis_labels,
                title=title)


    # Show the plot, if requested
    if show:
        fig.show()

    # return the figure
    return fig


def surface(x_values,
            y_values,
            calc_z:Callable,
            title=None,
            axis_title_x:Optional[str]=None,
            axis_title_y:Optional[str]=None,
            axis_title_z:Optional[str]=None,
            show:bool=False):
    '''
    Creates a surface plot using a function. Returns the figure for that plot.

    Note that if calling this from jupyter notebooks and not capturing the output
    it will appear on screen as though `.show()` has been called

    x_value: A numpy array of x values
    y_value: A numpy array of y values
    calc_z: A function to calculate z, given an x and a y value
    title: Plot title
    axis_title_x: Title for the x axis
    axis_title_y: Title for the y axis
    axis_title_z: Title for the z axis
    show:   appears on screen. NB that this is not needed if this is called from a
            notebook and the output is not captured
    '''

    # Check arguments
    assert len(x_values.shape) == 1, "Provide x_values as 1D"
    assert len(y_values.shape) == 1, "Provide y_values as 1D"


    # Calculate z for a range of x and y inputs
    # Note that z seems to be expected to be indexed [y,x] not [x,y] though this appears to
    # be counter to the documentation. If z is indexed [x, y] the result is flipped.
    # Potentially there is a bug here somewhere causing this issue or in plotly itself
    z = numpy.zeros((y_values.shape[0], x_values.shape[0]))
    for i_x in range(x_values.shape[0]):
        for i_y in range(y_values.shape[0]):
            z[i_y, i_x] = calc_z(x_values[i_x], y_values[i_y])
            
    # Create a graph of cost
    fig = graph_objects.Figure(data=[graph_objects.Surface(x=x_values, y=y_values, z=z)])
    fig.update_layout(title=title,
                      scene_xaxis_title=axis_title_x,
                      scene_yaxis_title=axis_title_y,
                      scene_zaxis_title=axis_title_z)

    #Add z-axis as colourbar title
    fig.update_traces(colorbar_title_text= axis_title_z, selector=dict(type='surface'))

    # Show the plot, if requested
    if show:
        fig.show()

    # return the figure
    return fig


def model_to_surface_plot(model, plot_features:List[str], data:pandas.DataFrame):
    '''Plots two features of a model as a surface. Other values are set at their means
    
    model:          A model that accepts a dataframe for prediction
    plot_features:  Two features to plot
    data:           A dataframe the model was trained or tested on
    '''

    # Give status as this can take several seconds to run
    print("Creating plot...")

    
    other_features = [f for f in data.columns if f not in plot_features]

    means = numpy.average(data[other_features], axis=0)
    mins = numpy.min(data[plot_features], axis=0)
    maxes = numpy.max(data[plot_features], axis=0)

    df = pandas.DataFrame()

    for f,m in zip(other_features, means):
        df[f] = [m]

    def predict(x, y):
        '''
        Makes a prediction using the model
        '''
        df[plot_features[0]] = [x]
        df[plot_features[1]] = [y]

        return model.predict(df)

    # Create a 3d plot of predictions
    x_vals = numpy.array(numpy.linspace(mins[plot_features[0]], maxes[plot_features[0]],20))
    y_vals = numpy.array(numpy.linspace(mins[plot_features[1]], maxes[plot_features[1]],20))

    return surface(x_vals, 
                    y_vals, 
                    predict, 
                    title="Model Prediction", 
                    axis_title_x=plot_features[0], 
                    axis_title_y=plot_features[1], 
                    axis_title_z="Probability")


def save_plot_as_image(fig, file="./plot.jpg", width=None, height="400", scale=1, format="jpg"):
    """
    Convert a figure to a static image and write it to a file or writeable object
    If "width" not set, plotly will set the aspect ration based on "hight"

    Parameters  

        fig – Figure object or dict representing a figure
        file (str or writeable) – A string representing a local file path or a writeable object (e.g. an open file descriptor)
        format (str or None) – The desired image format:

                ’png’
                ’jpg’ or ‘jpeg’
                ’webp’
                ’svg’
                ’pdf’
                ’eps’ (Requires the poppler library to be installed and on the PATH)

        width (int or None) – The width of the exported image in layout pixels. 
        height (int or None) – The height of the exported image in layout pixels. 

        scale (int or float or None) – The scale factor to use when exporting the figure. 
        A scale factor larger than 1.0 will increase the image resolution with respect to the 
        figure’s layout pixel dimensions. Whereas as scale factor of less than 1.0 will decrease 
        the image resolution.
    """
    pio.write_image(fig, 
                    file=file, 
                    width=width, 
                    height=height, 
                    scale=scale,
                    format=format, 
                    engine="kaleido",
                    )
    


# We use plotly to create plots and charts
import plotly.figure_factory as ff


def confussion_matrix_pretty(cm, labels):
    # Create the list of unique labels in the test set, to use in our plot
    # I.e., ['animal', 'hiker', 'rock', 'tree']
    x = y = sorted(labels)
    # Plot the matrix above as a heatmap with annotations (values) in its cells
    fig = ff.create_annotated_heatmap(cm, x, y)
    
    # Set titles and ordering
    fig.update_layout(  title_text="<b>Confusion matrix</b>", 
                        yaxis = dict(categoryorder = "category descending"))
    
    fig.add_annotation(dict(font=dict(color="black",size=14),
                            x=0.5,
                            y=-0.15,
                            showarrow=False,
                            text="Predicted label",
                            xref="paper",
                            yref="paper"))
    
    fig.add_annotation(dict(font=dict(color="black",size=14),
                            x=-0.1,
                            y=0.5,
                            showarrow=False,
                            text="Actual label",
                            textangle=-90,
                            xref="paper",
                            yref="paper"))
    
    # We need margins so the titles fit
    fig.update_layout(margin=dict(t=80, r=20, l=100, b=50))
    fig['data'][0]['showscale'] = True
    fig.show()

# =============================================================================
# ROC FUNCTIONS 
# =============================================================================
def calculate_tpr_fpr(prediction, actual):
    '''
    Calculates true positive rate and false positive rate

    prediction: the labels predicted by the model
    actual:     the correct labels we hope the model predicts
    '''

    # To calculate the true positive rate and true negative rate we need to know
    # TP - how many true positives (where the model predicts hiker, and it is a hiker)
    # TN - how many true negatives (where the model predicts tree, and it is a tree)
    # FP - how many false positives (where the model predicts hiker, but it was a tree)
    # FN - how many false negatives (where the model predicts tree, but it was a hiker)

    # First, make a note of which predictions were 'true' and which were 'false'
    prediction_true = numpy.equal(prediction, 1)
    prediction_false= numpy.equal(prediction, 0)

    # Now, make a note of which correct results were 'true' and which were 'false'
    actual_true = numpy.equal(actual, 1)
    actual_false = numpy.equal(actual, 0)

    # Calculate TP, TN, FP, and FN
    # The combination of sum and '&' counts the overlap
    # For example, TP calculates how many 'true' predictions 
    # overlapped with 'true' labels (correct answers)
    TP = numpy.sum(prediction_true  & actual_true)
    TN = numpy.sum(prediction_false & actual_false)
    FP = numpy.sum(prediction_true  & actual_false)
    FN = numpy.sum(prediction_false & actual_true)

    # Calculate the true positive rate
    # This is the proportion of 'hiker' labels that are identified as hikers
    tpr = TP / (TP + FN)

    # Calculate the false positive rate 
    # This is the proportion of 'tree' labels that are identified as hikers
    fpr = FP / (FP + TN)

    # Return both rates
    return tpr, fpr


def assess_model(model_predict, test, feature_name, threshold):
    '''
    Calculates the true positive rate and false positive rate of the model
    at a particular decision threshold

    model_predict: the model's predict function
    test: the test dataset
    feature_name: the feature the model is expecting
    threshold: the decision threshold to use 
    '''

    # Make model predictions for every sample in the test set
    # What we get back is a probability that the sample is a hiker
    # For example, if we had two samples in the test set, we might
    # get 0.45 and 0.65, meaning the model says there is a 45% chance
    # the first sample is a hiker, and 65% chance the second is a 
    # hiker
    probability_of_hiker = model_predict(test[feature_name])
    
    # See which predictions at this threshold would say hiker
    predicted_is_hiker = probability_of_hiker > threshold

    # calculate the true and false positives rates using our
    # handy method
    return calculate_tpr_fpr(predicted_is_hiker, test.is_hiker)


def create_roc_curve(model_predict, test, feature="motion"):
    '''
    This function creates a ROC curve for a given model by testing it
    on the test set for a range of decision thresholds. An ROC curve has
    the True Positive rate on the x-axis and False Positive rate on the 
    y-axis

    model_predict: The model's predict function, returning probabilities
    test: the test set
    feature: The feature to provide the model's predict function

    Returns the plotly figure and a dataframe of the results
    '''

    # Calculate what the true positive and false positive rate would be if
    # we had used different thresholds. 

    #  Make a list of thresholds to try
    # NB We need some specific values for our exercises but this is not typical
    thresholds = numpy.sort(numpy.hstack([[-1E-6], [1.0001], [0.5], [0.3], numpy.linspace(0,1,100)]))

    false_positive_rates = []
    true_positive_rates = []

    # Loop through all thresholds
    for threshold in thresholds:
        # calculate the true and false positives rates using our
        # handy method
        tpr, fpr = assess_model(model_predict, test, feature, threshold)

        # save the results
        true_positive_rates.append(tpr)
        false_positive_rates.append(fpr)


    # Graph the result
    # You don't need to understand this code, but essentially we are plotting
    # TPR versus FPR as a line plot
    # -- Prepare a dataframe, required by our graphinx code
    df_for_graphinx = pandas.DataFrame(dict(threshold=thresholds, fpr=false_positive_rates, tpr=true_positive_rates))
    # -- Generate the plot
    fig = scatter_2D(df_for_graphinx, label_x="fpr", label_y="tpr", x_range=[-0.05,1.05])
    fig.update_traces(mode='lines') # Comment our this line if you would like to see points rather than lines
    fig.update_yaxes(range=[-0.05, 1.05])

    # Return the graph
    return fig, df_for_graphinx
