import os
import subprocess
from loguru import logger
import shutil


class ModelOptimization:
    def __init__(self, num_class, weight_path, image_size=416):
        self.num_class = num_class
        self.image_size = image_size
        self.weight_path = weight_path

    def change_configurations(self):
        logger.info(f"Provided number of classes and image size are :  {self.num_class} and {self.image_size}")
        try:

            with open('yolov5/yololayer.h', 'r') as file:
                # read a list of lines into data
                data = file.readlines()

            data[19] = f"    static constexpr int CLASS_NUM = {self.num_class};\n"
            data[20] = f"    static constexpr int INPUT_H = {self.image_size};\n"
            data[21] = f"    static constexpr int INPUT_W = {self.image_size};\n"

            # and write everything back
            with open('yolov5/yololayer.h', 'w') as file:
                file.writelines(data)
            logger.info("Successfully changed configurations")
        except Exception as e:
            logger.info(f"Failed to change configurations :  {e}")

    def optimize_model(self):
        try:


            current_directory = os.getcwd()

            logger.info(f"Current directory is :  {current_directory}")
            build_path = os.path.join(current_directory, "yolov5", "build")
            if os.path.isdir('yolov5/build'):
                logger.info("build directory exists. Removing build directory!!")
                shutil.rmtree('yolov5/build')
            os.mkdir(build_path)
            weight_name_with_extension = os.path.basename(self.weight_path)
            weight_name, extension = os.path.splitext(weight_name_with_extension)
            src = f"{current_directory}/yolov5/build/"
            shutil.copy(self.weight_path, src)
            logger.info(f"Created build folder")
            os.chdir('yolov5/build')
            logger.info("Running CMake command")
            subprocess.run(['cmake', '..'])
            logger.info("Running Make command")
            subprocess.run(['make'])
            logger.info("Optimizing model")
            engine_name = weight_name + ".engine"
            subprocess.run(["sudo", "./yolov5", "-s", weight_name_with_extension, engine_name, "c", "0.33", "0.50"])
            engine_file_path = os.path.join(src, engine_name)

            shutil.copy(engine_file_path, "/app/data/")

        except Exception as e:
            logger.info(f"Failed to optimized model :  {e}")





