from ggplot import *
import matplotlib.pyplot as plt
import seaborn as sns
import statsmodels as stamo
import statsmodels.formula.api as sm
import patsy
from statsmodels.regression.linear_model import OLS

regressions_plot = {
    1: {
        'x': 'total_revenue',
        'y': 'spotify_playlist_streams'
    },
    2: {
        'x': 'total_revenue',
        'y': 'amazon_total_streams'
    },
    3: {
        'x': 'total_revenue',
        'y': 'pandora_total_streams'
    },
    4: {
        'x': 'total_revenue',
        'y': 'youtube_audio_streams'
    },
    5: {
        'x': 'total_revenue',
        'y': 'shazam_total_streams'
    },
    6: {
        'x': 'total_revenue',
        'y': 'alexa_total_streams'
    },
    7: {
        'x': 'total_revenue',
        'y': 'apple_playlist_streams'
    },
    8: {
        'x': 'total_revenue',
        'y': 'apple_active_streams'
    },
    9: {
        'x': 'total_revenue',
        'y': 'apple_passive_streams'
    },
    10: {
        'x': 'total_revenue',
        'y': 'spotify_passive_streams'
    },
    11: {
        'x': 'total_revenue',
        'y': 'spotify_active_streams'
    },
    11: {
        'x': 'total_revenue',
        'y': 'spotify_collection_streams'
    },
    12: {
        'x': 'total_revenue',
        'y': 'apple_passive_streams'
    },
    13: {
        'x': 'total_revenue',
        'y': 'apple_collection_streams'
    },
    14: {
        'x': 'total_revenue',
        'y': 'itunes_total_downloads'
    },
    15: {
        'x': 'total_revenue',
        'y': 'itunes_total_downloads'
    },
    16: {
        'x': 'total_revenue',
        'y': 'amazon_total_downloads'
    },
    17: {
        'x': 'total_revenue',
        'y': 'google_total_downloads'
    },
    18: {
        'x': 'total_revenue',
        'y': 'total_gross_revenue_last_year'
    },
    19: {
        'x': 'total_revenue',
        'y': 'apple_playlists_streams_last_year'
    },
    20: {
        'x': 'total_revenue',
        'y': 'spotify_playlists_streams_last_year'
    },
    21: {
        'x': 'total_revenue',
        'y': 'alexa_overall_streams_last_year'
    },
}


def generate_graph_with_categorical_var(
        x, y, data, categorical_variable, x_title, y_title, g_title):
    """ Generates a graph between two variables and a categorical variable.
    Args:
        x (str): The variable name for the x axis
        y (str): The variable name for the y axis
        data (DataFrame): A pandas dataframe
        categorical_variable (str): The categorical variable
        x_title (str): Title for x axis
        y_title (str): Title for y axis
        g_title (str): The general title of the graph
    """
    g = (
        ggplot(data, aes(
            x=x, y=y, color=categorical_variable)) + geom_point() + xlab(
            x_title) + ylab(y_title) + ggtitle(g_title))

    g.show()


def generate_graph_with_regression_line(x, y, data, intercept, slope):
    """ Generates a graph between two variables and a categorical variable.
    Args:
        x (str): The variable name for the x axis
        y (str): The variable name for the y axis
        data (DataFrame): A pandas dataframe
        intercept (float): The intercept of the line
        slope (float): The slope of the line
    """
    g = (ggplot(data, aes(x=x, y=y)) + geom_point(
        alpha=0.4) + geom_abline(intercept=intercept, slope=slope, size=3))

    g.show()


def generate_bar_plot(x, y, data):
    """ Generate a bar plot
    Args:
        x (str): The variable name for the x axis
        y (str): The variable name for the y axis
        data (DataFrame): A pandas dataframe
    """
    sns.barplot(x=x, y=y, data=data)
    plt.show()


def generate_pair_plot(data, variables_to_plot):
    """ Generate Pair Plot
    Args:
        data (DataFrame): A pandas dataframe
        variables_to_plot (list): List of strings with the variables that
        would be plotted
    """
    sns.pairplot(data_frame, vars=variables_to_plot)
    plt.show()