import os 
import sys 
import shutil 
from ultralytics import YOLO 
from src.logger import logging 
from src.exception import PlantException
from src.entity.config_entity import ModelTrainerConfig
from src.entity.artifact_entity import ModelTrainerArtifact
from src.entity.artifact_entity import DataIngestionArtifact


class ModelTrainer:
    
    try:
        def __init__(self, model_trainer_config:ModelTrainerConfig, data_ingestion_artifact:DataIngestionArtifact):
            self.model_trainer_config = model_trainer_config
            self.data_ingestion_artifact = data_ingestion_artifact

        def initiate_model_trainer(self) -> ModelTrainerArtifact:
            logging.info(f"Removing and existing runs directory from previous training")
            os.system("rm -rf runs")

            model_config_file_name = self.model_trainer_config.weight_name.split('.')[0]
            print(model_config_file_name)
            
            os.system(f"yolo task=classify mode=train model={self.model_trainer_config.weight_name} \
                data={self.data_ingestion_artifact.dataset_path} epochs={self.model_trainer_config.no_epochs} \
                    imgsz=128 batch={self.model_trainer_config.batch_size} patience={self.model_trainer_config.patience}")
            
            os.makedirs("custom_model_weights", exist_ok=True)
            os.system("cp runs/classify/train/weights/best.pt custom_model_weights/")
            
            os.makedirs(self.model_trainer_config.model_trainer_dir, exist_ok=True)

            
            os.system(f"cp runs/classify/train/weights/best.pt {self.model_trainer_config.model_trainer_dir}/")

            model_trainer_artifact = ModelTrainerArtifact(trained_model_file_path="custom_model_weights/best.pt",)

            logging.info("Exited initiate_model_trainer method of ModelTrainer class")
            logging.info(f"Model trainer artifact: {model_trainer_artifact}")

            return model_trainer_artifact

    except Exception as e:
            raise PlantException(e, sys)