import traceback
from pathlib import Path

log_path = Path('load_predict_log.txt')
with log_path.open('w', encoding='utf-8') as f:
    def log(*args, **kwargs):
        print(*args, file=f, **kwargs)
        f.flush()

    try:
        log('Starting model load & predict test')
        import tensorflow as tf
        import numpy as np
        from PIL import Image
        import os

        img_path = Path('data/UTKFace/53_1_1_20170110122449716.jpg.chip.jpg')
        log('Image path:', str(img_path))

        # Try HDF5 / .h5 first
        h5_path = Path('final_model.h5')
        keras_path = Path('saved_model_age_regressor.keras')
        saved_model_dir = Path('saved_model_age_regressor')

        if h5_path.exists():
            try:
                log('Attempting to load HDF5 model:', str(h5_path))
                m = tf.keras.models.load_model(str(h5_path), compile=False)
                log('Loaded HDF5 model:', type(m))
                img = Image.open(img_path).convert('RGB').resize((224,224))
                x = np.expand_dims(np.array(img, dtype=np.float32)/255.0, 0)
                log('Running predict on HDF5 model...')
                pred = m.predict(x)
                log('Prediction result (HDF5):', pred.tolist())
            except Exception:
                log('Exception while loading/predicting from HDF5:')
                traceback.print_exc(file=f)
        elif keras_path.exists():
            try:
                log('Attempting to load Keras native file:', str(keras_path))
                m = tf.keras.models.load_model(str(keras_path), compile=False)
                log('Loaded Keras native model:', type(m))
                img = Image.open(img_path).convert('RGB').resize((224,224))
                x = np.expand_dims(np.array(img, dtype=np.float32)/255.0, 0)
                pred = m.predict(x)
                log('Prediction result (KERAS):', pred.tolist())
            except Exception:
                log('Exception while loading/predicting from Keras file:')
                traceback.print_exc(file=f)
        elif saved_model_dir.exists():
            try:
                log('HDF5/.keras not found; attempting to wrap TF SavedModel using TFSMLayer...')
                try:
                    from keras.layers import TFSMLayer
                except Exception as e:
                    log('TFSMLayer import failed:', e)
                    raise
                # Build wrapper model
                inputs = tf.keras.Input(shape=(224,224,3))
                tfsml = TFSMLayer(str(saved_model_dir), call_endpoint='serving_default')
                outputs = tfsml(inputs)
                wrapper = tf.keras.Model(inputs, outputs)
                log('Wrapper model created; running predict...')
                img = Image.open(img_path).convert('RGB').resize((224,224))
                x = np.expand_dims(np.array(img, dtype=np.float32)/255.0, 0)
                pred = wrapper.predict(x)
                # The SavedModel serving signature can return a dict mapping names->arrays
                if isinstance(pred, dict):
                    log('Prediction returned a dict with keys:', list(pred.keys()))
                    import numpy as _np
                    for k, v in pred.items():
                        try:
                            arr = _np.array(v)
                            log(f"Output '{k}': shape={arr.shape} values={arr.flatten()[:10].tolist()}")
                        except Exception as _e:
                            log(f"Could not convert output '{k}' to numpy array:", _e)
                else:
                    try:
                        log('Prediction result (wrapped SavedModel):', pred.tolist())
                    except Exception:
                        log('Prediction result (wrapped SavedModel) type:', type(pred))
            except Exception:
                log('Exception while wrapping/using SavedModel:')
                traceback.print_exc(file=f)
        else:
            log('No model file found: looked for final_model.h5, saved_model_age_regressor.keras, or saved_model_age_regressor/')

        log('Finished load & predict test')
    except Exception:
        traceback.print_exc(file=f)
        log('Top-level exception')