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

from comparisons import histogram_comparison, structural_similarity_index, euclidean_distance
from load_test_data import load_data

# Allows imports from the style transfer submodule
sys.path.append('DCT-Net')

from source.cartoonize import Cartoonizer


def load_source(filename: str) -> np.ndarray:
  return cv2.imread(filename)[..., ::-1]


def find_and_crop_face(data: np.ndarray, classifier_file='haarcascades/haarcascade_frontalface_default.xml') -> np.ndarray:
  data_gray = cv2.cvtColor(data, cv2.COLOR_BGR2GRAY)
  face_cascade = cv2.CascadeClassifier(classifier_file)
  face = face_cascade.detectMultiScale(data_gray, 1.1, 3)
  face = max(face, key=len)
  x, y, w, h = face
  face = data[y:y + h, x:x + w]
  return face


def plot_two_images(a: np.ndarray, b: np.ndarray):
  plt.figure(figsize=[10, 10])
  plt.subplot(121)
  plt.imshow(a)
  plt.title("A")
  plt.subplot(122)
  plt.imshow(b)
  plt.title("B")
  plt.show()


def compare_with_anime_characters(source: np.ndarray, verbose=False) -> list[dict]:
  dataset = load_data('data/images')
  all_metrics = []
  for anime_image, label in zip(dataset['values'], dataset['labels']):
    current_result = {
      'name': label,
      'metrics': {}
    }
    # TODO: Use a different face detection method for anime images
    # anime_face = find_and_crop_face(anime_image, 'haarcascades/lbpcascade_animeface.xml')
    anime_face = anime_image
    source_rescaled = cv2.resize(source, anime_face.shape[:2])
    if verbose:
      plot_two_images(anime_face, source_rescaled)
    current_result['metrics'] = histogram_comparison(source_rescaled, anime_face)
    current_result['metrics']['structural-similarity'] = structural_similarity_index(source_rescaled, anime_face)
    current_result['metrics']['euclidean-distance'] = euclidean_distance(source_rescaled, anime_face)
    all_metrics.append(current_result)

  return all_metrics


def get_top_results(all_metrics: list[dict], metric='correlation', count=1):
  all_metrics.sort(reverse=True, key=lambda item: item['metrics'][metric])
  return list(map(lambda item: {'name': item['name'], 'score': item['metrics'][metric]}, all_metrics[:count]))


def transfer_to_anime(img: np.ndarray):
  algo = Cartoonizer(dataroot='DCT-Net/damo/cv_unet_person-image-cartoon_compound-models')
  return algo.cartoonize(img).astype(np.uint8)


if __name__ == '__main__':
  source = load_source('UAM-Andre.jpg')
  source_anime = transfer_to_anime(source)
  source_face_anime = find_and_crop_face(source_anime)
  results = compare_with_anime_characters(source_face_anime)
  print(get_top_results(results, count=5))