from functools import partial
import numpy
import os
import re
import random
import signal
from collections import OrderedDict
from scipy.misc import imread
from unicsv import DictUnicodeReader
from multiprocessing import Pool, cpu_count
from multiprocessing.pool import ThreadPool
from scipy.ndimage.interpolation import zoom

class AbstractSegmentation:
    def all_names(self, category, j):
        raise NotImplementedError
    def size(self, split=None):
        return 0
    def filename(self, i):
        raise NotImplementedError
    def metadata(self, i):
        return self.filename(i)
    @classmethod
    def resolve_segmentation(cls, m):
        return {}

    def name(self, category, i):
        '''
        Default implemtnation for segmentation_data,
        utilizing all_names.
        '''
        all_names = self.all_names(category, i)
        return all_names[0] if len(all_names) else ''

    def segmentation_data(self, category, i, c=0, full=False):
        '''
        Default implemtnation for segmentation_data,
        utilizing metadata and resolve_segmentation.
        '''
        segs = self.resolve_segmentation(
                self.metadata(i), categories=[category])
        if category not in segs:
            return 0
        data = segs[category]
        if not full and len(data.shape) >= 3:
            return data[0]
        return data

class SegmentationData(AbstractSegmentation):
    '''
    Represents and loads a multi-channel segmentation represented with
    a series of csv files: index.csv lists the images together with
    any label data avilable in each category; category.csv lists
    the categories of segmentations available; and label.csv lists the
    numbers used to describe each label class. In addition, the categories
    each have a separate c_*.csv file describing a dense coding of labels.
    '''
    def __init__(self, directory, categories=None, require_all=False):
        directory = os.path.expanduser(directory)
        self.directory = directory
        with open(os.path.join(directory, 'index.csv')) as f:
            self.image = [decode_index_dict(r) for r in DictUnicodeReader(f)]
        with open(os.path.join(directory, 'category.csv')) as f:
            self.category = OrderedDict()
            for row in DictUnicodeReader(f):
                self.category[row['name']] = row
        with open(os.path.join(directory, 'label.csv')) as f:
            label_data = [decode_label_dict(r) for r in DictUnicodeReader(f)]
        self.label = build_dense_label_array(label_data)
        if categories is not None:
            # Filter out unused categories
            categories = set([c for c in categories if c in self.category])
            for cat in self.category.keys():
                if cat not in categories:
                    del self.category[cat]
        else:
            categories = self.category.keys()
        # Filter out images with insufficient data
        filter_fn = partial(
                index_has_all_data if require_all else index_has_any_data,
                categories=categories)
        self.image = [row for row in self.image if filter_fn(row)]
        # Build dense remapping arrays for labels, so that you can
        # get dense ranges of labels for each category.
        self.category_map = {}
        self.category_unmap = {}
        self.category_label = {}
        for cat in self.category:
            with open(os.path.join(directory, 'c_%s.csv' % cat)) as f:
                c_data = [decode_label_dict(r) for r in DictUnicodeReader(f)]
            self.category_unmap[cat], self.category_map[cat] = (
                    build_numpy_category_map(c_data))
            self.category_label[cat] = build_dense_label_array(
                    c_data, key='code')

    def all_names(self, category, j):
        '''All English synonyms for the given label'''
        if category is not None:
            j = self.category_unmap[category][j]
        return [self.label[j]['name']] + self.label[i]['syns']

    def size(self, split=None):
        '''The number of images in this data set.'''
        if split is None:
            return len(self.image)
        return len([im for im in self.image if im['split'] == split])

    def filename(self, i):
        '''The filename of the ith jpeg (original image).'''
        return os.path.join(self.directory, 'images', self.image[i]['image'])

    def split(self, i):
        '''Which split contains item i.'''
        return self.image[i]['split']

    def metadata(self, i):
        '''Extract metadata for image i, For efficient data loading.'''
        return self.directory, self.image[i]

    meta_categories = ['image', 'split', 'ih', 'iw', 'sh', 'sw']

    @classmethod
    def resolve_segmentation(cls, m, categories=None):
        '''
        Resolves a full segmentation, potentially in a differenct process,
        for efficient multiprocess data loading.
        '''
        directory, row = m
        result = {}
        for cat, d in row.items():
            if cat in cls.meta_categories:
                continue
            if not wants(cat, categories):
                continue
            if all(isinstance(data, int) for data in d):
                result[cat] = d
                continue
            out = numpy.empty((len(d), row['sh'], row['sw']), dtype=numpy.int16)
            for i, channel in enumerate(d):
                if isinstance(channel, int):
                    out[i] = channel
                else:
                    rgb = imread(os.path.join(directory, 'images', channel))
                    out[i] = rgb[:,:,0] + rgb[:,:,1] * 256
            result[cat] = out
        return result, (row['sh'], row['sw'])

    def label_size(self, category=None):
        '''
        Returns the number of distinct labels (plus zero), i.e., one
        more than the maximum label number.  If a category is specified,
        returns the number of distinct labels within that category.
        '''
        if category is None:
            return len(self.label)
        else:
            return len(self.category_unmap[category])

    def name(self, category, j):
        '''
        Returns an English name for the jth label.  If a category is
        specified, returns the name for the category-specific nubmer j.
        If category=None, then treats j as a fully unified index number.
        '''
        if category is not None:
            j = self.category_unmap[category][j]
        return self.label[j]['name']

    def frequency(self, category, j):
        '''
        Returns the number of images for which the label appears.
        '''
        if category is not None:
            return self.category_label[category][j]['frequency']
        return self.label[j]['frequency']

    def coverage(self, category, j):
        '''
        Returns the pixel coverage of the label in units of whole-images.
        '''
        if category is not None:
            return self.category_label[category][j]['coverage']
        return self.label[j]['coverage']

    def category_names(self):
        '''
        Returns the set of category names.
        '''
        return self.category.keys()

    def category_frequency(self, category):
        '''
        Returns the number of images touched by a category.
        '''
        return float(self.category[category]['frequency'])

    def primary_categories_per_index(self, categories=None):
        '''
        Returns an array of primary category numbers for each label, where
        catagories are indexed according to the list of categories passed,
        or self.category_names() if none.
        '''
        if categories is None:
            categories = self.category_names()
        # Make lists which are nonzero for labels in a category
        catmap = {}
        for cat in categories:
            imap = self.category_index_map(cat)
            if len(imap) < self.label_size(None):
                imap = numpy.concatenate((imap, numpy.zeros(
                    self.label_size(None) - len(imap), dtype=imap.dtype)))
            catmap[cat] = imap
        # For each label, find the category with maximum coverage.
        result = []
        for i in range(self.label_size(None)):
            maxcov, maxcat = max(
                    (self.coverage(cat, catmap[cat][i])
                        if catmap[cat][i] else 0, ic)
                    for ic, cat in enumerate(categories))
            result.append(maxcat)
        # Return the max-coverage cateogry for each label.
        return numpy.array(result)

    def segmentation_data(self, category, i, c=0, full=False, out=None):
        '''
        Returns a 2-d numpy matrix with segmentation data for the ith image,
        restricted to the given category.  By default, maps all label numbers
        to the category-specific dense mapping described in the c_*.csv
        listing; but can be asked to expose the fully unique indexing by
        using full=True.
        '''
        row = self.image[i]
        data_channels = row.get(category, ())
        if c >= len(data_channels):
            channel = 0  # Deal with unlabeled data in this category
        else:
            channel = data_channels[c]
        if out is None:
            out = numpy.empty((row['sh'], row['sw']), dtype=numpy.int16)
        if isinstance(channel, int):
            if not full:
                channel = self.category_map[category][channel]
            out[:,:] = channel  # Single-label for the whole image
            return out
        png = imread(os.path.join(self.directory, 'images', channel))
        if full:
            # Full case: just combine png channels.
            out[...] = png[:,:,0] + png[:,:,1] * 256
        else:
            # Dense case: combine png channels and apply the category map.
            catmap = self.category_map[category]
            out[...] = catmap[png[:,:,0] + png[:,:,1] * 256]
        return out

    def full_segmentation_data(self, i,
            categories=None, max_depth=None, out=None):
        '''
        Returns a 3-d numpy tensor with segmentation data for the ith image,
        with multiple layers represnting multiple lables for each pixel.
        The depth is variable depending on available data but can be
        limited to max_depth.
        '''
        row = self.image[i]
        if categories:
            groups = [d for cat, d in row.items() if cat in categories and d]
        else:
            groups = [d for cat, d in row.items() if d and (
                cat not in self.meta_categories)]
        depth = sum(len(c) for c in groups)
        if max_depth is not None:
            depth = min(depth, max_depth)
        # Allocate an array if not already allocated.
        if out is None:
            out = numpy.empty((depth, row['sh'], row['sw']), dtype=numpy.int16)
        i = 0
        # Stack up the result segmentation one channel at a time
        for group in groups:
            for channel in group:
                if isinstance(channel, int):
                    out[i] = channel
                else:
                    png = imread(
                            os.path.join(self.directory, 'images', channel))
                    out[i] = png[:,:,0] + png[:,:,1] * 256
                i += 1
                if i == depth:
                    return out
        # Return above when we get up to depth
        assert False

    def category_index_map(self, category):
        return numpy.array(self.category_map[category])

def build_dense_label_array(label_data, key='number', allow_none=False):
    '''
    Input: set of rows with 'number' fields (or another field name key).
    Output: array such that a[number] = the row with the given number.
    '''
    result = [None] * (max([d[key] for d in label_data]) + 1)
    for d in label_data:
        result[d[key]] = d
    # Fill in none
    if not allow_none:
        example = label_data[0]
        def make_empty(k):
            return dict((c, k if c is key else type(v)())
                    for c, v in example.items())
        for i, d in enumerate(result):
            if d is None:
                result[i] = dict(make_empty(i))
    return result

def build_numpy_category_map(map_data, key1='code', key2='number'):
    '''
    Input: set of rows with 'number' fields (or another field name key).
    Output: array such that a[number] = the row with the given number.
    '''
    results = list(numpy.zeros((max([d[key] for d in map_data]) + 1),
            dtype=numpy.int16) for key in (key1, key2))
    for d in map_data:
        results[0][d[key1]] = d[key2]
        results[1][d[key2]] = d[key1]
    return results

def decode_label_dict(row):
    result = {}
    for key, val in row.items():
        if key == 'category':
            result[key] = dict((c, int(n))
                for c, n in [re.match('^([^(]*)\(([^)]*)\)$', f).groups()
                    for f in val.split(';')])
        elif key == 'name':
            result[key] = val
        elif key == 'syns':
            result[key] = val.split(';')
        elif re.match('^\d+$', val):
            result[key] = int(val)
        elif re.match('^\d+\.\d*$', val):
            result[key] = float(val)
        else:
            result[key] = val
    return result

def decode_index_dict(row):
    result = {}
    for key, val in row.items():
        if key in ['image', 'split']:
            result[key] = val
        elif key in ['sw', 'sh', 'iw', 'ih']:
            result[key] = int(val)
        else:
            item = [s for s in val.split(';') if s]
            for i, v in enumerate(item):
                if re.match('^\d+$', v):
                    item[i] = int(v)
            result[key] = item
    return result

def index_has_any_data(row, categories):
    for c in categories:
        for data in row[c]:
            if data: return True
    return False

def index_has_all_data(row, categories):
    for c in categories:
        cat_has = False
        for data in row[c]:
            if data:
                cat_has = True
                break
        if not cat_has:
            return False
    return True

class SegmentationPrefetcher:
    '''
    SegmentationPrefetcher will prefetch a bunch of segmentation
    images using a multiprocessing pool, so you do not have to wait
    around while the files get opened and decoded.  Just request
    batches of images and segmentations calling fetch_batch().
    '''
    def __init__(self, segmentation, split=None, randomize=False,
            segmentation_shape=None, categories=None, once=False,
            start=None, end=None, batch_size=4, ahead=24, thread=False):
        '''
        Constructor arguments:
        segmentation: The AbstractSegmentation to load.
        split: None for no filtering, or 'train' or 'val' etc.
        randomize: True to randomly shuffle order, or a random seed.
        categories: a list of categories to include in each batch.
        batch_size: number of data items for each batch.
        ahead: the number of data items to prefetch ahead.
        '''
        self.segmentation = segmentation
        self.split = split
        self.randomize = randomize
        self.random = random.Random()
        if randomize is not True:
            self.random.seed(randomize)
        self.categories = categories
        self.once = once
        self.batch_size = batch_size
        self.ahead = ahead
        # Initialize the multiprocessing pool
        n_procs = cpu_count()
        if thread:
            self.pool = ThreadPool(processes=n_procs)
        else:
            original_sigint_handler = setup_sigint()
            self.pool = Pool(processes=n_procs, initializer=setup_sigint)
            restore_sigint(original_sigint_handler)
        # Prefilter the image indexes of interest
        if start is None:
            start = 0
        if end is None:
            end = segmentation.size()
        self.indexes = range(start, end)
        if split:
            self.indexes = [i for i in self.indexes
                    if segmentation.split(i) == split]
        if self.randomize:
            self.random.shuffle(self.indexes)
        self.index = 0
        self.result_queue = []
        self.segmentation_shape = segmentation_shape
        # Get dense catmaps
        self.catmaps = [
                segmentation.category_index_map(cat) if cat != 'image' else None
                for cat in categories]

    def next_job(self):
        if self.index < 0:
            return None
        j = self.indexes[self.index]
        result = (j,
                self.segmentation.__class__,
                self.segmentation.metadata(j),
                self.segmentation.filename(j),
                self.categories,
                self.segmentation_shape)
        self.index += 1
        if self.index >= len(self.indexes):
            if self.once:
                self.index = -1
            else:
                self.index = 0
                if self.randomize:
                    # Reshuffle every time through
                    self.random.shuffle(self.indexes)
        return result

    def batches(self):
        '''Iterator for all batches'''
        while True:
            batch = self.fetch_batch()
            if batch is None:
                raise StopIteration
            yield batch

    def fetch_batch(self):
        '''Returns a single batch as an array of dictionaries.'''
        try:
            self.refill_tasks()
            if len(self.result_queue) == 0:
                return None
            result = self.result_queue.pop(0)
            return result.get(31536000)
        except KeyboardInterrupt:
            print("Caught KeyboardInterrupt, terminating workers")
            self.pool.terminate()
            raise

    def fetch_tensor_batch(self, bgr_mean=None, global_labels=False):
        '''Iterator for batches as arrays of tensors.'''
        batch = self.fetch_batch()
        return self.form_caffe_tensors(batch, bgr_mean, global_labels)

    def tensor_batches(self, bgr_mean=None, global_labels=False):
        '''Returns a single batch as an array of tensors, one per category.'''
        while True:
            batch = self.fetch_tensor_batch(
                    bgr_mean=bgr_mean, global_labels=global_labels)
            if batch is None:
                raise StopIteration
            yield batch

    def form_caffe_tensors(self, batch, bgr_mean=None, global_labels=False):
        # Assemble a batch in [{'cat': data,..},..] format into
        # an array of batch tensors, the first for the image, and the
        # remaining for each category in self.categories, in order.
        # This also applies a random flip if needed
        if batch is None:
            return None
        batches = [[] for c in self.categories]
        for record in batch:
            default_shape = (1, record['sh'], record['sw'])
            for c, cat in enumerate(self.categories):
                if cat == 'image':
                    # Normalize image with right RGB order and mean
                    batches[c].append(normalize_image(
                        record[cat], bgr_mean))
                elif global_labels:
                    batches[c].append(normalize_label(
                        record[cat], default_shape, flatten=True))
                else:
                    catmap = self.catmaps[c]
                    batches[c].append(catmap[normalize_label(
                        record[cat], default_shape, flatten=True)])
        return [numpy.concatenate(tuple(m[numpy.newaxis] for m in b))
                for b in batches]

    def refill_tasks(self):
        # It will call the sequencer to ask for a sequence
        # of batch_size jobs (indexes with categories)
        # Then it will call pool.map_async
        while len(self.result_queue) < self.ahead:
            data = []
            while len(data) < self.batch_size:
                job = self.next_job();
                if job is None:
                    break
                data.append(job)
            if len(data) == 0:
                return
            self.result_queue.append(self.pool.map_async(prefetch_worker, data))

    def close():
        while len(self.result_queue):
            result = self.result_queue.pop(0)
            if result is not None:
                result.wait(0.001)
        self.pool.close()
        self.poool.cancel_join_thread()

def prefetch_worker(d):
    if d is None:
        return None
    j, typ, m, fn, categories, segmentation_shape = d
    segs, shape = typ.resolve_segmentation(m, categories=categories)
    if segmentation_shape is not None:
        for k, v in segs.items():
            segs[k] = scale_segmentation(v, segmentation_shape)
        shape = segmentation_shape
    # Some additional metadata to provide
    segs['sh'], segs['sw'] = shape
    segs['i'] = j
    segs['fn'] = fn
    if categories is None or 'image' in categories:
        segs['image'] = imread(fn)
    return segs

def scale_segmentation(segmentation, dims, crop=False):
    '''
    Zooms a 2d or 3d segmentation to the given dims, using nearest neighbor.
    '''
    shape = numpy.shape(segmentation)
    if len(shape) < 2 or shape[-2:] == dims:
        return segmentation
    peel = (len(shape) == 2)
    if peel:
        segmentation = segmentation[numpy.newaxis]
    levels = segmentation.shape[0]
    result = numpy.zeros((levels, ) + dims,
            dtype=segmentation.dtype)
    ratio = (1,) + tuple(res / float(orig)
            for res, orig in zip(result.shape[1:], segmentation.shape[1:]))
    if not crop:
        safezoom(segmentation, ratio, output=result, order=0)
    else:
        ratio = max(ratio[1:])
        height = int(round(dims[0] / ratio))
        hmargin = (segmentation.shape[0] - height) // 2
        width = int(round(dims[1] / ratio))
        wmargin = (segmentation.shape[1] - height) // 2
        safezoom(segmentation[:, hmargin:hmargin+height,
            wmargin:wmargin+width],
            (1, ratio, ratio), output=result, order=0)
    if peel:
        result = result[0]
    return result

def safezoom(array, ratio, output=None, order=0):
    '''Like numpy.zoom, but does not crash when the first dimension
    of the array is of size 1, as happens often with segmentations'''
    dtype = array.dtype
    if array.dtype == numpy.float16:
        array = array.astype(numpy.float32)
    if array.shape[0] == 1:
        if output is not None:
            output = output[0,...]
        result = zoom(array[0,...], ratio[1:],
                output=output, order=order)
        if output is None:
            output = result[numpy.newaxis]
    else:
        result = zoom(array, ratio, output=output, order=order)
        if output is None:
            output = result
    return output.astype(dtype)

def setup_sigint():
    import threading
    if not isinstance(threading.current_thread(), threading._MainThread):
        return None
    return signal.signal(signal.SIGINT, signal.SIG_IGN)

def restore_sigint(original):
    import threading
    if not isinstance(threading.current_thread(), threading._MainThread):
        return
    if original is None:
        original = signal.SIG_DFL
    signal.signal(signal.SIGINT, original)

def wants(what, option):
    if option is None:
        return True
    return what in option

def normalize_image(rgb_image, bgr_mean):
    """
    Load input image and preprocess for Caffe:
    - cast to float
    - switch channels RGB -> BGR
    - subtract mean
    - transpose to channel x height x width order
    """
    img = numpy.array(rgb_image, dtype=numpy.float32)
    if (img.ndim == 2):
        img = numpy.repeat(img[:,:,None], 3, axis = 2)
    img = img[:,:,::-1]
    if bgr_mean is not None:
        img -= bgr_mean
    img = img.transpose((2,0,1))
    return img

def normalize_label(label_data, shape, flatten=False):
    """
    Given a 0, 1, 2, or 3-dimensional label_data and a default
    shape of the form (1, y, x), returns a 3d tensor by
    """
    dims = len(numpy.shape(label_data))
    if dims <= 2:
        # Scalar data on this channel: fill shape
        if dims == 1:
            if flatten:
                label_data = label_data[0] if len(label_data) else 0
            else:
                return (numpy.ones(shape, dtype=numpy.int16) *
                        numpy.asarray(label_data, dtype=numpy.int16)
                            [:, numpy.newaxis, numpy.newaxis])
        return numpy.full(shape, label_data, dtype=numpy.int16)
    else:
        if dims == 3:
            if flatten:
                label_data = label_data[0]
            else:
                return label_data
        return label_data[numpy.newaxis]

if __name__ == '__main__':
    import numpy
    ds = SegmentationData('~/bulk/small_test')
    print 'sample has size', ds.size()
    i = ds.size() - 1
    print 'last filename is', ds.filename(i)
    print 'image shape read was', imread(ds.filename(i)).shape
    sd = ds.segmentation_data('texture', i)
    print 'texture seg shape read was', sd.shape
    fs = ds.full_segmentation_data(i)
    print 'texture seg datatype', sd.dtype
    found = list(numpy.bincount(sd.ravel()).nonzero()[0])
    print 'unique textures seen', ';'.join(
            ds.name('texture', f) for f in found)
    fs, shape = ds.resolve_segmentation(ds.metadata(i))
    print 'full segmentation', shape, ', '.join(fs.keys())
    pf = SegmentationPrefetcher(ds, categories=['texture'],
            segmentation_shape=(96, 96), once=True)
    print ds.size()
    for i, dat in enumerate(pf.batches()):
        for d in dat:
            print i, d['i'], d.keys()
        if i >= 5:
            break
    for i, dat in enumerate(pf.tensor_batches()):
        for d in dat:
            print i, numpy.shape(d), numpy.max(d)
        if i >= 5:
            break
    print 'count of object labels', ds.label_size('object')
    for j in range(min(10, ds.label_size('object'))):
        print ds.name('object', j), 'freq', ds.frequency('object', j), (
                'cov'), ds.coverage('object', j)


# coding = utf-8
import sys
image_size = 48
sys.path.append('..\{0:d}net'.format(image_size))
import numpy as np
import cv2
import os
import numpy.random as npr
from utils import IoU
import _pickle as pickle
import matplotlib.pyplot as plt

anno_file = r"trainImageList.txt"
landmark_file = r"F:\YiChunYu_File\face\keras-mtcnn-master\data\landmark.txt"
im_dir = r"F:\YiChunYu_File\face\keras-mtcnn-master\data\lfw_5590"
# f_landmark = open("landmark.txt","w")

with open(anno_file, 'r') as f:
    annotations = f.readlines()
    '''提取Landmark数据'''
    # print(annotations[0].split(" "))
    # print(annotations[0].split()[-10:])
    # for annotation in annotations:
    #     landmark_val = annotation.split(" ")[-10:]
    #     land = annotation.split(" ")[0]
    #     for i in range(len(landmark_val)):
    #         land = land + " " + landmark_val[i]
    #         # print(land)
    #     f_landmark.write(land+ "\n")

# print(len(annotations[0].split()[5:]))
with open(landmark_file, 'r') as f2:
    landmark_positions_list = f2.readlines()
    # print(landmark_positions_list)
num = len(annotations)
print ("{:d} pics in total".format(num))
p_idx = 0 # positive
n_idx = 0 # negative
d_idx = 0 # dont care
idx = 0
box_idx = 0
pts_idx = 0
cls_list = []
roi_list = []

pos_cls_list = []
pos_roi_list = []
neg_cls_list = []
part_roi_list = []
pts_list = []
for id_annos in range (2,10000):  # too large, use only first 10000 data
    annotation = annotations[id_annos]
    # annotation = annotation.strip().split(' ')
    annotation = annotation.strip().split()
    annotation1 = annotation[:5]
    tmp = annotation1[5:]
    im_path = annotation1[0]
    bbox = list(map(float, annotation1[1:]) )
    boxes = np.array(bbox, dtype=np.float32).reshape(-1, 4)
    # print(landmark_positions_list)
    landmark_positions = landmark_positions_list[id_annos]
    landmark_positions = landmark_positions.strip().split()
    pts_list_raw = list(map(float, landmark_positions[1:]))   # need to convert to relative pts positions
    # print aa

    # file_basename = os.path.splitext(im_path)[0]
    print(annotation1[0])
    if not os.path.exists(annotation1[0]):
        continue
    img = cv2.imread(annotation1[0])  # the celebA database seems to have used wrong name
    idx += 1
    if idx % 100 == 0:
        print (idx, "images done")

    height, width, channel = img.shape
    print(str(height)+ "=================")

    for box in boxes:
        # box (x_left, y_top, x_right, y_bottom)
        x1, y1, w, h = box
        w_box = w
        h_box = h
        # w = x2 - x1 + 1
        # h = y2 - y1 + 1
        x2 = w + x1 -1
        y2 = h + y1 -1
        box = np.array([x1,y1,x2,y2])

        # ignore small faces
        # in case the ground truth boxes of small faces are not accurate
        if max(w, h) < 40 or x1 < 0 or y1 < 0:
            continue

        # generate positive examples and part faces
        for i in range(20):
            # size = npr.randint(int(min(w, h) * 0.9), np.ceil(1.25 * max(w, h)))# enlarge size here
            #
            # # delta here is the offset of box center
            # delta_x = npr.randint(int(-w * 0.1), int(w * 0.1))   # reduce uncertainty here
            # delta_y = npr.randint(int(-h * 0.1), int (h * 0.1))
            # nx1 = max( int(x1 + w / 2 + delta_x - size / 2 ), 0)   # position of selected box part
            # ny1 = max( int(y1 + h / 2 + delta_y - size / 2), 0)

            size = int(min(w,h))  # enlarge size here
            # delta here is the offset of box center

            nx1 = max(int (x1), 0)  # position of selected box part
            ny1 = max(int (y1), 0)

            nx2 = nx1 + size
            ny2 = ny1 + size

            if nx2 > width or ny2 > height:
                continue
            crop_box = np.array([nx1, ny1, nx2, ny2])

            offset_x1 = (x1 - nx1) / float(size)
            offset_y1 = (y1 - ny1) / float(size)
            offset_x2 = (x2 - nx2) / float(size)
            offset_y2 = (y2 - ny2) / float(size)
            cropped_im = img[ny1 : ny2, nx1 : nx2, :]
            resized_im = cv2.resize(cropped_im, (image_size, image_size), interpolation=cv2.INTER_LINEAR)

            box_ = box.reshape(1, -1)

            if IoU(crop_box, box_) >= 0.65:   # output both position vector and class label
                im = resized_im
                h, w, ch = resized_im.shape
                if h != image_size or w != image_size:
                    im = cv2.resize(im, (image_size, image_size))
                im = np.swapaxes(im, 0, 2)
                im = (im - 127.5) / 127.5
                label = 1
                roi = [-1, -1, -1, -1]
                pts = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1]
                pos_cls_list.append([im, label, roi])
                roi = [float(offset_x1), float(offset_y1), float(offset_x2), float(offset_y2)]
                pos_roi_list.append([im, label, roi])
                p_idx += 1

                (px1,py1,px2,py2,px3,py3,px4,py4,px5,py5) = pts_list_raw
                if min ([px1, px2, px3,px4, px5]) > nx1 and min ([py1, py2, py3,py4, py5]) > ny1 \
                        and max ([px1, px2, px3,px4, px5]) < nx2 and max([py1, py2, py3,py4, py5]) < ny2 :
                    # pts_list.append([im, label, roi, pts])
                    pts_xs = pts_list_raw[0::2]
                    pts_ys = pts_list_raw[1::2]
                    # pts_xs_real =(pts_xs +1 -x1 )/w   # not planning to use numpy here, so list comprehension here
                    # pts_xs_real = [(tmpval -x1 )/w_box for tmpval in pts_xs]
                    # # pts_ys_real = (pts_ys + 1 - y1) / h
                    # pts_ys_real = [(tmpval - y1) / h_box for tmpval in pts_ys]
                    pts_xs_real = [(tmpval - x1 )/size for tmpval in pts_xs]# mistake here last time
                    pts_ys_real = [(tmpval - y1 )/size for tmpval in pts_ys]

                    pts = [0]*10
                    pts[0::2] = pts_xs_real  # should be numbers between 0 and 1 .
                    pts[1::2] = pts_ys_real

                    # label = -1
                    # roi = [-1, -1, -1, -1]
                    label = 1
                    roi = [float(offset_x1), float(offset_y1), float(offset_x2), float(offset_y2)]
                    pts_list.append([im, label, roi, pts])
                    pts_idx += 1

            elif IoU(crop_box, box_) >= 0.4:   # output class label only

                im = resized_im
                h, w, ch = resized_im.shape
                if h != image_size or w != image_size:
                    im = cv2.resize(im, (image_size, image_size))
                im = np.swapaxes(im, 0, 2)
                im = (im - 127.5) / 127.5
                label = -1
                roi = [float(offset_x1), float(offset_y1), float(offset_x2), float(offset_y2)]
                pts = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1]
                part_roi_list.append([im, label, roi])
                d_idx += 1
        box_idx += 1
        print ("{:d} images done, Lanmark{:d} positive:  {:d} part: {:d} negative: {:d}".format(idx, pts_idx,p_idx, d_idx, n_idx))

    neg_num = 0
    while neg_num < 10:
        if int( min(width, height) / 2 )>40:
            size = npr.randint(40, int( min(width, height) / 2 ))
        else:
            size = npr.randint(int(min(width, height) / 2), 40)
        nx = npr.randint(0, width - size)
        ny = npr.randint(0, height - size)
        crop_box = np.array([nx, ny, nx + size, ny + size])

        Iou = IoU(crop_box, boxes)

        cropped_im = img[ny : ny + size, nx : nx + size, :]
        resized_im = cv2.resize(cropped_im, (image_size, image_size), interpolation=cv2.INTER_LINEAR)

        if np.max(Iou) < 0.3:   # output class label only
            # Iou with all gts must below 0.3
            im = resized_im
            h, w, ch = resized_im.shape
            if h != image_size or w != image_size:
                im = cv2.resize(im, (image_size, image_size))
            im = np.swapaxes(im, 0, 2)
            im = (im - 127.5) / 127.5
            label = 0
            roi = [-1, -1, -1, -1]
            pts = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1]
            neg_cls_list.append([im, label, roi])
            n_idx += 1
            neg_num += 1
            # success = cv2.imwrite(save_file, resized_im)
            # if not success:
            #     raise Exception('Not writing file!')

if p_idx < len(part_roi_list):
    part_keep = npr.choice(len(part_roi_list), size=p_idx * 1, replace=False)  # limit the amount of part sample to keep
    for i in part_keep:
        roi_list.append(part_roi_list[i])
else:
    roi_list.extend(part_roi_list)

if p_idx < len(neg_cls_list):
    neg_keep = npr.choice(len(neg_cls_list), size=p_idx*3, replace=False)   # limit the amount of negative sample to keep
    for i in neg_keep:
        cls_list.append(neg_cls_list[i])
else:
    cls_list.extend(neg_cls_list)

cls_list.extend(pos_cls_list)
roi_list.extend(pos_roi_list)

# before dumping data, check shape of data
for iline, line in enumerate(roi_list):
    if len(line) != 3:
        raise Exception('roi_list shape is wrong')
for iline, line in enumerate(pts_list):
    if len(line) != 4:
        raise Exception('pts_list shape is wrong')
for iline, line in enumerate(cls_list):
    if len(line) != 3:
        raise Exception('cls_list shape is wrong')

fid = open("{0:d}/roi.imdb".format(image_size),'wb')
pickle.dump(roi_list, fid)
fid.close()
fid = open("{0:d}/pts.imdb".format(image_size),'wb')
pickle.dump(pts_list, fid)
fid.close()
fid = open("{0:d}/cls.imdb".format(image_size),'wb')
pickle.dump(cls_list, fid)
fid.close()


import os
import numpy as np
import pandas as pd
import requests
import json
import pyarrow.parquet as parquet
from sklearn.linear_model import LogisticRegression

input_path = 'https://cloud.mail.ru/public/A9cF/bVSWJCJgt/collabTest/'
# print(os.listdir(input_path))
result = (-8)*48+6*2+3*147
print(result)


import numpy as np
import torch
import torch.nn as nn
from mmcv.cnn import normal_init

from mmdet.models.registry import HEADS
from mmdet.models.utils import ConvModule, bias_init_with_prob
from mmdet.models.anchor_heads.anchor_head import AnchorHead
from mmdet.models.builder import build_loss
from mmdet.core import (AnchorGenerator, delta2bbox, force_fp32,
                        multi_apply, multiclass_nms)
from mmdet.core.bbox import PseudoSampler, assign_and_sample, bbox2delta, build_assigner

from .max_iou_assigner import MaxIoUAssigner

@HEADS.register_module
class CIRHead(nn.Module):

    def __init__(self,
                 num_classes,
                 in_channels,
                 feat_channels=256,
                 stacked_convs=4,
                 octave_base_scale=4,
                 scales_per_octave=3,
                 conv_cfg=None,
                 norm_cfg=dict(type='GN', num_groups=32, requires_grad=True),
                 anchor_ratios=[0.5, 1.0, 2.0],
                 anchor_strides=[4, 8, 16, 32, 64],
                 anchor_base_sizes=None,
                 target_means=(.0, .0, .0, .0),
                 target_stds=(1.0, 1.0, 1.0, 1.0),
                 loss_cls=dict(
                     type='FocalLoss',
                     use_sigmoid=True,
                     gamma=2.0,
                     alpha=0.25,
                     loss_weight=1.0),
                 loss_bbox=dict(
                     type='SmoothL1Loss', beta=0.11, loss_weight=1.0),
                 loss_IoUness=dict(
                     type='CrossEntropyLoss',
                     use_sigmoid=True,
                     loss_weight=1.0),
                 loss_softcls=dict(
                     type='MSELoss',
                     loss_weight=1.0)):

        super(CIRHead, self).__init__()
        self.stacked_convs = stacked_convs
        self.octave_base_scale = octave_base_scale
        self.scales_per_octave = scales_per_octave
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        octave_scales = np.array(
            [2**(i / scales_per_octave) for i in range(scales_per_octave)])
        anchor_scales = octave_scales * octave_base_scale
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.feat_channels = feat_channels
        self.anchor_scales = anchor_scales
        self.anchor_ratios = anchor_ratios
        self.anchor_strides = anchor_strides
        self.anchor_base_sizes = list(
            anchor_strides) if anchor_base_sizes is None else anchor_base_sizes
        self.target_means = target_means
        self.target_stds = target_stds

        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
        self.sampling = loss_cls['type'] not in ['FocalLoss', 'GHMC']
        if self.use_sigmoid_cls:
            self.cls_out_channels = num_classes - 1
        else:
            self.cls_out_channels = num_classes

        if self.cls_out_channels <= 0:
            raise ValueError('num_classes={} is too small'.format(num_classes))

        self.loss_cls = build_loss(loss_cls)
        self.loss_bbox = build_loss(loss_bbox)
        self.loss_iou = build_loss(loss_IoUness)
        self.loss_softcls = build_loss(loss_softcls)
        self.fp16_enabled = False

        self.anchor_generators = []
        for anchor_base in self.anchor_base_sizes:
            self.anchor_generators.append(
                AnchorGenerator(anchor_base, anchor_scales, anchor_ratios))

        self.num_anchors = len(self.anchor_ratios) * len(self.anchor_scales)
        self._init_layers()

    def _init_layers(self):
        self.relu = nn.ReLU(inplace=True)
        self.cls_convs = nn.ModuleList()
        self.reg_convs = nn.ModuleList()
        for i in range(self.stacked_convs):
            chn = self.in_channels if i == 0 else self.feat_channels
            self.cls_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    bias=self.norm_cfg is None))
            self.reg_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    bias=self.norm_cfg is None))
        self.CIR_cls = nn.Conv2d(
            self.feat_channels,
            self.num_anchors * self.cls_out_channels,
            3,
            padding=1)
        self.CIR_reg = nn.Conv2d(
            self.feat_channels, self.num_anchors * 4, 3, padding=1)
        self.CIR_IoUness = nn.Conv2d(self.feat_channels, self.num_anchors, 3, padding=1)

    def init_weights(self):
        for m in self.cls_convs:
            normal_init(m.conv, std=0.01)
        for m in self.reg_convs:
            normal_init(m.conv, std=0.01)
        bias_cls = bias_init_with_prob(0.01)
        normal_init(self.CIR_cls, std=0.01, bias=bias_cls)
        normal_init(self.CIR_reg, std=0.01)
        normal_init(self.CIR_IoUness, std=0.01)

    def forward_single(self, x):
        cls_feat = x
        reg_feat = x
        for cls_conv in self.cls_convs:
            cls_feat = cls_conv(cls_feat)
        for reg_conv in self.reg_convs:
            reg_feat = reg_conv(reg_feat)
        cls_score = self.CIR_cls(cls_feat)
        bbox_pred = self.CIR_reg(reg_feat)
        IoU_feat = cls_feat+reg_feat
        IoUness_pred = self.CIR_IoUness(IoU_feat)
        return cls_score, bbox_pred, IoUness_pred

    def forward(self, feats):
        return multi_apply(self.forward_single, feats)

    def get_anchors(self, featmap_sizes, img_metas, device='cuda'):
        """Get anchors according to feature map sizes.

        Args:
            featmap_sizes (list[tuple]): Multi-level feature map sizes.
            img_metas (list[dict]): Image meta info.
            device (torch.device | str): device for returned tensors

        Returns:
            tuple: anchors of each image, valid flags of each image
        """
        num_imgs = len(img_metas)
        num_levels = len(featmap_sizes)

        # since feature map sizes of all images are the same, we only compute
        # anchors for one time
        multi_level_anchors = []
        for i in range(num_levels):
            anchors = self.anchor_generators[i].grid_anchors(
                featmap_sizes[i], self.anchor_strides[i], device=device)
            multi_level_anchors.append(anchors)
        anchor_list = [multi_level_anchors for _ in range(num_imgs)]

        # for each image, we compute valid flags of multi level anchors
        valid_flag_list = []
        for img_id, img_meta in enumerate(img_metas):
            multi_level_flags = []
            for i in range(num_levels):
                anchor_stride = self.anchor_strides[i]
                feat_h, feat_w = featmap_sizes[i]
                h, w, _ = img_meta['pad_shape']
                valid_feat_h = min(int(np.ceil(h / anchor_stride)), feat_h)
                valid_feat_w = min(int(np.ceil(w / anchor_stride)), feat_w)
                flags = self.anchor_generators[i].valid_flags(
                    (feat_h, feat_w), (valid_feat_h, valid_feat_w),
                    device=device)
                multi_level_flags.append(flags)
            valid_flag_list.append(multi_level_flags)

        return anchor_list, valid_flag_list

    def loss_single(self, cls_score, bbox_pred, iou_pred, labels, label_weights,
                    bbox_targets, bbox_weights, iou_targets, iou_weights, softcls,
                    softcls_weights,num_total_samples, cfg):
        # classification loss
        labels = labels.reshape(-1)
        label_weights = label_weights.reshape(-1)
        cls_score = cls_score.permute(0, 2, 3,
                                      1).reshape(-1, self.cls_out_channels)
        loss_cls = self.loss_cls(
            cls_score, labels, label_weights, avg_factor=num_total_samples)
        # regression loss
        bbox_targets = bbox_targets.reshape(-1, 4)
        bbox_weights = bbox_weights.reshape(-1, 4)
        bbox_pred = bbox_pred.permute(0, 2, 3, 1).reshape(-1, 4)
        loss_bbox = self.loss_bbox(
            bbox_pred,
            bbox_targets,
            bbox_weights,
            avg_factor=num_total_samples)
        iou_targets = iou_targets.reshape(-1)
        iou_weights = iou_weights.reshape(-1)
        iou_pred = iou_pred.permute(0, 2, 3, 1).reshape(-1)
        loss_iou = self.loss_iou(iou_pred, iou_targets, iou_weights, avg_factor=num_total_samples)

        softcls = softcls.reshape(-1, self.cls_out_channels)
        softcls_weights = softcls_weights.reshape(-1, self.cls_out_channels)
        loss_softcls = self.loss_softcls(cls_score, softcls, softcls_weights, avg_factor = num_total_samples)

        return loss_cls, loss_bbox, loss_iou, loss_softcls

    def loss(self,
             cls_scores,
             bbox_preds,
             iou_preds,
             gt_bboxes,
             gt_labels,
             img_metas,
             cfg,
             gt_bboxes_ignore=None):
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
        assert len(featmap_sizes) == len(self.anchor_generators)

        device = cls_scores[0].device

        anchor_list, valid_flag_list = self.get_anchors(
            featmap_sizes, img_metas, device=device)
        label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1
        cls_reg_iou_targets = self.anchor_target(
            anchor_list,
            valid_flag_list,
            gt_bboxes,
            img_metas,
            self.target_means,
            self.target_stds,
            cfg,
            gt_bboxes_ignore_list=gt_bboxes_ignore,
            gt_labels_list=gt_labels,
            label_channels=label_channels,
            sampling=self.sampling)
        if cls_reg_iou_targets is None:
            return None
        (labels_list, label_weights_list, bbox_targets_list, bbox_weights_list,
         IoU_targets_list, IoU_weights_list, num_total_pos, num_total_neg, softcls_list,
         softcls_weights_list) = cls_reg_iou_targets
        num_total_samples = (
            num_total_pos + num_total_neg if self.sampling else num_total_pos)
        losses_cls, losses_bbox, losses_iou, losses_softcls = multi_apply(
            self.loss_single,
            cls_scores,
            bbox_preds,
            iou_preds,
            labels_list,
            label_weights_list,
            bbox_targets_list,
            bbox_weights_list,
            IoU_targets_list,
            IoU_weights_list,
            softcls_list,
            softcls_weights_list,
            num_total_samples=num_total_samples,
            cfg=cfg)
        # return dict(loss_cls=losses_cls, loss_bbox=losses_bbox,
        #             loss_iou=losses_iou, loss_softcls=losses_softcls)
        return dict(loss_cls=losses_cls, loss_bbox=losses_bbox)

    def anchor_target(self,
                      anchor_list,
                      valid_flag_list,
                      gt_bboxes_list,
                      img_metas,
                      target_means,
                      target_stds,
                      cfg,
                      gt_bboxes_ignore_list=None,
                      gt_labels_list=None,
                      label_channels=1,
                      sampling=True,
                      unmap_outputs=True):
        num_imgs = len(img_metas)
        assert len(anchor_list) == len(valid_flag_list) == num_imgs

        # anchor number of multi levels
        num_level_anchors = [anchors.size(0) for anchors in anchor_list[0]]
        # concat all level anchors and flags to a single tensor
        for i in range(num_imgs):
            assert len(anchor_list[i]) == len(valid_flag_list[i])
            anchor_list[i] = torch.cat(anchor_list[i])
            valid_flag_list[i] = torch.cat(valid_flag_list[i])
        if gt_bboxes_ignore_list is None:
            gt_bboxes_ignore_list = [None for _ in range(num_imgs)]
        if gt_labels_list is None:
            gt_labels_list = [None for _ in range(num_imgs)]
        (all_labels, all_label_weights, all_bbox_targets, all_bbox_weights, all_IoU_targets, all_IoU_weights,
         pos_inds_list, neg_inds_list, all_softcls, all_softcls_weights) = multi_apply(
            self.anchor_target_single,
            anchor_list,
            valid_flag_list,
            gt_bboxes_list,
            gt_bboxes_ignore_list,
            gt_labels_list,
            img_metas,
            target_means=target_means,
            target_stds=target_stds,
            cfg=cfg,
            label_channels=label_channels,
            sampling=sampling,
            unmap_outputs=unmap_outputs)
        # no valid anchors
        if any([labels is None for labels in all_labels]):
            return None
        # sampled anchors of all images
        num_total_pos = sum([max(inds.numel(), 1) for inds in pos_inds_list])
        num_total_neg = sum([max(inds.numel(), 1) for inds in neg_inds_list])
        # split targets to a list w.r.t. multiple levels
        labels_list = images_to_levels(all_labels, num_level_anchors)
        label_weights_list = images_to_levels(all_label_weights, num_level_anchors)
        bbox_targets_list = images_to_levels(all_bbox_targets, num_level_anchors)
        bbox_weights_list = images_to_levels(all_bbox_weights, num_level_anchors)
        # new add
        IoU_targets_list = images_to_levels(all_IoU_targets, num_level_anchors)
        IoU_weights_list = images_to_levels(all_IoU_weights, num_level_anchors)

        softcls_list = images_to_levels(all_softcls, num_level_anchors)
        softcls_weights_list = images_to_levels(all_softcls_weights, num_level_anchors)
        return (labels_list, label_weights_list, bbox_targets_list,
                bbox_weights_list, IoU_targets_list, IoU_weights_list,
                num_total_pos, num_total_neg, softcls_list, softcls_weights_list)

    def anchor_target_single(self,
                             flat_anchors,
                             valid_flags,
                             gt_bboxes,
                             gt_bboxes_ignore,
                             gt_labels,
                             img_meta,
                             target_means,
                             target_stds,
                             cfg,
                             label_channels=1,
                             sampling=True,
                             unmap_outputs=True):
        inside_flags = self.anchor_inside_flags(flat_anchors, valid_flags,
                                           img_meta['img_shape'][:2],
                                           cfg.allowed_border)
        if not inside_flags.any():
            return (None,) * 6
        # assign gt and sample anchors
        anchors = flat_anchors[inside_flags, :]

        if sampling:
            assign_result, sampling_result = assign_and_sample(
                anchors, gt_bboxes, gt_bboxes_ignore, None, cfg)
        else:
            bbox_assigner = MaxIoUAssigner()
            assign_result = bbox_assigner.assign(anchors, gt_bboxes,
                                                 gt_bboxes_ignore, gt_labels)
            bbox_sampler = PseudoSampler()
            sampling_result = bbox_sampler.sample(assign_result, anchors,
                                                  gt_bboxes)
        num_valid_anchors = anchors.shape[0]
        bbox_targets = torch.zeros_like(anchors)
        bbox_weights = torch.zeros_like(anchors)
        labels = anchors.new_zeros(num_valid_anchors, dtype=torch.long)
        label_weights = anchors.new_zeros(num_valid_anchors, dtype=torch.float)
        # new add
        IoU_targets = anchors.new_zeros(num_valid_anchors, dtype=torch.float)
        IoU_weights = anchors.new_zeros(num_valid_anchors, dtype=torch.float)

        softcls_targets = anchors.new_zeros((num_valid_anchors, self.cls_out_channels), dtype=torch.float)
        softcls_weights = anchors.new_zeros((num_valid_anchors, self.cls_out_channels), dtype=torch.float)
        soft_gt_labels = sampling_result.soft_gt_labels
        unassigned_bboxes_inds = sampling_result.unassigned_bboxes_inds
        softcls_targets[unassigned_bboxes_inds, soft_gt_labels[unassigned_bboxes_inds, 0]-1] = \
            soft_gt_labels[unassigned_bboxes_inds, 1]
        softcls_weights[unassigned_bboxes_inds, :] = 1

        pos_inds = sampling_result.pos_inds
        neg_inds = sampling_result.neg_inds
        if len(pos_inds) > 0:
            pos_bbox_targets = bbox2delta(sampling_result.pos_bboxes,
                                          sampling_result.pos_gt_bboxes,
                                          target_means, target_stds)
            bbox_targets[pos_inds, :] = pos_bbox_targets
            bbox_weights[pos_inds, :] = 1.0
            if gt_labels is None:
                labels[pos_inds] = 1
            else:
                labels[pos_inds] = gt_labels[sampling_result.pos_assigned_gt_inds]
            if cfg.pos_weight <= 0:
                label_weights[pos_inds] = 1.0
            else:
                label_weights[pos_inds] = cfg.pos_weight
            # new add
            IoU_targets[pos_inds] = sampling_result.pos_gt_IoU
            IoU_weights[pos_inds] = 1.0

        if len(neg_inds) > 0:
            label_weights[neg_inds] = 1.0

        # map up to original set of anchors
        if unmap_outputs:
            num_total_anchors = flat_anchors.size(0)
            labels = unmap(labels, num_total_anchors, inside_flags)
            label_weights = unmap(label_weights, num_total_anchors, inside_flags)
            bbox_targets = unmap(bbox_targets, num_total_anchors, inside_flags)
            bbox_weights = unmap(bbox_weights, num_total_anchors, inside_flags)
            # new add
            IoU_targets = unmap(IoU_targets, num_total_anchors, inside_flags)
            IoU_weights = unmap(IoU_weights, num_total_anchors, inside_flags)

            softcls_targets = unmap(softcls_targets, num_total_anchors, inside_flags)
            softcls_weights = unmap(softcls_weights, num_total_anchors, inside_flags)

        return (labels, label_weights, bbox_targets, bbox_weights, IoU_targets, IoU_weights, pos_inds,
                neg_inds, softcls_targets, softcls_weights)

    def anchor_inside_flags(self, flat_anchors, valid_flags, img_shape,
                            allowed_border=0):
        img_h, img_w = img_shape[:2]
        if allowed_border >= 0:
            inside_flags = valid_flags & \
                           (flat_anchors[:, 0] >= -allowed_border).type(torch.uint8) & \
                           (flat_anchors[:, 1] >= -allowed_border).type(torch.uint8) & \
                           (flat_anchors[:, 2] < img_w + allowed_border).type(torch.uint8) & \
                           (flat_anchors[:, 3] < img_h + allowed_border).type(torch.uint8)
        else:
            inside_flags = valid_flags
        return inside_flags

def unmap(data, count, inds, fill=0):
    """ Unmap a subset of item (data) back to the original set of items (of
    size count) """
    if data.dim() == 1:
        ret = data.new_full((count, ), fill)
        ret[inds] = data
    else:
        new_size = (count, ) + data.size()[1:]
        ret = data.new_full(new_size, fill)
        ret[inds, :] = data
    return ret

def images_to_levels(target, num_level_anchors):
    """Convert targets by image to targets by feature level.

    [target_img0, target_img1] -> [target_level0, target_level1, ...]
    """
    target = torch.stack(target, 0)
    level_targets = []
    start = 0
    for n in num_level_anchors:
        end = start + n
        level_targets.append(target[:, start:end].squeeze(0))
        start = end
    return level_targets


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import sys
import time
import numpy as np
import tensorflow as tf
import os
import re

import data_utils
from data_utils import *
import argparse
from model import TILGAN
import collections
from gensim.models import KeyedVectors
FLAGS = None

# tf.enable_eager_execution()
def add_arguments(parser):
    parser.register("type", "bool", lambda v: v.lower() == "true")
    parser.add_argument("--data_dir", type=str, default="data/", help="Data directory")
    parser.add_argument("--model_dir", type=str, default="inference_ckpt_example/", help="Model directory")
    parser.add_argument("--out_dir", type=str, default="output/", help="Out directory")
    parser.add_argument("--train_dir", type=str, default="tilgan/", help="Training directory")
    parser.add_argument("--gpu_device", type=str, default="0", help="which gpu to use")
    parser.add_argument("--train_data", type=str, default="training",
                        help="Training data path")
    parser.add_argument("--valid_data", type=str, default="dev",
                        help="Valid data path")
    parser.add_argument("--test_data", type=str, default="test",
                        help="Test data path")
    parser.add_argument("--from_vocab", type=str, default="data/vocab_20000",
                        help="from vocab path")
    parser.add_argument("--to_vocab", type=str, default="data/vocab_20000",
                        help="to vocab path")
    parser.add_argument("--output_dir", type=str, default="tfm/")
    parser.add_argument("--max_train_data_size", type=int, default=0, help="Limit on the size of training data (0: no limit)")
    parser.add_argument("--from_vocab_size", type=int, default=20000, help="source vocabulary size")
    parser.add_argument("--to_vocab_size", type=int, default=20000, help="target vocabulary size")
    parser.add_argument("--num_layers", type=int, default=6, help="Number of layers in the model")
    parser.add_argument("--num_units", type=int, default=512, help="Size of each model layer")
    parser.add_argument("--num_heads", type=int, default=8, help="Number of heads in attention")
    parser.add_argument("--emb_dim", type=int, default=300, help="Dimension of word embedding")
    parser.add_argument("--latent_dim", type=int, default=64, help="Dimension of latent variable")
    parser.add_argument("--batch_size", type=int, default=1, help="Batch size to use during training")
    parser.add_argument("--max_gradient_norm", type=float, default=3.0, help="Clip gradients to this norm")
    parser.add_argument("--learning_rate_decay_factor", type=float, default=0.5, help="Learning rate decays by this much")
    parser.add_argument("--learning_rate", type=float, default=1, help="Learning rate")
    parser.add_argument("--dropout_rate", type=float, default=0.15, help="Dropout rate")
    parser.add_argument("--epoch_num", type=int, default=100, help="Number of epoch")

def create_hparams(flags):
    return tf.contrib.training.HParams(
        # dir path
        data_dir=flags.data_dir,
        train_dir=flags.train_dir,
        output_dir=flags.output_dir,

        # data params
        batch_size=flags.batch_size,
        from_vocab_size=flags.from_vocab_size,
        to_vocab_size=flags.to_vocab_size,
        GO_ID=data_utils.GO_ID,
        EOS_ID=data_utils.EOS_ID,
        PAD_ID=data_utils.PAD_ID,

        train_data=flags.train_data,
        valid_data=flags.valid_data,
        test_data=flags.test_data,

        from_vocab=flags.from_vocab,
        to_vocab=flags.to_vocab,

        dropout_rate=flags.dropout_rate,
        init_weight=0.1,
        emb_dim=flags.emb_dim,
        latent_dim=flags.latent_dim,
        num_units=flags.num_units,
        num_heads=flags.num_heads,
        num_layers=flags.num_layers,
        learning_rate=flags.learning_rate,
        clip_value=flags.max_gradient_norm,
        decay_factor=flags.learning_rate_decay_factor,
        epoch_num=flags.epoch_num,
    )

def get_config_proto(log_device_placement=False, allow_soft_placement=True):
  config_proto = tf.ConfigProto(
      log_device_placement=log_device_placement,
      allow_soft_placement=allow_soft_placement)
  config_proto.gpu_options.allow_growth = True
  return config_proto

class TrainModel(
    collections.namedtuple("TrainModel",
                           ("graph", "model"))):
  pass

class EvalModel(
    collections.namedtuple("EvalModel",
                           ("graph", "model"))):
  pass

class InferModel(
    collections.namedtuple("InferModel",
                           ("graph", "model"))):
  pass

def create_model(hparams, model, length=22):
    train_graph = tf.Graph()
    with train_graph.as_default():
        train_model = model(hparams, tf.contrib.learn.ModeKeys.TRAIN)

    eval_graph = tf.Graph()
    with eval_graph.as_default():
        eval_model = model(hparams, tf.contrib.learn.ModeKeys.EVAL)

    infer_graph = tf.Graph()
    with infer_graph.as_default():
        infer_model = model(hparams, tf.contrib.learn.ModeKeys.INFER)

    return TrainModel(graph=train_graph, model=train_model), EvalModel(graph=eval_graph, model=eval_model), InferModel(
        graph=infer_graph, model=infer_model)

def read_data(src_path):
    data_set = []
    counter = 0
    max_length1 = 0
    with tf.gfile.GFile(src_path, mode="r") as src_file:
        src = src_file.readline()
        while src:
            if counter % 100000 == 0:
                print("  reading data line %d" % counter)
                sys.stdout.flush()

            sentences = []
            s = []
            for x in src.split(" "):
                id = int(x)
                if id != -1:
                    s.append(id)
                else:
                    if len(s) > max_length1:
                        max_length1 = len(s)
                    sentences.append(s)
                    s = []

            data_set.append(sentences)
            counter += 1
            src = src_file.readline()
    print(counter)
    print(max_length1)
    return data_set

def safe_exp(value):
  """Exponentiation with catching of overflow error."""
  try:
    ans = math.exp(value)
  except OverflowError:
    ans = float("inf")
  return ans

def get_inference_position(path):
    with open(path,'r') as f:
        position_list = []
        for line in f:
            sentences = line.strip().split(".")
            if (sentences[0].strip() == "<pending_infer>"):
                position = 0
            elif (sentences[1].strip() == "<pending_infer>"):
                position = 1
            elif (sentences[2].strip() == "<pending_infer>"):
                position = 2
            elif (sentences[3].strip() == "<pending_infer>"):
                position = 3
            else:
                position = 4
            position_list.append(position)
    return position_list
def train(hparams):
    embeddings = init_embedding(hparams)
    hparams.add_hparam(name="embeddings", value=embeddings)
    print("Vocab load over.")
    train_model, eval_model, infer_model = create_model(hparams, TILGAN)
    config = get_config_proto(
        log_device_placement=False)
    train_sess = tf.Session(config=config, graph=train_model.graph)
    eval_sess = tf.Session(config=config, graph=eval_model.graph)
    infer_sess = tf.Session(config=config, graph=infer_model.graph)
    print("Model create over.")
    #train_data = read_data("data/train.ids")
    valid_data = read_data("data/valid.ids")
    #test_data = read_data("data/test.ids")
    infer_data = read_data("inference_data/infer.ids")
    position_list = get_inference_position("inference_data/infer.txt")
    print("begin")
    print(position_list)
    ckpt = tf.train.get_checkpoint_state(hparams.train_dir)
    ckpt_path = os.path.join(hparams.train_dir, "ckpt")
    with train_model.graph.as_default():
        if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
            print("Reading model parameters from %s" % ckpt.model_checkpoint_path)
            train_model.model.saver.restore(train_sess, ckpt.model_checkpoint_path)
            #eval_model.model.saver.restore(eval_sess, ckpt.model_checkpoint_path)
            infer_model.model.saver.restore(infer_sess, ckpt.model_checkpoint_path)
            global_step = train_model.model.global_step.eval(session=train_sess)
        else:
            raise ValueError("Inference ckpt does not exist")
    to_vocab, rev_to_vocab = data_utils.initialize_vocabulary(hparams.from_vocab)

    step_loss, step_time, total_predict_count, total_loss, total_time, avg_loss, avg_time = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    total_loss_disc, total_loss_gen, total_loss_gan_ae,avg_disc_loss, avg_gen_loss ,avg_gan_ae_loss= 0.0, 0.0, 0.0, 0.0,0.0,0.0

    inference_out_f = open("inference/" + "predict2_file" + str(global_step),"w", encoding="utf-8")

    for id in range(0, int(len(infer_data) / hparams.batch_size)):
        given, answer, predict = infer_model.model.infer_step(infer_sess, infer_data, no_random=True,
                                                                id=id* hparams.batch_size,position=position_list[id])
        for i in range(hparams.batch_size):
            sample_output = predict[i]
            if hparams.EOS_ID in sample_output:
                sample_output = sample_output[:sample_output.index(hparams.EOS_ID)]
            pred = []
            for output in sample_output:
                pred.append(tf.compat.as_str(rev_to_vocab[output]))

            sample_output = answer[i]
            if hparams.EOS_ID in sample_output[:]:
                if sample_output[0] == hparams.GO_ID:
                    sample_output = sample_output[1:sample_output.index(hparams.EOS_ID)]
                else:
                    sample_output = sample_output[0:sample_output.index(hparams.EOS_ID)]
            ans = []
            for output in sample_output:
                ans.append(tf.compat.as_str(rev_to_vocab[output]))
            if id < 8:
                #print("answer: ", " ".join(ans))
                print("predict: ", " ".join(pred))

            inference_out_f.write(" ".join(pred) + "\n")
    inference_out_f.close()
    print("infer done.")

def init_embedding(hparams):
    f = open("data/vocab_20000", "r", encoding="utf-8")
    vocab = []
    for line in f:
        vocab.append(line.rstrip("\n"))
    # word_vectors = KeyedVectors.load_word2vec_format("GoogleNews-vectors-negative300.bin", binary=True)

    word_vectors = KeyedVectors.load_word2vec_format("data/roc_vector.txt")
    emb = []
    num = 0
    for i in range(0, len(vocab)):
        word = vocab[i]
        if word in word_vectors:
            num += 1
            emb.append(word_vectors[word])
        else:
            emb.append((0.1 * np.random.random([hparams.emb_dim]) - 0.05).astype(np.float32))

    print(" init embedding finished")
    emb = np.array(emb)
    print(num)
    print(emb.shape)
    return emb

def main(_):
    hparams = create_hparams(FLAGS)
    train(hparams)

if __name__ == "__main__":
    my_parser = argparse.ArgumentParser()
    add_arguments(my_parser)
    FLAGS, remaining = my_parser.parse_known_args()
    FLAGS.train_dir = FLAGS.model_dir + FLAGS.train_dir
    FLAGS.output_dir = FLAGS.out_dir + FLAGS.output_dir
    print(FLAGS)
    os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu_device
    tf.app.run()


import argparse
import json
import os

import faiss
import numpy as np
from tqdm import tqdm
from transformers import DPRContextEncoderTokenizer, DPRContextEncoder

def encode_passage(titles, texts, tokenizer, model, device='cuda:0'):
    max_length = 256  # hardcode for now
    inputs = tokenizer(
        titles,
        text_pair=texts,
        max_length=max_length,
        padding='longest',
        truncation=True,
        add_special_tokens=True,
        return_tensors='pt'
    )
    inputs.to(device)
    embeddings = model(inputs["input_ids"]).pooler_output.detach().cpu().numpy()
    return embeddings

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--encoder', type=str, help='encoder name or path', required=True)
    parser.add_argument('--dimension', type=int, help='dimension of passage embeddings', required=False, default=768)
    parser.add_argument('--corpus', type=str,
                        help='directory that contains corpus files to be encoded, in jsonl format.', required=True)
    parser.add_argument('--index', type=str, help='directory to store brute force index of corpus', required=True)
    parser.add_argument('--batch', type=int, help='batch size', default=8)
    parser.add_argument('--device', type=str, help='device cpu or cuda [cuda:0, cuda:1...]', default='cuda:0')
    args = parser.parse_args()

    tokenizer = DPRContextEncoderTokenizer.from_pretrained(args.encoder)
    model = DPRContextEncoder.from_pretrained(args.encoder)
    model.to(args.device)

    index = faiss.IndexFlatIP(args.dimension)

    if not os.path.exists(args.index):
        os.mkdir(args.index)

    titles = []
    texts = []
    with open(os.path.join(args.index, 'docid'), 'w') as id_file:
        for file in sorted(os.listdir(args.corpus)):
            file = os.path.join(args.corpus, file)
            if file.endswith('json') or file.endswith('jsonl'):
                print(f'Encoding {file}')
                with open(file, 'r') as corpus:
                    for idx, line in enumerate(tqdm(corpus.readlines())):
                        info = json.loads(line)
                        docid = info['id']
                        title, text = info['contents'].split('\n')
                        title = title.replace('"', '')
                        text = text.replace('""', '"')
                        id_file.write(f'{docid}\n')
                        titles.append(title)
                        texts.append(text)
    for idx in tqdm(range(0, len(titles), args.batch)):
        title_batch = titles[idx: idx+args.batch]
        text_batch = texts[idx: idx+args.batch]
        embeddings = encode_passage(title_batch, text_batch, tokenizer, model, args.device)
        index.add(np.array(embeddings))
    faiss.write_index(index, os.path.join(args.index, 'index'))