"""
Ashley Wu 12/4/25
Creates visualizations for frontend of app 
"""
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr
import math
import utils


def time_series(df, xcol, ycol):
    """
    Creates a time series graph 
    
    Args: 
        df: dataframe 
        xcol: column of x calues
        ycoL: column of y values
    Returns: 
        fig: time series graph
    """
    fig, ax = plt.subplots()

    if df[ycol].dtype.name == 'object':
        raise gr.Error("The Y value can not be a string or a category!", duration=10)
    
    
    x = df[xcol].to_numpy() 
    y = df[ycol].to_numpy()

    ax.plot(x, y, color='red', linewidth=2, marker='o')
    ax.set_xlabel(xcol)
    ax.set_ylabel(ycol)
    ax.set_title(f"Time Series of {xcol} and {ycol}")

    return fig

def bar_chart(df, xcol, ycol, aggmethod):
    """
    Creates bar chart of information 
    
    Args: 
        df: dataframe
        xcol: column of x values
        ycol: column of y values
        aggmethod: aggregation method 
    Return: 
        fig: bar chart 
    """

    df = utils.detect_categories(df)
    fig, ax = plt.subplots()

    if df[ycol].dtype.name == 'category' or df[ycol].dtype.name == 'object':
        raise gr.Error("The Y variable can not be a string or a category!", duration=10)
    if df[xcol].dtype.name != 'category':
        raise gr.Error("The X variable must be a category!")
    
    if aggmethod == 'sum':
        y = df.groupby(xcol)[ycol].sum()
    elif aggmethod == 'mean':
        y = df.groupby(xcol)[ycol].mean()
    elif aggmethod == 'count':
        y = df[xcol].value_counts().tolist()
    elif aggmethod == 'median':
        y = df.groupby(xcol)[ycol].median()
    else: 
        y = df[xcol].tolist() 

    x = df[xcol].unique().tolist()

    ax.bar(x, y)
    
    ax.set_xlabel('Category')
    ax.set_ylabel(ycol)
    ax.set_title(f'{aggmethod} by Category')
    ax.set_xticks(ticks=x)
    return fig

def histogram(df, col, bins):
    """
    Creates a histogram 
    
    Args: 
        df: dataframe
        col: data for histogram 
        bins: number of bins to put histogram data into
    Returns: 
        fig: histogram created 
    """
    fig, ax = plt.subplots()

    bins = math.floor(bins)

    if df[col].dtype.name == 'object':
        raise gr.Error("The Y value can not be a string or a category!", duration=10)
    
    ax.hist(df[col], bins)
    ax.set_xlabel(f'{col}')
    ax.set_title(f'Histogram of {col}')
    return fig

def scatter(df, xcol, ycol):
    """
    Creates a scatter plot
    
    Args: 
        df: dataframe 
        xcol: column of x values
        ycol: column of y values
    Returns: 
        fig: scatter plot 
    """
    fig, ax = plt.subplots()

    if df[ycol].dtype.name == 'object' or df[xcol].dtype.name == 'object':
        raise gr.Error("The X or Y values can not be a string or a category!", duration=10)
    
    x = df[xcol].to_numpy() 
    y = df[ycol].to_numpy()

    ax.plot(x, y, color='red', marker='o')
    ax.set_xlabel(xcol)
    ax.set_ylabel(ycol)
    ax.set_title(f"Scatter Plot of {xcol} and {ycol}")

    return fig