import gradio as gr
import tensorflow as tf
import joblib
import numpy as np
import zipfile
import os
import re

# Step 1: Unzip models only once
unzip_dir = "unzipped_models"
zip_file = "Models.zip"  # Ensure this matches exactly the file name uploaded in your Space repo

if not os.path.exists(unzip_dir):
    print("Extracting model zip file...")
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        zip_ref.extractall(unzip_dir)
    print("Extraction complete.")

# Step 2: Parse folders to dynamically populate dropdowns

model_root = os.path.join(unzip_dir, 'Models')  # Adjust if ZIP structure is different

activations = []
seeds_dict = dict()
neurons_dict = dict()

for act in os.listdir(model_root):
    act_path = os.path.join(model_root, act)
    if os.path.isdir(act_path) and not act.startswith("linear_models"):
        activations.append(act)
        seeds = []
        for seed_folder in os.listdir(act_path):
            seed_path = os.path.join(act_path, seed_folder)
            if os.path.isdir(seed_path):
                seeds.append(seed_folder)
                neuron_list = []
                for model_file in os.listdir(seed_path):
                    match = re.match(r"model_(\d+)\.keras", model_file)
                    if match:
                        neuron_list.append(int(match.group(1)))
                neurons_dict[(act, seed_folder)] = sorted(neuron_list)
        seeds_dict[act] = sorted(seeds)

activations = sorted(activations)

# Step 3: Prediction function
def predict(r, g, b, activation, seed, neurons):
    try:
        # Normalise R G B 
        r = r/255
        g = g/255
        b = b/255
        
        X = np.array([[r, g, b]])
        
        # Linear prediction (you can replace this with your actual linear model)
        lin_pred_rgb = (1.9221 * r) - (1.3817 * g) + (1.4058 * b) - 0.1318

        # ANN prediction
        keras_path = os.path.join(model_root, activation, seed, f"model_{neurons}.keras")
        if not os.path.exists(keras_path):
            raise FileNotFoundError(f"Model not found: {keras_path}")
        
        model = tf.keras.models.load_model(keras_path)
        ann_pred = model.predict(X)[0][0]

        # Rescale cholestrol concentration prediction in mM
        return ann_pred*50, lin_pred_rgb*50

    except Exception as e:
        return f"Error: {str(e)}", ""

# Step 4: Dynamic UI update functions (Gradio 4.x compliant)
def update_seeds(activation):
    return gr.update(choices=seeds_dict[activation], value=seeds_dict[activation][0])

def update_neurons(activation, seed):
    neurons = neurons_dict[(activation, seed)]
    return gr.update(choices=neurons, value=neurons[0])

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# ANN vs Linear Model Predictor")
    gr.Markdown("Dynamically select models and predict cholesterol concentration.")
    
    with gr.Row():
        r = gr.Number(label="R (0 -255)")
        g = gr.Number(label="G (0 -255)")
        b = gr.Number(label="B (0 -255)")

    with gr.Row():
        activation = gr.Dropdown(choices=activations, label="Activation Function", interactive=True)
        seed = gr.Dropdown(choices=seed, label="Seed", interactive=True)
        neurons = gr.Dropdown(choices=neurons, label="Neurons", interactive=True)

    activation.change(update_seeds, inputs=[activation], outputs=[seed])
    seed.change(update_neurons, inputs=[activation, seed], outputs=[neurons])

    with gr.Row():
        btn = gr.Button("Predict")
    
    with gr.Row():
        ann_output = gr.Text(label="Cholestrol Conentration (mM) - ANN Model Prediction ")
        lin_rgb_output = gr.Text(label="Cholestrol Conentration (mM) - Linear Model Prediction")

    btn.click(
        fn=predict, 
        inputs=[r, g, b, activation, seed, neurons],
        outputs=[ann_output, lin_rgb_output]
    )

if __name__ == "__main__":
    demo.launch()