import tensorflow as tf
import numpy as np
import os
import cv2

# 1. Path setup
saved_model_path = 'best_saved_model'
# Try to find where saved_model.pb actually is
target_saved_model = saved_model_path
if not os.path.exists(os.path.join(saved_model_path, 'saved_model.pb')):
    for root, dirs, files in os.walk(saved_model_path):
        if 'saved_model.pb' in files:
            target_saved_model = root
            break

print(f"Using SavedModel at: {target_saved_model}")

# 2. Representative Dataset
def representative_dataset():
    img_dir = 'unified_dataset/images/val'
    count = 0
    for f in os.listdir(img_dir):
        if f.endswith(('.jpg', '.jpeg', '.png')) and count < 50:
            img = cv2.imread(os.path.join(img_dir, f))
            if img is None: continue
            img = cv2.resize(img, (640, 640))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float32) / 255.0
            img = np.expand_dims(img, axis=0)
            yield [img]
            count += 1

# 3. Converter
converter = tf.lite.TFLiteConverter.from_saved_model(target_saved_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

try:
    tflite_model = converter.convert()
    output_path = 'best_int8.tflite'
    with open(output_path, 'wb') as f:
        f.write(tflite_model)
    print(f"Success: {output_path} generated.")
except Exception as e:
    print(f"Conversion failed: {e}")