#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.

import os
import numpy as np
import sys

sys.path.insert(0, 'third_party/CenterNet2/projects/CenterNet2/')
sys.path.insert(0, 'third_party/Deformable-DETR')
from detic.data.tar_dataset import _TarDataset, DiskTarDataset
import pickle
import io
import gzip
import time


class _RawTarDataset(object):

    def __init__(self, filename, indexname, preload=False):
        self.filename = filename
        self.names = []
        self.offsets = []

        for l in open(indexname):
            ll = l.split()
            a, b, c = ll[:3]
            offset = int(b[:-1])
            if l.endswith('** Block of NULs **\n'):
                self.offsets.append(offset)
                break
            else:
                if c.endswith('JPEG'):
                    self.names.append(c)
                    self.offsets.append(offset)
                else:
                    # ignore directories
                    pass
        if preload:
            self.data = np.memmap(filename, mode='r', dtype='uint8')
        else:
            self.data = None

    def __len__(self):
        return len(self.names)

    def __getitem__(self, idx):
        if self.data is None:
            self.data = np.memmap(self.filename, mode='r', dtype='uint8')
        ofs = self.offsets[idx] * 512
        fsize = 512 * (self.offsets[idx + 1] - self.offsets[idx])
        data = self.data[ofs:ofs + fsize]

        if data[:13].tostring() == '././@LongLink':
            data = data[3 * 512:]
        else:
            data = data[512:]

        # just to make it more fun a few JPEGs are GZIP compressed...
        # catch this case
        if tuple(data[:2]) == (0x1f, 0x8b):
            s = io.StringIO(data.tostring())
            g = gzip.GzipFile(None, 'r', 0, s)
            sdata = g.read()
        else:
            sdata = data.tostring()
        return sdata



def preprocess():
    # Follow https://github.com/Alibaba-MIIL/ImageNet21K/blob/main/dataset_preprocessing/processing_script.sh
    # Expect 12358684 samples with 11221 classes
    # ImageNet folder has 21841 classes (synsets)

    i22kdir = '/datasets01/imagenet-22k/062717/'
    i22ktarlogs = '/checkpoint/imisra/datasets/imagenet-22k/tarindex'
    class_names_file = '/checkpoint/imisra/datasets/imagenet-22k/words.txt'

    output_dir = '/checkpoint/zhouxy/Datasets/ImageNet/metadata-22k/'
    i22knpytarlogs = '/checkpoint/zhouxy/Datasets/ImageNet/metadata-22k/tarindex_npy'
    print('Listing dir')
    log_files = os.listdir(i22ktarlogs)
    log_files = [x for x in log_files if x.endswith(".tarlog")]
    log_files.sort()
    chunk_datasets = []
    dataset_lens = []
    min_count = 0
    create_npy_tarlogs = True
    print('Creating folders')
    if create_npy_tarlogs:
        os.makedirs(i22knpytarlogs, exist_ok=True)
        for log_file in log_files:
            syn = log_file.replace(".tarlog", "")
            dataset = _RawTarDataset(os.path.join(i22kdir, syn + ".tar"),
                                    os.path.join(i22ktarlogs, syn + ".tarlog"),
                                    preload=False)
            names = np.array(dataset.names)
            offsets = np.array(dataset.offsets, dtype=np.int64)
            np.save(os.path.join(i22knpytarlogs, f"{syn}_names.npy"), names)
            np.save(os.path.join(i22knpytarlogs, f"{syn}_offsets.npy"), offsets)

    os.makedirs(output_dir, exist_ok=True)

    start_time = time.time()
    for log_file in log_files:
        syn = log_file.replace(".tarlog", "")
        dataset = _TarDataset(os.path.join(i22kdir, syn + ".tar"), i22knpytarlogs)
        # dataset = _RawTarDataset(os.path.join(i22kdir, syn + ".tar"),
        #                             os.path.join(i22ktarlogs, syn + ".tarlog"),
        #                             preload=False)
        dataset_lens.append(len(dataset))
    end_time = time.time()
    print(f"Time {end_time - start_time}")


    dataset_lens = np.array(dataset_lens)
    dataset_valid = dataset_lens > min_count

    syn2class = {}
    with open(class_names_file) as fh:
        for line in fh:
            line = line.strip().split("\t")
            syn2class[line[0]] = line[1]

    tarlog_files = []
    class_names = []
    tar_files = []
    for k in range(len(dataset_valid)):
        if not dataset_valid[k]:
            continue
        syn = log_files[k].replace(".tarlog", "")
        tarlog_files.append(os.path.join(i22ktarlogs, syn + ".tarlog"))
        tar_files.append(os.path.join(i22kdir, syn + ".tar"))
        class_names.append(syn2class[syn])

    tarlog_files = np.array(tarlog_files)
    tar_files = np.array(tar_files)
    class_names = np.array(class_names)
    print(f"Have {len(class_names)} classes and {dataset_lens[dataset_valid].sum()} samples")

    np.save(os.path.join(output_dir, "tarlog_files.npy"), tarlog_files)
    np.save(os.path.join(output_dir, "tar_files.npy"), tar_files)
    np.save(os.path.join(output_dir, "class_names.npy"), class_names)
    np.save(os.path.join(output_dir, "tar_files.npy"), tar_files)


if __name__ == "__main__":
    preprocess()