"""
Construct TFRecord serialized datasets for pose estimation models
"""
import cv2
import numpy as np
import tensorflow as tf
from random import shuffle, sample
from sklearn.model_selection import train_test_split

from itertools import chain
from functools import partial

from modules import interhand_utils, wflw_utils, coco_utils, halpe_utils, mwb_utils, rhd_utils, hands_utils
from modules.config import CFG


def interhand():
    """Construct Interhand TFRecord files
    """
    # for each subset
    for subset in ['train', 'val', 'test']:
        # load annotation files
        annots_data = hands_utils.load_json(CFG.INTERHAND_DATASET_PATHS['annotations'][subset]['data'])
        annots_joint = hands_utils.load_json(CFG.INTERHAND_DATASET_PATHS['annotations'][subset]['joint'])
        annots_camera = hands_utils.load_json(CFG.INTERHAND_DATASET_PATHS['annotations'][subset]['camera'])
        images_path = CFG.INTERHAND_DATASET_PATHS['images'][subset]
        # open TFRecord file
        with tf.io.TFRecordWriter(CFG.INTERHAND_TFRECORD_PATHS[subset]) as writer:
            # for every record in annotation
            for frame_annot, data_annot in zip(annots_data['images'], annots_data['annotations']):
                assert frame_annot['id'] == data_annot['image_id'], 'Annotation missmatch'
                # read data from record
                data = interhand_utils.data_from_annotation(frame_annot, data_annot, annots_joint, annots_camera, images_path)
                img, bbox, joints, joints_valid, handness, handness_valid = data
                # preprocess example
                img, trans, _ = interhand_utils.preprocess_image(img, bbox, CFG.INTERHAND_PARAMS.IMG_SHAPE)
                # transform keypoints to image space
                joints[:, :2] = np.dot(trans, np.concatenate((joints[:, :2], np.ones((joints.shape[0], 1))), axis=1).T).T
                if img is None:
                    print(f'Invalid bbox - Frame ID:{frame_annot["id"]}')
                    continue
                # encode image in jpeg as BGR
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img_encoded = cv2.imencode('.jpg', img)[1].tobytes()
                joints = joints.reshape(-1).tolist()
                # construct example object and write it
                example = tf.train.Example(features=tf.train.Features(feature={
                    'image':
                        tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_encoded])),
                    'keypoints/points':
                        tf.train.Feature(float_list=tf.train.FloatList(value=joints)),
                    'keypoints/validity':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=joints_valid)),
                    'hand/handness':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=handness)),
                    'hand/handness_validity':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=[handness_valid])),
                }))
                writer.write(example.SerializeToString())


def random_chain(*iters):
    """Random sample generator chain

    Yields:
        list: list of generators to chain
    """
    while True:
        # randomly order generators and iterate over them
        for iterable in sample(iters, len(iters)):
            # try to get a generator yield
            try:
                # if a generator yields a value, break and yield it
                val = next(iterable)
                break
            # if generator ended, try the next one
            except StopIteration:
                continue
        # if all generators stopped iterating, terminate generation
        else:
            return
        yield val


def hands():
    """Construct hand pose estimation TFRecord files
    """
    print('Loading data...')
    # load coco training subset
    coco_train_data = coco_utils.load_annotation(CFG.COCO_DATASET_PATHS['train']['annotations'])
    keys = list(coco_train_data.keys())
    # split dataset to train-validation splits
    coco_train_keys, coco_val_keys = train_test_split(keys, train_size=0.9, random_state=42)
    # get coco train and val generators
    coco_train_generator = hands_utils.data_generator(coco_utils.data_from_annotation, coco_train_data, coco_train_keys, CFG.COCO_DATASET_PATHS['train']['images'], scale=1.3)
    coco_val_generator = hands_utils.data_generator(coco_utils.data_from_annotation, coco_train_data, coco_val_keys, CFG.COCO_DATASET_PATHS['train']['images'], scale=1.3)

    # load halpe training subset
    halpe_train_data = halpe_utils.load_annotation(CFG.HALPE_DATASET_PATHS['annotations'])
    keys = list(halpe_train_data.keys())
    # split dataset to train-validation splits
    halpe_train_keys, halpe_val_keys = train_test_split(keys, train_size=0.9, random_state=42)
    # get halpe train and val generators
    halpe_train_generator = hands_utils.data_generator(halpe_utils.data_from_annotation, halpe_train_data, halpe_train_keys, CFG.HALPE_DATASET_PATHS['images'], scale=1.3)
    halpe_val_generator = hands_utils.data_generator(halpe_utils.data_from_annotation, halpe_train_data, halpe_val_keys, CFG.HALPE_DATASET_PATHS['images'], scale=1.3)

    # load mwb training subset
    mwb_train_data = mwb_utils.load_annotation(CFG.MWB_DATASET_PATHS['manual']['train'])
    keys = list(mwb_train_data.keys())
    # split dataset to train-validation splits
    mwb_train_keys, mwb_val_keys = train_test_split(keys, train_size=0.9, random_state=42)
    # get mwb train and val generators
    mwb_train_generator = hands_utils.data_generator(mwb_utils.data_from_annotation, mwb_train_data, mwb_train_keys, CFG.MWB_DATASET_PATHS['manual']['train'], scale=1.3)
    mwb_val_generator = hands_utils.data_generator(mwb_utils.data_from_annotation, mwb_train_data, mwb_val_keys, CFG.MWB_DATASET_PATHS['manual']['train'], scale=1.3)

    # load rhd training subset
    rhd_train_data = rhd_utils.load_annotation(CFG.RHD_DATASET_PATHS['train']['annotations'])
    keys = list(rhd_train_data.keys())
    # split dataset to train-validation splits
    rhd_train_keys, rhd_val_keys = train_test_split(keys, train_size=0.9, random_state=42)
    # get rhd train and val generators
    rhd_train_generator = hands_utils.data_generator(rhd_utils.data_from_annotation, rhd_train_data, rhd_train_keys, CFG.RHD_DATASET_PATHS['train']['images'], scale=1.3)
    rhd_val_generator = hands_utils.data_generator(rhd_utils.data_from_annotation, rhd_train_data, rhd_val_keys, CFG.RHD_DATASET_PATHS['train']['images'], scale=1.3)

    # randomly chain train generators and chain validation generators
    train_generator = random_chain(coco_train_generator, halpe_train_generator, mwb_train_generator, rhd_train_generator)
    val_generator = chain(coco_val_generator, halpe_val_generator, mwb_val_generator, rhd_val_generator)

    # load coco testing subset
    coco_test_data = coco_utils.load_annotation(CFG.COCO_DATASET_PATHS['test']['annotations'])
    # make two instances of the same generator - one for combined testing subset and one for isolated coco testing subset
    coco_test_generator_partial = partial(hands_utils.data_generator, coco_utils.data_from_annotation, coco_test_data, coco_test_data.keys(), CFG.COCO_DATASET_PATHS['test']['images'], scale=1.3)
    coco_test_generator = coco_test_generator_partial()
    coco_test_generator_chain = coco_test_generator_partial()

    # load mwb testing subset
    mwb_test_data = mwb_utils.load_annotation(CFG.MWB_DATASET_PATHS['manual']['test'])
    # make two instances of the same generator - one for combined testing subset and one for isolated mwb testing subset
    mwb_test_generator_partial = partial(hands_utils.data_generator, mwb_utils.data_from_annotation, mwb_test_data, mwb_test_data.keys(), CFG.MWB_DATASET_PATHS['manual']['test'], scale=1.3)
    mwb_test_generator_chain = mwb_test_generator_partial()
    mwb_test_generator = mwb_test_generator_partial()

    # load rhd testing subset
    rhd_test_data = rhd_utils.load_annotation(CFG.RHD_DATASET_PATHS['test']['annotations'])
    # make two instances of the same generator - one for combined testing subset and one for isolated rhd testing subset
    rhd_test_generator_partial = partial(hands_utils.data_generator, rhd_utils.data_from_annotation, rhd_test_data, rhd_test_data.keys(), CFG.RHD_DATASET_PATHS['test']['images'], scale=1.3)
    rhd_test_generator_chain = rhd_test_generator_partial()
    rhd_test_generator = rhd_test_generator_partial()

    # chain one of the two testing generator instances into a unified testing subset
    test_generator = chain(coco_test_generator_chain, mwb_test_generator_chain, rhd_test_generator_chain)

    # data distribution counters
    num_classes = {'train': [0, 0], 'val': [0, 0], 'test': [0, 0]}
    dataset_classes = {'train': {'coco': [0, 0], 'halpe': [0, 0], 'mwb': [0, 0], 'rhd': [0, 0]},
                       'val': {'coco': [0, 0], 'halpe': [0, 0], 'mwb': [0, 0], 'rhd': [0, 0]},
                       'test': {'coco': [0, 0], 'halpe': [0, 0], 'mwb': [0, 0], 'rhd': [0, 0]}}
    samples = {'train': 0, 'val': 0, 'test': 0}
    dataset_samples = {'train': {'coco': 0, 'halpe': 0, 'mwb': 0, 'rhd': 0},
                       'val': {'coco': 0, 'halpe': 0, 'mwb': 0, 'rhd': 0},
                       'test': {'coco': 0, 'halpe': 0, 'mwb': 0, 'rhd': 0}}
    rejects = {'train': {'coco': 0, 'halpe': 0, 'mwb': 0, 'rhd': 0},
               'val': {'coco': 0, 'halpe': 0, 'mwb': 0, 'rhd': 0},
               'test': {'coco': 0, 'halpe': 0, 'mwb': 0, 'rhd': 0}}
    # for every subset
    for generator, subset in zip([train_generator, val_generator, test_generator, coco_test_generator, mwb_test_generator, rhd_test_generator], ['train', 'val', 'test', 'coco_test', 'mwb_test', 'rhd_test']):
        print(f'Parsing {subset} subset...')
        # open TFRecord file
        with tf.io.TFRecordWriter(CFG.HANDS_TFRECORD_PATHS[subset]) as writer:
            # for every annotation entry
            for image, bbox, kpts, valid, cls, dataset in generator:
                # check min hand bbox size
                if (bbox[2] * bbox[3]) <= CFG.HANDS_PARAMS.MIN_BBOX_AREA:
                    if subset in ['train', 'test', 'val']:
                        rejects[subset][dataset] += 1
                    continue
                # preprocess image and keypoints
                img, trans, _ = hands_utils.preprocess_image(image, bbox, CFG.HANDS_PARAMS.IMG_SHAPE)
                kpts = np.dot(trans, np.concatenate((kpts[:, :2], np.ones((kpts.shape[0], 1))), axis=1).T).T
                # encode image in jpeg as BGR
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img_encoded = cv2.imencode('.jpg', img)[1].tobytes()
                kpts = kpts.reshape(-1).tolist()
                valid = (valid > 0).reshape(-1).astype(np.int64).tolist()
                # construct example and write it
                example = tf.train.Example(features=tf.train.Features(feature={
                    'image':
                        tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_encoded])),
                    'keypoints':
                        tf.train.Feature(float_list=tf.train.FloatList(value=kpts)),
                    'validity':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=valid)),
                    'class':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=[int(cls)])),
                }))
                writer.write(example.SerializeToString())
                if subset in ['train', 'test', 'val']:
                    num_classes[subset][int(cls)] += 1
                    dataset_classes[subset][dataset][int(cls)] += 1
                    samples[subset] += 1
                    dataset_samples[subset][dataset] += 1
    print(f'Examples: {num_classes}')
    print('Examples per dataset:')
    for key, val in dataset_classes.items():
        print(f'\t{key}: {val}')
    print(f'Samples: {samples}')
    print('Samples per dataset:')
    for key, val in dataset_samples.items():
        print(f'\t{key}: {val}')
    print(f'Class weighting: {sum(num_classes["train"]) / (2 * np.array(num_classes["train"]))}')
    print('Rejected samples:')
    for key, val in dataset_samples.items():
        print(f'\t{key}: {val}')


def wflw():
    """Construct face pose estimation TFRecord files
    """
    train_data = wflw_utils.read_annotation(CFG.WFLW_DATASET_PATHS['annotations']['train'])
    keys = list(train_data.keys())
    train_keys, val_keys = train_test_split(keys, train_size=0.9, random_state=42)
    shuffle(train_keys)
    test_data = wflw_utils.read_annotation(CFG.WFLW_DATASET_PATHS['annotations']['test'])
    test_keys = test_data.keys()
    # data distribution counters
    samples = {'train': 0, 'val': 0, 'test': 0}
    # for every subset
    for keys, subset in zip([train_keys, val_keys, test_keys], ['train', 'val', 'test']):
        # validation entries are from training dictionary, testing subset has its own dict
        data_dict = test_data if subset == 'test' else train_data
        # open TFRecord file
        with tf.io.TFRecordWriter(CFG.WFLW_TFRECORD_PATHS[subset]) as writer:
            # for every annotation entry
            for key in keys:
                samples[subset] += 1
                annot = data_dict[key]
                # preprocess image and keypoints
                img, bbox, keypoints, cls = wflw_utils.data_from_annotation(annot, CFG.WFLW_DATASET_PATHS['images'], scale=1.25)
                img, trans, _ = wflw_utils.preprocess_image(img, bbox, CFG.WFLW_PARAMS.IMG_SHAPE)
                keypoints = np.dot(trans, np.concatenate((keypoints, np.ones((keypoints.shape[0], 1))), axis=1).T).T
                # encode image in jpeg as BGRi
                mg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img_encoded = cv2.imencode('.jpg', img)[1].tobytes()
                keypoints = keypoints.reshape(-1).tolist()
                cls = [int(c) for c in cls]
                # construct example and write it
                example = tf.train.Example(features=tf.train.Features(feature={
                    'image':
                        tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_encoded])),
                    'keypoints':
                        tf.train.Feature(float_list=tf.train.FloatList(value=keypoints)),
                    'attribute':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=cls)),
                }))
                writer.write(example.SerializeToString())
    print(samples)


if __name__ == '__main__':
    # interhand()
    # wflw()
    hands()
