report.py

"""
Report generation module for productivity tracker.

This module handles generating various reports including summaries,
analytics, and visualizations using matplotlib.
"""

import os
from datetime import datetime, date, timedelta
from typing import List, Dict, Tuple, Optional
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import numpy as np

from tracker import Task, ProductivityTracker
from utils import format_duration, format_date, get_week_start_date, get_short_weekday_name


class ReportGenerator:
    """Generates various reports and visualizations for productivity tracking."""
    
    def __init__(self, productivity_tracker: ProductivityTracker):
        """
        Initialize the report generator.
        
        Args:
            productivity_tracker: ProductivityTracker instance
        """
        self.tracker = productivity_tracker
        self._ensure_reports_directory()
    
    def _ensure_reports_directory(self) -> None:
        """Ensure the reports directory exists."""
        os.makedirs("reports", exist_ok=True)
    
    def generate_daily_summary(self, target_date: str, save_to_file: bool = True) -> str:
        """
        Generate a daily summary report.
        
        Args:
            target_date: Date in YYYY-MM-DD format
            save_to_file: Whether to save the report to a file
            
        Returns:
            str: Generated report content
        """
        tasks = self.tracker.get_tasks_by_date(target_date)
        total_time = self.tracker.get_total_time_by_date(target_date)
        avg_focus = self.tracker.get_average_focus_by_date(target_date)
        
        report_lines = []
        report_lines.append("=" * 60)
        report_lines.append(f"DAILY PRODUCTIVITY SUMMARY - {format_date(target_date)}")
        report_lines.append("=" * 60)
        report_lines.append(f"Generated on: {datetime.now().strftime('%B %d, %Y at %I:%M %p')}")
        report_lines.append("")
        
        # Overview section
        report_lines.append("OVERVIEW")
        report_lines.append("-" * 20)
        report_lines.append(f"Total Tasks: {len(tasks)}")
        report_lines.append(f"Total Time: {format_duration(total_time)}")
        report_lines.append(f"Average Focus Score: {avg_focus:.1f}/10")
        report_lines.append("")
        
        # Task breakdown
        if tasks:
            report_lines.append("TASK BREAKDOWN")
            report_lines.append("-" * 20)
            for task in tasks:
                report_lines.append(f"• {task.task_name} ({task.category})")
                report_lines.append(f"  Duration: {format_duration(task.duration_minutes)}")
                report_lines.append(f"  Focus Score: {task.focus_score}/10")
                if task.notes:
                    report_lines.append(f"  Notes: {task.notes}")
                report_lines.append("")
        else:
            report_lines.append("No tasks recorded for this date.")
        
        # Category breakdown
        if tasks:
            category_times = {}
            category_focus = {}
            for task in tasks:
                if task.category not in category_times:
                    category_times[task.category] = 0
                    category_focus[task.category] = []
                category_times[task.category] += task.duration_minutes
                category_focus[task.category].append(task.focus_score)
            
            report_lines.append("CATEGORY BREAKDOWN")
            report_lines.append("-" * 25)
            for category, time in sorted(category_times.items(), key=lambda x: x[1], reverse=True):
                avg_cat_focus = sum(category_focus[category]) / len(category_focus[category])
                percentage = (time / total_time * 100) if total_time > 0 else 0
                report_lines.append(f"{category}: {format_duration(time)} ({percentage:.1f}%) - Avg Focus: {avg_cat_focus:.1f}")
        
        report_content = "\n".join(report_lines)
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/daily_summary_{target_date}_{timestamp}.txt"
            with open(filename, 'w') as f:
                f.write(report_content)
            print(f"Daily summary saved to: {filename}")
        
        return report_content
    
    def generate_weekly_summary(self, start_date: str, save_to_file: bool = True) -> str:
        """
        Generate a weekly summary report.
        
        Args:
            start_date: Start date of the week in YYYY-MM-DD format
            save_to_file: Whether to save the report to a file
            
        Returns:
            str: Generated report content
        """
        weekly_data = self.tracker.get_weekly_summary(start_date)
        
        report_lines = []
        report_lines.append("=" * 60)
        report_lines.append(f"WEEKLY PRODUCTIVITY SUMMARY - Week of {format_date(start_date)}")
        report_lines.append("=" * 60)
        report_lines.append(f"Generated on: {datetime.now().strftime('%B %d, %Y at %I:%M %p')}")
        report_lines.append("")
        
        # Weekly totals
        total_week_time = sum(day['total_time_minutes'] for day in weekly_data.values())
        total_week_tasks = sum(day['task_count'] for day in weekly_data.values())
        avg_week_focus = sum(day['average_focus'] for day in weekly_data.values()) / 7
        
        report_lines.append("WEEKLY OVERVIEW")
        report_lines.append("-" * 20)
        report_lines.append(f"Total Time: {format_duration(total_week_time)}")
        report_lines.append(f"Total Tasks: {total_week_tasks}")
        report_lines.append(f"Average Daily Focus: {avg_week_focus:.1f}/10")
        report_lines.append(f"Average Daily Time: {format_duration(total_week_time // 7)}")
        report_lines.append("")
        
        # Daily breakdown
        report_lines.append("DAILY BREAKDOWN")
        report_lines.append("-" * 20)
        for date_str, data in weekly_data.items():
            weekday = get_short_weekday_name(date_str)
            formatted_date = format_date(date_str, output_format='%b %d')
            report_lines.append(f"{weekday} {formatted_date}: {format_duration(data['total_time_minutes'])} "
                              f"({data['task_count']} tasks, Focus: {data['average_focus']:.1f})")
        
        report_content = "\n".join(report_lines)
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/weekly_summary_{start_date}_{timestamp}.txt"
            with open(filename, 'w') as f:
                f.write(report_content)
            print(f"Weekly summary saved to: {filename}")
        
        return report_content
    
    def generate_category_summary(self, save_to_file: bool = True) -> str:
        """
        Generate a category-wise summary report.
        
        Args:
            save_to_file: Whether to save the report to a file
            
        Returns:
            str: Generated report content
        """
        category_data = self.tracker.get_category_summary()
        
        report_lines = []
        report_lines.append("=" * 60)
        report_lines.append("CATEGORY-WISE PRODUCTIVITY SUMMARY")
        report_lines.append("=" * 60)
        report_lines.append(f"Generated on: {datetime.now().strftime('%B %d, %Y at %I:%M %p')}")
        report_lines.append("")
        
        if not category_data:
            report_lines.append("No task data available.")
            return "\n".join(report_lines)
        
        # Sort categories by total time
        sorted_categories = sorted(category_data.items(), key=lambda x: x[1]['total_time_minutes'], reverse=True)
        total_time = sum(data['total_time_minutes'] for data in category_data.values())
        
        report_lines.append("CATEGORY BREAKDOWN")
        report_lines.append("-" * 25)
        for category, data in sorted_categories:
            percentage = (data['total_time_minutes'] / total_time * 100) if total_time > 0 else 0
            report_lines.append(f"{category}:")
            report_lines.append(f"  Total Time: {format_duration(data['total_time_minutes'])} ({percentage:.1f}%)")
            report_lines.append(f"  Tasks: {data['task_count']}")
            report_lines.append(f"  Average Focus: {data['average_focus']:.1f}/10")
            report_lines.append(f"  Average Task Duration: {format_duration(data['total_time_minutes'] // data['task_count'])}")
            report_lines.append("")
        
        report_content = "\n".join(report_lines)
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/category_summary_{timestamp}.txt"
            with open(filename, 'w') as f:
                f.write(report_content)
            print(f"Category summary saved to: {filename}")
        
        return report_content
    
    def generate_weekly_time_chart(self, start_date: str, save_to_file: bool = True) -> Optional[str]:
        """
        Generate a bar chart showing daily time distribution for the week.
        
        Args:
            start_date: Start date of the week in YYYY-MM-DD format
            save_to_file: Whether to save the chart to a file
            
        Returns:
            Optional[str]: Path to saved file if save_to_file is True
        """
        weekly_data = self.tracker.get_weekly_summary(start_date)
        
        if not weekly_data:
            print("No data available for chart generation.")
            return None
        
        # Prepare data
        dates = []
        times = []
        labels = []
        
        for date_str, data in weekly_data.items():
            dates.append(date_str)
            times.append(data['total_time_hours'])
            weekday = get_short_weekday_name(date_str)
            labels.append(f"{weekday}\n{format_date(date_str, output_format='%m/%d')}")
        
        # Create bar chart
        plt.figure(figsize=(12, 6))
        bars = plt.bar(range(len(dates)), times, color=plt.cm.viridis(np.linspace(0, 1, len(dates))))
        
        # Add value labels on bars
        for bar, time in zip(bars, times):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1,
                    f'{time:.1f}h', ha='center', va='bottom', fontweight='bold')
        
        # Customize the chart
        plt.title(f'Weekly Time Distribution - Week of {format_date(start_date)}', 
                 fontsize=16, fontweight='bold')
        plt.xlabel('Day of Week', fontsize=12)
        plt.ylabel('Hours Spent', fontsize=12)
        plt.xticks(range(len(dates)), labels, rotation=0)
        plt.grid(True, alpha=0.3, axis='y')
        
        # Add total time annotation
        total_time = sum(times)
        plt.text(0.02, 0.98, f'Total: {total_time:.1f}h', 
                transform=plt.gca().transAxes, fontsize=12, fontweight='bold',
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
        
        plt.tight_layout()
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/weekly_time_chart_{start_date}_{timestamp}.png"
            plt.savefig(filename, dpi=300, bbox_inches='tight')
            print(f"Weekly time chart saved to: {filename}")
            plt.close()
            return filename
        else:
            plt.show()
            return None
    
    def generate_category_time_chart(self, save_to_file: bool = True) -> Optional[str]:
        """
        Generate a pie chart showing time distribution by category.
        
        Args:
            save_to_file: Whether to save the chart to a file
            
        Returns:
            Optional[str]: Path to saved file if save_to_file is True
        """
        category_data = self.tracker.get_category_summary()
        
        if not category_data:
            print("No data available for chart generation.")
            return None
        
        # Prepare data
        categories = list(category_data.keys())
        times = [data['total_time_hours'] for data in category_data.values()]
        
        # Create pie chart
        plt.figure(figsize=(10, 8))
        colors = plt.cm.Set3(np.linspace(0, 1, len(categories)))
        
        wedges, texts, autotexts = plt.pie(
            times,
            labels=categories,
            autopct='%1.1f%%',
            colors=colors,
            startangle=90
        )
        
        # Customize the chart
        plt.title('Time Distribution by Category', fontsize=16, fontweight='bold')
        
        # Improve text readability
        for autotext in autotexts:
            autotext.set_color('white')
            autotext.set_fontweight('bold')
        
        plt.axis('equal')
        plt.tight_layout()
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/category_time_chart_{timestamp}.png"
            plt.savefig(filename, dpi=300, bbox_inches='tight')
            print(f"Category time chart saved to: {filename}")
            plt.close()
            return filename
        else:
            plt.show()
            return None
    
    def generate_focus_trend_chart(self, days: int = 30, save_to_file: bool = True) -> Optional[str]:
        """
        Generate a line chart showing focus score trends over time.
        
        Args:
            days: Number of days to show in the trend
            save_to_file: Whether to save the chart to a file
            
        Returns:
            Optional[str]: Path to saved file if save_to_file is True
        """
        # Get data for the last N days
        end_date = date.today()
        start_date = end_date - timedelta(days=days-1)
        
        dates = []
        focus_scores = []
        
        current_date = start_date
        while current_date <= end_date:
            date_str = current_date.strftime('%Y-%m-%d')
            avg_focus = self.tracker.get_average_focus_by_date(date_str)
            
            if avg_focus > 0:  # Only include days with data
                dates.append(current_date)
                focus_scores.append(avg_focus)
            
            current_date += timedelta(days=1)
        
        if not dates:
            print("No focus score data available for chart generation.")
            return None
        
        # Create line chart
        plt.figure(figsize=(12, 6))
        plt.plot(dates, focus_scores, marker='o', linewidth=2, markersize=6, color='#2E8B57')
        
        # Customize the chart
        plt.title(f'Focus Score Trend - Last {days} Days', fontsize=16, fontweight='bold')
        plt.xlabel('Date', fontsize=12)
        plt.ylabel('Average Focus Score', fontsize=12)
        plt.ylim(0, 10)
        plt.grid(True, alpha=0.3)
        
        # Format x-axis
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m/%d'))
        plt.gca().xaxis.set_major_locator(mdates.DayLocator(interval=max(1, len(dates)//10)))
        plt.xticks(rotation=45)
        
        # Add average line
        if focus_scores:
            avg_focus = sum(focus_scores) / len(focus_scores)
            plt.axhline(y=avg_focus, color='red', linestyle='--', alpha=0.7, 
                       label=f'Average: {avg_focus:.1f}')
            plt.legend()
        
        plt.tight_layout()
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/focus_trend_{days}days_{timestamp}.png"
            plt.savefig(filename, dpi=300, bbox_inches='tight')
            print(f"Focus trend chart saved to: {filename}")
            plt.close()
            return filename
        else:
            plt.show()
            return None
    
    def generate_productivity_heatmap(self, weeks: int = 4, save_to_file: bool = True) -> Optional[str]:
        """
        Generate a heatmap showing productivity patterns by day of week and time.
        
        Args:
            weeks: Number of weeks to include
            save_to_file: Whether to save the chart to a file
            
        Returns:
            Optional[str]: Path to saved file if save_to_file is True
        """
        # This is a simplified version - in a real app, you'd track hourly data
        end_date = date.today()
        start_date = end_date - timedelta(weeks=weeks)
        
        # Create a 7x4 grid (7 days, 4 weeks)
        heatmap_data = np.zeros((7, weeks))
        weekdays = ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun']
        
        current_date = start_date
        week_idx = 0
        
        while current_date <= end_date and week_idx < weeks:
            date_str = current_date.strftime('%Y-%m-%d')
            total_time = self.tracker.get_total_time_by_date(date_str)
            
            # Convert to hours and normalize (max 8 hours = 1.0)
            normalized_time = min(total_time / 480, 1.0)  # 480 minutes = 8 hours
            day_idx = current_date.weekday()  # 0 = Monday
            
            heatmap_data[day_idx, week_idx] = normalized_time
            
            # Move to next week
            if current_date.weekday() == 6:  # Sunday
                week_idx += 1
            
            current_date += timedelta(days=1)
        
        # Create heatmap
        plt.figure(figsize=(10, 6))
        im = plt.imshow(heatmap_data, cmap='YlOrRd', aspect='auto')
        
        # Customize the chart
        plt.title(f'Productivity Heatmap - Last {weeks} Weeks', fontsize=16, fontweight='bold')
        plt.xlabel('Week', fontsize=12)
        plt.ylabel('Day of Week', fontsize=12)
        
        # Set ticks
        plt.xticks(range(weeks), [f'W{i+1}' for i in range(weeks)])
        plt.yticks(range(7), weekdays)
        
        # Add colorbar
        cbar = plt.colorbar(im)
        cbar.set_label('Productivity Level (0-8 hours)', fontsize=10)
        
        # Add text annotations
        for i in range(7):
            for j in range(weeks):
                text = plt.text(j, i, f'{heatmap_data[i, j]:.1f}',
                              ha="center", va="center", color="black", fontsize=8)
        
        plt.tight_layout()
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/productivity_heatmap_{weeks}weeks_{timestamp}.png"
            plt.savefig(filename, dpi=300, bbox_inches='tight')
            print(f"Productivity heatmap saved to: {filename}")
            plt.close()
            return filename
        else:
            plt.show()
            return None
    
    def generate_comprehensive_report(self, save_to_file: bool = True) -> str:
        """
        Generate a comprehensive report with all available data and charts.
        
        Args:
            save_to_file: Whether to save the report and charts to files
            
        Returns:
            str: Generated report content
        """
        print("Generating comprehensive productivity report...")
        
        # Get current week start
        current_week_start = get_week_start_date(date.today().strftime('%Y-%m-%d'))
        
        # Generate text reports
        daily_report = self.generate_daily_summary(date.today().strftime('%Y-%m-%d'), save_to_file)
        weekly_report = self.generate_weekly_summary(current_week_start, save_to_file)
        category_report = self.generate_category_summary(save_to_file)
        
        # Generate charts
        if save_to_file:
            print("Generating charts...")
            self.generate_weekly_time_chart(current_week_start, save_to_file=True)
            self.generate_category_time_chart(save_to_file=True)
            self.generate_focus_trend_chart(save_to_file=True)
            self.generate_productivity_heatmap(save_to_file=True)
        
        # Get productivity stats
        stats = self.tracker.get_productivity_stats()
        
        # Create comprehensive summary
        comprehensive_content = []
        comprehensive_content.append("COMPREHENSIVE PRODUCTIVITY REPORT")
        comprehensive_content.append("=" * 50)
        comprehensive_content.append("")
        comprehensive_content.append("OVERALL STATISTICS")
        comprehensive_content.append("-" * 25)
        comprehensive_content.append(f"Total Tasks: {stats['total_tasks']}")
        comprehensive_content.append(f"Total Time: {format_duration(int(stats['total_time_hours'] * 60))}")
        comprehensive_content.append(f"Average Focus Score: {stats['average_focus']:.1f}/10")
        comprehensive_content.append(f"Average Task Duration: {format_duration(int(stats['average_task_duration']))}")
        comprehensive_content.append(f"Most Productive Category: {stats['most_productive_category']}")
        comprehensive_content.append(f"Most Productive Day: {stats['most_productive_day']}")
        comprehensive_content.append("")
        comprehensive_content.append("=" * 50)
        comprehensive_content.append("")
        comprehensive_content.append(daily_report)
        comprehensive_content.append("")
        comprehensive_content.append("=" * 50)
        comprehensive_content.append("")
        comprehensive_content.append(weekly_report)
        comprehensive_content.append("")
        comprehensive_content.append("=" * 50)
        comprehensive_content.append("")
        comprehensive_content.append(category_report)
        
        if save_to_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"reports/comprehensive_report_{timestamp}.txt"
            with open(filename, 'w') as f:
                f.write("\n".join(comprehensive_content))
            print(f"Comprehensive report saved to: {filename}")
        
        return "\n".join(comprehensive_content)


# TODO: Add Pomodoro Timer Integration (countdowns, breaks, notifications)
# TODO: Add Productivity Scoring Algorithm (based on focus + time spent)
# TODO: Fix occasional bug: chart not updating if data added after chart generation
# TODO: Add Export Reports to Excel or PDF feature
# TODO: Integrate Rich library for colored CLI output
# TODO: Add Daily Goal Tracker (target hours per day)
# TODO: Implement Unit Tests for utils.py functions