#!/usr/bin/env python
if __name__ != '__main__': raise Exception("Do not import me!")
import sys
sys.path.insert(0, "..")

import logging
import numpy as np
import matplotlib.pyplot as plt

from argparse import ArgumentParser

from cvdatasets import AnnotationType
from utils import parser

def plot_crops(crops, spec, spec_offset, scatter_mid=False, names=None):

	n_crops = len(crops)
	if n_crops == 0: return
	rows = int(np.ceil(np.sqrt(n_crops)))
	cols = int(np.ceil(n_crops / rows))
	dx, dy = spec_offset

	for i, crop in enumerate(crops):
		x, y = np.unravel_index(i, (rows, cols))

		ax = plt.subplot(spec[x+dx, y+dy])

		if names is not None:
			ax.set_title(names[i])

		ax.imshow(crop)
		ax.axis("off")

		if scatter_mid:
			middle_h, middle_w = crop.shape[0] / 2, crop.shape[1] / 2
			ax.scatter(middle_w, middle_h, marker="x")


def main(args):
	# assert args.dataset in AnnotationType, \
	# 	f"AnnotationType is not known: \"{args.dataset}\""

	annot = AnnotationType.new_annotation(args)

	kwargs = {}
	if annot.info is None:
		kwargs = dict(
			part_rescale_size=args.rescale_size,
			uniform_parts=args.uniform_parts,
			ratio=args.ratio,
		)

	data = annot.new_dataset(
		args.subset,

		center_cropped=not args.no_center_crop,
		crop_to_bb=args.crop_to_bb,
		crop_uniform=args.crop_uniform,

		parts_in_bb=args.parts_in_bb,

		rnd_select=args.rnd,
		seed=args.seed,

		**kwargs
	)

	logging.info(f"Loaded {len(data)} {args.subset} images")

	if args.only_class >= 0:
		mask = data.labels == args.only_class
		logging.info(f"Showing only {mask.sum()} images from class {args.only_class}")
		idxs = np.where(mask)[0]
	else:
		start = max(args.start, 0)
		n_images = min(args.n_images, len(data) - start)
		end = max(start, start + n_images)
		logging.info(f"Showing only images {start} - {end}")
		idxs = range(start, end)


	for i in idxs:
		im, parts, label = data[i]
		n_parts = len(parts)

		if args.no_parts:
			cols, rows = 1, 1
			factor = 1
		else:
			assert n_parts != 0
			rows = int(np.ceil(np.sqrt(n_parts)))
			cols = int(np.ceil(n_parts / rows))
			factor = 3 if args.rnd else 2

		grid_spec = plt.GridSpec(rows, factor * cols)

		fig = plt.figure()

		im_ax = plt.subplot(grid_spec[:, :cols])

		im_ax.axis("off")
		im_ax.set_title("Visible Parts")
		im_ax.imshow(im)

		if not args.crop_to_bb and not args.no_bboxes:
			data.plot_bounding_box(i, im_ax)

		# axs[1].axis("off")
		# axs[1].set_title("{}selected parts".format("randomly " if args.rnd else ""))
		# axs[1].imshow(parts.reveal(im, ratio=data.ratio))

		if not args.no_parts:
			parts.plot(im=im, ax=im_ax, ratio=data.ratio, linewidth=3)
			if data.uniform_parts:
				crop_names = None
			else:
				crop_names = list(data._annot.part_names.values())

			part_crops = parts.visible_crops(im, ratio=data.ratio)
			if args.rnd:
				parts.invert_selection()
				action_crops = parts.visible_crops(im, ratio=data.ratio)

			plot_crops(part_crops, grid_spec, spec_offset=(0, cols), names=crop_names)

			if args.rnd:
				plot_crops(action_crops, grid_spec, spec_offset=(0, 2*cols), names=crop_names)

		plt.show()
		plt.close()

main(parser.parse_args())