import gradio as gr
import pandas as pd
from plotly import graph_objects as go


def process_dataset():
    """
    Process the dataset and perform the following operations:
    1. Read the file_counts_and_sizes, repo_by_size_df, unique_files_df, and file_extensions data from parquet files.
    2. Convert the total size to petabytes and format it to two decimal places.
    3. Capitalize the 'type' column in the file_counts_and_sizes dataframe.
    4. Rename the columns in the file_counts_and_sizes dataframe.
    5. Sort the file_counts_and_sizes dataframe by total size in descending order.
    6. Drop rows with missing values in the 'extension' column of the file_extensions dataframe.
    7. Return the repo_by_size_df, unique_files_df, file_counts_and_sizes, and file_extensions dataframes.
    """

    file_counts_and_sizes = pd.read_parquet(
        "hf://datasets/jsulz/lfs-analysis-data/file_counts_and_sizes.parquet"
    )
    repo_by_size_df = pd.read_parquet(
        "hf://datasets/jsulz/lfs-analysis-data/repo_by_size.parquet"
    )
    unique_files_df = pd.read_parquet(
        "hf://datasets/jsulz/lfs-analysis-data/repo_by_size_file_dedupe.parquet"
    )
    file_extensions = pd.read_parquet(
        "hf://datasets/jsulz/lfs-analysis-data/file_extensions.parquet"
    )

    # Convert the total size to petabytes and format to two decimal places
    file_counts_and_sizes = format_dataframe_size_column(
        file_counts_and_sizes, "total_size"
    )

    file_counts_and_sizes["type"] = file_counts_and_sizes["type"].str.capitalize()
    # update the column name to 'total size (PB)'
    file_counts_and_sizes = file_counts_and_sizes.rename(
        columns={
            "type": "Repository Type",
            "num_files": "Number of Files",
            "total_size": "Total Size (PBs)",
        }
    )
    # sort the dataframe by total size in descending order
    file_counts_and_sizes = file_counts_and_sizes.sort_values(
        by="Total Size (PBs)", ascending=False
    )

    # drop nas from the extension column
    file_extensions = file_extensions.dropna(subset=["extension"])

    return repo_by_size_df, unique_files_df, file_counts_and_sizes, file_extensions


def format_dataframe_size_column(_df, column_name):
    """
    Format the size to petabytes and return the formatted size.
    """
    _df[column_name] = _df[column_name] / 1e15
    _df[column_name] = _df[column_name].map("{:.2f}".format)
    return _df


def cumulative_growth_plot_analysis(df, df_compressed):
    """
    Calculates the cumulative growth of models, spaces, and datasets over time and generates a plot and dataframe from the analysis.

    Args:
        df (DataFrame): The input dataframe containing the data.
        df_compressed (DataFrame): The input dataframe containing the compressed data.

    Returns:
        tuple: A tuple containing two elements:
            - fig (Figure): The Plotly figure showing the cumulative growth of models, spaces, and datasets over time.
            - last_10_months (DataFrame): The last 10 months of data showing the month-to-month growth in petabytes.

    Raises:
        None
    """
    # Convert year and month into a datetime column
    df["date"] = pd.to_datetime(df[["year", "month"]].assign(day=1))
    df_compressed["date"] = pd.to_datetime(
        df_compressed[["year", "month"]].assign(day=1)
    )

    # Sort by date to ensure correct cumulative sum
    df = df.sort_values(by="date")
    df_compressed = df_compressed.sort_values(by="date")

    # Pivot the dataframe to get the totalsize for each type
    pivot_df = df.pivot_table(
        index="date", columns="type", values="totalsize", aggfunc="sum"
    ).fillna(0)
    pivot_df_compressed = df_compressed.pivot_table(
        index="date", columns="type", values="totalsize", aggfunc="sum"
    ).fillna(0)

    # Calculate cumulative sum for each type
    cumulative_df = pivot_df.cumsum()
    cumulative_df_compressed = pivot_df_compressed.cumsum()

    last_10_months = cumulative_df.tail(10).copy()
    last_10_months["total"] = last_10_months.sum(axis=1)
    last_10_months["total_change"] = last_10_months["total"].diff()
    last_10_months = format_dataframe_size_column(last_10_months, "total_change")
    last_10_months["date"] = cumulative_df.tail(10).index
    # drop the dataset, model, and space
    last_10_months = last_10_months.drop(columns=["model", "space", "dataset"])
    # pretiffy the date column to not have 00:00:00
    last_10_months["date"] = last_10_months["date"].dt.strftime("%Y-%m")
    # drop the first row
    last_10_months = last_10_months.drop(last_10_months.index[0])
    # order the columns date, total, total_change
    last_10_months = last_10_months[["date", "total_change"]]
    # rename the columns
    last_10_months = last_10_months.rename(
        columns={"date": "Date", "total_change": "Month-to-Month Growth (PBs)"}
    )

    # Create a Plotly figure
    fig = go.Figure()

    # Define a color map for each type
    color_map = {"model": "blue", "space": "green", "dataset": "red"}

    # Add a scatter trace for each type
    for column in cumulative_df.columns:
        fig.add_trace(
            go.Scatter(
                x=cumulative_df.index,
                y=cumulative_df[column] / 1e15,  # Convert to petabytes
                mode="lines",
                name=column.capitalize(),
                line=dict(color=color_map.get(column, "black")),  # Use color map
            )
        )

    # Add a scatter trace for each type
    for column in cumulative_df_compressed.columns:
        fig.add_trace(
            go.Scatter(
                x=cumulative_df_compressed.index,
                y=cumulative_df_compressed[column] / 1e15,  # Convert to petabytes
                mode="lines",
                name=column.capitalize() + " (Compressed)",
                line=dict(color=color_map.get(column, "black"), dash="dash"),
            )
        )

    # Update layout
    fig.update_layout(
        title="Cumulative Growth of Models, Spaces, and Datasets Over Time",
        xaxis_title="Date",
        yaxis_title="Cumulative Size (PBs)",
        legend_title="Type",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
    )
    return fig, last_10_months


def plot_total_sum(by_type_arr):
    # Sort the array by size in decreasing order
    by_type_arr = sorted(by_type_arr, key=lambda x: x[1], reverse=True)

    # Create a Plotly figure
    fig = go.Figure()

    # Add a bar trace for each type
    for type, size in by_type_arr:
        fig.add_trace(
            go.Bar(
                x=[type],
                y=[size / 1e15],  # Convert to petabytes
                name=type.capitalize(),
            )
        )

    # Update layout
    fig.update_layout(
        title="Top 20 File Extensions by Total Size",
        xaxis_title="File Extension",
        yaxis_title="Total Size (PBs)",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
    )
    return fig


# Create a gradio blocks interface and launch a demo
with gr.Blocks() as demo:
    df, file_df, by_type, by_extension = process_dataset()

    # Add a heading
    gr.Markdown("# Git LFS Analysis Across the Hub")
    with gr.Row():
        # scale so that
        # group the data by month and year and compute a cumulative sum of the total_size column
        fig, last_10_months = cumulative_growth_plot_analysis(df, file_df)
        with gr.Column(scale=1):
            gr.Markdown("# Repository Growth")
            gr.Markdown(
                "The cumulative growth of models, spaces, and datasets over time can be seen in the adjacent chart. Beside that is a view of the total change, month to month, of LFS files stored on the hub over 2024. We're averaging nearly 2.3 PBs uploaded to LFS per month!"
            )
            gr.Dataframe(last_10_months, height=250)
        with gr.Column(scale=3):
            gr.Plot(fig)
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                "This table shows the total number of files and cumulative size of those files across all repositories on the Hub. These numbers might be hard to grok, so let's try to put them in context. The last [Common Crawl](https://commoncrawl.org/) download was [451 TBs](https://github.com/commoncrawl/cc-crawl-statistics/blob/master/stats/crawler/CC-MAIN-2024-38.json#L31). The Spaces repositories alone outpaces that. Meanwhile, between Datasets and Model repos, the Hub stores 64 Common Crawls."
            )
        with gr.Column(scale=3):
            gr.Dataframe(by_type)

    # Add a heading
    gr.Markdown("## File Extension Analysis")
    gr.Markdown(
        "Breaking this down by file extension, some interesting trends emerge. [Safetensors](https://huggingface.co/docs/safetensors/en/index) are quickly becoming the defacto standard on the hub, accounting for over 7PBs (25%) of LFS storage. The top 20 file extensions seen here and in the table below account for 82% of all LFS storage on the hub."
    )
    # Get the top 10 file extnesions by size
    by_extension_size = by_extension.sort_values(by="size", ascending=False).head(22)
    # get the top 10 file extensions by count
    # by_extension_count = by_extension.sort_values(by="count", ascending=False).head(20)

    # make a pie chart of the by_extension_size dataframe
    gr.Plot(plot_total_sum(by_extension_size[["extension", "size"]].values))
    # drop the unnamed: 0 column
    by_extension_size = by_extension_size.drop(columns=["Unnamed: 0"])
    # format the size column
    by_extension_size = format_dataframe_size_column(by_extension_size, "size")
    # Rename the other columns
    by_extension_size = by_extension_size.rename(
        columns={
            "extension": "File Extension",
            "count": "Number of Files",
            "size": "Total Size (PBs)",
        }
    )
    gr.Dataframe(by_extension_size)


demo.launch()