report_generator.py

import os
import io
import tempfile
from datetime import datetime
from PIL import Image
from reportlab.lib.pagesizes import A4, letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as RLImage, Table, TableStyle, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER, TA_LEFT, TA_JUSTIFY
from utils import format_figure_type, get_figure_type_emoji

class PDFReportGenerator:
    """Generate comprehensive PDF reports of figure analysis results."""
    
    def __init__(self):
        self.styles = getSampleStyleSheet()
        self._setup_custom_styles()
    
    def _setup_custom_styles(self):
        """Setup custom paragraph styles for the report."""
        # Title style
        self.title_style = ParagraphStyle(
            'CustomTitle',
            parent=self.styles['Heading1'],
            fontSize=24,
            spaceAfter=30,
            alignment=TA_CENTER,
            textColor=colors.darkblue
        )
        
        # Section header style
        self.section_style = ParagraphStyle(
            'CustomSection',
            parent=self.styles['Heading2'],
            fontSize=16,
            spaceAfter=12,
            spaceBefore=20,
            textColor=colors.darkblue,
            borderWidth=1,
            borderColor=colors.lightblue,
            borderPadding=5
        )
        
        # Subsection style
        self.subsection_style = ParagraphStyle(
            'CustomSubsection',
            parent=self.styles['Heading3'],
            fontSize=14,
            spaceAfter=8,
            spaceBefore=10,
            textColor=colors.darkgreen
        )
        
        # Body text style
        self.body_style = ParagraphStyle(
            'CustomBody',
            parent=self.styles['Normal'],
            fontSize=11,
            spaceAfter=6,
            alignment=TA_JUSTIFY,
            firstLineIndent=20
        )
        
        # Caption style
        self.caption_style = ParagraphStyle(
            'CustomCaption',
            parent=self.styles['Normal'],
            fontSize=9,
            spaceAfter=10,
            alignment=TA_CENTER,
            textColor=colors.darkgrey,
            italic=True
        )
    
    def generate_report(self, figures, classifications, source_info="PDF Document", output_path=None):
        """
        Generate a comprehensive PDF report of the analysis.
        
        Args:
            figures (list): List of extracted figure data
            classifications (list): List of classification results
            source_info (str): Information about the source document
            output_path (str): Optional path for the output file
            
        Returns:
            str: Path to the generated PDF report
        """
        if output_path is None:
            output_path = tempfile.mktemp(suffix='_analysis_report.pdf')
        
        # Create PDF document
        doc = SimpleDocTemplate(
            output_path,
            pagesize=A4,
            rightMargin=72,
            leftMargin=72,
            topMargin=72,
            bottomMargin=18
        )
        
        # Build the story (content)
        story = []
        
        # Title page
        story.extend(self._create_title_page(source_info, len(figures)))
        
        # Executive summary
        story.extend(self._create_executive_summary(figures, classifications))
        
        # Statistical analysis
        story.extend(self._create_statistical_analysis(classifications))
        
        # Figure type distribution
        story.extend(self._create_type_distribution(classifications))
        
        # Detailed figure analysis
        story.extend(self._create_detailed_analysis(figures, classifications))
        
        # Build PDF
        doc.build(story)
        
        return output_path
    
    def _create_title_page(self, source_info, total_figures):
        """Create the title page of the report."""
        story = []
        
        # Main title
        story.append(Paragraph("PDF Figure Analysis Report", self.title_style))
        story.append(Spacer(1, 30))
        
        # Report info
        report_info = f"""
        <b>Document Source:</b> {source_info}<br/>
        <b>Analysis Date:</b> {datetime.now().strftime('%B %d, %Y at %I:%M %p')}<br/>
        <b>Total Figures Extracted:</b> {total_figures}<br/>
        <b>Generated by:</b> AI-Powered PDF Figure Extraction Tool
        """
        story.append(Paragraph(report_info, self.body_style))
        story.append(Spacer(1, 50))
        
        # Description
        description = """
        This report provides a comprehensive analysis of all figures, charts, diagrams, 
        and images extracted from the source PDF document. Each figure has been automatically 
        classified using advanced AI technology to identify its type, purpose, and characteristics.
        
        The analysis includes statistical summaries, type distributions, confidence assessments, 
        and detailed descriptions of each identified figure.
        """
        story.append(Paragraph(description, self.body_style))
        
        story.append(PageBreak())
        return story
    
    def _create_executive_summary(self, figures, classifications):
        """Create executive summary section."""
        story = []
        
        story.append(Paragraph("Executive Summary", self.section_style))
        
        # Calculate key metrics
        total_figures = len(figures)
        unique_types = len(set(c['classification'] for c in classifications))
        avg_confidence = sum(c['confidence'] for c in classifications) / len(classifications) if classifications else 0
        high_confidence_count = sum(1 for c in classifications if c['confidence'] > 0.8)
        
        # Create summary text
        summary = f"""
        This analysis processed <b>{total_figures}</b> figures extracted from the source document. 
        The AI classification system identified <b>{unique_types}</b> distinct figure types with an 
        average confidence score of <b>{avg_confidence:.1%}</b>.
        
        <b>{high_confidence_count}</b> figures ({high_confidence_count/total_figures:.1%}) 
        were classified with high confidence (>80%), indicating reliable identification of their types and purposes.
        
        The extracted figures span multiple categories including data visualizations, technical diagrams, 
        photographs, and specialized scientific or engineering illustrations.
        """
        
        story.append(Paragraph(summary, self.body_style))
        story.append(Spacer(1, 20))
        
        return story
    
    def _create_statistical_analysis(self, classifications):
        """Create statistical analysis section."""
        story = []
        
        story.append(Paragraph("Statistical Analysis", self.section_style))
        
        # Confidence distribution
        confidence_ranges = {
            "High (80-100%)": sum(1 for c in classifications if c['confidence'] > 0.8),
            "Medium (60-80%)": sum(1 for c in classifications if 0.6 < c['confidence'] <= 0.8),
            "Low (0-60%)": sum(1 for c in classifications if c['confidence'] <= 0.6)
        }
        
        story.append(Paragraph("Confidence Score Distribution", self.subsection_style))
        
        # Create confidence table
        confidence_data = [['Confidence Range', 'Count', 'Percentage']]
        total = len(classifications)
        
        for range_name, count in confidence_ranges.items():
            percentage = f"{count/total:.1%}" if total > 0 else "0%"
            confidence_data.append([range_name, str(count), percentage])
        
        confidence_table = Table(confidence_data)
        confidence_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (-1, 0), colors.lightblue),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
            ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
            ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
            ('FONTSIZE', (0, 0), (-1, 0), 12),
            ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
            ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
            ('GRID', (0, 0), (-1, -1), 1, colors.black)
        ]))
        
        story.append(confidence_table)
        story.append(Spacer(1, 20))
        
        return story
    
    def _create_type_distribution(self, classifications):
        """Create figure type distribution section."""
        story = []
        
        story.append(Paragraph("Figure Type Distribution", self.section_style))
        
        # Calculate type counts
        type_counts = {}
        for classification in classifications:
            fig_type = classification['classification']
            type_counts[fig_type] = type_counts.get(fig_type, 0) + 1
        
        # Sort by count (descending)
        sorted_types = sorted(type_counts.items(), key=lambda x: x[1], reverse=True)
        
        # Create distribution table
        distribution_data = [['Figure Type', 'Count', 'Percentage', 'Examples']]
        total = len(classifications)
        
        for fig_type, count in sorted_types:
            percentage = f"{count/total:.1%}"
            formatted_type = format_figure_type(fig_type)
            emoji = get_figure_type_emoji(fig_type)
            
            # Get example descriptions
            examples = [c.get('description', '') for c in classifications 
                       if c['classification'] == fig_type][:2]
            example_text = '; '.join(examples[:2]) if examples else 'N/A'
            if len(example_text) > 50:
                example_text = example_text[:47] + "..."
            
            distribution_data.append([
                f"{emoji} {formatted_type}",
                str(count),
                percentage,
                example_text
            ])
        
        distribution_table = Table(distribution_data, colWidths=[2*inch, 0.8*inch, 0.8*inch, 2.4*inch])
        distribution_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (-1, 0), colors.darkblue),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
            ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
            ('ALIGN', (1, 0), (2, -1), 'CENTER'),
            ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
            ('FONTSIZE', (0, 0), (-1, 0), 10),
            ('FONTSIZE', (0, 1), (-1, -1), 9),
            ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
            ('BACKGROUND', (0, 1), (-1, -1), colors.lightgrey),
            ('GRID', (0, 0), (-1, -1), 1, colors.black),
            ('VALIGN', (0, 0), (-1, -1), 'TOP')
        ]))
        
        story.append(distribution_table)
        story.append(Spacer(1, 20))
        
        return story
    
    def _create_detailed_analysis(self, figures, classifications):
        """Create detailed analysis section with figure thumbnails."""
        story = []
        
        story.append(Paragraph("Detailed Figure Analysis", self.section_style))
        
        for i, (figure_data, classification) in enumerate(zip(figures, classifications)):
            # Figure header
            fig_num = i + 1
            emoji = get_figure_type_emoji(classification['classification'])
            formatted_type = format_figure_type(classification['classification'])
            
            figure_title = f"Figure {fig_num}: {emoji} {formatted_type}"
            story.append(Paragraph(figure_title, self.subsection_style))
            
            # Create thumbnail with better error handling
            try:
                # Create a smaller version of the image for the report
                img = figure_data['image'].copy()
                
                # Ensure image is in RGB mode
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Create thumbnail
                img.thumbnail((300, 300), Image.Resampling.LANCZOS)
                
                # Save to temporary file with unique name
                import uuid
                temp_img_path = tempfile.mktemp(suffix=f'_{uuid.uuid4().hex[:8]}.png')
                
                try:
                    img.save(temp_img_path, 'PNG')
                    
                    # Add image to report with proper sizing
                    actual_width, actual_height = img.size
                    max_width = min(actual_width, 300)
                    max_height = min(actual_height, 300)
                    
                    rl_img = RLImage(temp_img_path, width=max_width, height=max_height)
                    story.append(rl_img)
                    
                except Exception as save_error:
                    story.append(Paragraph(f"[Figure {fig_num}: Unable to include image - {str(save_error)}]", self.caption_style))
                
                finally:
                    # Clean up temp file safely
                    try:
                        if os.path.exists(temp_img_path):
                            os.unlink(temp_img_path)
                    except:
                        pass  # Ignore cleanup errors
                
            except Exception as e:
                # If image processing fails, just note it
                story.append(Paragraph(f"[Figure {fig_num}: Image processing failed - {str(e)}]", self.caption_style))
            
            # Figure details
            confidence_indicator = "🟢" if classification['confidence'] > 0.8 else "🟡" if classification['confidence'] > 0.6 else "🔴"
            
            details_text = f"""
            <b>Page:</b> {classification['page']}<br/>
            <b>Confidence:</b> {confidence_indicator} {classification['confidence']:.1%}<br/>
            <b>Description:</b> {classification.get('description', 'No description available')}<br/>
            """
            
            if classification.get('reasoning'):
                details_text += f"<b>AI Reasoning:</b> {classification['reasoning']}<br/>"
            
            if classification.get('details', {}).get('visual_elements'):
                elements = ', '.join(classification['details']['visual_elements'])
                details_text += f"<b>Visual Elements:</b> {elements}<br/>"
            
            story.append(Paragraph(details_text, self.body_style))
            story.append(Spacer(1, 15))
            
            # Add page break every 3 figures to avoid overcrowding
            if (i + 1) % 3 == 0 and i < len(figures) - 1:
                story.append(PageBreak())
        
        return story
    
    def create_summary_buffer(self, figures, classifications, source_info="PDF Document"):
        """
        Create a PDF report and return it as a BytesIO buffer.
        
        Args:
            figures (list): List of extracted figure data
            classifications (list): List of classification results
            source_info (str): Information about the source document
            
        Returns:
            io.BytesIO: Buffer containing the PDF report
        """
        # Generate PDF to temporary file
        temp_path = self.generate_report(figures, classifications, source_info)
        
        # Read into buffer
        with open(temp_path, 'rb') as f:
            pdf_buffer = io.BytesIO(f.read())
        
        # Clean up temporary file
        os.unlink(temp_path)
        
        return pdf_buffer