import os
import sys
from pathlib import Path
import importlib
import torchvision
import wget
from ultralytics import YOLO
import torchvision

import torch    
import cv2
import numpy as np
import matplotlib.pyplot as plt

import torchvision.transforms as transforms
from PIL import Image
import io


FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv8 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
# ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative


sys.path.insert(0, os.path.abspath(
    os.path.join(os.path.dirname(__file__), './../../')
))
from ecos_core.base_model import BaseModel
from yolo_cam.eigen_cam import EigenCAM
from yolo_cam.utils.image import show_cam_on_image

class Model(BaseModel):
    NAME="yolo8"
    def __init__(self,opt):
        self.opt=opt
        # self.check_weight_path = os.path.join(str(ROOT), self.opt.weights)

    def get_transform(self) -> torchvision.transforms.transforms:
        """Get image data transform function for preprocessing

        Returns:
            transform (func): transform function.
                    This function will take input as image path and output 
                    - raw_image: image numpy mat
                    - image_transform: image tensor (pytorch)
                    Example: torchvision.transforms.transforms: data transforms function
        """
        pass
    
    def build_model(self,model_path=None):
        """
        Load Yolo model for retraining 
        Args:
            model_path(str): Custom path for loading weight. Default: model_path="model/yolov8{version}.pt"

        Returns:
            model(YOLO): created YOLO model

        """

        if model_path is not None:
            self.path = model_path
            print("Using custom model weights:", self.path)
        else:
            version=self.opt.version

            # Define a dictionary mapping YOLO version letters to their corresponding model names
            #model_names_pt = {'n': 'yolov8n.pt', 's': 'yolov8s.pt', 'm': 'yolov8m.pt', 'l':'yolov8l.pt','x':'yolov8x.pt'}
            #Construct the path to the corresponding  YOLO model file
            #self.path = f'model/yolov8{version}.pt'
            self.path = self.opt.weights
            print("Using default model weights:", self.path)

        if os.path.getsize(self.path) > 0:
          path = self.path
          print("Loading model weights from file:", path)
        else:
          print("Downloading model weights")
          URL = "https://ecos-ai-test-upload.s3.ap-northeast-1.amazonaws.com/models/yolo8/pretrained_weight/model_weight.pt"
          path = wget.download(URL)
          print("Downloaded model weights to file:", path)

        self.model=YOLO(self.path)

        return self.model



    def fit(self):
        """
        Args:
          ****minimun value******
          data(str): input data.yaml
          epochs(int): number of epochs
          imgsz(int): Image size (640)
          device:0,1,2,3 for GPU or device:cpu
          batch_size: a number of samples processed before the model is updated
          **************************************************************

        Returns:
            model_path(str): default "runs/detect/train/weights/best.pt"
        """

        return self.model.train(data=self.opt.data, epochs=self.opt.epochs,
                                imgsz=self.opt.imgsz, device=self.opt.device,
                                batch=self.opt.batch_size)

    def validate(self):
        """
        Args:
        Returns:
            metrics(list): Validated results
        """
        
        # Validate the model
        metrics =self.model.val()   # no arguments needed, dataset and settings remembered
        metrics.box.map    # map50-95
        metrics.box.map50  # map50
        metrics.box.map75  # map75
        metrics.box.maps   # a list contains map50-95 of each category
        if self.opt.task == "segment":
            metrics.seg.map    # map50-95(M)
            metrics.seg.map50  # map50(M)
            metrics.seg.map75  # map75(M)
            metrics.seg.maps   # a list contains map50-95(M) of each category
        print("validation succeed")
        return metrics

    def predict(self, source=None):
        """
          Args:
              source(str): file or folder for testing
          Returns:
              folder_path(str): runs/detect/predict
        """
        if source is not None:
          self.source = source
          print("custom Image or image folder ", self.source )
        else:
          self.source  = os.path.join(str(ROOT), self.opt.source)
          print("Using default model weights:", self.source )
        return self.model(self.source,save=self.opt.save)


    def get_explanation(self,source=None):
       
       
        if source is not None:
            self.source = source
            print("custom Image or image folder ", self.source )
        else:
            self.source  = os.path.join(str(ROOT), self.opt.source)
            print("Using default model weights:", self.source )

        file_name = os.path.basename(source)  # Get the base name (including extension)
        img_name, extension = os.path.splitext(file_name)
            
        model=self.model
        #target layer [-4] or [-2]
        target_layers =[model.model.model[-4]]

        img = cv2.imread(source)
        img = cv2.resize(img, (640, 640))
        rgb_img = img.copy()
        img = np.float32(img) / 255
        cam = EigenCAM(model, target_layers, task='od')
        grayscale_cam = cam(rgb_img)[0, :, :]
        cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True)
        
        # Save the CAM image as a picture
        output_path=img_name+'_explain'+extension
        
        Image.fromarray(cam_image).save(output_path)
        return print("success")

    def process_predictions(self, net_output, raw_image_path, raw_image_mat, image_transform, save_image_path):
        """post process the output of the net

        Args:
            net_output (_type_): output of detection net
            raw_image_path (str): raw image path
            raw_image_mat (numpy.Array) raw image numpy mat type
            image_transform (func) image_transform function
            save_image_path (str) save image path
            
        Returns:
            save_image_path
                inform that the annotation image has been written successfully
                in the same directory contain the annotation image, the annotation text file will be named "annotated_image.txt"
                each line format (yolo): 
                    class, x, y, w, h, confidence, class_name
        """

        pass   
    
    

