app.py

import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
# Assuming TKAN and TKAT are available after installing the respective packages
from tkan import TKAN
# If TKAT is from a different library, import it similarly
try:
    from tkat import TKAT
except ImportError:
    print("TKAT library not found. If your model uses TKAT, make sure the library is installed.")
    TKAT = None # Set to None if TKAT is not available

from tensorflow.keras.utils import custom_object_scope
import pickle # Used for saving/loading the scaler

# --- Your MinMaxScaler Class (Copied from Notebook) ---
class MinMaxScaler:
    def __init__(self, feature_axis=None, minmax_range=(0, 1)):
        """
        Initialize the MinMaxScaler.
        Args:
        feature_axis (int, optional): The axis that represents the feature dimension if applicable.
                                      Use only for 3D data to specify which axis is the feature axis.
                                      Default is None, automatically managed based on data dimensions.
        """
        self.feature_axis = feature_axis
        self.min_ = None
        self.max_ = None
        self.scale_ = None
        self.minmax_range = minmax_range # Default range for scaling (min, max)

    def fit(self, X):
        """
        Fit the scaler to the data based on its dimensionality.
        Args:
        X (np.array): The data to fit the scaler on.
        """
        if X.ndim == 3 and self.feature_axis is not None:  # 3D data
            axis = tuple(i for i in range(X.ndim) if i != self.feature_axis)
            self.min_ = np.min(X, axis=axis)
            self.max_ = np.max(X, axis=axis)
        elif X.ndim == 2:  # 2D data
            self.min_ = np.min(X, axis=0)
            self.max_ = np.max(X, axis=0)
        elif X.ndim == 1:  # 1D data
            self.min_ = np.min(X)
            self.max_ = np.max(X)
        else:
            raise ValueError("Data must be 1D, 2D, or 3D.")

        self.scale_ = self.max_ - self.min_
        return self

    def transform(self, X):
        """
        Transform the data using the fitted scaler.
        Args:
        X (np.array): The data to transform.
        Returns:
        np.array: The scaled data.
        """
        X_scaled = (X - self.min_) / self.scale_
        X_scaled = X_scaled * (self.minmax_range[1] - self.minmax_range[0]) + self.minmax_range[0]
        return X_scaled

    def fit_transform(self, X):
        """
        Fit to data, then transform it.
        Args:
        X (np.array): The data to fit and transform.
        Returns:
        np.array: The scaled data.
        """
        return self.fit(X).transform(X)

    def inverse_transform(self, X_scaled):
        """
        Inverse transform the scaled data to original data.
        Args:
        X_scaled (np.array): The scaled data to inverse transform.
        Returns:
        np.array: The original data scale.
        """
        X = (X_scaled - self.minmax_range[0]) / (self.minmax_range[1] - self.minmax_range[0])
        X = X * self.scale_ + self.min_
        return X
# --- End of MinMaxScaler Class ---


# --- Configuration ---
MODEL_PATH = "best_model_TKAN_nahead_1 (2).keras"
INPUT_SCALER_PATH = "input_scaler.pkl" # You need to save your X_scaler to this file
# TARGET_SCALER_PATH = "target_scaler.pkl" # You might also need this if you predict scaled target
SEQUENCE_LENGTH = 24 # Matches the notebook
NUM_INPUT_FEATURES = 5 # ['calculated_aqi', 'temp', 'pm25', 'pm10', 'co']
N_AHEAD = 1 # Matches the notebook

# --- Load Model and Scalers ---
custom_objects = {"TKAN": TKAN}
if TKAT is not None:
    custom_objects["TKAT"] = TKAT
# Also add your custom MinMaxScaler to custom_objects for loading the scaler object
custom_objects["MinMaxScaler"] = MinMaxScaler


model = None
input_scaler = None
# target_scaler = None # Load if needed

try:
    with custom_object_scope(custom_objects):
        model = load_model(MODEL_PATH)
    print("Model loaded successfully!")
    model.summary()

except Exception as e:
    print(f"Error loading model from {MODEL_PATH}: {e}")
    import traceback
    traceback.print_exc()
    import sys
    sys.exit("Failed to load the model. Exiting.")

try:
    # When loading the scaler, you need custom_objects if it's your custom class
    with custom_object_scope(custom_objects):
        with open(INPUT_SCALER_PATH, 'rb') as f:
            input_scaler = pickle.load(f)
    print(f"Input scaler loaded successfully from {INPUT_SCALER_PATH}")

    # If you also scaled your target variable and need to inverse transform the prediction,
    # load the target scaler here as well.
    # with custom_object_scope(custom_objects):
    #     with open(TARGET_SCALER_PATH, 'rb') as f:
    #         target_scaler = pickle.load(f)
    # print(f"Target scaler loaded successfully from {TARGET_SCALER_PATH}")

except FileNotFoundError as e:
    print(f"Error loading scaler: {e}. Make sure your scaler files are in the correct path.")
    import sys
    sys.exit("Failed to load scaler(s). Exiting.")
except Exception as e:
    print(f"Error loading scaler: {e}")
    import traceback
    traceback.print_exc()
    import sys
    sys.exit("Failed to load scaler(s). Exiting.")


# --- Data Preparation (get_latest_data_sequence needs implementation) ---

def get_latest_data_sequence(sequence_length, num_features):
    """
    Retrieves the latest sequence of data for the required features.

    This function needs to be implemented based on your data source in the
    deployment environment. It should return a numpy array with shape
    (sequence_length, num_features).

    Args:
        sequence_length (int): The length of the historical sequence required.
        num_features (int): The number of features in each time step.

    Returns:
        np.ndarray: A numpy array containing the historical data sequence.
                    Shape: (sequence_length, num_features)
    """
    print("WARNING: Using dummy data sequence. Implement get_latest_data_sequence.")
    # --- REPLACE THIS WITH YOUR ACTUAL DATA RETRIEVAL LOGIC ---
    # Example: Load from a database, fetch from an API, read from a file.
    # The data should be in the correct order (oldest to newest time step).
    # The columns should be in the order ['calculated_aqi', 'temp', 'pm25', 'pm10', 'co'].

    # For now, returning a placeholder with the correct shape.
    dummy_data = np.zeros((sequence_length, num_features))
    # Populate dummy_data with some values for testing if you can load historical data
    # For example, load a small sample of your training data's X_test_unscaled
    # and use it here temporarily to verify the scaling and prediction pipeline.
    return dummy_data
    # --- END OF PLACEHOLDER ---


# --- Define Predict Function ---

def predict(): # Modify inputs as needed based on how you get data
    """
    Retrieves the latest data sequence, preprocesses it, and makes a prediction.

    The Gradio interface will need to trigger this function.
    The input parameters to this function might change depending on how you
    provide the necessary historical data sequence via the Gradio interface.
    """
    if model is None or input_scaler is None:
         return "Model or scaler not loaded. Check logs."

    # 1. Get the latest historical data sequence
    latest_data_sequence = get_latest_data_sequence(SEQUENCE_LENGTH, NUM_INPUT_FEATURES)

    # Ensure the retrieved data has the correct shape
    if latest_data_sequence.shape != (SEQUENCE_LENGTH, NUM_INPUT_FEATURES):
        return f"Error: Retrieved data has incorrect shape {latest_data_sequence.shape}. Expected ({SEQUENCE_LENGTH}, {NUM_INPUT_FEATURES})."

    # 2. Scale the data sequence using the loaded input scaler
    # Your MinMaxScaler from the notebook had feature_axis=2 for 3D data (samples, sequence, features).
    # So, for a single sequence (2D array), you'll likely need to add a batch dimension (1) before scaling.
    latest_data_sequence_with_batch = latest_data_sequence[np.newaxis, :, :]
    scaled_input_data = input_scaler.transform(latest_data_sequence_with_batch)

    # 3. Perform prediction
    # The model expects input shape (batch_size, sequence_length, num_features)
    output = model.predict(scaled_input_data)

    # 4. Process the output
    # The output shape is (batch_size, n_ahead). Since n_ahead=1, shape is (batch_size, 1).
    predicted_scaled_value = output[0][0] # Get the first prediction for the first sample

    # 5. Inverse transform the prediction if the target was scaled
    # If you scaled the target variable (calculated_aqi) before training,
    # you need to inverse transform the prediction back to the original scale.
    # This requires saving and loading the target_scaler as well and using it here.

    # Example if you need to inverse transform the target:
    # if target_scaler is not None:
    #     predicted_original_scale = target_scaler.inverse_transform(np.array([[predicted_scaled_value]]))[0][0]
    # else:
    #     predicted_original_scale = predicted_scaled_value
    # predicted_value = predicted_original_scale

    # For now, assuming the model outputs directly in the desired scale or
    # you handle inverse transformation elsewhere if needed.
    predicted_value = predicted_scaled_value # Adjust this if inverse transformation is needed

    return float(predicted_value)


# --- Gradio Interface ---
# The Gradio interface needs to allow the user to provide the necessary
# historical data for the `predict` function. The current inputs (pm25, pm10, co, temp)
# are NOT used by the predict function as written, since predict calls `get_latest_data_sequence`.
# You need to decide how the user will provide the historical data.

# Option 1: No direct inputs to `predict` function, it fetches data internally.
# In this case, the Gradio interface might just have a button to trigger the prediction.
interface = gr.Interface(
    fn=predict,
    inputs=None, # `predict` function doesn't take direct inputs from Gradio
    outputs=gr.Number(label=f"Predicted AQI (Next {N_AHEAD} Hour(s))")
)

# Option 2: Modify `predict` to accept some parameters that help retrieve the data.
# For example, if you pass the current time, and `get_latest_data_sequence` uses that
# to find the last 24 hours ending at that time.

# Option 3: Design a more complex Gradio interface to input the full sequence.
# This is more complex but directly aligns with the model input.

# Choose the Gradio interface definition that matches how your `predict` function
# will receive the data it needs. Option 1 is shown below.

# --- Launch Gradio Interface ---
if __name__ == "__main__":
    interface.launch()