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):
    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"""
        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"""
        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')