import pickle
import pandas as pd
from datetime import datetime, timedelta

class ModelManager:
    """Manages ML model loading and inference"""
    
    def __init__(self, model_path='model_ai_bta.pkl', laju_path='laju_penipisan.pkl'):
        self.model = None
        self.laju_penipisan = None
        self.load_models(model_path, laju_path)
    
    def load_models(self, model_path, laju_path):
        """Load ML model dan laju penipisan dari pickle files"""
        try:
            with open(model_path, 'rb') as f:
                self.model = pickle.load(f)
            print(f"✓ Model loaded successfully from {model_path}")
        except Exception as e:
            print(f"✗ Error loading model: {e}")
            raise
        
        try:
            with open(laju_path, 'rb') as f:
                self.laju_penipisan = pickle.load(f)
            print(f"✓ Laju penipisan loaded successfully: {self.laju_penipisan}")
        except Exception as e:
            print(f"✗ Error loading laju penipisan: {e}")
            raise
    
    def predict(self, cone_depan, bodi_tengah, cone_belakang):
        """Predict ketebalan BTA given temperature inputs
        
        Returns: ketebalan_prediksi (float)
        """
        suhu_avg = (cone_depan + bodi_tengah + cone_belakang) / 3
        diff_std = suhu_avg - 400
        
        # Prepare input features untuk model
        input_features = pd.DataFrame([{
            'cone_depan_smooth': cone_depan,
            'body_tengah_smooth': bodi_tengah,
            'cone_belakang_smooth': cone_belakang,
            'suhu_avg_smooth': suhu_avg,
            'diff_std': diff_std
        }])
        
        # Predict
        ketebalan_prediksi = self.model.predict(input_features)[0]
        return ketebalan_prediksi
    
    def predict_batch(self, df):
        """Predict ketebalan untuk batch data
        
        Args:
            df: DataFrame dengan columns [cone_depan, bodi_tengah, cone_belakang, suhu_avg]
        
        Returns: array of predictions
        """
        # Ensure numeric columns
        for col in ['cone_depan', 'bodi_tengah', 'cone_belakang']:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        if 'suhu_avg' not in df.columns:
            df['suhu_avg'] = (df['cone_depan'] + df['bodi_tengah'] + df['cone_belakang']) / 3
        
        df['diff_std'] = df['suhu_avg'] - 400
        
        input_features = pd.DataFrame({
            'cone_depan_smooth': df['cone_depan'],
            'body_tengah_smooth': df['bodi_tengah'],
            'cone_belakang_smooth': df['cone_belakang'],
            'suhu_avg_smooth': df['suhu_avg'],
            'diff_std': df['diff_std']
        })
        
        predictions = self.model.predict(input_features)
        return predictions
    
    def calculate_rul(self, ketebalan_prediksi, min_safe_thickness=115):
        """Calculate Remaining Useful Life (RUL)
        
        Args:
            ketebalan_prediksi: predicted thickness (mm)
            min_safe_thickness: minimum safe thickness (mm)
        
        Returns: dict with keys: sisa_tebal, sisa_hari, tgl_target, status
        """
        sisa_tebal = ketebalan_prediksi - min_safe_thickness
        
        if sisa_tebal <= 0:
            status = "⚠️ CRITICAL (Perlu Perbaikan Segera)"
            sisa_hari = 0
            tgl_target = datetime.now()
            sisa_hari_str = "0 Hari"
            tgl_target_str = "SEGERA AKSI"
        else:
            status = "✅ AMAN"
            sisa_hari = int(sisa_tebal / self.laju_penipisan)
            tgl_target = datetime.now() + timedelta(days=sisa_hari)
            sisa_hari_str = f"± {sisa_hari} Hari"
            tgl_target_str = tgl_target.strftime('%d %B %Y')
        
        return {
            'sisa_tebal': sisa_tebal,
            'sisa_hari': sisa_hari,
            'tgl_target': tgl_target_str,
            'status': status,
            'sisa_hari_str': sisa_hari_str
        }