import os

# Quieter TensorFlow C++ logs: 0=all, 1=warn, 2=error, 3=fatal
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
# Disable oneDNN custom ops to avoid the startup info line
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

import json
import numpy as np
import pandas as pd
import streamlit as st
from PIL import Image, ImageOps
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import preprocess_input

tf.get_logger().setLevel("ERROR")

# ---------------- Streamlit page config ----------------
st.set_page_config(page_title="Weld Defect Classifier", layout="centered")

# ---- Mixed precision off on CPU to be safe
tf.keras.mixed_precision.set_global_policy("float32")

# ---- Session state
if "upload" not in st.session_state:
    st.session_state.upload = None
if "probs" not in st.session_state:
    st.session_state.probs = None

# ---- Local model file paths (inside THIS Space repo) --- #
MODEL_PATH = "model/final_single_phase.h5"
CONFIG_PATH = "model/training_config.json"
IMG_SIZE = (224, 224)

# ---- Pretty display labels
DISPLAY_LABELS = {
    "PO": "PO (Porosity)",
    "CR": "CR (Crack)",
    "LP": "LP (Lack of Penetration)",
    "ND": "ND (No Defect)",
}


def pretty_label(code: str) -> str:
    return DISPLAY_LABELS.get(code, code)


# ---- Confidence threshold for displaying the prediction
THRESHOLD = 0.65


# ---------- Custom layer to handle unknown "Cast" ----------
class CastLayer(tf.keras.layers.Layer):
    """
    Minimal custom layer used to replace the unknown 'Cast' layer
    when loading the saved model from H5.

    If the original object was effectively just casting to float32,
    this reproduces that behavior.
    """

    def __init__(self, dtype="float32", **kwargs):
        super().__init__(**kwargs)
        self.target_dtype = tf.dtypes.as_dtype(dtype)

    def call(self, inputs):
        return tf.cast(inputs, self.target_dtype)

    def get_config(self):
        config = super().get_config()
        config.update({"dtype": self.target_dtype.name})
        return config


@st.cache_resource
def load_model_and_config():
    """Loads model + config from local files inside the Space."""
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"Model file not found at: {MODEL_PATH}")
    if not os.path.exists(CONFIG_PATH):
        raise FileNotFoundError(f"Config file not found at: {CONFIG_PATH}")

    # Load the Keras model with custom_objects so that 'Cast' is known
    custom_objects = {
        "Cast": CastLayer,
    }

    model = tf.keras.models.load_model(
        MODEL_PATH,
        compile=False,
        custom_objects=custom_objects,
    )

    # Load class names from the config file
    with open(CONFIG_PATH, "r") as f:
        cfg = json.load(f)
    class_names = cfg.get("class_names", ["CR", "LP", "ND", "PO"])  # Fallback

    return model, class_names


def prepare_image(pil_img: Image.Image, target_size=(224, 224)) -> np.ndarray:
    """
    Letterbox (resize-with-pad) to target_size, fix EXIF orientation,
    convert to RGB, and apply ResNet50 preprocess_input.
    """
    # 1) Honor camera EXIF orientation
    img = ImageOps.exif_transpose(pil_img)

    # 2) Convert to RGB (handles grayscale seamlessly)
    img = img.convert("RGB")

    # 3) Resize with aspect ratio preserved + pad to target (letterbox)
    img = ImageOps.pad(
        img,
        target_size,
        method=Image.Resampling.BILINEAR,
        color=(0, 0, 0),
    )

    # 4) To array, add batch dimension, preprocess like training
    x = np.asarray(img, dtype=np.float32)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    return x


def upload_cb():
    st.session_state.upload = st.session_state.upload_k
    st.session_state.probs = None  # reset because the user has new input


def weld():
    st.title("🔎 Weld Defect Classifier")

    # Load resources from local files
    try:
        model, class_names = load_model_and_config()
    except Exception as e:
        st.error(f"Error loading model/config: {str(e)}")
        st.stop()
        return

    st.file_uploader(
        "Upload an image",
        type=["jpg", "jpeg", "png", "bmp", "webp"],
        accept_multiple_files=False,
        on_change=upload_cb,
        key="upload_k",
    )

    if st.session_state.upload and model is not None and class_names:
        pil_img = Image.open(st.session_state.upload)
        st.image(pil_img, caption="Input image")

        image_batch = prepare_image(pil_img, IMG_SIZE)

        if st.session_state.probs is None:
            with st.spinner("Running inference..."):
                probs = model.predict(image_batch, verbose=0)[0].astype(float)
                st.session_state.probs = probs

        # Build DataFrame and add pretty labels
        df = pd.DataFrame(
            {"class": class_names, "probability": st.session_state.probs}
        )
        df["label"] = df["class"].map(pretty_label)
        df = df.sort_values("probability", ascending=False).reset_index(drop=True)

        # Top-1 with thresholding
        top_prob = float(df.loc[0, "probability"])
        top_label = df.loc[0, "label"]
        display_label = "Unclear" if top_prob < THRESHOLD else top_label

        st.subheader("Prediction")
        st.markdown(f"**{display_label}** — Confidence: {top_prob:.3f}")

        # All probabilities
        st.subheader("All class probabilities")
        st.dataframe(
            df[["label", "probability"]]
            .rename(columns={"label": "Class"})
            .style.format({"probability": "{:.3f}"})
        )


def credits():
    st.title("Credits")
    st.markdown(
        """
[1] Benito Totino, Fanny Spagnolo, Stefania Perri,  
"RIAWELC: A Novel Dataset of Radiographic Images for Automatic Weld Defects Classification",  
ICMECE 2022, Barcelona, Spain.

[2] Stefania Perri, Fanny Spagnolo, Fabio Frustaci, Pasquale Corsonello,  
"Welding Defects Classification Through a Convolutional Neural Network",  
Manufacturing Letters, Elsevier.

[3] [Github RIAWELC](https://github.com/stefyste/RIAWELC)
"""
    )


# --- Main app navigation ---
weld_page = st.Page(weld, title="Weld Defect Classifier", default=True)
credit_page = st.Page(credits, title="Credits")

pg = st.navigation([weld_page, credit_page])
pg.run()