src/utils/scheduler.py

"""
Production Scheduling Adjustment Module
Author: Antigravity AI
Date: August 2026

This module handles closed-loop production schedule adjustments.
It evaluates a production schedule against ML predictions (failure probabilities
and Remaining Useful Life), identifies at-risk assets, and dynamically reschedules/reallocates
jobs to healthy backup machines of the same model.
"""

import os
import logging
import pandas as pd
try:
    from utils.project_paths import LOGS_DIR
except ImportError:
    from project_paths import LOGS_DIR

# Setup logging
os.makedirs(LOGS_DIR, exist_ok=True)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOGS_DIR, "scheduler.log"))
    ]
)
logger = logging.getLogger(__name__)


def generate_mock_schedule(output_path: str):
    """Generate a realistic baseline production schedule for an FMCG manufacturing plant."""
    logger.info(f"Generating mock baseline production schedule at {output_path}...")
    
    # 50 production jobs running on machines 1 to 20
    jobs = []
    for i in range(1, 51):
        mid = (i % 20) + 1  # machines 1 to 20
        job_id = f"JOB_FMCG_{i:03d}"
        duration = float((i * 7) % 24 + 4)  # 4 to 27 hours
        priority = "High" if (i % 3 == 0) else "Medium" if (i % 3 == 1) else "Low"
        
        jobs.append({
            'job_id': job_id,
            'machine_id': mid,
            'duration_hours': duration,
            'priority': priority,
            'status': 'Scheduled',
            'notes': 'Original schedule assignment.'
        })
        
    df = pd.DataFrame(jobs)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    df.to_csv(output_path, index=False)
    logger.info("Mock schedule created successfully.")


class ProductionScheduler:
    def __init__(self, schedule_path: str, machines_metadata_path: str):
        self.schedule_path = schedule_path
        self.machines_path = machines_metadata_path
        
        if not os.path.exists(self.schedule_path):
            generate_mock_schedule(self.schedule_path)
            
        if not os.path.exists(self.machines_path):
            raise FileNotFoundError(f"Missing machine metadata file: {self.machines_path}")
            
        self.schedule = pd.read_csv(self.schedule_path)
        self.machines = pd.read_csv(self.machines_path)
        
    def adjust_schedule(self, predictions: dict, output_adjusted_path: str,
                        failure_threshold: float = 0.7, min_rul_hours: float = 48.0) -> pd.DataFrame:
        """
        Adjust schedule by rerouting jobs from at-risk machines to healthy backup machines.
        
        Parameters:
        - predictions: dict like { machine_id: { 'failure_probability': float, 'rul': float } }
        - output_adjusted_path: file path to save the modified schedule.
        """
        logger.info("Starting production schedule adjustment run...")
        
        adjusted_df = self.schedule.copy()
        
        # Build machine health lookup maps
        at_risk_machines = set()
        healthy_machines = {}  # model -> list of (machine_id, failure_prob, RUL)
        
        # Initialize model groups from machines metadata
        model_groups = {}
        for _, row in self.machines.iterrows():
            mid = row['machineID']
            model = row['model']
            model_groups[mid] = model
            
        for _, row in self.machines.iterrows():
            mid = row['machineID']
            model = row['model']
            
            # Retrieve predictions (default to healthy if no predictions available)
            pred = predictions.get(mid, {'failure_probability': 0.0, 'rul': 500.0})
            f_prob = pred.get('failure_probability', 0.0)
            rul = pred.get('rul', 500.0)
            
            # Determine if machine is at risk
            is_risk = False
            reasons = []
            
            if f_prob > failure_threshold:
                is_risk = True
                reasons.append(f"high failure probability ({f_prob*100:.1f}%)")
            if rul is not None and rul < min_rul_hours:
                is_risk = True
                reasons.append(f"low remaining useful life ({rul:.1f} hours)")
                
            if is_risk:
                at_risk_machines.add(mid)
                logger.warning(f"Machine {mid} ({model}) is AT RISK: {', '.join(reasons)}.")
            else:
                if model not in healthy_machines:
                    healthy_machines[model] = []
                healthy_machines[model].append((mid, f_prob, rul))
                
        # Sort healthy backup machines by lowest failure probability, then highest RUL
        for model in healthy_machines:
            healthy_machines[model].sort(key=lambda x: (x[1], -x[2] if x[2] is not None else 0))
            
        reassigned_count = 0
        
        # Round-robin counters per model for load balancing
        rr_counters = {model: 0 for model in healthy_machines}
        
        # Adjust assignments
        for idx, row in adjusted_df.iterrows():
            mid = row['machine_id']
            if mid in at_risk_machines:
                model = model_groups.get(mid, 'Unknown')
                backups = healthy_machines.get(model, [])
                
                if backups:
                    # Round-robin: pick next healthy machine in rotation
                    rr_idx = rr_counters.get(model, 0) % len(backups)
                    backup_id, backup_f_prob, backup_rul = backups[rr_idx]
                    rr_counters[model] = rr_idx + 1
                    
                    adjusted_df.at[idx, 'machine_id'] = backup_id
                    adjusted_df.at[idx, 'status'] = 'Rescheduled'
                    adjusted_df.at[idx, 'notes'] = (
                        f"Reallocated from at-risk Machine {mid} to healthy backup Machine {backup_id} "
                        f"(Model: {model}). Risk mitigation active."
                    )
                    
                    logger.info(f"Job {row['job_id']} reassigned from Machine {mid} to Machine {backup_id}.")
                    reassigned_count += 1
                else:
                    adjusted_df.at[idx, 'status'] = 'Risk Alert'
                    adjusted_df.at[idx, 'notes'] = (
                        f"Machine {mid} is at risk, but no healthy backup machines of model {model} were available!"
                    )
                    logger.error(f"Failed to reassign Job {row['job_id']}: No healthy backup machines available for model {model}!")
                    
        os.makedirs(os.path.dirname(output_adjusted_path), exist_ok=True)
        adjusted_df.to_csv(output_adjusted_path, index=False)
        
        logger.info(f"Production schedule adjustment complete. Reallocated {reassigned_count} jobs.")
        logger.info(f"Adjusted schedule saved to {output_adjusted_path}")
        
        return adjusted_df


if __name__ == '__main__':
    # Test block
    base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
    sch_path = os.path.join(base_dir, 'data', 'raw', 'production_schedule.csv')
    mach_path = os.path.join(base_dir, 'data', 'raw', 'PdM_machines.csv')
    adj_path = os.path.join(base_dir, 'data', 'processed', 'adjusted_schedule.csv')
    
    # Mock some machine risk predictions
    mock_preds = {
        3: {'failure_probability': 0.85, 'rul': 12.0},  # At-risk
        8: {'failure_probability': 0.12, 'rul': 350.0}, # Healthy
        12: {'failure_probability': 0.92, 'rul': 4.0}   # At-risk
    }
    
    scheduler = ProductionScheduler(sch_path, mach_path)
    scheduler.adjust_schedule(mock_preds, adj_path)