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.pkl'):
                with open('threshold_model_83.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.pkl'):
                with open('threshold_model_90.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)
            
            # Parse input data
            date_obj = datetime.strptime(data['data'], '%Y-%m-%d')
            
            # Color mapping
            color_mapping = {0: 'incolor', 1: 'verde', 2: 'cinza', 3: 'bronze'}
            cor_str = color_mapping.get(data['cor'], 'incolor')
            
            # Create input features
            input_data = {
                'boosting': data['pot_boost'],
                'espessura': data['espessura'],
                'extracao_forno': data['extracao_forno'],
                'porcentagem_caco': data['porcentagem_caco'],
                'cor': cor_str,
                'prod_e': data['Prod_E'],
                'prod_l': data['Prod_L'],
                '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
            }
            
            # Convert to DataFrame
            input_df = pd.DataFrame([input_data])
            
            # Preprocess
            X_processed = self.threshold_preprocessor.transform(input_df)
            
            # Make predictions
            prob_83 = self.threshold_model_83.predict_proba(X_processed)[0][1] if len(self.threshold_model_83.classes_) > 1 else 0.0
            pred_83 = int(prob_83 > 0.5)
            
            prob_90 = self.threshold_model_90.predict_proba(X_processed)[0][1] if len(self.threshold_model_90.classes_) > 1 else 0.0
            pred_90 = int(prob_90 > 0.5)
            
            # Format response
            next_date = (date_obj + pd.Timedelta(days=1)).strftime('%Y-%m-%d')
            
            result = {
                "predictions": {
                    "prediction_1": [
                        {
                            "datetime": data['data'],
                            "probabilidade_de_estouro": float(prob_83),
                            "estouro_previsto": pred_83
                        },
                        {
                            "datetime": next_date,
                            "probabilidade_de_estouro": float(prob_83 * 0.98),
                            "estouro_previsto": int(prob_83 * 0.98 > 0.5)
                        }
                    ],
                    "prediction_2": [
                        {
                            "datetime": data['data'],
                            "probabilidade_de_estouro": float(prob_90),
                            "estouro_previsto": pred_90
                        },
                        {
                            "datetime": next_date,
                            "probabilidade_de_estouro": float(prob_90 * 0.99),
                            "estouro_previsto": int(prob_90 * 0.99 > 0.5)
                        }
                    ]
                }
            }
            
            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')
                boosting_val = float(item['boosting'].replace(',', '.'))
                extracao_val = float(item['extracao_forno'].replace(',', '.'))
                
                # Create features
                input_data = {
                    'boosting': boosting_val,
                    'espessura': item['espessura'],
                    'extracao_forno': extracao_val,
                    'porcentagem_caco': item['porcentagem_caco'],
                    'cor': item['cor'].lower(),
                    'prod_e': 1,
                    'prod_l': 1,
                    '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')
                boosting_val = float(item['boosting'].replace(',', '.'))
                extracao_val = float(item['extracao_forno'].replace(',', '.'))
                
                # Create features
                input_data = {
                    'boosting': boosting_val,
                    'espessura': item['espessura'],
                    'extracao_forno': extracao_val,
                    'porcentagem_caco': item['porcentagem_caco'],
                    'cor': item['cor'].lower(),
                    'prod_e': 1,
                    'prod_l': 1,
                    '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
                }
                
                # 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
  }
]"""

# 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"
            )
            
            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
    
    model_choice.change(update_example, inputs=[model_choice], 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
    )