import gradio as gr
import pandas as pd
from plotly import graph_objects as go
import plotly.io as pio
import plotly.express as px

# @TODO: Add a custom template to the plotly figure
"""
pio.templates["custom"] = go.layout.Template()
pio.templates["custom"].layout = dict(
    plot_bgcolor="#bde5ec", paper_bgcolor="#bbd5da"
)

# Set the default theme to "plotly_dark"
pio.templates.default = "custom"
"""


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/xet-team/lfs-analysis-data/transformed/file_counts_and_sizes.parquet"
    )
    repo_by_size_df = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/repo_by_size.parquet"
    )
    unique_files_df = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/repo_by_size_file_dedupe.parquet"
    )
    file_extensions = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/file_extensions.parquet"
    )

    # read the file_extensions_by_month.parquet file
    file_extensions_by_month = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/file_extensions_by_month.parquet"
    )
    # drop any nas
    file_extensions_by_month = file_extensions_by_month.dropna()

    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)",
        }
    )
    file_counts_and_sizes = file_counts_and_sizes.drop(columns=["Number of Files"])

    # 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,
        file_extensions_by_month,
    )


def cumulative_growth_df(_df):
    # Sort by date to ensure correct cumulative sum
    _df = _df.sort_values(by="date")
    # Pivot the dataframe to get the totalsize
    pivot_df = _df.pivot_table(
        index="date", columns="type", values="totalsize", aggfunc="sum"
    ).fillna(0)
    # Calculate cumulative sum
    cumulative_df = pivot_df.cumsum()
    return cumulative_df


def compare_last_10_months(_cumulative_df, _cumulative_df_compressed):
    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["compressed_change"] = (
        _cumulative_df_compressed.tail(10).sum(axis=1).diff()
    )
    last_10_months["savings"] = (
        last_10_months["total_change"] - last_10_months["compressed_change"]
    )

    last_10_months = format_dataframe_size_column(
        last_10_months, ["total_change", "compressed_change", "savings"]
    )

    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", "compressed_change", "savings"]
    ]
    # rename the columns
    last_10_months = last_10_months.rename(
        columns={
            "date": "Date",
            "total_change": "Month-to-Month Growth (PBs)",
            "compressed_change": "Growth with File-Level Deduplication (PBs)",
            "savings": "Dedupe Savings (PBs)",
        }
    )
    return last_10_months


def tabular_analysis(repo_sizes, cumulative_df, cumulative_df_compressed):
    # create a new column in the repository sizes dataframe for "compressed size" and set it to empty atif rist
    repo_sizes["Compressed Size (PBs)"] = ""
    repo_sizes["Dedupe Savings (PBs)"] = ""

    for column in cumulative_df.columns:
        cum_repo_size = cumulative_df[column].iloc[-1]
        comp_repo_size = cumulative_df_compressed[column].iloc[-1]
        repo_size_diff = cum_repo_size - comp_repo_size
        repo_sizes.loc[
            repo_sizes["Repository Type"] == column.capitalize(),
            "Compressed Size (PBs)",
        ] = comp_repo_size
        repo_sizes.loc[
            repo_sizes["Repository Type"] == column.capitalize(), "Dedupe Savings (PBs)"
        ] = repo_size_diff

    # add a row that sums the total size and compressed size
    repo_sizes.loc["Total"] = repo_sizes.sum()
    repo_sizes.loc["Total", "Repository Type"] = "Total"
    return repo_sizes


def cumulative_growth_plot_analysis(cumulative_df, cumulative_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
    """

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

    # Define a color map for each type
    color_map = {
        "model": px.colors.qualitative.Alphabet[3],
        "space": px.colors.qualitative.Alphabet[2],
        "dataset": px.colors.qualitative.Alphabet[9],
    }

    # 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() + " (File-Level Deduplication)",
                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<br><sup>Dotted lines represent growth with file-level deduplication</sup>",
        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


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])

    # 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 (in PBs)",
        xaxis_title="File Extension",
        yaxis_title="Total Size (PBs)",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
        colorway=px.colors.qualitative.Alphabet,  # Use Plotly color palette
    )
    return fig


def filter_by_extension_month(_df, _extension):
    """
    Filters the given DataFrame (_df) by the specified extension and creates a line plot using Plotly.

    Parameters:
    _df (DataFrame): The input DataFrame containing the data.
    extension (str): The extension to filter the DataFrame by. If None, no filtering is applied.

    Returns:
    fig (Figure): The Plotly figure object representing the line plot.
    """
    # Filter the DataFrame by the specified extension or extensions
    if _extension is None:
        pass
    elif len(_extension) == 0:
        pass
    else:
        _df = _df[_df["extension"].isin(_extension)].copy()

    # Convert year and month into a datetime column and sort by date
    _df["date"] = pd.to_datetime(_df[["year", "month"]].assign(day=1))
    _df = _df.sort_values(by="date")

    # Pivot the DataFrame to get the total size for each extension and make this plotable as a time series
    pivot_df = _df.pivot_table(
        index="date", columns="extension", values="total_size"
    ).fillna(0)

    # Plot!!
    fig = go.Figure()
    for i, column in enumerate(pivot_df.columns):
        if column != "":
            fig.add_trace(
                go.Scatter(
                    x=pivot_df.index,
                    y=pivot_df[column] / 1e12,  # Convert to TBs
                    mode="lines",
                    name=column,
                    line=dict(color=px.colors.qualitative.Alphabet[i]),
                )
            )

    # Update layout
    fig.update_layout(
        title="Monthly Additions of LFS Files by Extension (in TBs)",
        xaxis_title="Date",
        yaxis_title="Size (TBs)",
        legend_title="Type",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
    )

    return fig


def area_plot_by_extension_month(_df):
    _df["total_size"] = _df["total_size"] / 1e15
    _df["date"] = pd.to_datetime(_df[["year", "month"]].assign(day=1))
    # make a plotly area chart with data and extension
    fig = px.area(_df, x="date", y="total_size", color="extension")
    # Update layout
    fig.update_layout(
        title="File Extension Monthly Additions (in PBs) Over Time",
        xaxis_title="Date",
        yaxis_title="Size (PBs)",
        legend_title="Type",
        # format y-axis to be PBs (currently bytes) with two decimal places
        yaxis=dict(tickformat=".2f"),
    )

    return fig


## Utility functions
def div_px(height):
    """
    Returns a string representing a div element with the specified height in pixels.
    """
    return f"<div style='height: {height}px;'></div>"


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


def month_year_to_date(_df):
    """
    Converts the 'year' and 'month' columns in the given DataFrame to a single 'date' column.
    """
    _df["date"] = pd.to_datetime(_df[["year", "month"]].assign(day=1))
    return _df


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

    # Convert year and month into a datetime column
    df = month_year_to_date(df)
    df_compressed = month_year_to_date(file_df)

    # Calculate the cumulative growth of models, spaces, and datasets over time
    cumulative_df = cumulative_growth_df(df)
    cumulative_df_compressed = cumulative_growth_df(df_compressed)

    last_10_months = compare_last_10_months(cumulative_df, cumulative_df_compressed)

    by_repo_type = tabular_analysis(
        by_repo_type, cumulative_df, cumulative_df_compressed
    )

    # get the figure for the cumulative growth plot and the last 10 months dataframe
    fig = cumulative_growth_plot_analysis(cumulative_df, cumulative_df_compressed)

    # Add top level heading and introduction text
    gr.Markdown("# Git LFS Usage Across the Hub")
    gr.Markdown(
        "The Hugging Face Hub has just crossed 1,000,000 models - but where is all that data stored? Most of it is stored in Git LFS. This analysis dives into the LFS storage on the Hub, breaking down the data by repository type, file extension, and growth over time. The data is based on a snapshot of the Hub's LFS storage, starting in March 2022 and ending September 20th, 2024 (meaning the data is incomplete for September 2024). Right now, this is a one-time analysis, but as we do our work we hope to revisit and update the underlying data to provide more insights."
    )

    gr.Markdown(
        "Now, you might ask yourself, 'Why are you doing this?' Well, the [Xet Team](https://huggingface.co/xet-team) is a [new addition to Hugging Face](https://huggingface.co/blog/xethub-joins-hf), bringing a new way to store massive datasets and models to enable ML teams to operate like software teams: Quickly and without friction. Because this story all starts with storage, that's where we've begun with our own deep dives into what the Hub holds. As part of this, we've included a look at what happens with just one simple deduplication strategy - deduplicating at the file level. Read on to see more!"
    )
    gr.HTML(div_px(25))
    # Cumulative growth analysis
    gr.Markdown("## Repository Growth")
    gr.Markdown(
        "The plot below shows the growth of Git LFS storage on the Hub over the past two years. The solid lines represent the cumulative growth of models, spaces, and datasets, while the dashed lines represent the growth with file-level deduplication."
    )
    gr.Plot(fig)

    gr.HTML(div_px(5))
    # @TODO Talk to Allison about variant="panel"
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                "In this table, we can see what the final picture looks like as of September 20th, 2024, along with the potential file-level deduplication savings."
            )
            gr.Markdown(
                "To put this 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** 🤯. Current estimates put file deduplication savings at approximately 3.24 PBs (7.2 Common Crawls)!"
            )
        with gr.Column(scale=3):
            # Convert the total size to petabytes and format to two decimal places
            by_repo_type = format_dataframe_size_column(
                by_repo_type,
                ["Total Size (PBs)", "Compressed Size (PBs)", "Dedupe Savings (PBs)"],
            )
            gr.Dataframe(by_repo_type)

    gr.HTML(div_px(5))
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                "The month-to-month growth of models, spaces, can be seen in the adjacent table. In 2024, the Hub has averaged nearly **2.3 PBs uploaded to LFS per month!** By the same token, the monthly file deduplication savings are nearly 225TBs. "
            )

            gr.Markdown(
                "Borrowing from the Common Crawl analogy, that's about *5 crawls* uploaded every month, with an _easy savings of half a crawl every month_ by deduplicating at the file-level!"
            )
        with gr.Column(scale=3):
            gr.Dataframe(last_10_months)

    gr.HTML(div_px(25))
    # File Extension analysis
    gr.Markdown("## File Extensions on the Hub")
    gr.Markdown(
        "Breaking this down by file extension, some interesting trends emerge. The following sections filter the analysis to the top 20 file extensions stored (in bytes) using LFS (which accounts for 82% of storage consumption)."
    )
    gr.Markdown(
        "As is evident in the chart below, [Safetensors](https://huggingface.co/docs/safetensors/en/index) is quickly becoming the defacto standard on the Hub for storing tensor files, accounting for over 7PBs (25%) of LFS storage. If you want to know why you'd want to check out YAF (yet another format), this explanation from the [Safetensors docs](https://github.com/huggingface/safetensors?tab=readme-ov-file#yet-another-format-) is a good place to start. Speaking of YAF, [GGUF (GPT-Generated Unified Format)](https://github.com/ggerganov/ggml/blob/master/docs/gguf.md) is also on the rise, accounting for 3.2 PBs (11%) of LFS storage. GGUF, like Safetensors, is a format for storing tensor files, with a different set of optimizations. The Hub has a few [built-in tools](https://huggingface.co/docs/hub/en/gguf) for working with GGUF."
    )
    # Get the top 10 file extnesions by size
    by_extension_size = by_extension.sort_values(by="size", ascending=False).head(22)

    # make a bar 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"])
    # average size
    by_extension_size["Average File Size (MBs)"] = (
        by_extension_size["size"].astype(float) / by_extension_size["count"]
    )
    by_extension_size["Average File Size (MBs)"] = (
        by_extension_size["Average File Size (MBs)"] / 1e6
    )
    by_extension_size["Average File Size (MBs)"] = by_extension_size[
        "Average File Size (MBs)"
    ].map("{:.2f}".format)
    # 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.HTML(div_px(5))
    gr.Markdown(
        "Below, we have a more detailed tabular view of the same top 20 file extensions by total size, number of files, and average file size."
    )
    gr.Dataframe(by_extension_size)

    gr.HTML(div_px(5))
    gr.Markdown("### File Extension Monthly Additions (in PBs)")
    gr.Markdown(
        "What if we want to see trends over time? The following area chart shows the number of bytes added to LFS storage each month, faceted by the most popular file extensions."
    )
    gr.Plot(area_plot_by_extension_month(by_extension_month))

    gr.HTML(div_px(5))
    gr.Markdown(
        "To dig a little deeper, the following dropdown allows you to filter the area chart by file extension. Because we're dealing with individual file extensions, the data is presented in terabytes (TBs)."
    )

    # build a dropdown using the unique values in the extension column
    extension = gr.Dropdown(
        choices=by_extension["extension"].unique().tolist(),
        multiselect=True,
        label="File Extension",
    )
    _by_extension_month = gr.State(by_extension_month)
    gr.Plot(filter_by_extension_month, inputs=[_by_extension_month, extension])

# launch the dang thing
demo.launch()