import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from .basePlotAnalysis import BasePlotAnalysis

import matplotlib
matplotlib.use('Agg')

class PlotActivityHeatmap(BasePlotAnalysis):
    """
    Class for analyzing user activity trends over multiple days and generating a heatmap.

    Attributes:
    name (str): The name of the analysis.
    description (str): A brief description of the analysis.
    plot_filename (str): The filename for the output plot.
    note (str): Additional notes for the analysis.
    """
    name = "Activity Heatmap"
    description = "Displays user activity trends over multiple days using a heatmap. Generates a downloadable PNG image."
    plot_filename = "activity_heatmap.png"
    note = ""
    
    def transform_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Transform data for the heatmap.

        Parameters:
        df (pd.DataFrame): The input DataFrame containing user activity data.

        Returns:
        pd.DataFrame: The transformed DataFrame with activity counts by hour.
        """
        active_counts = df[df['was_active']].pivot_table(
            index='name', 
            columns='hour', 
            values='was_active', 
            aggfunc='sum', 
            fill_value=0
        )
        active_counts['total_active_minutes'] = active_counts.sum(axis=1)
        return active_counts.sort_values(by='total_active_minutes', ascending=False)

    def plot_data(self, df: pd.DataFrame):
        """
        Generate heatmap plot.

        Parameters:
        df (pd.DataFrame): The transformed DataFrame containing activity counts by hour.
        """
        plt.figure(figsize=(12, 8))
        sns.heatmap(df.loc[:, df.columns != 'total_active_minutes'], cmap='viridis', cbar_kws={'label': 'Count of was_active == True'})
        plt.xlabel('Hour of Day')
        plt.ylabel('User ID')
        plt.title('User Activity Heatmap')