import random
import shutil
import json
import math
from glob import glob
import os
from pathlib import Path
import sys
import cv2
import numpy as np
from PIL import Image, ImageDraw
import logica.controladores.feature_extractor_controller as controlador

from datos.Config import ConfigManager

config = ConfigManager()

def unir_imagenes(imagen1, imagen2, matches, keypoints):
	#              0    1     2
    # imagenX -> (id, name, image)
    image1 = cv2.cvtColor(imagen1[2], cv2.COLOR_BGR2RGB)
    image2 = cv2.cvtColor(imagen2[2], cv2.COLOR_BGR2RGB)

    image1 = Image.fromarray(image1)
    image2 = Image.fromarray(image2)

    # Obtener las dimensiones de ambas imágenes
    ancho_combinado = image1.width + image2.width
    alto_combinado = max(image1.height, image2.height)
    imagen_combinada = Image.new('RGB', (ancho_combinado, alto_combinado))
    
	# Pegar la primera imagen en la esquina superior izquierda
    imagen_combinada.paste(image1, (0, 0))
    
    # Pegar la segunda imagen al lado de la primera
    imagen_combinada.paste(image2, (image1.width, 0))

    draw = ImageDraw.Draw(imagen_combinada)
    for row in matches:
        keypoint1_idx, keypoint2_idx = row

        keypoint1 = keypoints[imagen1[1]][keypoint1_idx]
        keypoint2 = keypoints[imagen2[1]][keypoint2_idx]

        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        draw_point(draw, keypoint1, color)
        draw_point(draw, (keypoint2[0] + image1.width, keypoint2[1]), color)
        draw_line(draw, keypoint1, (keypoint2[0] + image1.width, keypoint2[1]))

    # Guardar la imagen combinada en el path de salida
    if os.path.exists(config.DIR_MATCHES) and config.is_matches_folder_set():
        try:
            nombre_salida = f"{imagen1[0]}_{imagen2[0]}_combinada.jpg"
            ruta_salida = os.path.join(config.DIR_MATCHES, nombre_salida)
            imagen_combinada.save(ruta_salida)
            controlador.set_output("Combined image saved correctly.")

        except OSError as e:
            controlador.set_output(f"Error saving the combined image: {e}")
            return

def draw_point(draw, keypoint, color):
    x, y = int(keypoint[0]) , int(keypoint[1]) 
    # Definir las coordenadas del cuadrado que circunscribe el círculo
    x1 = x - 6
    y1 = y - 6
    x2 = x + 6
    y2 = y + 6
    # Dibujar el círculo inscrito en el cuadrado
    draw.ellipse([x1, y1, x2, y2], fill=color)

def draw_line(imagen, keypoint1, keypoint2):
    x1, y1 = int(keypoint1[0]) , int(keypoint1[1])
    x2, y2 = int(keypoint2[0]) , int(keypoint2[1])
    imagen.line([x1, y1, x2, y2], fill=(255, 255, 0), width=2)

def feature_extraction_callback():
    colmap_binary = "colmap"
    colmap_camera_params = ""
    text = "colmap_text"

    current_path = os.path.dirname(os.path.realpath(__file__))
    father_path = os.path.dirname(current_path)
    script_path = os.path.join(father_path, "Scripts")

    # On Windows, if FFmpeg isn't found, try automatically downloading it from the internet
    if os.name == "nt" and os.system(f"where {colmap_binary} >nul 2>nul") != 0:
        colmap_glob = os.path.join(script_path, "external", "colmap", "*", "COLMAP.bat")
        candidates = glob(colmap_glob)
        
        if not candidates:
            command = os.path.join(script_path, "download_colmap.bat")
            # Ejecuta el comando y captura su output y estado de finalización
            estado = controlador.init_process(command)
            
            if estado != 0:
                controlador.set_output("Error downloading COLMAP.")
                return
			
            else:
                controlador.set_output("COLMAP successfully downloaded.")
                candidates = glob(colmap_glob)

        if candidates:
            colmap_binary = candidates[0]

    images = config.get_output_folder()
    db_noext = str(Path(config.DATABASE_NAME).with_suffix(""))

    sparse = db_noext + "_sparse"
    controlador.set_output("Extracting features...")

    feature_extract_result = controlador.init_process(f"{colmap_binary} feature_extractor --ImageReader.camera_model {config.MODEL_CAMERA} --ImageReader.camera_params \"{colmap_camera_params}\" --SiftExtraction.estimate_affine_shape=true --SiftExtraction.domain_size_pooling=true --ImageReader.single_camera 1 --database_path {config.DATABASE_NAME} --image_path {images}")

    if feature_extract_result != 0:
        controlador.set_output("Error while extracting features.")
        return
    else:
        controlador.set_output("Features extracted correctly.")
    
    match_cmd = f"{colmap_binary} {controlador.get_model_matcher()}_matcher --SiftMatching.guided_matching=true --database_path {config.DATABASE_NAME}"   
    matche_result = controlador.init_process(match_cmd)
    if matche_result != 0:
        controlador.set_output("Error while matching.")
        return
    else:
        controlador.set_output("Matching done correctly.")

    try:
        shutil.rmtree(sparse)
    except:
        pass
    
    os.system(f"mkdir \"{os.path.join(config.DIR_ACT, sparse)}\"")

    mapper_result = controlador.init_process(f"{colmap_binary} mapper --database_path {config.DATABASE_NAME} --image_path {images} --output_path {sparse}")
    if mapper_result != 0:
        controlador.set_output("Mapping error.")
        return
    else:
        controlador.set_output("Mapping done correctly.")

    bundle_result = controlador.init_process(f"{colmap_binary} bundle_adjuster --input_path {sparse}/0 --output_path {sparse}/0 --BundleAdjustment.refine_principal_point 1")
    if bundle_result != 0:
        controlador.set_output("Error while adjusting the bundle.")
        return
    else:
        controlador.set_output("Bundle adjusted correctly.")

    try:
        shutil.rmtree(text)
    except:
        pass

    os.system(f"mkdir \"{os.path.join(config.DIR_ACT, text)}\"")

    convert_result = controlador.init_process(f"{colmap_binary} model_converter --input_path {sparse}/0 --output_path {text} --output_type TXT")
    if convert_result != 0:
        controlador.set_output("Error while converting the model.")
        return
    else:
        controlador.set_output("Correctly converted model.")
        create_json()

def variance_of_laplacian(image):
	return cv2.Laplacian(image, cv2.CV_64F).var()

def sharpness(imagePath):
	image = cv2.imread(imagePath)
	gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
	fm = variance_of_laplacian(gray)
	return fm

def qvec2rotmat(qvec):
	return np.array([
		[
			1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
			2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
			2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]
		], [
			2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
			1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
			2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]
		], [
			2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
			2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
			1 - 2 * qvec[1]**2 - 2 * qvec[2]**2
		]
	])

def rotmat(a, b):
	a, b = a / np.linalg.norm(a), b / np.linalg.norm(b)
	v = np.cross(a, b)
	c = np.dot(a, b)
	# handle exception for the opposite direction input
	if c < -1 + 1e-10:
		return rotmat(a + np.random.uniform(-1e-2, 1e-2, 3), b)
	s = np.linalg.norm(v)
	kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
	return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2 + 1e-10))

def closest_point_2_lines(oa, da, ob, db): # returns point closest to both rays of form o+t*d, and a weight factor that goes to 0 if the lines are parallel
	da = da / np.linalg.norm(da)
	db = db / np.linalg.norm(db)
	c = np.cross(da, db)
	denom = np.linalg.norm(c)**2
	t = ob - oa
	ta = np.linalg.det([t, db, c]) / (denom + 1e-10)
	tb = np.linalg.det([t, da, c]) / (denom + 1e-10)
	if ta > 0:
		ta = 0
	if tb > 0:
		tb = 0
	return (oa+ta*da+ob+tb*db) * 0.5, denom

def create_json():
    # Check that we can save the output before we do a lot of work
	try:
		open(config.DIR_ACT, "a").close()
	except Exception as e:
		controlador.set_output(f"Could not save transforms JSON to {config.DIR_ACT}: {e}")

	controlador.set_output(f"outputting to {config.DIR_ACT}...")
	cameras = {}
	with open(os.path.join(config.DIR_ACT, "colmap_text", "cameras.txt"), "r") as f:
		camera_angle_x = math.pi / 2
		for line in f:
			# 1 SIMPLE_RADIAL 2048 1536 1580.46 1024 768 0.0045691
			# 1 OPENCV 3840 2160 3178.27 3182.09 1920 1080 0.159668 -0.231286 -0.00123982 0.00272224
			# 1 RADIAL 1920 1080 1665.1 960 540 0.0672856 -0.0761443
			if line[0] == "#":
				continue
			els = line.split(" ")
			camera = {}
			camera_id = int(els[0])
			camera["w"] = float(els[2])
			camera["h"] = float(els[3])
			camera["fl_x"] = float(els[4])
			camera["fl_y"] = float(els[4])
			camera["k1"] = 0
			camera["k2"] = 0
			camera["k3"] = 0
			camera["k4"] = 0
			camera["p1"] = 0
			camera["p2"] = 0
			camera["cx"] = camera["w"] / 2
			camera["cy"] = camera["h"] / 2
			camera["is_fisheye"] = False
			if els[1] == "SIMPLE_PINHOLE":
				camera["cx"] = float(els[5])
				camera["cy"] = float(els[6])
			elif els[1] == "PINHOLE":
				camera["fl_y"] = float(els[5])
				camera["cx"] = float(els[6])
				camera["cy"] = float(els[7])
			elif els[1] == "SIMPLE_RADIAL":
				camera["cx"] = float(els[5])
				camera["cy"] = float(els[6])
				camera["k1"] = float(els[7])
			elif els[1] == "RADIAL":
				camera["cx"] = float(els[5])
				camera["cy"] = float(els[6])
				camera["k1"] = float(els[7])
				camera["k2"] = float(els[8])
			elif els[1] == "OPENCV":
				camera["fl_y"] = float(els[5])
				camera["cx"] = float(els[6])
				camera["cy"] = float(els[7])
				camera["k1"] = float(els[8])
				camera["k2"] = float(els[9])
				camera["p1"] = float(els[10])
				camera["p2"] = float(els[11])
			elif els[1] == "SIMPLE_RADIAL_FISHEYE":
				camera["is_fisheye"] = True
				camera["cx"] = float(els[5])
				camera["cy"] = float(els[6])
				camera["k1"] = float(els[7])
			elif els[1] == "RADIAL_FISHEYE":
				camera["is_fisheye"] = True
				camera["cx"] = float(els[5])
				camera["cy"] = float(els[6])
				camera["k1"] = float(els[7])
				camera["k2"] = float(els[8])
			elif els[1] == "OPENCV_FISHEYE":
				camera["is_fisheye"] = True
				camera["fl_y"] = float(els[5])
				camera["cx"] = float(els[6])
				camera["cy"] = float(els[7])
				camera["k1"] = float(els[8])
				camera["k2"] = float(els[9])
				camera["k3"] = float(els[10])
				camera["k4"] = float(els[11])
			else:
				controlador.set_output("Unknown camera model " + els[1])
			# fl = 0.5 * w / tan(0.5 * angle_x);
			camera["camera_angle_x"] = math.atan(camera["w"] / (camera["fl_x"] * 2)) * 2
			camera["camera_angle_y"] = math.atan(camera["h"] / (camera["fl_y"] * 2)) * 2
			camera["fovx"] = camera["camera_angle_x"] * 180 / math.pi
			camera["fovy"] = camera["camera_angle_y"] * 180 / math.pi

			controlador.set_output(f"camera {camera_id}:\n\tres={camera['w'],camera['h']}\n\tcenter={camera['cx'],camera['cy']}\n\tfocal={camera['fl_x'],camera['fl_y']}\n\tfov={camera['fovx'],camera['fovy']}\n\tk={camera['k1'],camera['k2']} p={camera['p1'],camera['p2']} ")
			cameras[camera_id] = camera

	if len(cameras) == 0:
		controlador.set_output("No cameras found!")

	with open(os.path.join(config.DIR_ACT, "colmap_text", "images.txt"), "r") as f:
		i = 0
		bottom = np.array([0.0, 0.0, 0.0, 1.0]).reshape([1, 4])
		if len(cameras) == 1:
			camera = cameras[camera_id]
			out = {
				"camera_angle_x": camera["camera_angle_x"],
				"camera_angle_y": camera["camera_angle_y"],
				"fl_x": camera["fl_x"],
				"fl_y": camera["fl_y"],
				"k1": camera["k1"],
				"k2": camera["k2"],
				"k3": camera["k3"],
				"k4": camera["k4"],
				"p1": camera["p1"],
				"p2": camera["p2"],
				"is_fisheye": camera["is_fisheye"],
				"cx": camera["cx"],
				"cy": camera["cy"],
				"w": camera["w"],
				"h": camera["h"],
				"aabb_scale": controlador.get_aabb_scale(),
				"frames": [],
			}
		else:
			out = {
				"frames": [],
				"aabb_scale": controlador.get_aabb_scale(),
			}

		up = np.zeros(3)
		for line in f:
			line = line.strip()
			if line[0] == "#":
				continue
			i = i + 1
			if  i % 2 == 1:
				elems=line.split(" ") # 1-4 is quat, 5-7 is trans, 9ff is filename (9, if filename contains no spaces)
				#name = str(PurePosixPath(Path(IMAGE_FOLDER, elems[9])))
				# why is this requireing a relitive path while using ^
				image_rel = config.get_output_folder()
				name = str(f"{image_rel}/{'_'.join(elems[9:])}")
				b = sharpness(name)
				controlador.set_output(name + "  sharpness=" + str(b))
				image_id = int(elems[0])
				qvec = np.array(tuple(map(float, elems[1:5])))
				tvec = np.array(tuple(map(float, elems[5:8])))
				R = qvec2rotmat(-qvec)
				t = tvec.reshape([3,1])
				m = np.concatenate([np.concatenate([R, t], 1), bottom], 0)
				c2w = np.linalg.inv(m)
				if not controlador.get_keep_colmap_coords():
					c2w[0:3,2] *= -1 # flip the y and z axis
					c2w[0:3,1] *= -1
					c2w = c2w[[1,0,2,3],:]
					c2w[2,:] *= -1 # flip whole world upside down

					up += c2w[0:3,1]

				frame = {"file_path":name,"sharpness":b,"transform_matrix": c2w}
				if len(cameras) != 1:
					frame.update(cameras[int(elems[8])])
				out["frames"].append(frame)
	nframes = len(out["frames"])

	if controlador.get_keep_colmap_coords():
		flip_mat = np.array([
			[1, 0, 0, 0],
			[0, -1, 0, 0],
			[0, 0, -1, 0],
			[0, 0, 0, 1]
		])

		for f in out["frames"]:
			f["transform_matrix"] = np.matmul(f["transform_matrix"], flip_mat) # flip cameras (it just works)
	else:
		# don't keep colmap coords - reorient the scene to be easier to work with

		up = up / np.linalg.norm(up)
		controlador.set_output("up vector was " + str(up))
		R = rotmat(up,[0,0,1]) # rotate up vector to [0,0,1]
		R = np.pad(R,[0,1])
		R[-1, -1] = 1

		for f in out["frames"]:
			f["transform_matrix"] = np.matmul(R, f["transform_matrix"]) # rotate up to be the z axis

		# find a central point they are all looking at
		controlador.set_output("computing center of attention...")
		totw = 0.0
		totp = np.array([0.0, 0.0, 0.0])
		for f in out["frames"]:
			mf = f["transform_matrix"][0:3,:]
			for g in out["frames"]:
				mg = g["transform_matrix"][0:3,:]
				p, w = closest_point_2_lines(mf[:,3], mf[:,2], mg[:,3], mg[:,2])
				if w > 0.00001:
					totp += p*w
					totw += w
		if totw > 0.0:
			totp /= totw
		controlador.set_output(str(totp)) # the cameras are looking at totp
		for f in out["frames"]:
			f["transform_matrix"][0:3,3] -= totp

		avglen = 0.
		for f in out["frames"]:
			avglen += np.linalg.norm(f["transform_matrix"][0:3,3])
		avglen /= nframes
		controlador.set_output("avg camera distance from origin " + str(avglen))
		for f in out["frames"]:
			f["transform_matrix"][0:3,3] *= 4.0 / avglen # scale to "nerf sized"

	for f in out["frames"]:
		f["transform_matrix"] = f["transform_matrix"].tolist()
	controlador.set_output(str(nframes) +" frames")
	controlador.set_output(f"writing {config.DIR_ACT}")
	with open(os.path.join(config.DIR_ACT, "transforms.json"), "w") as outfile:
		json.dump(out, outfile, indent=2)

	if len(controlador.get_mask_categories()) > 0:
		# Check if detectron2 is installed. If not, install it.
		try:
			import detectron2
		except ModuleNotFoundError:
			try:
				import torch
			except ModuleNotFoundError:
				controlador.set_output("PyTorch is not installed. For automatic masking, install PyTorch from https://pytorch.org/")

			controlador.set_output("Detectron2 is not installed. Press enter to install it.")
			
			import subprocess
			package = 'git+https://github.com/facebookresearch/detectron2.git'
			subprocess.check_call([sys.executable, "-m", "pip", "install", package])
			import detectron2

		import torch
		from pathlib import Path
		from detectron2.config import get_cfg
		from detectron2 import model_zoo
		from detectron2.engine import DefaultPredictor

		category2id = json.load(open(os.path.join(".\Scripts", "category2id.json"), "r"))
		mask_ids = [category2id[c] for c in controlador.get_mask_categories()]

		cfg = get_cfg()
		# Add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
		cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
		cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
		# Find a model from detectron2's model zoo.
		cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
		predictor = DefaultPredictor(cfg)

		for frame in out["frames"]:
			img = cv2.imread(frame["file_path"])
			outputs = predictor(img)

			output_mask = np.zeros((img.shape[0], img.shape[1]))
			for i in range(len(outputs["instances"])):
				if outputs["instances"][i].pred_classes.cpu().numpy()[0] in mask_ids:
					pred_mask = outputs["instances"][i].pred_masks.cpu().numpy()[0]
					output_mask = np.logical_or(output_mask, pred_mask)

			rgb_path = Path(frame["file_path"])
			mask_name = str(rgb_path.parents[0] / Path("dynamic_mask_" + rgb_path.name.replace(".jpg", ".png")))
			cv2.imwrite(mask_name, (output_mask*255).astype(np.uint8))

	controlador.set_output("Transforms JSON saved successfully.")
	controlador.set_output("Features created correctly.")


def draw_keypoints(image, keypoints):
	
	# Dibujar cada punto clave en la imagen
    for kp in keypoints.tolist():
        x, y = int(kp[0]) , int(kp[1]) 
        cv2.circle(image, (x, y), 5, (0, 255, 0), 1)  # Dibujar un círculo verde en el punto clave

def save_images_with_keypoints():
    cargar_imagenes_cv2(True)

	# keyPoints es un diccionario que contiene los keypoints de cada imagen id_imagen: keypoints np.array
    keypoints, num_images = controlador.extract_keypoints()

    # Asegurarse de que el directorio de salida exista
    if not os.path.exists(config.DIR_KEYPOINTS):
        try:
            os.makedirs(config.DIR_KEYPOINTS)
            controlador.set_output("Output directory created correctly.")

        except OSError as e:
            controlador.set_output("Error creating the output directory:", e)

    # Iterar sobre todas las imágenes y sus keypoints
    for i, image in enumerate(controlador.get_keypoints_images()):
        # Crear una copia de la imagen para no modificar la original
        image_with_keypoints = image.copy()

        # Dibujar los keypoints en la imagen
        draw_keypoints(image_with_keypoints, keypoints[controlador.get_name_file(i)])

        # Guardar la imagen con los keypoints
        output_path = os.path.join(config.DIR_KEYPOINTS, f"image_{i+1}_with_keypoints.jpg")
        cv2.imwrite(output_path, image_with_keypoints)

        controlador.set_output(f"Image with keypoints saved in {output_path}")

    controlador.set_output("Images with keypoints saved correctly.")

	# Eliminem dades en memoria
    controlador.delete_keypoints_images()
    controlador.delete_name_files()

	# Inicializar el visor para la posterior carga de imágenes con keypoints
    controlador.load_images_with_keypoints()

def save_images_with_matches():
	# Cargar en el sistema las imagenes
    cargar_imagenes_cv2(False)

	# Matches es un diccionario que contiene los match de cada par de imagenes (id1, id2): Matches np.array(2, x)
	# cada elemento matches[(id1, id2)], es una lista de indices de keypoints que hacen match
    matches, num_images = controlador.extract_matches()
	# keyPoints es un diccionario que contiene los keypoints, keypoint['name_image']: keypoints (puntos, dim(6)) np.array
    keypoints, num_images = controlador.extract_keypoints()
    
    # Asegurarse de que el directorio de salida exista
    if not os.path.exists(config.DIR_MATCHES):
        try:
            os.makedirs(config.DIR_MATCHES)
            controlador.set_output("Output directory created correctly.")

        except OSError as e:
            controlador.set_output(f"Error creating the output directory: {e}")

	# Seleccionar 30 filas de forma aleatoria de cada array en el diccionario
    matches_aleatorios = {}
	# clave = (id1, id2), array = np.array(num_mat, 2)
    for clave, (array, num_mat) in matches.items():
        if num_mat >= 10:
			# Obtener un índice aleatorio para seleccionar 30 filas
            indices_aleatorios = np.random.choice(array.shape[0], size=30, replace=False)
			
			# Seleccionar las filas aleatorias del array
            filas_aleatorias = array[indices_aleatorios]
			# Nuevo diccionario con las filas aleatorias
            matches_aleatorios[clave] = filas_aleatorias
        else:
            # Si no hay suficientes filas, simplemente asigna el array original
            matches_aleatorios[clave] = array

	# Creamos los matches de los pares de imagenes
    for clave, array in matches_aleatorios.items():
        idx1, idx2 = clave  # Suponiendo que clave es una tupla de dos índices
		# Convertir idx a entero
        idx1 = int(idx1)    
        idx2 = int(idx2)
        imagen1 = controlador.obtener_data_imagen_por_id(idx1)
        imagen2 = controlador.obtener_data_imagen_por_id(idx2)
        unir_imagenes(imagen1, imagen2, array, keypoints)

	# Eliminem dades en memoria
    controlador.delete_matches_images()
    controlador.delete_matches_name_files()

	# Inicializar el visor para la posterior carga de imágenes con keypoints
    controlador.load_images_with_matches()

def cargar_imagenes_cv2(keypoints):
    # Obtener todas las imágenes en la carpeta
    output_path = config.get_output_folder()

    controlador.cargar_imagenes_id()
    images_id = controlador.get_images_id()

    images = os.listdir(output_path)
    if len(images) == 0:
        controlador.set_output("No images were found in the specified folder.")
        return
	
    for archivo in images:
        if archivo.endswith((".png", ".jpg", ".jpeg")):
            if keypoints:
                controlador.push_name_file(archivo)
                controlador.push_keypoint_image(cv2.imread(os.path.join(output_path, archivo)))

            else:
                indice = next((id for id, name in images_id.items() if name == archivo), None)
                full_index = (indice, archivo)
                controlador.push_matches_image(full_index, cv2.imread(os.path.join(output_path, archivo)))

            controlador.set_output(f"Images {archivo} loaded correctly.")
    
def generate_image_matches():
	controlador.extract_matches()