import os
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import os
import pickle
import numpy as np
import tensorflow as tf

from tensorflow.keras import Model
from tensorflow.keras.layers import Input, Embedding, Dropout, LSTM, Dense, Add
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.xception import Xception, preprocess_input

MODEL_WEIGHTS = r"D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\Exercise-7\Xception_LSTM_OCT_DATASET.hdf5"
TOKENIZER_PKL = r"D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\Exercise-7\tokenizer.pkl"
ONE_IMAGE     = r"D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\OCT_Dataset\OCT_Dataset\val\DME\dme_val_1001.jpg"

MAX_LEN    = 107
EMB_DIM    = 256
LSTM_UNITS = 256
PDROP      = 0.3
FEAT_DIM   = 2048
IMG_SIZE   = (299, 299)

with open(TOKENIZER_PKL, "rb") as f:
    tokenizer = pickle.load(f)
VOCAB_SIZE = len(tokenizer.word_index) + 1
index_word = {v: k for k, v in tokenizer.word_index.items()}

def build_caption_model(max_len: int,
                        vocab_size: int,
                        feat_dim: int = FEAT_DIM,
                        emb_dim: int = EMB_DIM,
                        lstm_units: int = LSTM_UNITS,
                        pdrop: float = PDROP) -> Model:
    img_in = Input(shape=(feat_dim,), name="image_features")
    x_img = Dropout(pdrop, name="img_dropout")(img_in)
    x_img = Dense(lstm_units, activation="relu", name="img_dense")(x_img)

    txt_in = Input(shape=(max_len,), name="seq_input")
    x_txt = Embedding(vocab_size, emb_dim, mask_zero=True, name="embedding")(txt_in)
    x_txt = Dropout(pdrop, name="txt_dropout")(x_txt)
    x_txt = LSTM(lstm_units, name="lstm")(x_txt)

    merged = Add(name="add")([x_img, x_txt])
    x = Dense(lstm_units, activation="relu", name="fc")(merged)
    out = Dense(vocab_size, activation="softmax", name="softmax")(x)
    return Model(inputs=[img_in, txt_in], outputs=out, name="caption_model")

model = build_caption_model(MAX_LEN, VOCAB_SIZE)
model.load_weights(MODEL_WEIGHTS)

xception = Xception(weights="imagenet", include_top=False, pooling="avg")

def extract_features_xception(path: str) -> np.ndarray:
    im = load_img(path, target_size=IMG_SIZE)
    x = img_to_array(im)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    feat = xception.predict(x, verbose=0)
    return feat

def word_for_id(integer: int) -> str | None:
    return index_word.get(integer)

def generate_caption(photo_feat_batched: np.ndarray, max_length: int) -> str:
    words = ["startseq"]
    for _ in range(max_length):
        seq = tokenizer.texts_to_sequences([" ".join(words)])[0]
        seq = pad_sequences([seq], maxlen=max_length, padding="post")
        yhat = model.predict([photo_feat_batched, seq], verbose=0)
        next_id = int(np.argmax(yhat))
        w = word_for_id(next_id)
        if w is None:
            break
        words.append(w)
        if w == "endseq":
            break
    return " ".join([w for w in words if w not in ("startseq", "endseq")])

if not os.path.exists(ONE_IMAGE):
    raise FileNotFoundError(f"Image not found: {ONE_IMAGE}")

feat = extract_features_xception(ONE_IMAGE)
caption = generate_caption(feat, MAX_LEN)
print("Image:", ONE_IMAGE)
print("Caption:", caption if caption else "[empty]")

