from PIL import Image
import os
import random
import numpy as np
import json

pororo_source_frame_paths = {
    'Pororo': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH1_2/Pororo_ENGLISH1_2_ep6/12.png',
    'Loopy': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH1_1/Pororo_ENGLISH1_1_ep12/26.png',
    'Crong': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH1_1/Pororo_ENGLISH1_1_ep12/10.png',
    'Poby': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH1_1/Pororo_ENGLISH1_1_ep9/34.png',
    'Eddy': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH1_1/Pororo_ENGLISH1_1_ep12/46.png',
    'Petty': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH2_1/Pororo_ENGLISH2_1_ep1/34.png',
    'Tongtong': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH3_1/Pororo_ENGLISH3_1_ep7/8.png',
    'Rody': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH3_1/Pororo_ENGLISH3_1_ep6/66.png',
    'Harry': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH3_1/Pororo_ENGLISH3_1_ep7/39.png',
}


flintstones_source_frame_paths = {
    "Wilma": '',
    "Fred": '',
    "Betty": '',
    "Barney": '',
    "Dino": '',
    "Pebbles": '',
    "Mr Slate": ''
}


def sample_image(im):
    shorter, longer = min(im.size[0], im.size[1]), max(im.size[0], im.size[1])
    video_len = int(longer / shorter)
    se = np.random.randint(0, video_len, 1)[0]
    return im.crop((0, se * shorter, shorter, (se + 1) * shorter))


def get_pororo_source_frames():

    # sample_image(Image.open(os.path.join(img_folder, tgt_img_path)).convert('RGB'))
    # labels = np.load('../../StoryGAN/pororo_png/labels.npy', allow_pickle=True, encoding='latin1').item()
    # for i in range(9):
    #     print(i)
    #     individual_frames = [(k, v) for k, v in labels.items() if v[i] == 1 and not any([v[j] == 1 for j in range(9) if j!=i])]
    #     print(random.sample(individual_frames, k=10))

    for k, v in pororo_source_frame_paths.items():

        img = sample_image(Image.open(v).convert('RGB'))
        img.save(k + '.png')


def get_flintstones_source_frames():

    dir_path = '../../StoryGAN/flintstones'
    annotations = json.load(open('../../StoryGAN/flintstones/flintstones_annotations_v1-0.json', 'r'))
    for k in flintstones_source_frame_paths.keys():

        if k != "Barney":
            continue

        character_frames = []
        for sample in annotations:
            sample_characters = [c["entityLabel"].strip().lower() for c in sample["characters"]]
            if sample_characters[0] == k.lower():
                character_frames.append(sample["globalID"])

        globalID = random.choice(character_frames)
        arr = np.load(os.path.join(dir_path, 'video_frames_sampled', globalID + '.npy'))
        n_frames = arr.shape[0]
        im = arr[random.randrange(n_frames)]
        im = Image.fromarray(im)
        im.save(k.replace(' ', '') + '.png')

get_flintstones_source_frames()