import gradio as gr
import json
import pickle
import pandas as pd
import numpy as np
from datetime import datetime
import os

class EnergyMLPredictor:
    def __init__(self):
        self.rf_model = None
        self.rf_preprocessor = None
        self.xgb_model = None
        self.xgb_encoders = None
        self.threshold_model_83 = None
        self.threshold_model_90 = None
        self.threshold_preprocessor = None
        self.models_loaded = False
        
    def load_models(self):
        """Load all models from pickle files"""
        try:
            # Load Random Forest Energy Model
            if os.path.exists('rf_energy_model.pkl'):
                with open('rf_energy_model.pkl', 'rb') as f:
                    rf_data = pickle.load(f)
                    self.rf_model = rf_data['model']
                    self.rf_preprocessor = rf_data['preprocessor']
            
            # Load XGBoost Energy Model
            if os.path.exists('xgboost_energy_model.pkl'):
                with open('xgboost_energy_model.pkl', 'rb') as f:
                    xgb_data = pickle.load(f)
                    self.xgb_model = xgb_data['model']
                    self.xgb_encoders = xgb_data['label_encoders']
            
            # Load Threshold Models
            if os.path.exists('threshold_model_83_autoclave.pkl'):
                with open('threshold_model_83_autoclave.pkl', 'rb') as f:
                    threshold_data = pickle.load(f)
                    self.threshold_model_83 = threshold_data['model']
                    self.threshold_preprocessor = threshold_data['preprocessor']
            
            if os.path.exists('threshold_model_90_autoclave.pkl'):
                with open('threshold_model_90_autoclave.pkl', 'rb') as f:
                    threshold_data = pickle.load(f)
                    self.threshold_model_90 = threshold_data['model']
            
            self.models_loaded = True
            return "Models loaded successfully"
        
        except Exception as e:
            return f"Error loading models: {str(e)}"
    
    def predict_threshold(self, json_input):
        """Predict threshold exceedance"""
        try:
            if not self.models_loaded:
                return "Error: Models not loaded"
            
            if not self.threshold_model_83 or not self.threshold_model_90:
                return "Error: Threshold models not available"
            
            data = json.loads(json_input)
            
            # Handle both single object and array formats
            if not isinstance(data, list):
                data = [data]
            
            # Process all items
            results_83 = []
            results_90 = []
            
            for item in data:
                # Parse input data
                date_obj = datetime.strptime(item['data'], '%Y-%m-%d')
                
                # Color mapping
                color_mapping = {0: 'incolor', 1: 'verde', 2: 'cinza', 3: 'bronze'}
                if isinstance(item['cor'], str):
                    cor_str = item['cor'].lower()
                else:
                    cor_str = color_mapping.get(item['cor'], 'incolor')
                
                # Handle different field name formats for threshold
                boosting_val = item.get('pot_boost', item.get('ext_boosting', 3.0))
                
                # Create input features (with autoclave)
                input_data = {
                    'boosting': boosting_val,
                    'espessura': item['espessura'],
                    'extracao_forno': item['extracao_forno'],
                    'porcentagem_caco': item['porcentagem_caco'],
                    'cor': cor_str,
                    'prod_e': item.get('Prod_E', item.get('prod_e', 1)),
                    'prod_l': item.get('Prod_L', item.get('prod_l', 1)),
                    'autoclave': item.get('autoclave', 1),
                    'week_day': date_obj.weekday(),
                    'month': date_obj.month,
                    'quarter': (date_obj.month - 1) // 3 + 1,
                    'is_weekend': int(date_obj.weekday() >= 5),
                    'week_of_year': date_obj.isocalendar()[1]
                }
                
                # Convert to DataFrame
                input_df = pd.DataFrame([input_data])
                
                # Preprocess
                X_processed = self.threshold_preprocessor.transform(input_df)
                
                # Make predictions with error handling
                try:
                    prob_83_raw = self.threshold_model_83.predict_proba(X_processed)
                    prob_83 = prob_83_raw[0][1] if len(prob_83_raw[0]) > 1 else prob_83_raw[0][0]
                    # Ensure probability is between 0 and 1
                    prob_83 = max(0.0, min(1.0, float(prob_83)))
                except Exception as e:
                    print(f"Error with threshold_83 prediction: {e}")
                    prob_83 = 0.0
                
                pred_83 = int(prob_83 > 0.5)
                
                try:
                    prob_90_raw = self.threshold_model_90.predict_proba(X_processed)
                    prob_90 = prob_90_raw[0][1] if len(prob_90_raw[0]) > 1 else prob_90_raw[0][0]
                    # Ensure probability is between 0 and 1
                    prob_90 = max(0.0, min(1.0, float(prob_90)))
                except Exception as e:
                    print(f"Error with threshold_90 prediction: {e}")
                    prob_90 = 0.0
                
                pred_90 = int(prob_90 > 0.5)
                
                # Add to results
                results_83.append({
                    "datetime": item['data'],
                    "probabilidade_de_estouro": round(prob_83, 4),
                    "estouro_previsto": pred_83
                })
                
                results_90.append({
                    "datetime": item['data'],
                    "probabilidade_de_estouro": round(prob_90, 4),
                    "estouro_previsto": pred_90
                })
            
            # Format response
            result = {
                "predictions": {
                    "prediction_1": results_83,
                    "prediction_2": results_90
                }
            }
            
            return json.dumps(result, indent=2)
            
        except json.JSONDecodeError:
            return "Error: Invalid JSON format"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def predict_energy_rf(self, json_input):
        """Predict energy using Random Forest"""
        try:
            if not self.models_loaded or not self.rf_model:
                return "Error: Random Forest model not available"
            
            data = json.loads(json_input)
            if not isinstance(data, list):
                data = [data]
            
            results = []
            
            for item in data:
                # Parse input
                date_obj = datetime.strptime(item['data'], '%Y-%m-%d')
                
                # Handle different field name formats (boosting vs ext_boosting)
                if 'boosting' in item:
                    boosting_val = float(str(item['boosting']).replace(',', '.'))
                elif 'ext_boosting' in item:
                    boosting_val = float(str(item['ext_boosting']).replace(',', '.'))
                else:
                    boosting_val = 0.0
                
                # Handle extracao_forno field
                if 'extracao_forno' in item:
                    extracao_val = float(str(item['extracao_forno']).replace(',', '.'))
                else:
                    extracao_val = 800.0
                
                # Create features
                input_data = {
                    'boosting': boosting_val,
                    'espessura': item['espessura'],
                    'extracao_forno': extracao_val,
                    'porcentagem_caco': item['porcentagem_caco'],
                    'cor': str(item['cor']).lower() if isinstance(item['cor'], str) else {0: 'incolor', 1: 'verde', 2: 'cinza', 3: 'bronze'}.get(item['cor'], 'incolor'),
                    'prod_e': item.get('prod_e', item.get('Prod_E', 1)),
                    'prod_l': item.get('prod_l', item.get('Prod_L', 1)),
                    'autoclave': item.get('autoclave', 1),
                    'week_day': date_obj.weekday(),
                    'month': date_obj.month,
                    'quarter': (date_obj.month - 1) // 3 + 1,
                    'is_weekend': int(date_obj.weekday() >= 5),
                    'week_of_year': date_obj.isocalendar()[1],
                    'day_of_month': date_obj.day,
                    'day_of_year': date_obj.timetuple().tm_yday
                }
                
                # Predict
                input_df = pd.DataFrame([input_data])
                X_processed = self.rf_preprocessor.transform(input_df)
                prediction = self.rf_model.predict(X_processed)[0]
                
                results.append({
                    "data": date_obj.strftime('%d-%m-%Y'),
                    "predictions": float(prediction)
                })
            
            return json.dumps(results, indent=2)
            
        except json.JSONDecodeError:
            return "Error: Invalid JSON format"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def predict_energy_xgb(self, json_input):
        """Predict energy using XGBoost"""
        try:
            if not self.models_loaded or not self.xgb_model:
                return "Error: XGBoost model not available"
            
            data = json.loads(json_input)
            if not isinstance(data, list):
                data = [data]
            
            results = []
            
            for item in data:
                # Parse input
                date_obj = datetime.strptime(item['data'], '%Y-%m-%d')
                
                # Handle different field name formats (boosting vs ext_boosting)
                if 'boosting' in item:
                    boosting_val = float(str(item['boosting']).replace(',', '.'))
                elif 'ext_boosting' in item:
                    boosting_val = float(str(item['ext_boosting']).replace(',', '.'))
                else:
                    boosting_val = 0.0
                
                # Handle extracao_forno field
                if 'extracao_forno' in item:
                    extracao_val = float(str(item['extracao_forno']).replace(',', '.'))
                else:
                    extracao_val = 800.0
                
                # Create features (match training: numerical + categorical + boolean, NO day_of_month/day_of_year)
                input_data = {
                    'boosting': boosting_val,
                    'espessura': item['espessura'],
                    'extracao_forno': extracao_val,
                    'porcentagem_caco': item['porcentagem_caco'],
                    'cor': str(item['cor']).lower() if isinstance(item['cor'], str) else {0: 'incolor', 1: 'verde', 2: 'cinza', 3: 'bronze'}.get(item['cor'], 'incolor'),
                    'week_day': date_obj.weekday(),
                    'month': date_obj.month,
                    'quarter': (date_obj.month - 1) // 3 + 1,
                    'week_of_year': date_obj.isocalendar()[1],
                    'prod_e': item.get('prod_e', 1),
                    'prod_l': item.get('prod_l', 1),
                    'is_weekend': int(date_obj.weekday() >= 5),
                    'autoclave': item.get('autoclave', 1)
                }
                
                # Encode categorical features
                input_df = pd.DataFrame([input_data])
                
                for col in input_df.columns:
                    if col in self.xgb_encoders:
                        try:
                            input_df[col] = self.xgb_encoders[col].transform(input_df[col].astype(str))
                        except ValueError:
                            # Handle unknown categories
                            input_df[col] = 0
                
                # Predict
                prediction = self.xgb_model.predict(input_df.values)[0]
                
                results.append({
                    "data": date_obj.strftime('%d-%m-%Y'),
                    "predictions": float(prediction)
                })
            
            return json.dumps(results, indent=2)
            
        except json.JSONDecodeError:
            return "Error: Invalid JSON format"
        except Exception as e:
            return f"Error: {str(e)}"

# Initialize predictor
predictor = EnergyMLPredictor()

def make_prediction(model_choice, json_input):
    """Make prediction based on model choice"""
    if not predictor.models_loaded:
        load_msg = predictor.load_models()
        if "Error" in load_msg:
            return load_msg
    
    if model_choice == "Threshold Detection":
        return predictor.predict_threshold(json_input)
    elif model_choice == "Energy Prediction (Random Forest)":
        return predictor.predict_energy_rf(json_input)
    elif model_choice == "Energy Prediction (XGBoost)":
        return predictor.predict_energy_xgb(json_input)
    else:
        return "Error: Please select a model"

# Default examples
threshold_example = """{
  "data": "2023-01-01",
  "cor": 0,
  "espessura": 8.0,
  "ext_boosting": 65.0,
  "extracao_forno": 851.1,
  "porcentagem_caco": 15.0,
  "pot_boost": 3.0,
  "Prod_E": 1,
  "Prod_L": 1
}"""

energy_example = """[
  {
    "data": "2023-01-01",
    "boosting": "0,0",
    "cor": "incolor",
    "espessura": 10,
    "extracao_forno": "651,6",
    "porcentagem_caco": 10.0,
    "prod_e": 1,
    "prod_l": 1,
    "autoclave": 1
  }
]"""

# Test data from holdout period (last 2 months used in training)
week_test_data = """[
  {"data": "2025-04-19", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2007, "extracao_forno": 699.561202512973, "porcentagem_caco": 10.0062724674475, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-20", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2026, "extracao_forno": 699.169485837721, "porcentagem_caco": 9.99757589767354, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-04-21", "cor": 0, "espessura": 8.0, "ext_boosting": 1.201, "extracao_forno": 699.134346519477, "porcentagem_caco": 9.99807838764974, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-04-22", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2074, "extracao_forno": 701.318973743488, "porcentagem_caco": 9.99545180216949, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-23", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2028, "extracao_forno": 702.765143096952, "porcentagem_caco": 9.97488288777139, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 1},
  {"data": "2025-04-24", "cor": 0, "espessura": 8.0, "ext_boosting": 1.3973, "extracao_forno": 700.8439481142, "porcentagem_caco": 10.002226628142, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 1},
  {"data": "2025-04-25", "cor": 0, "espessura": 8.0, "ext_boosting": 1.6005, "extracao_forno": 702.032548397562, "porcentagem_caco": 9.98529201530728, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1}
]"""

month_test_data = """[
  {"data": "2025-04-19", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2007, "extracao_forno": 699.561202512973, "porcentagem_caco": 10.0062724674475, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-20", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2026, "extracao_forno": 699.169485837721, "porcentagem_caco": 9.99757589767354, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-04-21", "cor": 0, "espessura": 8.0, "ext_boosting": 1.201, "extracao_forno": 699.134346519477, "porcentagem_caco": 9.99807838764974, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-04-22", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2074, "extracao_forno": 701.318973743488, "porcentagem_caco": 9.99545180216949, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-23", "cor": 0, "espessura": 8.0, "ext_boosting": 1.2028, "extracao_forno": 702.765143096952, "porcentagem_caco": 9.97488288777139, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 1},
  {"data": "2025-04-24", "cor": 0, "espessura": 8.0, "ext_boosting": 1.3973, "extracao_forno": 700.8439481142, "porcentagem_caco": 10.002226628142, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 1},
  {"data": "2025-04-25", "cor": 0, "espessura": 8.0, "ext_boosting": 1.6005, "extracao_forno": 702.032548397562, "porcentagem_caco": 9.98529201530728, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-26", "cor": 0, "espessura": 8.0, "ext_boosting": 1.7549, "extracao_forno": 703.33718364331, "porcentagem_caco": 9.96677008271902, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-27", "cor": 0, "espessura": 8.0, "ext_boosting": 1.8022, "extracao_forno": 698.519152270116, "porcentagem_caco": 10.0355158154479, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-04-28", "cor": 0, "espessura": 8.0, "ext_boosting": 1.8023, "extracao_forno": 699.802291106822, "porcentagem_caco": 10.0171149610168, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-04-29", "cor": 0, "espessura": 8.0, "ext_boosting": 1.803, "extracao_forno": 702.213883737496, "porcentagem_caco": 9.98271347568585, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 0},
  {"data": "2025-04-30", "cor": 0, "espessura": 8.0, "ext_boosting": 1.801, "extracao_forno": 701.164091438783, "porcentagem_caco": 9.99765972843181, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 0},
  {"data": "2025-05-01", "cor": 0, "espessura": 8.0, "ext_boosting": 1.7999, "extracao_forno": 701.096395285213, "porcentagem_caco": 9.99862507800837, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 0},
  {"data": "2025-05-02", "cor": 0, "espessura": 8.0, "ext_boosting": 1.8016, "extracao_forno": 701.004721690124, "porcentagem_caco": 9.99993264396119, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 1, "autoclave": 0},
  {"data": "2025-05-03", "cor": 0, "espessura": 8.0, "ext_boosting": 1.8023, "extracao_forno": 699.505291072901, "porcentagem_caco": 10.021368086077, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-05-04", "cor": 0, "espessura": 8.0, "ext_boosting": 1.8036, "extracao_forno": 700.073447985429, "porcentagem_caco": 10.0132350686523, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-05-05", "cor": 0, "espessura": 8.0, "ext_boosting": 0.689, "extracao_forno": 700.60585295748, "porcentagem_caco": 10.0056258028798, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-06", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 699.123418185867, "porcentagem_caco": 10.026841924692, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-07", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 699.086556585488, "porcentagem_caco": 10.0273706223712, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-08", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 698.120389195209, "porcentagem_caco": 10.0412480547676, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-09", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 697.228099576186, "porcentagem_caco": 9.9680434627127, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-05-10", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 697.37935572186, "porcentagem_caco": 9.96588147179382, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 0},
  {"data": "2025-05-11", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 699.563378916139, "porcentagem_caco": 10.0205359675357, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-12", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 698.733542903546, "porcentagem_caco": 10.0324366436888, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-13", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 699.509702244859, "porcentagem_caco": 10.0213048904162, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-14", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 701.657766576732, "porcentagem_caco": 9.99062553558067, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-15", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 674.645706945424, "porcentagem_caco": 10.0052515424159, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-16", "cor": 0, "espessura": 8.0, "ext_boosting": 0.0, "extracao_forno": 653.148421891636, "porcentagem_caco": 9.95179622600148, "pot_boost": 3.0, "Prod_E": 0, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-17", "cor": 0, "espessura": 6.0, "ext_boosting": 0.0, "extracao_forno": 611.090907286899, "porcentagem_caco": 9.98214819965588, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1},
  {"data": "2025-05-18", "cor": 0, "espessura": 6.0, "ext_boosting": 0.0, "extracao_forno": 599.399563235682, "porcentagem_caco": 10.0100173040013, "pot_boost": 3.0, "Prod_E": 1, "Prod_L": 0, "autoclave": 1}
]"""

# Create Gradio interface
with gr.Blocks(title="Energy ML Cloud", theme=gr.themes.Default()) as app:
    
    gr.Markdown("# Energy ML Prediction System")
    gr.Markdown("Cloud deployment with embedded models")
    
    with gr.Row():
        with gr.Column():
            model_choice = gr.Radio(
                choices=[
                    "Threshold Detection",
                    "Energy Prediction (Random Forest)",
                    "Energy Prediction (XGBoost)"
                ],
                label="Select Model",
                value="Threshold Detection"
            )
            
            # Quick test data options
            gr.Markdown("### Quick Test Options")
            with gr.Row():
                load_week_btn = gr.Button("Load Week Test Data", size="sm")
                load_month_btn = gr.Button("Load Month Test Data", size="sm")
                clear_btn = gr.Button("Clear", size="sm")
            
            json_input = gr.Textbox(
                label="JSON Input",
                placeholder="Enter JSON data here...",
                lines=15,
                value=threshold_example
            )
            
            predict_btn = gr.Button("Make Prediction", variant="primary")
            
        with gr.Column():
            output = gr.Textbox(
                label="Prediction Result",
                lines=20,
                interactive=False
            )
    
    def update_example(choice):
        if "Threshold" in choice:
            return threshold_example
        else:
            return energy_example
    
    def load_week_data():
        return week_test_data
        
    def load_month_data():
        return month_test_data
        
    def clear_data():
        return ""
    
    model_choice.change(update_example, inputs=[model_choice], outputs=[json_input])
    load_week_btn.click(fn=load_week_data, outputs=json_input)
    load_month_btn.click(fn=load_month_data, outputs=json_input)  
    clear_btn.click(fn=clear_data, outputs=json_input)
    predict_btn.click(make_prediction, inputs=[model_choice, json_input], outputs=[output])
    
    with gr.Accordion("Model Information", open=False):
        gr.Markdown("""
        ## Available Models
        - **Threshold Detection**: Predict probability of exceeding 8.3 and 9.0 MWh
        - **Random Forest**: Energy consumption prediction (R² = 0.72)
        - **XGBoost**: Energy consumption prediction (R² = 0.56, winner model)
        
        ## Input Formats
        See examples that change when you select different models.
        """)

if __name__ == "__main__":
    app.launch(
        auth=("admin", "energy123"),
        share=True,
        ssr_mode=False
    )