import os
import torch
import sys
from pathlib import Path
import importlib
from typing import List
from easydict import EasyDict
import json
import numpy as np
import groundingdino.datasets.transforms as T
from groundingdino.util.inference import load_model, load_image, predict, annotate
import supervision as sv
import wget
import cv2
from torchvision.ops import box_convert
from PIL import Image

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
sys.path.insert(0, os.path.abspath(
	os.path.join(os.path.dirname(__file__), './../../')
))

if str(ROOT) not in sys.path:
	sys.path.append(str(ROOT))  # add ROOT to PATH	


sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from ecos_core.base_model import BaseModel

class Model(BaseModel):
	NAME = "GroundingDINO"

	def __init__(self, opt):
		self.opt = opt
		self.opt_path = "./opt.json"
		self.DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
		self.wh_error_ratio = 15 if not hasattr(self.opt, "wh_error_ratio") else self.opt.wh_error_ratio

	def __call__(self,x):
		return None

	def get_transform(self):
		"""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,custom_weight=None):
		"""Build model in GroundingDINO model

		Returns:
						_type_: _description_
		"""
		self.custom_weight = custom_weight
		if self.custom_weight is not None:
			self.CHECKPOINT_PATH = self.custom_weight
		else:
			self.CHECKPOINT_PATH =  os.path.join(str(ROOT), self.opt.weights)

			if os.path.exists(self.CHECKPOINT_PATH):

				print("have model weight")

			else:

				print("downloading model weight")
				# create folder
				os.makedirs(os.path.dirname(self.CHECKPOINT_PATH), exist_ok=True)
				URL = self.opt.weights_url

				wget.download(URL, self.CHECKPOINT_PATH)

				print("downloaded model weight")

			


		self.CONFIG_PATH = os.path.join(str(ROOT), 'config','GroundingDINO_SwinT_OGC.py')

		self.model = load_model(self.CONFIG_PATH, self.CHECKPOINT_PATH)

		return self

	def reload_param(self, opt_path="./opt.json"):
		"""reload parameters

		Args:
			opt_path (str, optional): Opt path. Defaults to "./opt.json".
		"""
		if not os.path.isfile(opt_path):
			opt_path = self.opt_path
		
		with open(opt_path, 'r') as f:
			data = json.load(f)
			for key, value in data.items():
				setattr(self.opt, key, value)
	
	def inference(self, image, caption, box_threshold, text_threshold):
		"""Run inference grounding dino given image, caption, box_threshold and text_threshold

		Args:
			image (torch.Tensor): input image
			caption (str): prompt to detect
			box_threshold (float): box threshold for filtering
			text_threshold (float): text threshold for filtering

		Returns:
			boxes (torch.Tensor): tensor containing boxes, shape(n, 4)
			logits (torch.Tensor): tensor containing confidence scores for each box
			phrases (list): list of class names
		"""
		channel, height, width = image.shape
		# If width and height differ greatly, ignore it
		if width==0 or height == 0 or height/width > self.wh_error_ratio\
				or width/height > self.wh_error_ratio:

			boxes = torch.empty(0,4)
			phrases = []
			logits = torch.empty(0)
			return boxes, logits, phrases

		boxes, logits, phrases = predict(
			model=self.model,
			image=image,
			caption=caption,
			box_threshold=box_threshold,
			text_threshold=text_threshold,
			device=self.DEVICE)
		# filter boxes
		filter_box = []
		filter_logit = []
		filter_phrase = []
		for box,logit,phrase in zip(boxes,logits,phrases):
			S = box[2]*box[3]
			if S < self.opt.area_percent : 
				filter_box.append(box)
				filter_logit.append(logit)
				filter_phrase.append(phrase)
		if len(filter_box) != 0 :
			boxes = torch.stack(filter_box)
			phrases = filter_phrase
			logits = torch.stack(filter_logit)
		else :
			boxes = torch.empty(0,4)
			phrases = []
			logits = torch.empty(0)

		return boxes, logits, phrases

	def predict_with_classes(
        self,
        image: np.ndarray,
        classes: List[str],
        box_threshold: float,
        text_threshold: float
    ) -> sv.Detections:

		caption = ". ".join(classes)
		processed_image = Model.preprocess_image(image_bgr=image).to(self.DEVICE)
		boxes, logits, phrases = self.inference(processed_image, caption, box_threshold, text_threshold)
		source_h, source_w, _ = image.shape
		detections = Model.post_process_result(
			source_h=source_h,
			source_w=source_w,
			boxes=boxes,
			logits=logits)
		class_id = Model.phrases2classes(phrases=phrases, classes=classes)
		detections.class_id = class_id
		return detections

	def process_image(self, image_path):
		"""Process image in groundingDINO model

		Args:
					image_path (_type_):  img path

		Returns:

		"""
		TEXT_PROMPT = self.opt.TEXT_PROMPT
		BOX_TRESHOLD = self.opt.BOX_TRESHOLD
		TEXT_TRESHOLD = self.opt.TEXT_TRESHOLD

		self.image_source, self.image = load_image(image_path)
		boxes, logits, phrases = self.inference(self.image, TEXT_PROMPT, BOX_TRESHOLD, TEXT_TRESHOLD)
			
		self.annotated_frame = annotate(image_source=self.image_source, boxes=boxes, logits=logits, phrases=phrases)

		return self.annotated_frame

	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

	@staticmethod
	def preprocess_image(image_bgr: np.ndarray) -> torch.Tensor:
		transform = T.Compose(
			[
				T.RandomResize([800], max_size=1333),
				T.ToTensor(),
				T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
			]
		)
		image_pillow = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
		image_transformed, _ = transform(image_pillow, None)
		return image_transformed

	@staticmethod
	def post_process_result(
			source_h: int,
			source_w: int,
			boxes: torch.Tensor,
			logits: torch.Tensor
	) -> sv.Detections:
		boxes = boxes * torch.Tensor([source_w, source_h, source_w, source_h])
		xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
		confidence = logits.numpy()
		return sv.Detections(xyxy=xyxy, confidence=confidence)

	@staticmethod
	def phrases2classes(phrases: List[str], classes: List[str]) -> np.ndarray:
		class_ids = []
		for phrase in phrases:
			try:
				class_ids.append(classes.index(phrase))
			except ValueError:
				class_ids.append(None)
		return np.array(class_ids)