import argparse
import os
from random import shuffle

import numpy as np
from matplotlib import pyplot as plt
from PIL import Image

parser = argparse.ArgumentParser(description="Inspect image versions")
parser.add_argument("-f", "--base_folder", type=str, required=True, help="Base folder to inspect")
parser.add_argument("-opt_fg_ratio", type=float, default=0.3, help="Optimal foreground ratio")

args = parser.parse_args()

bg_folder = os.path.join(args.base_folder, "backgrounds")
fg_folder = os.path.join(args.base_folder, "foregrounds")

classes = os.listdir(fg_folder)
classes = sorted(classes, key=lambda x: int(x[1:]))
assert len(classes) in [200, 1_000], f"Expected 200 or 1_000 classes, got {len(classes)}"

total_images = set()
for in_cls in classes:
    cls_images = {
        os.path.join(in_cls, "_".join(img.split(".")[0].split("_")[:3]))
        for img in os.listdir(os.path.join(fg_folder, in_cls))
    }
    total_images.update(cls_images)
total_images = list(total_images)
shuffle(total_images)
print(total_images[:5], "...")

# in_cls to print name/lemma
with open("wordnet_data/tinyimagenet_synset_names.txt", "r") as f:
    in_cls_to_name = {line.split(":")[0].strip(): line.split(":")[1].strip() for line in f.readlines() if len(line) > 2}

for image_name in total_images:
    in_cls, img_name = image_name.split("/")
    versions = set()
    for img in os.listdir(os.path.join(fg_folder, in_cls)):
        if img.startswith(img_name):
            versions.add(img)
    if len(versions) <= 1:
        print(f"Image {image_name} has only one version")
        continue
    versions = sorted(list(versions))

    bg_img = Image.open(os.path.join(bg_folder, in_cls, f"{versions[0].split('.')[0]}.JPEG"))

    # plot all versions at once
    fig, axs = plt.subplots(1, len(versions) + 1, figsize=(15, 5))
    axs[0].imshow(bg_img)
    axs[0].axis("off")
    axs[0].set_title("Background")
    for version, ax in zip(versions, axs[1:]):
        img = Image.open(os.path.join(fg_folder, in_cls, version))
        img_mask = np.array(img.convert("RGBA").split()[-1])
        fg_dens_fact = np.sum(img_mask) / (255 * img_mask.size)
        fg_ratio = np.sum(img_mask) / (255 * bg_img.size[0] * bg_img.size[1])
        ax.imshow(img)
        ax.axis("off")
        ax.set_title(f"{version}\nfg_rat: {fg_ratio:.2f} dist to optimal: {abs(fg_ratio - args.opt_fg_ratio):.2f}")
    fig.suptitle(in_cls_to_name[in_cls])
    plt.show()
    plt.close()