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

from rule_based_quality import RuleBasedQualityEvaluator
from preprocess_cnn import preprocess_for_cnn
from preprocess_quality import GeometricFeatureExtractor

cnn_model = tf.keras.models.load_model("custom_cnn_shape.keras")
xgb_quality = joblib.load("xgb_quality.pkl")
scaler = joblib.load("scaler.pkl")
le_quality = joblib.load("le_quality.pkl")
le_shape = joblib.load("le_shape.pkl")
feature_extractor = GeometricFeatureExtractor()
rule_evaluator = RuleBasedQualityEvaluator()
SHAPES = ['Triangle', 'Square', 'Circle', 'Rectangle']

def predict(image):
    response = {}

    x = preprocess_for_cnn(image)
    shape_probs = cnn_model.predict(x, verbose=0)[0]

    shape_idx = int(np.argmax(shape_probs))
    shape_label = le_shape.inverse_transform([shape_idx])[0]

    response["shape"] = shape_label
    response["shape_confidence"] = float(shape_probs[shape_idx])

    features = feature_extractor.extract_features(image)

    if features is None:
        response["quality"] = "unknown"
        response["quality_confidence"] = 0.0
        response["rl_quality"] = "unknown"
        response["rl_confidence"] = 0.0
        return response

    features_scaled = scaler.transform(features.reshape(1, -1))
    quality_probs = xgb_quality.predict_proba(features_scaled)[0]
    q_idx = int(np.argmax(quality_probs))

    response["quality"] = le_quality.inverse_transform([q_idx])[0]
    response["quality_confidence"] = float(quality_probs[q_idx])

    rl_label, rl_conf = rule_evaluator.evaluate(
        feature_extractor.last_feature_dict,  
        shape_label
    )

    response["rl_quality"] = rl_label
    response["rl_confidence"] = float(rl_conf)

    return response


gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="Shape & Quality Recognition System",
    description="CNN (Keras) for shape recognition + XGBoost for drawing quality assessment"
).launch()