app.py

import os
import sys
import time
import re
import math
import tempfile
import uuid
from io import BytesIO
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from pptx import Presentation
from pptx.util import Pt, Inches
from pptx.enum.text import PP_ALIGN
from pptx.dml.color import RGBColor
from pptx.chart.data import CategoryChartData, ChartData
from pptx.enum.chart import XL_CHART_TYPE, XL_LEGEND_POSITION

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Model configuration
MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
MAX_LENGTH = 4000
TEMPERATURE = 0.7

# ===== ENHANCED PROPOSAL GENERATION FUNCTIONS ===== #

def generate_proposal(model, tokenizer, description, temperature=0.7, max_length=4000):
    """Generate a detailed, well-structured proposal from a description"""
    system_prompt = """You are an expert proposal writer with years of experience crafting successful business and project proposals.
Your task is to create a comprehensive, professional proposal that is detailed, specific, and tailored to the project description.
Follow a structured format with clear sections, use professional language, and include concrete details that make the proposal compelling.
Be specific with numbers, timelines, and methodologies where appropriate."""

    user_prompt = f"""Create a detailed project proposal for the following project description:

"{description}"

Your proposal must include these sections:
1. Executive Summary - Concise overview of the entire proposal (1-2 paragraphs)
2. Project Background - Context, history, and need for the project (2-3 paragraphs)
3. Goals and Objectives - Clear, measurable outcomes using SMART criteria (specific, measurable, achievable, relevant, time-bound)
4. Methodology and Approach - Detailed explanation of how the project will be executed, including specific techniques and frameworks
5. Timeline - Realistic schedule with key milestones and deliverables (presented in chronological order with approximate dates)
6. Budget Considerations - Cost breakdown by category with justifications 
7. Expected Outcomes - Tangible and intangible benefits with metrics for measuring success
8. Team and Resources - Key personnel, their qualifications, and resource requirements
9. Risk Assessment - Potential challenges and mitigation strategies
10. Conclusion - Compelling closing with clear next steps

For each section, include specific details that would convince stakeholders to approve this project.
Use professional language and maintain a confident, authoritative tone throughout.
Format each section with clear headings and organized paragraphs."""

    # Format the prompt according to the model's expected format
    # Zephyr uses a specific format with system and user messages
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    
    # Convert messages to the format expected by the model
    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
    
    # Tokenize and generate
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=max_length,
            temperature=temperature,
            do_sample=True,
            top_p=0.9,  # Add nucleus sampling for better quality
            top_k=50,   # Limit vocab to top 50 tokens at each step
            repetition_penalty=1.2  # Reduce repetition
        )
    
    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract just the generated proposal without the prompt
    proposal = full_output[len(prompt):].strip()
    
    # Post-processing to ensure clean formatting
    proposal = clean_proposal_format(proposal)
    
    return proposal

def clean_proposal_format(proposal_text):
    """Clean up the proposal text for better formatting"""
    
    # Remove any assistant prefixes that might be generated
    proposal_text = re.sub(r'^(Assistant:|A:|Response:)\s*', '', proposal_text)
    
    # Ensure section titles are properly formatted
    section_titles = [
        "Executive Summary", "Project Background", "Goals and Objectives", 
        "Methodology and Approach", "Timeline", "Budget Considerations", 
        "Expected Outcomes", "Team and Resources", "Risk Assessment", "Conclusion"
    ]
    
    for title in section_titles:
        # Replace various formats of section titles with consistent formatting
        proposal_text = re.sub(
            rf'(?i)(^|\n)[\d\.\s]*{title}[\s\:]*(\n|\r)',
            f'\n\n{title}\n\n',
            proposal_text
        )
    
    # Ensure consistent paragraph breaks
    proposal_text = re.sub(r'\n{3,}', '\n\n', proposal_text)
    
    # Ensure bullet points are consistent
    proposal_text = re.sub(r'\n\s*•\s*', '\n- ', proposal_text)
    proposal_text = re.sub(r'\n\s*\*\s*', '\n- ', proposal_text)
    
    return proposal_text

# ===== ENHANCED POWERPOINT GENERATION FUNCTIONS ===== #

def create_slides(proposal, project_title="Project Proposal"):
    """Create professional PowerPoint slides from the proposal"""
    
    # Create presentation with widescreen dimensions
    prs = Presentation()
    prs.slide_width = Inches(13.33)
    prs.slide_height = Inches(7.5)
    
    # Define a professional color scheme
    colors = {
        'primary': RGBColor(0, 112, 192),     # Blue
        'secondary': RGBColor(0, 176, 80),    # Green
        'accent1': RGBColor(255, 102, 0),     # Orange
        'accent2': RGBColor(112, 48, 160),    # Purple
        'dark': RGBColor(54, 54, 54),         # Dark Gray
        'light': RGBColor(244, 244, 244)      # Light Gray
    }
    
    # Add title slide with professional styling
    title_slide = prs.slides.add_slide(prs.slide_layouts[0])
    title_slide.shapes.title.text = project_title
    title_slide.shapes.title.text_frame.paragraphs[0].font.size = Pt(44)
    title_slide.shapes.title.text_frame.paragraphs[0].font.color.rgb = colors['primary']
    title_slide.shapes.title.text_frame.paragraphs[0].alignment = PP_ALIGN.CENTER
    
    # Add subtitle
    subtitle = title_slide.placeholders[1]
    subtitle.text = "Professional Project Proposal"
    subtitle.text_frame.paragraphs[0].font.size = Pt(28)
    subtitle.text_frame.paragraphs[0].font.color.rgb = colors['dark']
    
    # Add background shape for visual interest
    left = top = 0
    width = prs.slide_width
    height = prs.slide_height
    
    # Add decorative element to title slide
    shape = title_slide.shapes.add_shape(
        1, Inches(0), Inches(6.5), prs.slide_width, Inches(1)
    )
    shape.fill.solid()
    shape.fill.fore_color.rgb = colors['primary']
    shape.line.color.rgb = colors['primary']
    
    # List of sections to look for
    sections = [
        "Executive Summary", 
        "Project Background", 
        "Goals and Objectives", 
        "Methodology and Approach", 
        "Timeline", 
        "Budget Considerations", 
        "Expected Outcomes", 
        "Team and Resources", 
        "Risk Assessment", 
        "Conclusion"
    ]
    
    # Create a slide for table of contents
    toc_slide = prs.slides.add_slide(prs.slide_layouts[2])
    toc_slide.shapes.title.text = "Table of Contents"
    toc_slide.shapes.title.text_frame.paragraphs[0].font.size = Pt(40)
    toc_slide.shapes.title.text_frame.paragraphs[0].font.color.rgb = colors['primary']
    
    # Add decorative element to TOC slide
    shape = toc_slide.shapes.add_shape(
        1, Inches(0), Inches(0), Inches(1), prs.slide_height
    )
    shape.fill.solid()
    shape.fill.fore_color.rgb = colors['primary']
    shape.line.color.rgb = colors['primary']
    
    # Add TOC content
    toc_content = toc_slide.placeholders[1].text_frame
    for i, section in enumerate(sections, 1):
        p = toc_content.add_paragraph()
        p.text = f"{i}. {section}"
        p.font.size = Pt(24)
        p.font.color.rgb = colors['dark']
        p.space_after = Pt(12)
    
    # Split text into paragraphs and identify sections
    paragraphs = proposal.split('\n\n')
    
    # Process each paragraph
    current_section = None
    current_content = []
    found_sections = []
    
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
            
        # Check if this is a section header
        is_header = False
        for section in sections:
            # More robust section detection
            if (section.lower() in para.lower() and len(para) < 100) or \
               re.match(r'^[0-9]+\.?\s*' + section, para, re.IGNORECASE):
                # Save previous section
                if current_section and current_content:
                    found_sections.append((current_section, current_content))
                    
                # Start new section
                current_section = section  # Use standard section name
                current_content = []
                is_header = True
                break
                
        if not is_header:
            current_content.append(para)
    
    # Add the last section
    if current_section and current_content:
        found_sections.append((current_section, current_content))
    
    # Create slides for each section
    for section_index, (title, content_paras) in enumerate(found_sections):
        # Section title slide with visual distinction
        section_slide = prs.slides.add_slide(prs.slide_layouts[2])
        section_slide.shapes.title.text = title
        section_slide.shapes.title.text_frame.paragraphs[0].font.size = Pt(40)
        section_slide.shapes.title.text_frame.paragraphs[0].font.color.rgb = colors['primary']
        
        # Add decorative accent based on section type
        accent_color = colors['primary']
        if "goal" in title.lower() or "objective" in title.lower():
            accent_color = colors['secondary']
        elif "risk" in title.lower():
            accent_color = colors['accent1']
        elif "conclusion" in title.lower():
            accent_color = colors['accent2']
        
        # Add section number
        subtitle = section_slide.placeholders[1]
        subtitle.text = f"Section {section_index + 1}"
        subtitle.text_frame.paragraphs[0].font.size = Pt(28)
        subtitle.text_frame.paragraphs[0].font.color.rgb = accent_color
        
        # Decorative shape for section divider
        shape = section_slide.shapes.add_shape(
            1, Inches(0), Inches(6.5), prs.slide_width, Inches(1)
        )
        shape.fill.solid()
        shape.fill.fore_color.rgb = accent_color
        shape.line.color.rgb = accent_color
        
        # Content slides with better formatting
        current_slide = None
        text_frame = None
        paragraphs_on_slide = 0
        
        for para_index, para in enumerate(content_paras):
            # Detect if paragraph is a bullet point
            is_bullet = para.strip().startswith("-") or para.strip().startswith("*")
            is_numbered = bool(re.match(r'^\d+\.', para.strip()))
            
            # Start a new slide if needed
            # Fewer paragraphs per slide for better readability
            max_paragraphs = 3 if len(para) > 200 else 4
            if current_slide is None or paragraphs_on_slide >= max_paragraphs:
                current_slide = prs.slides.add_slide(prs.slide_layouts[1])
                current_slide.shapes.title.text = title
                current_slide.shapes.title.text_frame.paragraphs[0].font.size = Pt(36)
                current_slide.shapes.title.text_frame.paragraphs[0].font.color.rgb = colors['primary']
                
                # Add decorative element
                shape = current_slide.shapes.add_shape(
                    1, Inches(0), Inches(0), Inches(0.3), prs.slide_height
                )
                shape.fill.solid()
                shape.fill.fore_color.rgb = accent_color
                shape.line.color.rgb = accent_color
                
                text_frame = current_slide.placeholders[1].text_frame
                text_frame.word_wrap = True
                paragraphs_on_slide = 0
                
                # Add subtitle for content continuation if not the first content slide
                if para_index > 0:
                    p = text_frame.add_paragraph()
                    p.text = "Continued..."
                    p.font.italic = True
                    p.font.size = Pt(14)
                    p.font.color.rgb = colors['dark']
                    paragraphs_on_slide += 0.5  # Count as half a paragraph
            
            # Add the paragraph with appropriate formatting
            p = text_frame.add_paragraph()
            
            # Clean up bullet points
            if is_bullet:
                clean_text = para.strip()[1:].strip()
                p.text = clean_text
                p.level = 1
            elif is_numbered:
                p.text = para.strip()
                p.level = 1
            else:
                p.text = para.strip()
            
            # Apply formatting
            p.font.size = Pt(20)
            p.font.color.rgb = colors['dark']
            p.space_after = Pt(12)
            
            # Highlight key phrases with color
            text_frame.fit_text(max_size=Pt(20))
            
            paragraphs_on_slide += 1
    
    # Add a closing slide
    closing_slide = prs.slides.add_slide(prs.slide_layouts[5])
    title_shape = closing_slide.shapes.title
    title_shape.text = "Thank You"
    title_shape.text_frame.paragraphs[0].font.size = Pt(54)
    title_shape.text_frame.paragraphs[0].font.color.rgb = colors['primary']
    title_shape.text_frame.paragraphs[0].alignment = PP_ALIGN.CENTER
    
    subtitle_shape = closing_slide.shapes.placeholders[1]
    subtitle_shape.text = "Questions & Discussion"
    subtitle_shape.text_frame.paragraphs[0].font.size = Pt(32)
    subtitle_shape.text_frame.paragraphs[0].font.color.rgb = colors['dark']
    subtitle_shape.text_frame.paragraphs[0].alignment = PP_ALIGN.CENTER
    
    # Add large decorative shape to closing slide
    shape = closing_slide.shapes.add_shape(
        1, Inches(0), Inches(6.5), prs.slide_width, Inches(1)
    )
    shape.fill.solid()
    shape.fill.fore_color.rgb = colors['primary']
    shape.line.color.rgb = colors['primary']
    
    # Generate a unique filename
    unique_id = str(uuid.uuid4())[:8]
    output_path = f"proposal_slides_{unique_id}.pptx"
    
    # Save the presentation
    prs.save(output_path)
    return output_path

# ===== CHARTS AND VISUALS FUNCTIONS ===== #

def add_charts_and_visuals(prs, proposal_text, include_visuals=True):
    """Add professional charts and diagrams to the presentation based on proposal content"""
    if not include_visuals:
        return prs
    
    # Extract potential data for charts from the proposal text
    timeline_data = extract_timeline_data(proposal_text)
    budget_data = extract_budget_data(proposal_text)
    risk_data = extract_risk_data(proposal_text)
    
    # Add a timeline slide if data available
    if timeline_data and len(timeline_data) >= 2:
        timeline_slide = prs.slides.add_slide(prs.slide_layouts[6])  # Blank layout
        
        # Add title
        title_shape = timeline_slide.shapes.add_textbox(
            Inches(0.5), Inches(0.5), Inches(9), Inches(1)
        )
        title_frame = title_shape.text_frame
        title_frame.text = "Project Timeline"
        title_frame.paragraphs[0].font.size = Pt(40)
        title_frame.paragraphs[0].font.bold = True
        title_frame.paragraphs[0].font.color.rgb = RGBColor(0, 112, 192)
        
        # Create chart data
        chart_data = ChartData()
        chart_data.categories = [item[0] for item in timeline_data]  # Milestone names
        
        # Convert timeline data to numeric values (months from start)
        def month_to_number(month_name):
            months = {'January': 1, 'February': 2, 'March': 3, 'April': 4, 
                     'May': 5, 'June': 6, 'July': 7, 'August': 8, 
                     'September': 9, 'October': 10, 'November': 11, 'December': 12}
            for month, num in months.items():
                if month.lower() in month_name.lower():
                    return num
            return 0
        
        # Normalize timeline values based on extracted data
        start_month = min(month_to_number(item[1]) for item in timeline_data) if timeline_data else 1
        timeline_values = [max(1, month_to_number(item[1]) - start_month + 1) for item in timeline_data]
        
        chart_data.add_series('Duration (months)', timeline_values)
        
        # Add chart to the slide
        x, y, cx, cy = Inches(1), Inches(2), Inches(11), Inches(5)
        chart = timeline_slide.shapes.add_chart(
            XL_CHART_TYPE.BAR_CLUSTERED, x, y, cx, cy, chart_data
        ).chart
        
        # Style the chart
        chart.has_legend = True
        chart.legend.position = XL_LEGEND_POSITION.BOTTOM
        chart.legend.include_in_layout = False
        
        plot = chart.plots[0]
        plot.has_data_labels = True
        data_labels = plot.data_labels
        data_labels.font.size = Pt(12)
        data_labels.font.color.rgb = RGBColor(0, 0, 0)
        data_labels.position = 2  # Inside End
    
    # Add a budget breakdown slide if data available
    if budget_data and len(budget_data) >= 2:
        budget_slide = prs.slides.add_slide(prs.slide_layouts[6])  # Blank layout
        
        # Add title
        title_shape = budget_slide.shapes.add_textbox(
            Inches(0.5), Inches(0.5), Inches(9), Inches(1)
        )
        title_frame = title_shape.text_frame
        title_frame.text = "Budget Allocation"
        title_frame.paragraphs[0].font.size = Pt(40)
        title_frame.paragraphs[0].font.bold = True
        title_frame.paragraphs[0].font.color.rgb = RGBColor(0, 112, 192)
        
        # Create pie chart data
        chart_data = CategoryChartData()
        chart_data.categories = [item[0] for item in budget_data]  # Category names
        chart_data.add_series('Budget', [item[1] for item in budget_data])  # Values
        
        # Add chart to the slide
        x, y, cx, cy = Inches(2), Inches(2), Inches(9), Inches(5)
        chart = budget_slide.shapes.add_chart(
            XL_CHART_TYPE.PIE, x, y, cx, cy, chart_data
        ).chart
        
        # Style the chart
        chart.has_legend = True
        chart.legend.position = XL_LEGEND_POSITION.RIGHT
        chart.legend.font.size = Pt(14)
        
        plot = chart.plots[0]
        plot.has_data_labels = True
        data_labels = plot.data_labels
        data_labels.font.size = Pt(12)
        data_labels.position = 1  # Outside End
        data_labels.number_format = '0%'
        
    # Add a risk assessment matrix if data available
    if risk_data and len(risk_data) >= 2:
        risk_slide = prs.slides.add_slide(prs.slide_layouts[6])  # Blank layout
        
        # Add title
        title_shape = risk_slide.shapes.add_textbox(
            Inches(0.5), Inches(0.5), Inches(9), Inches(1)
        )
        title_frame = title_shape.text_frame
        title_frame.text = "Risk Assessment Matrix"
        title_frame.paragraphs[0].font.size = Pt(40)
        title_frame.paragraphs[0].font.bold = True
        title_frame.paragraphs[0].font.color.rgb = RGBColor(0, 112, 192)
        
        # Create a table for the risk matrix
        rows, cols = len(risk_data) + 1, 3  # +1 for header
        table_width, table_height = Inches(10), Inches(5)
        table = risk_slide.shapes.add_table(
            rows, cols, 
            Inches(1.5), Inches(2), 
            table_width, table_height
        ).table
        
        # Set column widths
        table.columns[0].width = Inches(5)  # Risk description
        table.columns[1].width = Inches(2.5)  # Probability
        table.columns[2].width = Inches(2.5)  # Impact
        
        # Add header row
        cell = table.cell(0, 0)
        cell.text = "Risk Description"
        cell.text_frame.paragraphs[0].font.bold = True
        cell.text_frame.paragraphs[0].font.size = Pt(16)
        
        cell = table.cell(0, 1)
        cell.text = "Probability"
        cell.text_frame.paragraphs[0].font.bold = True
        cell.text_frame.paragraphs[0].font.size = Pt(16)
        
        cell = table.cell(0, 2)
        cell.text = "Impact"
        cell.text_frame.paragraphs[0].font.bold = True
        cell.text_frame.paragraphs[0].font.size = Pt(16)
        
        # Add risk data rows
        for i, (risk, probability, impact) in enumerate(risk_data, 1):
            # Risk description
            cell = table.cell(i, 0)
            cell.text = risk
            cell.text_frame.paragraphs[0].font.size = Pt(14)
            
            # Probability
            cell = table.cell(i, 1)
            cell.text = probability
            cell.text_frame.paragraphs[0].font.size = Pt(14)
            
            # Impact
            cell = table.cell(i, 2)
            cell.text = impact
            cell.text_frame.paragraphs[0].font.size = Pt(14)
            
            # Color code based on overall risk level
            prob_level = get_risk_level(probability)
            impact_level = get_risk_level(impact)
            overall_risk = prob_level * impact_level
            
            # Set background color based on risk level
            fill_color = RGBColor(255, 255, 255)  # Default white
            if overall_risk >= 9:
                fill_color = RGBColor(255, 153, 153)  # Red for high risk
            elif overall_risk >= 4:
                fill_color = RGBColor(255, 204, 153)  # Orange for medium risk
            else:
                fill_color = RGBColor(198, 239, 206)  # Green for low risk
            
            for col in range(3):
                table.cell(i, col).fill.solid()
                table.cell(i, col).fill.fore_color.rgb = fill_color
    
    # Add a SMART objectives slide - a generic visual we can add regardless of content
    objectives_slide = prs.slides.add_slide(prs.slide_layouts[6])  # Blank layout
    
    # Add title
    title_shape = objectives_slide.shapes.add_textbox(
        Inches(0.5), Inches(0.5), Inches(9), Inches(1)
    )
    title_frame = title_shape.text_frame
    title_frame.text = "SMART Objectives Framework"
    title_frame.paragraphs[0].font.size = Pt(40)
    title_frame.paragraphs[0].font.bold = True
    title_frame.paragraphs[0].font.color.rgb = RGBColor(0, 112, 192)
    
    # Create a SmartArt-like diagram for SMART objectives
    # Since python-pptx doesn't directly support SmartArt, we'll simulate it with shapes
    
    # Center point for the circular layout
    center_x, center_y = Inches(6.5), Inches(4)
    radius = Inches(2.5)
    
    # SMART components with colors
    smart_components = [
        ("Specific", RGBColor(91, 155, 213)),   # Blue
        ("Measurable", RGBColor(112, 173, 71)), # Green
        ("Achievable", RGBColor(237, 125, 49)), # Orange
        ("Relevant", RGBColor(165, 105, 189)),  # Purple
        ("Time-bound", RGBColor(68, 114, 196))  # Dark Blue
    ]
    
    # Add central circle
    central_shape = objectives_slide.shapes.add_shape(
        1, center_x - Inches(1), center_y - Inches(1), Inches(2), Inches(2)
    )
    central_shape.fill.solid()
    central_shape.fill.fore_color.rgb = RGBColor(0, 112, 192)
    central_shape.line.color.rgb = RGBColor(255, 255, 255)
    central_shape.line.width = Pt(2)
    
    # Add text to central circle
    text_frame = central_shape.text_frame
    text_frame.text = "SMART\nObjectives"
    text_frame.paragraphs[0].alignment = 1  # Center
    text_frame.paragraphs[0].font.color.rgb = RGBColor(255, 255, 255)
    text_frame.paragraphs[0].font.size = Pt(18)
    text_frame.paragraphs[0].font.bold = True
    
    # Add surrounding circles for each SMART component
    for i, (component, color) in enumerate(smart_components):
        angle = (2 * 3.14159 * i) / len(smart_components)
        x = center_x + radius * 0.8 * math.cos(angle) - Inches(1)
        y = center_y + radius * 0.8 * math.sin(angle) - Inches(0.75)
        
        # Component circle
        component_shape = objectives_slide.shapes.add_shape(
            1, x, y, Inches(2), Inches(1.5)
        )
        component_shape.fill.solid()
        component_shape.fill.fore_color.rgb = color
        component_shape.line.color.rgb = RGBColor(255, 255, 255)
        component_shape.line.width = Pt(2)
        
        # Add component text
        text_frame = component_shape.text_frame
        text_frame.text = component
        text_frame.paragraphs[0].alignment = 1  # Center
        text_frame.paragraphs[0].font.color.rgb = RGBColor(255, 255, 255)
        text_frame.paragraphs[0].font.size = Pt(16)
        text_frame.paragraphs[0].font.bold = True
        
        # Add connecting line
        line = objectives_slide.shapes.add_connector(
            3, # Straight connector
            central_shape.left + Inches(1),
            central_shape.top + Inches(1),
            component_shape.left + Inches(1),
            component_shape.top + Inches(0.75)
        )
        line.line.color.rgb = RGBColor(0, 0, 0)
        line.line.width = Pt(1.5)
    
    return prs

def extract_timeline_data(proposal_text):
    """Extract timeline data from the proposal text using regex patterns"""
    
    # Look for timeline section
    timeline_section = re.search(r'(?i)Timeline.*?(?=\n\n[A-Z]|$)', proposal_text, re.DOTALL)
    if not timeline_section:
        # Generate example data if no real data found
        return [
            ("Project Initiation", "Month 1"),
            ("Requirements Analysis", "Month 2"),
            ("Design Phase", "Month 3"),
            ("Implementation", "Month 4-5"),
            ("Testing", "Month 6"),
            ("Deployment", "Month 7"),
            ("Post-Implementation Review", "Month 8")
        ]
    
    timeline_text = timeline_section.group(0)
    
    # Look for milestone patterns like "Phase 1: Project Initiation (Month 1)"
    # or "Project Initiation - Month 1" or "Week 1-2: Requirements Gathering"
    milestone_patterns = [
        r'([^:]+):\s*([^(]+)\s*\(([^)]+)\)',  # Phase 1: Project Initiation (Month 1)
        r'([^-]+)\s*-\s*([^(]+)',             # Project Initiation - Month 1
        r'((?:Week|Month)[^:]+):\s*([^(]+)',  # Week 1-2: Requirements Gathering
        r'([^:]+):\s*([^(]+)'                 # Any other pattern with colon
    ]
    
    extracted_data = []
    for pattern in milestone_patterns:
        matches = re.finditer(pattern, timeline_text)
        for match in matches:
            if len(match.groups()) >= 2:
                milestone = match.group(1).strip()
                timeframe = match.group(2).strip()
                
                # Clean up the milestone and timeframe
                milestone = re.sub(r'^[0-9]+\.\s*', '', milestone)  # Remove leading numbers
                
                # Add to extracted data
                extracted_data.append((milestone, timeframe))
    
    # If no data extracted, create example data
    if not extracted_data:
        # Generate some example data
        return [
            ("Project Initiation", "Month 1"),
            ("Requirements Analysis", "Month 2"),
            ("Design Phase", "Month 3"),
            ("Implementation", "Month 4-5"),
            ("Testing", "Month 6"),
            ("Deployment", "Month 7"),
            ("Post-Implementation Review", "Month 8")
        ]
    
    return extracted_data

def extract_budget_data(proposal_text):
    """Extract budget data from the proposal text using regex patterns"""
    
    # Look for budget section
    budget_section = re.search(r'(?i)Budget.*?(?=\n\n[A-Z]|$)', proposal_text, re.DOTALL)
    if not budget_section:
        # Generate example data if no real data found
        return [
            ("Hardware", 30),
            ("Software", 25),
            ("Personnel", 35),
            ("Training", 10),
            ("Contingency", 10)
        ]
    
    budget_text = budget_section.group(0)
    
    # Pattern for budget items with percentages or amounts
    # e.g. "Hardware: $50,000 (20%)" or "Personnel: 35% of total budget"
    budget_patterns = [
        r'([^:]+):\s*\$?[\d,]+\s*\((\d+)%\)',   # Hardware: $50,000 (20%)
        r'([^:]+):\s*(\d+)%',                   # Personnel: 35%
        r'([^-]+)\s*-\s*(\d+)%'                 # Software - 25%
    ]
    
    extracted_data = []
    for pattern in budget_patterns:
        matches = re.finditer(pattern, budget_text)
        for match in matches:
            if len(match.groups()) >= 2:
                category = match.group(1).strip()
                percentage = int(match.group(2).strip())
                
                # Clean up the category
                category = re.sub(r'^[0-9]+\.\s*', '', category)  # Remove leading numbers
                
                # Add to extracted data
                extracted_data.append((category, percentage))
    
    # If no data extracted, look for dollar amounts instead
    if not extracted_data:
        # Pattern for dollar amounts: "Hardware: $50,000"
        amount_pattern = r'([^:]+):\s*\$?([\d,]+)'
        matches = re.finditer(amount_pattern, budget_text)
        
        total_amount = 0
        temp_data = []
        
        for match in matches:
            if len(match.groups()) >= 2:
                category = match.group(1).strip()
                try:
                    amount = int(match.group(2).replace(',', ''))
                    total_amount += amount
                    temp_data.append((category, amount))
                except ValueError:
                    continue
        
        # Convert absolute amounts to percentages
        if total_amount > 0:
            for category, amount in temp_data:
                percentage = round((amount / total_amount) * 100)
                extracted_data.append((category, percentage))
    
    # If still no data extracted, create example data
    if not extracted_data:
        return [
            ("Hardware", 30),
            ("Software", 25),
            ("Personnel", 35),
            ("Training", 10),
            ("Contingency", 10)
        ]
    
    return extracted_data

def extract_risk_data(proposal_text):
    """Extract risk assessment data from the proposal text"""
    
    # Look for risk section
    risk_section = re.search(r'(?i)Risk Assessment.*?(?=\n\n[A-Z]|$)', proposal_text, re.DOTALL)
    if not risk_section:
        # Generate example data if no real data found
        return [
            ("Resource availability constraints", "Medium", "High"),
            ("Technology integration issues", "High", "Medium"),
            ("Budget overruns", "Medium", "High"),
            ("Schedule delays", "High", "Medium"),
            ("Stakeholder resistance", "Medium", "Medium")
        ]
    
    risk_text = risk_section.group(0)
    
    # Split into lines
    lines = risk_text.split('\n')
    extracted_data = []
    
    # Look for risk items in various formats
    for line in lines:
        # Skip empty lines
        if not line.strip():
            continue
            
        # Look for risk items with probability and impact
        # Patterns:
        # 1. "Risk: Resource constraints - Probability: Medium, Impact: High"
        # 2. "Resource constraints (Medium probability, High impact)"
        # 3. "- Resource constraints: Medium probability, High impact"
        
        # Pattern 1
        match = re.search(r'(?i)(?:Risk:)?\s*([^-]+)\s*-\s*Probability:\s*(\w+),\s*Impact:\s*(\w+)', line)
        if match:
            risk = match.group(1).strip()
            probability = match.group(2).strip()
            impact = match.group(3).strip()
            extracted_data.append((risk, probability, impact))
            continue
            
        # Pattern 2
        match = re.search(r'([^(]+)\s*\((\w+)\s*probability,\s*(\w+)\s*impact\)', line)
        if match:
            risk = match.group(1).strip()
            probability = match.group(2).strip()
            impact = match.group(3).strip()
            extracted_data.append((risk, probability, impact))
            continue
            
        # Pattern 3
        match = re.search(r'(?:-|\*)\s*([^:]+):\s*(\w+)\s*probability,\s*(\w+)\s*impact', line)
        if match:
            risk = match.group(1).strip()
            probability = match.group(2).strip()
            impact = match.group(3).strip()
            extracted_data.append((risk, probability, impact))
            continue
    
    # If we couldn't extract structured data, look for bullet points and guess
    if not extracted_data:
        bullet_items = re.findall(r'(?:^|\n)(?:-|\*|•|\d+\.)\s*([^\n]+)', risk_text)
        for item in bullet_items:
            # Clean the item text
            item = item.strip()
            
            # Guess probability and impact based on keywords
            probability = "Medium"
            impact = "Medium"
            
            # Check for probability indicators
            if re.search(r'(?i)\b(?:high probability|likely|frequently|often|high chance)\b', item):
                probability = "High"
            elif re.search(r'(?i)\b(?:low probability|unlikely|rarely|seldom|small chance)\b', item):
                probability = "Low"
                
            # Check for impact indicators
            if re.search(r'(?i)\b(?:high impact|severe|critical|major|significant)\b', item):
                impact = "High"
            elif re.search(r'(?i)\b(?:low impact|minor|minimal|negligible)\b', item):
                impact = "Low"
                
            # Extract the risk description (removing any probability/impact text)
            risk = re.sub(r'(?i)\b(?:high|medium|low)(?:\s+(?:probability|impact|chance|risk))\b', '', item)
            risk = re.sub(r'(?i)\b(?:likely|unlikely|critical|severe|major|minor)\b', '', risk)
            risk = risk.strip()
            
            if risk:
                extracted_data.append((risk, probability, impact))
    
    # If still no data extracted, create example data
    if not extracted_data:
        return [
            ("Resource availability constraints", "Medium", "High"),
            ("Technology integration issues", "High", "Medium"),
            ("Budget overruns", "Medium", "High"),
            ("Schedule delays", "High", "Medium"),
            ("Stakeholder resistance", "Medium", "Medium")
        ]
    
    return extracted_data

def get_risk_level(text):
    """Convert text risk level to numeric value"""
    text = text.lower()
    if "high" in text:
        return 3
    elif "medium" in text:
        return 2
    else:
        return 1  # Low or default

# ===== ANALYSIS AND UTILITY FUNCTIONS ===== #

def analyze_proposal(proposal_text):
    """Analyze the generated proposal for metrics and structure"""
    # Count words
    word_count = len(proposal_text.split())
    
    # Count sections
    sections = [
        "Executive Summary", "Project Background", "Goals and Objectives", 
        "Methodology", "Timeline", "Budget", "Expected Outcomes", 
        "Team and Resources", "Risk Assessment", "Conclusion"
    ]
    section_count = 0
    for section in sections:
        if re.search(rf'\b{section}\b', proposal_text, re.IGNORECASE):
            section_count += 1
    
    # Calculate reading time (average 200 words per minute)
    reading_time_min = word_count / 200
    
    # Calculate readability (simple algorithm based on word and sentence length)
    sentences = re.split(r'[.!?]+', proposal_text)
    sentence_count = len([s for s in sentences if s.strip()])
    avg_sentence_length = word_count / max(1, sentence_count)
    
    # Simplified Flesch-Kincaid calculation
    readability_score = 206.835 - (1.015 * avg_sentence_length) - (84.6 * 1.5 / avg_sentence_length)
    readability_score = max(0, min(100, readability_score))
    
    # Analyze keyword presence for common business terms
    keywords = {
        "Strategic": proposal_text.lower().count("strategic"),
        "Innovation": proposal_text.lower().count("innovat"),
        "Efficiency": proposal_text.lower().count("efficien"),
        "ROI": proposal_text.lower().count("roi") + proposal_text.lower().count("return on investment"),
        "Stakeholder": proposal_text.lower().count("stakeholder"),
        "Sustainable": proposal_text.lower().count("sustainab"),
    }
    
    # Return analysis results
    return {
        "word_count": word_count,
        "section_count": section_count,
        "readability_score": round(readability_score),
        "estimated_reading_time": f"{reading_time_min:.1f} minutes",
        "keyword_analysis": keywords,
        "sections_present": {section: (1 if re.search(rf'\b{section}\b', proposal_text, re.IGNORECASE) else 0) for section in sections},
    }

def create_slide_preview(pptx_path):
    """Create a preview image of the first slide of the presentation"""
    try:
        # If using Windows with COM support, try this method
        try:
            from pptx import Presentation
            import win32com.client
            import os
            
            # Get absolute path
            abs_path = os.path.abspath(pptx_path)
            
            # Export first slide as image using PowerPoint COM object
            ppt_app = win32com.client.Dispatch('PowerPoint.Application')
            presentation = ppt_app.Presentations.Open(abs_path)
            
            # Save the first slide as PNG
            output_path = abs_path.replace('.pptx', '_preview.png')
            presentation.Slides[1].Export(output_path, 'PNG')
            
            # Clean up
            presentation.Close()
            ppt_app.Quit()
            
            return output_path
        except:
            # Fallback: return the PPT file path itself as we can't generate a preview
            return pptx_path
        
    except Exception as e:
        print(f"Error creating slide preview: {e}")
        # Return the PPT file path itself as we can't generate a preview
        return pptx_path

def load_model():
    """Load the model and tokenizer"""
    print(f"Loading model: {MODEL_NAME}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        low_cpu_mem_usage=True,
        device_map="auto"
    )
    return model, tokenizer

def process_input(title, description, temperature=0.7, max_length=4000, 
                  color_scheme="Professional Blue", include_visuals=True, 
                  focus_areas=None):
    """Process the input and generate both proposal and slides with advanced options"""
    # Load model if not already loaded
    global model, tokenizer
    if 'model' not in globals() or model is None:
        model, tokenizer = load_model()
    
    # Set parameters based on input
    temperature_val = float(temperature)
    max_length_val = int(max_length)
    
    # Modify prompt based on focus areas
    focus_instruction = ""
    if focus_areas and len(focus_areas) > 0:
        focus_instruction = f"Pay special attention to these aspects: {', '.join(focus_areas)}."
    
    # Generate the proposal
    proposal = generate_proposal(model, tokenizer, description + " " + focus_instruction, 
                                temperature=temperature_val, max_length=max_length_val)
    
    # Create the slides
    ppt_path = create_slides(proposal, title)
    
    # Add charts and visuals if requested
    if include_visuals:
        prs = Presentation(ppt_path)
        prs = add_charts_and_visuals(prs, proposal, include_visuals)
        prs.save(ppt_path)
    
    # Create a preview image of the first slide
    preview_path = create_slide_preview(ppt_path)
    
    # Analyze the proposal
    analysis = analyze_proposal(proposal)
    
    return proposal, ppt_path, preview_path, analysis

# ===== GRADIO INTERFACE FUNCTIONS ===== #

def create_interface():
    """Create an enhanced Gradio interface with advanced options"""
    
    # Define color schemes for PowerPoint
    color_schemes = {
        "Professional Blue": {"primary": "#0070C0", "secondary": "#00B050"},
        "Elegant Gray": {"primary": "#404040", "secondary": "#7030A0"},
        "Bold Impact": {"primary": "#C00000", "secondary": "#FFC000"},
        "Modern Green": {"primary": "#00B050", "secondary": "#5B9BD5"},
        "Corporate Purple": {"primary": "#7030A0", "secondary": "#ED7D31"}
    }
    
    with gr.Blocks(title="Advanced Project Proposal Generator", theme=gr.themes.Soft()) as app:
        gr.Markdown("# Professional Project Proposal Generator")
        gr.Markdown("#### Transform your idea into a complete project proposal with presentation slides")
        
        with gr.Row():
            with gr.Column(scale=2):
                # Input section
                project_title = gr.Textbox(
                    label="Project Title",
                    placeholder="Enter your project title",
                    value="New Project Proposal"
                )
                
                description_input = gr.Textbox(
                    label="Project Description",
                    placeholder="Describe your project in detail. Include the purpose, scope, stakeholders, and any specific requirements or challenges...",
                    lines=10
                )
                
                with gr.Accordion("Advanced Options", open=False):
                    with gr.Row():
                        temperature_slider = gr.Slider(
                            minimum=0.1, 
                            maximum=1.0, 
                            value=0.7,
                            step=0.1,
                            label="Creativity Level"
                        )
                        
                        max_length_slider = gr.Slider(
                            minimum=2000,
                            maximum=6000,
                            value=4000,
                            step=500,
                            label="Maximum Length"
                        )
                    
                    with gr.Row():
                        color_scheme = gr.Dropdown(
                            choices=list(color_schemes.keys()),
                            value="Professional Blue",
                            label="Presentation Style"
                        )
                        
                        include_visuals = gr.Checkbox(
                            label="Include Sample Charts/Diagrams",
                            value=True
                        )
                    
                    with gr.Row():
                        focus_areas = gr.CheckboxGroup(
                            choices=[
                                "Technical Details", 
                                "Business Impact", 
                                "Implementation Plan",
                                "Cost Analysis",
                                "Risk Management"
                            ],
                            value=["Business Impact", "Implementation Plan"],
                            label="Focus Areas (Emphasize these aspects)"
                        )
                
                with gr.Row():
                    clear_button = gr.Button("Clear All", variant="secondary")
                    example_button = gr.Button("Load Example", variant="secondary")
                    generate_button = gr.Button("Generate Proposal & Slides", variant="primary")
                
            with gr.Column(scale=3):
                # Output tabs
                with gr.Tabs():
                    with gr.TabItem("Proposal Text"):
                        proposal_output = gr.Textbox(
                            label="Generated Proposal", 
                            lines=25,
                            show_copy_button=True
                        )
                    
                    with gr.TabItem("PowerPoint Preview"):
                        gr.Markdown("#### PowerPoint Slides Preview")
                        slides_preview = gr.Image(
                            label="Preview (first slide)",
                            type="filepath",
                            height=400
                        )
                        slides_output = gr.File(
                            label="Download Complete Presentation"
                        )
                    
                    with gr.TabItem("Proposal Analysis"):
                        analysis_output = gr.JSON(
                            label="Proposal Structure Analysis"
                        )
        
        # Example project description
        example_description = """
        Our company needs to implement a new customer relationship management (CRM) system to replace our outdated solution. The current system is 8 years old and lacks modern features like cloud integration, mobile access, and AI-powered analytics. We have approximately 5,000 customer records that need to be migrated, and 75 employees across sales, marketing, and customer service departments who will use the system. The project should include software selection, data migration, staff training, and integration with our existing ERP system. Budget constraints are significant, with a maximum allocation of $250,000. The implementation needs to be completed within 6 months to align with our fiscal year planning.
        """
        
        # Set up event handlers
        def load_example():
            return "Enterprise CRM Implementation Project", example_description
        
        example_button.click(
            load_example,
            outputs=[project_title, description_input]
        )
        
        clear_button.click(
            lambda: ("", ""),
            outputs=[project_title, description_input]
        )
        
        generate_button.click(
            process_input,
            inputs=[
                project_title,
                description_input,
                temperature_slider,
                max_length_slider,
                color_scheme,
                include_visuals,
                focus_areas
            ],
            outputs=[
                proposal_output,
                slides_output,
                slides_preview,
                analysis_output
            ]
        )
    
    return app

# ===== MAIN SCRIPT CODE ===== #

# Global model and tokenizer
model = None
tokenizer = None

# Only load model at startup when running as main program
if __name__ == "__main__":
    print("Loading initial model...")
    model, tokenizer = load_model()
    print("Starting Gradio interface...")
    app = create_interface()
    app.launch(share=True)