src/schedule_ai.py

import google.generativeai as genai
import os
from dotenv import load_dotenv
import re
import joblib
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler

class SimpleScheduler:
    def __init__(self):
        load_dotenv()
        api_key = os.getenv('GEMINI_API_KEY')
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel('gemini-2.0-flash-exp')
        
        # Load ML model dan scaler
        try:
            self.rf_model = joblib.load('Models/random_forest_model.pkl')
            self.scaler = joblib.load('Models/scaler.pkl')
        except Exception as e:
            print(f"Warning: Could not load ML models: {str(e)}")
            self.rf_model = None
            self.scaler = None

    def generate_schedule(self, duration, goals, available_hours, considerations):
        """Generate schedule using Gemini with improved formatting"""
        duration_lower = duration.lower()
        
        # Select appropriate prompt based on duration
        if "year" in duration_lower:
            prompt = self._create_yearly_prompt(goals, available_hours, considerations)
        elif "month" in duration_lower:
            prompt = self._create_monthly_prompt(goals, available_hours, considerations)
        else:  # weekly
            prompt = self._create_weekly_prompt(goals, available_hours, considerations)

        try:
            response = self.model.generate_content(prompt)
            return self._format_output(response.text, duration, goals, available_hours, considerations)
        except Exception as e:
            raise Exception(f"Error generating schedule: {str(e)}")

    def _get_priority(self, hour):
        """Determine priority based on hour of day"""
        if hour < 12:
            return "HIGH"
        elif hour < 15:
            return "MEDIUM"
        else:
            return "LOW"

    def _format_time(self, time_str):
        """Format time string consistently"""
        # Remove any spaces and AM/PM
        time_str = time_str.replace(" ", "").upper()
        
        # Convert to 24-hour format if needed
        if "AM" in time_str or "PM" in time_str:
            hour = int(time_str.split(":")[0])
            if "PM" in time_str and hour != 12:
                hour += 12
            return f"{hour:02d}:00"
        
        return time_str

    def _format_output(self, schedule, duration, goals, available_hours, considerations):
        """Format the output with clean and consistent structure"""
        
        # Create header
        header = f"""SCHEDULE OVERVIEW
-------------------------------------------
Duration: {duration}
Goals: {goals}
Available Hours: {available_hours} hours/day
Considerations: {considerations}
-------------------------------------------\n"""

        # Process schedule content
        formatted_schedule = ""
        current_section = None
        
        for line in schedule.strip().split('\n'):
            line = line.strip()
            
            # Skip unnecessary lines
            if not line or "[Continue" in line or "rest" in line.lower():
                continue
                
            # Process QUARTER sections for yearly schedule
            if "QUARTER" in line.upper():
                formatted_schedule += f"\n{line}\n{'=' * len(line)}\n"
                continue
                
            # Process Week sections
            elif "WEEK" in line.upper() and "-" in line:
                week_info = line.replace('**', '').strip()
                formatted_schedule += f"\n{week_info}\n{'-' * len(week_info)}\n"
                continue
                
            # Process day headers for weekly schedule
            elif any(day in line.upper() for day in ['MONDAY', 'TUESDAY', 'WEDNESDAY', 'THURSDAY', 'FRIDAY', 'SATURDAY', 'SUNDAY']):
                day = line.split(':')[0].strip()
                formatted_schedule += f"\n{day}:\n"
                continue
            
            # Process time blocks
            elif ':' in line:
                try:
                    # Split into time and task
                    parts = line.split(':', 1)
                    if len(parts) < 2:
                        continue
                        
                    time_str = self._format_time(parts[0])
                    task = parts[1].strip()
                    
                    # Clean up task text
                    task = task.split('[')[0].strip()  # Remove existing priority if any
                    
                    # Get hour for priority
                    hour = int(time_str.split(':')[0])
                    priority = self._get_priority(hour)
                    
                    # Format schedule entry
                    formatted_schedule += f"  {time_str}: {task} [{priority}]\n"
                except Exception:
                    # If processing fails, add line as is
                    formatted_schedule += f"  {line}\n"
            
            # Add other lines that might be important
            elif line.strip() and not line.startswith('*'):
                formatted_schedule += f"{line}\n"

        return header + formatted_schedule

    def _create_weekly_prompt(self, goals, available_hours, considerations):
        return f"""
        Create a WEEKLY schedule for: {goals}
        Available Hours: {available_hours}/day
        Consider: {considerations}
        
        Format exactly as follows:
        
        MONDAY:
        09:00: [Morning activity]
        10:00: [Morning activity]
        11:00: [Morning activity]
        12:00: Lunch Break
        13:00: [Afternoon activity]
        14:00: [Afternoon activity]
        15:00: [Late activity]
        16:00: [Break/Gym if needed]
        
        TUESDAY:
        [Same format]
        
        Continue for each day including SATURDAY and SUNDAY.
        Use 24-hour format (09:00, 14:00).
        Keep descriptions brief and specific.
        Include breaks and considerations.
        """

    def _create_monthly_prompt(self, goals, available_hours, considerations):
        return f"""
        Create a MONTHLY schedule for: {goals}
        Available Hours: {available_hours}/day
        Consider: {considerations}
        
        Format as follows:
        
        WEEK 1:
        Monday-Friday:
        09:00: [Morning activity]
        10:00: [Morning activity]
        11:00: [Morning activity]
        12:00: Lunch Break
        13:00: [Afternoon activity]
        14:00: [Afternoon activity]
        15:00: [Late activity]
        16:00: [Break/Gym if needed]
        
        WEEK 2:
        [Same format]
        
        Continue for 4 weeks.
        Use 24-hour format.
        Keep descriptions brief and specific.
        Include breaks and considerations.
        """

    def _create_yearly_prompt(self, goals, available_hours, considerations):
        return f"""
        Create a YEARLY schedule for: {goals}
        Available Hours: {available_hours}/day
        Consider: {considerations}
        
        Format exactly as:
        
        QUARTER 1 (January-March)
        
        Week 1-4 - [Phase Name]
        09:00: [Primary morning activity]
        11:00: [Late morning activity]
        14:00: [Afternoon activity]
        
        Week 5-8 - [Next Phase]
        [Same format]
        
        Continue for all quarters.
        Use 24-hour format.
        Focus on major milestones.
        Keep descriptions brief and specific.
        """