import gradio as gr
import tensorflow as tf
import numpy as np
import joblib
import logging

# Load the model and the scaler
model = tf.keras.models.load_model('my_lstm_model.h5')
scaler = joblib.load('scaler.pkl')

# Configure logging
logging.basicConfig(level=logging.DEBUG)

def predict(el_access_urban, el_demand, el_access_rural, population, net_imports, el_demand_pc, fin_support, el_from_gas, pop_no_el_access_total, urban_share, income_group_num, year, el_access_total, gdp_pc):
    try:
        # Extract the features from the form data
        features = {
            'el_access_urban': float(el_access_urban),
            'el_demand': float(el_demand),
            'el_access_rural': float(el_access_rural),
            'population': float(population),
            'net_imports': float(net_imports),
            'el_demand_pc': float(el_demand_pc),
            'fin_support': float(fin_support),
            'el_from_gas': float(el_from_gas),
            'pop_no_el_access_total': float(pop_no_el_access_total),
            'urban_share': float(urban_share),
            'income_group_num': float(income_group_num),
            'year': float(year),
            'el_access_total': float(el_access_total),
            'gdp_pc': float(gdp_pc)
        }

        # Calculate supply_rate and t_demand
        features['supply_rate'] = features['el_demand'] / features['el_access_total']
        features['t_demand'] = 100 * features['supply_rate']

        logging.debug(f"Extracted features: {features}")

        # Prepare features for prediction
        feature_values = [
            features['el_access_urban'],
            features['el_demand'],
            features['el_access_rural'],
            features['population'],
            features['net_imports'],
            features['el_demand_pc'],
            features['fin_support'],
            features['el_from_gas'],
            features['pop_no_el_access_total'],
            features['urban_share'],
            features['income_group_num'],
            features['year'],
            features['el_access_total'],
            features['gdp_pc'],
            features['supply_rate'],
            features['t_demand']
        ]

        logging.debug(f"Feature values before scaling: {feature_values}")

        # Ensure feature_values has the correct number of features
        if len(feature_values) != 16:
            logging.error("Incorrect number of features provided.")
            return {'error': 'Incorrect number of features provided.'}

        # Scale the features
        feature_values = np.array(feature_values).reshape(1, -1)
        feature_values_scaled = scaler.transform(feature_values)

        logging.debug(f"Scaled feature values: {feature_values_scaled}")

        # Reshape for LSTM
        feature_values_scaled = feature_values_scaled.reshape((feature_values_scaled.shape[0], 1, feature_values_scaled.shape[1]))

        logging.debug(f"Reshaped feature values for LSTM: {feature_values_scaled}")

        # Make the prediction
        prediction = model.predict(feature_values_scaled)
        prediction = np.abs(prediction[0][0])

        logging.debug(f"Prediction: {prediction}")

        return prediction
    except Exception as e:
        logging.error(f"An error occurred: {e}")
        return {'error': str(e)}

def fill_example():
    return [23, 12, 12, 2334354, 5, 2, 23, -1, 234567, 2, 2, 2018, 56, 6]

with gr.Blocks() as iface:
    with gr.Row():
        gr.Markdown("# Electricity Gap Prediction")
        gr.Markdown("Predict the electricity demand and supply gap for sub-Saharan African countries using socio-economic and energy data.")
    
    with gr.Row():
        with gr.Column():
            el_access_urban = gr.Number(label="Electricity Access Urban")
            el_demand = gr.Number(label="Electricity Demand")
            el_access_rural = gr.Number(label="Electricity Access Rural")
            population = gr.Number(label="Population")
            net_imports = gr.Number(label="Net Imports")
            el_demand_pc = gr.Number(label="Electricity Demand per Capita")
            fin_support = gr.Number(label="Financial Support")
            el_from_gas = gr.Number(label="Electricity from Gas")
        
        with gr.Column():
            pop_no_el_access_total = gr.Number(label="Population with No Electricity Access Total")
            urban_share = gr.Number(label="Urban Share")
            income_group_num = gr.Number(label="Income Group Number")
            year = gr.Number(label="Year")
            el_access_total = gr.Number(label="Electricity Access Total")
            gdp_pc = gr.Number(label="GDP per Capita")

    prediction_output = gr.Textbox(label="Predicted Electricity Gap")

    predict_btn = gr.Button("Predict")
    example_btn = gr.Button("Fill with Example Values")

    predict_btn.click(
        fn=predict,
        inputs=[el_access_urban, el_demand, el_access_rural, population, net_imports, el_demand_pc, fin_support, el_from_gas, pop_no_el_access_total, urban_share, income_group_num, year, el_access_total, gdp_pc],
        outputs=prediction_output
    )

    example_btn.click(
        fn=fill_example,
        inputs=[],
        outputs=[el_access_urban, el_demand, el_access_rural, population, net_imports, el_demand_pc, fin_support, el_from_gas, pop_no_el_access_total, urban_share, income_group_num, year, el_access_total, gdp_pc]
    )

iface.launch()