workflow.py

"""
Workflow controller for the paper polishing system.
Orchestrates the iterative loop between Author and Reviewer agents.
"""

import logging
from typing import Dict, Any, Optional
from agents import AuthorAgent, ReviewerAgent
from models import LLMModel
from latex_utils import LaTeXValidator
from io_handler import IOHandler


logger = logging.getLogger(__name__)


class WorkflowController:
    """Controls the iterative paper polishing workflow."""
    
    def __init__(self, 
                 model: LLMModel,
                 max_rounds: int = 5,
                 output_dir: str = "output",
                 validate_latex: bool = True,
                 project_dir: Optional[str] = None,
                 main_tex_file: Optional[str] = None):
        """
        Initialize workflow controller.
        
        Args:
            model: LLM model instance
            max_rounds: Maximum iterations
            output_dir: Output directory for results
            validate_latex: Whether to validate LaTeX syntax
            project_dir: Path to project folder containing all dependencies
            main_tex_file: Path to main LaTeX file
        """
        self.model = model
        self.max_rounds = max_rounds
        self.validate_latex = validate_latex
        self.project_dir = project_dir
        self.main_tex_file = main_tex_file
        
        # Initialize agents
        self.author = AuthorAgent(model)
        # Create three reviewers with different attitudes
        self.reviewers = {
            "harsh": ReviewerAgent(model, attitude="harsh"),
            "kind": ReviewerAgent(model, attitude="kind"),
            "median": ReviewerAgent(model, attitude="median"),
        }
        
        # Initialize I/O handler
        self.io_handler = IOHandler(output_dir, project_dir)
        
        # Initialize LaTeX validator
        self.latex_validator = LaTeXValidator() if validate_latex else None
        
        # Metadata tracking
        self.metadata = {
            "total_iterations": 0,
            "status": "PENDING",
            "accepted_at_iteration": None,
            "iterations": [],
        }
        
        logger.info(f"Workflow initialized (max_rounds={max_rounds})")
    
    def run(self, input_file: str = None) -> Dict[str, Any]:
        """
        Run the paper polishing workflow.
        
        Args:
            input_file: Path to input .tex file (optional if project_dir and main_tex_file set)
            
        Returns:
            Results dictionary with final paper and metadata
            
        Raises:
            FileNotFoundError: If input file not found
            RuntimeError: If workflow fails
        """
        try:
            # Determine which tex file to use
            if input_file is None:
                if self.main_tex_file is None:
                    raise ValueError("No input file specified")
                input_file = self.main_tex_file
            
            # Load initial paper
            logger.info(f"Starting workflow with input: {input_file}")
            initial_paper = self.io_handler.load_paper(input_file)
            
            # Create session directory
            self.io_handler.create_session(input_file)
            
            # Validate initial LaTeX
            if self.validate_latex:
                is_valid, errors = self.latex_validator.validate_syntax(initial_paper)
                if not is_valid:
                    logger.warning(f"Initial LaTeX validation failed: {errors}")
            
            current_paper = initial_paper
            
            # Main iteration loop
            for round_num in range(self.max_rounds):
                logger.info(f"\n{'='*60}")
                logger.info(f"ROUND {round_num + 1}/{self.max_rounds}")
                logger.info(f"{'='*60}")
                
                # Step 1: Get reviews from all three reviewers
                logger.info("Getting reviews from 3 reviewers (harsh, kind, median)...")
                reviews = {}
                try:
                    for attitude in ["harsh", "kind", "median"]:
                        logger.info(f"  → {attitude.upper()} reviewer analyzing paper...")
                        feedback, is_acceptable = self.reviewers[attitude].process(current_paper)
                        reviews[attitude] = {
                            "feedback": feedback,
                            "acceptable": is_acceptable
                        }
                except Exception as e:
                    logger.error(f"Review process failed: {str(e)}")
                    raise RuntimeError(f"Reviewer error: {str(e)}")
                
                # Step 2: Aggregate feedback from all reviewers
                aggregated_feedback = self._aggregate_feedback(reviews)
                is_acceptable = self._determine_acceptance(reviews)
                
                # Store iteration metadata
                iter_metadata = {
                    "round": round_num + 1,
                    "accepted": is_acceptable,
                    "reviews": {k: v["acceptable"] for k, v in reviews.items()},
                    "feedback": aggregated_feedback[:500],  # Store first 500 chars
                }
                self.metadata["iterations"].append(iter_metadata)
                
                # Save iteration
                logger.info(f"Saving iteration {round_num + 1}...")
                self.io_handler.save_iteration(
                    round_num + 1,
                    current_paper,
                    aggregated_feedback,
                    original=initial_paper if round_num == 0 else None
                )
                
                # Check if paper is acceptable (all reviewers agree OR median accepts)
                if is_acceptable:
                    logger.info(f"✓ Paper ACCEPTED at iteration {round_num + 1}")
                    self.metadata["status"] = "ACCEPTED"
                    self.metadata["accepted_at_iteration"] = round_num + 1
                    self.metadata["total_iterations"] = round_num + 1
                    break
                
                # If last round, stop here
                if round_num == self.max_rounds - 1:
                    logger.warning(
                        f"Maximum rounds ({self.max_rounds}) reached. "
                        "Paper not accepted."
                    )
                    self.metadata["status"] = "MAX_ROUNDS_REACHED"
                    self.metadata["total_iterations"] = self.max_rounds
                    break
                
                # Step 3: Author aggregates feedback and rewrites
                logger.info("Author aggregating feedback from 3 reviewers...")
                logger.info("Author rewriting paper...")
                try:
                    rewritten_paper = self.author.process(current_paper, aggregated_feedback)
                except Exception as e:
                    logger.error(f"Author failed: {str(e)}")
                    raise RuntimeError(f"Author error: {str(e)}")
                
                # Validate LaTeX preservation
                if self.validate_latex:
                    logger.info("Validating LaTeX preservation...")
                    is_safe, issues = self.latex_validator.check_preservation(
                        current_paper,
                        rewritten_paper
                    )
                    
                    if not is_safe:
                        logger.warning(f"LaTeX preservation issues: {issues}")
                        # Continue anyway but log the issues
                
                current_paper = rewritten_paper
                logger.info(f"Author rewrote paper ({len(current_paper)} chars)")
            
            # Save final results
            logger.info("Saving final results...")
            final_paper_path = self.io_handler.save_final_results(
                current_paper,
                self.metadata
            )
            
            # Return results
            results = {
                "success": True,
                "status": self.metadata["status"],
                "total_iterations": self.metadata["total_iterations"],
                "final_paper_path": str(final_paper_path),
                "session_dir": str(self.io_handler.session_dir),
                "metadata": self.metadata,
            }
            
            logger.info(f"\nWorkflow completed: {results['status']}")
            return results
            
        except Exception as e:
            logger.error(f"Workflow failed: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "status": "FAILED",
                "metadata": self.metadata,
            }
    
    def get_status(self) -> Dict[str, Any]:
        """Get current workflow status."""
        return {
            "total_iterations": self.metadata["total_iterations"],
            "status": self.metadata["status"],
            "iterations_completed": len(self.metadata["iterations"]),
            "accepted_at_iteration": self.metadata["accepted_at_iteration"],
        }
    
    def _aggregate_feedback(self, reviews: Dict[str, Dict[str, Any]]) -> str:
        """
        Aggregate feedback from all three reviewers.
        
        Args:
            reviews: Dictionary with keys 'harsh', 'kind', 'median'
                    Each contains 'feedback' and 'acceptable' keys
            
        Returns:
            Aggregated feedback text combining all three perspectives
        """
        aggregated = "AGGREGATED FEEDBACK FROM THREE REVIEWERS\n"
        aggregated += "=" * 60 + "\n\n"
        
        # Add feedback from each reviewer with clear attribution
        for attitude in ["harsh", "kind", "median"]:
            feedback = reviews[attitude]["feedback"]
            acceptable = reviews[attitude]["acceptable"]
            status = "ACCEPTABLE" if acceptable else "NEEDS REVISION"
            
            aggregated += f"[{attitude.upper()} REVIEWER - {status}]\n"
            aggregated += "-" * 60 + "\n"
            aggregated += feedback
            aggregated += "\n\n"
        
        aggregated += "=" * 60 + "\n"
        aggregated += "AUTHOR SHOULD CONSIDER:\n"
        aggregated += "1. Common points raised by multiple reviewers require priority attention\n"
        aggregated += "2. Understand constructive criticism from the kind reviewer\n"
        aggregated += "3. Address rigorous demands from the harsh reviewer where valid\n"
        aggregated += "4. Use the median reviewer's balanced perspective to prioritize improvements\n"
        aggregated += "5. Synthesize feedback to improve the paper comprehensively\n"
        
        return aggregated
    
    def _determine_acceptance(self, reviews: Dict[str, Dict[str, Any]]) -> bool:
        """
        Determine if paper is acceptable based on reviews from all three reviewers.
        
        Acceptance criteria:
        - If all 3 reviewers accept: ACCEPT
        - If median (balanced) reviewer accepts and harsh reviewer requests revision: 
          Need consensus (median alone is not enough)
        - If only kind reviewer accepts: NEEDS REVISION
        - Default: NEEDS REVISION unless clear acceptance
        
        Args:
            reviews: Dictionary with acceptance status from each reviewer
            
        Returns:
            True if paper should be accepted, False if needs revision
        """
        harsh_accepts = reviews["harsh"]["acceptable"]
        kind_accepts = reviews["kind"]["acceptable"]
        median_accepts = reviews["median"]["acceptable"]
        
        logger.info(f"Review decisions - Harsh: {harsh_accepts}, Kind: {kind_accepts}, Median: {median_accepts}")
        
        # Acceptance logic: majority or consensus
        acceptance_count = sum([harsh_accepts, kind_accepts, median_accepts])
        
        if acceptance_count >= 2:  # At least 2 out of 3 reviewers accept
            logger.info("✓ Paper accepted (2 or more reviewers agree)")
            return True
        else:
            logger.info("Paper needs revision (majority of reviewers request changes)")
            return False