pdf_generator.py

from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from datetime import datetime
import os
from data_loader import DataLoader

class PDFGenerator:
    def __init__(self, df, plot_files, plots_dir, csv_path):
        self.df = df
        self.plot_files = plot_files
        self.plots_dir = plots_dir
        self.csv_path = csv_path
        
        self.styles = getSampleStyleSheet()
        self.setup_styles()
    
    def setup_styles(self):
        self.styles.add(ParagraphStyle(
            name='CompactTitle',
            parent=self.styles['Title'],
            fontSize=14,
            spaceAfter=5,
            textColor=colors.HexColor('#2c3e50'),
            alignment=1
        ))
        
        self.styles.add(ParagraphStyle(
            name='CompactHeading1',
            parent=self.styles['Heading1'],
            fontSize=11,
            spaceAfter=2,
            spaceBefore=2,
            textColor=colors.HexColor('#3498db')
        ))
        
        self.styles.add(ParagraphStyle(
            name='CompactHeading2',
            parent=self.styles['Heading2'],
            fontSize=9,
            spaceAfter=1,
            spaceBefore=1,
            textColor=colors.HexColor('#2c3e50')
        ))
        
        self.styles.add(ParagraphStyle(
            name='CompactBody',
            parent=self.styles['Normal'],
            fontSize=8,
            leading=9,
            spaceAfter=1
        ))
        
        self.styles.add(ParagraphStyle(
            name='NoSpace',
            parent=self.styles['Normal'],
            fontSize=8,
            leading=9,
            spaceAfter=0
        ))
        
        self.styles.add(ParagraphStyle(
            name='SmallText',
            parent=self.styles['Normal'],
            fontSize=7,
            leading=8,
            spaceAfter=0
        ))
        
        self.styles.add(ParagraphStyle(
            name='TinyText',
            parent=self.styles['Normal'],
            fontSize=6,
            leading=7,
            spaceAfter=0
        ))
    
    def get_top_ligands(self):
        loader = DataLoader(self.csv_path)
        return loader.get_top_ligands(self.df)
    
    def truncate_ligand_name(self, name, max_length=40):
        if len(name) <= max_length:
            return name
        
        if '_' in name:
            parts = name.split('_')
            if len(parts) > 3:
                truncated = '_'.join(parts[:2]) + '...' + parts[-1]
                if len(truncated) <= max_length:
                    return truncated
        
        return name[:max_length-3] + "..."
    
    def generate_pdf(self, pdf_path, task_queue=None):
        try:
            if task_queue:
                task_queue.put(('update_progress', 90, "Creating PDF structure..."))
            
            doc = SimpleDocTemplate(pdf_path, pagesize=A4,
                                  rightMargin=15, leftMargin=15,
                                  topMargin=15, bottomMargin=15)
            story = []
                        
            if task_queue:
                task_queue.put(('update_progress', 92, "Creating cover page..."))
            
            story.append(Paragraph("DOCKING RESULTS ANALYSIS", self.styles['CompactTitle']))
            story.append(Spacer(1, 5))
            
            cover_info = [
                ["Analysis Date:", datetime.now().strftime("%Y-%m-%d")],
                ["Input File:", os.path.basename(self.csv_path)],
                ["Total Records:", f"{len(self.df):,}"],
                ["Generated By:", "Docking Analyzer"]
            ]
            
            cover_table = Table(cover_info, colWidths=[0.8*inch, 4*inch])
            cover_table.setStyle(TableStyle([
                ('FONTNAME', (0, 0), (-1, -1), 'Helvetica'),
                ('FONTSIZE', (0, 0), (-1, -1), 7),
                ('BOTTOMPADDING', (0, 0), (-1, -1), 2),
                ('TOPPADDING', (0, 0), (-1, -1), 1),
            ]))
            
            story.append(cover_table)
            story.append(Spacer(1, 8))
            
            if 'binding_affinity' in self.df.columns:
                stats_text = f"""<b>QUICK STATISTICS:</b>
                • Best Binding: {self.df['binding_affinity'].min():.2f} kcal/mol
                • Worst Binding: {self.df['binding_affinity'].max():.2f} kcal/mol
                • Average: {self.df['binding_affinity'].mean():.2f} ± {self.df['binding_affinity'].std():.2f} kcal/mol
                • Visualizations: {len(self.plot_files)} plots analyzed"""
                story.append(Paragraph(stats_text, self.styles['SmallText']))
            
            story.append(Spacer(1, 10))
            
            story.append(Paragraph("1. EXECUTIVE SUMMARY", self.styles['CompactHeading1']))
            
            summary = f"""This report analyzes {len(self.df):,} docking poses. Key metrics show binding affinities ranging from {self.df['binding_affinity'].min():.2f} to {self.df['binding_affinity'].max():.2f} kcal/mol with a mean of {self.df['binding_affinity'].mean():.3f} ± {self.df['binding_affinity'].std():.3f} kcal/mol."""
            
            story.append(Paragraph(summary, self.styles['CompactBody']))
            story.append(Spacer(1, 5))
                        
            if task_queue:
                task_queue.put(('update_progress', 94, "Adding top ligands analysis..."))
            
            story.append(Paragraph("2. TOP LIGANDS ANALYSIS", self.styles['CompactHeading1']))
            
            top_ligands = self.get_top_ligands()
            
            if not top_ligands.empty:
                story.append(Paragraph("Top 10 Ligands (Minimum Binding Affinity)", self.styles['CompactHeading2']))
                
                zero_rmsd_count = len(top_ligands[(top_ligands['rmsd_ub'] == 0) & (top_ligands['rmsd_lb'] == 0)])
                
                selection_note = f"<b>Selection:</b> Minimum binding affinity, RMSD ub/lb = 0 ({zero_rmsd_count} ligands)"
                story.append(Paragraph(selection_note, self.styles['SmallText']))
                story.append(Spacer(1, 2))
                
                top_data = [["#", "Ligand Name", "Binding", "RMSD ub", "RMSD lb"]]
                
                for i, (_, row) in enumerate(top_ligands.head(10).iterrows(), 1):
                    ligand_name = str(row['ligand'])
                    
                    top_data.append([
                        str(i),
                        ligand_name,  
                        f"{row['binding_affinity']:.3f}",
                        f"{row['rmsd_ub']:.2f}",
                        f"{row['rmsd_lb']:.2f}"
                    ])
                
                top_table = Table(top_data, colWidths=[0.25*inch, 3.5*inch, 0.6*inch, 0.6*inch, 0.6*inch])
                top_table.setStyle(TableStyle([
                    ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#27ae60')),
                    ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                    ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                    ('ALIGN', (1, 0), (1, -1), 'LEFT'), 
                    ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
                    ('FONTSIZE', (0, 0), (-1, -1), 6),
                    ('FONTSIZE', (1, 0), (1, -1), 5),  
                    ('BOTTOMPADDING', (0, 0), (-1, -1), 1),
                    ('TOPPADDING', (0, 0), (-1, -1), 1),
                    ('BACKGROUND', (0, 1), (-1, -1), colors.lightgreen),
                    ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
                    ('WORDWRAP', (1, 0), (1, -1), True),  
                ]))
                
                story.append(top_table)
                story.append(Spacer(1, 8))
                        
            from reportlab.platypus import PageBreak
            story.append(PageBreak())
            
            if task_queue:
                task_queue.put(('update_progress', 95, "Adding recommended ligands..."))
            
            if not top_ligands.empty and len(top_ligands) >= 3:
                top3 = top_ligands.head(3)
                
                story.append(Paragraph("3. TOP 3 RECOMMENDED LIGANDS", self.styles['CompactHeading1']))
                
                rec_text = "These ligands show the most promise for further study:"
                story.append(Paragraph(rec_text, self.styles['CompactBody']))
                story.append(Spacer(1, 3))
                
                top3_data = [["Rank", "Ligand Name", "Binding Affinity", "RMSD ub", "RMSD lb"]]
                
                for i, (_, row) in enumerate(top3.iterrows(), 1):
                    ligand_name = str(row['ligand'])
                    
                    top3_data.append([
                        str(i),
                        ligand_name,  
                        f"{row['binding_affinity']:.4f}",
                        f"{row['rmsd_ub']:.3f}",
                        f"{row['rmsd_lb']:.3f}"
                    ])
                
                top3_table = Table(top3_data, colWidths=[0.3*inch, 3.7*inch, 0.7*inch, 0.6*inch, 0.6*inch])
                top3_table.setStyle(TableStyle([
                    ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#f39c12')),
                    ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                    ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                    ('ALIGN', (1, 0), (1, -1), 'LEFT'),  
                    ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
                    ('FONTSIZE', (0, 0), (-1, -1), 7),
                    ('FONTSIZE', (1, 0), (1, -1), 6),  
                    ('BOTTOMPADDING', (0, 0), (-1, -1), 2),
                    ('BACKGROUND', (0, 1), (-1, -1), colors.whitesmoke),
                    ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
                    ('WORDWRAP', (1, 0), (1, -1), True),  
                ]))
                
                story.append(top3_table)
                story.append(Spacer(1, 8))
                
                story.append(Paragraph("<b>Additional Details:</b>", self.styles['CompactBody']))
                
                details_text = ""
                for i, (_, row) in enumerate(top3.iterrows(), 1):
                    ligand_name = str(row['ligand'])
                    details_text += f"{i}. <b>{ligand_name}</b>: Binding = {row['binding_affinity']:.4f} kcal/mol, RMSD ub/lb = {row['rmsd_ub']:.3f}/{row['rmsd_lb']:.3f}<br/>"
                
                story.append(Paragraph(details_text, self.styles['TinyText']))
                story.append(Spacer(1, 5))
                        
            if task_queue:
                task_queue.put(('update_progress', 97, "Adding statistics..."))
            
            story.append(Paragraph("4. STATISTICAL ANALYSIS", self.styles['CompactHeading1']))
            
            if 'binding_affinity' in self.df.columns:
                story.append(Paragraph("Binding Affinity Statistics", self.styles['CompactHeading2']))
                
                stats = self.df['binding_affinity'].describe()
                stats_data = [
                    ["Metric", "Value (kcal/mol)"],
                    ["Mean", f"{stats['mean']:.4f}"],
                    ["Std Deviation", f"{stats['std']:.4f}"],
                    ["Minimum (Best)", f"{stats['min']:.4f}"],
                    ["Maximum", f"{stats['max']:.4f}"],
                    ["25th Percentile", f"{stats['25%']:.4f}"],
                    ["50th Percentile", f"{stats['50%']:.4f}"],
                    ["75th Percentile", f"{stats['75%']:.4f}"],
                    ["Count", f"{int(stats['count']):,}"]
                ]
                
                stats_table = Table(stats_data, colWidths=[1.5*inch, 1.2*inch])
                stats_table.setStyle(TableStyle([
                    ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#3498db')),
                    ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                    ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                    ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
                    ('FONTSIZE', (0, 0), (-1, -1), 7),
                    ('BOTTOMPADDING', (0, 0), (-1, -1), 2),
                    ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
                    ('GRID', (0, 0), (-1, -1), 0.5, colors.grey)
                ]))
                
                story.append(stats_table)
                story.append(Spacer(1, 8))
                                    
            top_ligand_names = []
            if not top_ligands.empty:
                for i, (_, row) in enumerate(top_ligands.head(3).iterrows(), 1):
                    ligand_name = str(row['ligand'])
                    if len(ligand_name) > 30:
                        ligand_name = self.truncate_ligand_name(ligand_name, 30)
                    top_ligand_names.append(f"{i}. {ligand_name}")
                        
            story.append(Spacer(1, 10))
            story.append(Paragraph(f"Report generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}", 
                                  ParagraphStyle(name='Footer', fontSize=6, alignment=1, textColor=colors.gray)))
            
            if task_queue:
                task_queue.put(('update_progress', 98, "Finalizing PDF..."))
            
            doc.build(story)
            
            if task_queue:
                task_queue.put(('update_progress', 99, "PDF completed!"))
            
        except Exception as e:
            import traceback
            raise Exception(f"PDF creation failed: {str(e)}\n{traceback.format_exc()}")