import csv
import os
import pathlib
import warnings

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
warnings.filterwarnings("ignore", message=".*tf.lite.Interpreter is deprecated.*")

import numpy as np
import tensorflow as tf

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

SAVEDMODEL_PATH = pathlib.Path("support_score_savedmodel")
REPRESENTATIVE_PATH = pathlib.Path("support_score_representative.csv")
FLOAT_TFLITE_PATH = pathlib.Path("support_score_float.tflite")
INT8_TFLITE_PATH = pathlib.Path("support_score_int8.tflite")
SAMPLE = np.array([[0.72, 0.24, 0.76, 0.82]], dtype=np.float32)


def load_representative_rows():
    rows = []

    with REPRESENTATIVE_PATH.open("r", newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            rows.append(
                [
                    float(row["feature_a"]),
                    float(row["feature_b"]),
                    float(row["feature_c"]),
                    float(row["feature_d"]),
                ]
            )

    return np.asarray(rows, dtype=np.float32)


representative_rows = load_representative_rows()

float_converter = tf.lite.TFLiteConverter.from_saved_model(str(SAVEDMODEL_PATH))
FLOAT_TFLITE_PATH.write_bytes(float_converter.convert())


def representative_dataset():
    for row in representative_rows:
        yield [row.reshape(1, 4)]


int8_converter = tf.lite.TFLiteConverter.from_saved_model(str(SAVEDMODEL_PATH))
int8_converter.optimizations = [tf.lite.Optimize.DEFAULT]
int8_converter.representative_dataset = representative_dataset
int8_converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
int8_converter.inference_input_type = tf.int8
int8_converter.inference_output_type = tf.int8
INT8_TFLITE_PATH.write_bytes(int8_converter.convert())

interpreter = tf.lite.Interpreter(model_path=str(INT8_TFLITE_PATH))
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()[0]
output_details = interpreter.get_output_details()[0]
input_scale, input_zero_point = input_details["quantization"]
output_scale, output_zero_point = output_details["quantization"]

quantized_sample = np.round(SAMPLE / input_scale + input_zero_point).astype(np.int8)
interpreter.set_tensor(input_details["index"], quantized_sample)
interpreter.invoke()
raw_output = interpreter.get_tensor(output_details["index"])
int8_score = (raw_output.astype(np.float32) - output_zero_point) * output_scale

float_interpreter = tf.lite.Interpreter(model_path=str(FLOAT_TFLITE_PATH))
float_interpreter.allocate_tensors()
float_input = float_interpreter.get_input_details()[0]
float_output = float_interpreter.get_output_details()[0]
float_interpreter.set_tensor(float_input["index"], SAMPLE)
float_interpreter.invoke()
float_score = float_interpreter.get_tensor(float_output["index"])

print(f"TensorFlow {tf.__version__}")
print(f"saved_model_dir={SAVEDMODEL_PATH.name}")
print(f"representative_rows={len(representative_rows)}")
print(f"float_model_file={FLOAT_TFLITE_PATH.name}")
print(f"int8_model_file={INT8_TFLITE_PATH.name}")
print(f"float_model_bytes={FLOAT_TFLITE_PATH.stat().st_size}")
print(f"int8_model_bytes={INT8_TFLITE_PATH.stat().st_size}")
print(f"input_dtype={np.dtype(input_details['dtype']).name}")
print(f"output_dtype={np.dtype(output_details['dtype']).name}")
print(f"input_quantization=scale:{input_scale:.8f},zero_point:{input_zero_point}")
print(f"output_quantization=scale:{output_scale:.8f},zero_point:{output_zero_point}")
print(f"float_sample_output={float_score[0][0]:.6f}")
print(f"int8_sample_output={int8_score[0][0]:.6f}")
print(f"absolute_difference={abs(float_score[0][0] - int8_score[0][0]):.6f}")