import cv2
import numpy as np
from importlib import import_module
from typing import Any, List

from darcyai.config_registry import ConfigRegistry
from darcyai.utils import validate_not_none, validate_type, validate

from .coral_perceptor_base import CoralPerceptorBase

class ImageClassificationPerceptor(CoralPerceptorBase):
    """
    ImageClassificationPerceptor is a class that implements the Perceptor interface for
    image classification.

    Arguments:
        threshold (float): The threshold for object detection.
        top_k (int): The number of top predictions to return.
        labels_file (str): The path to the labels file.
        labels (dict): A dictionary of labels.
        mean (float): The mean of the image.
        std (float): The standard deviation of the image.
        **kwargs: Keyword arguments to pass to Perceptor.
    """

    def __init__(self,
                 threshold:float,
                 top_k:int=None,
                 labels_file:str=None,
                 labels:dict=None,
                 mean:float=128.0,
                 std:float=128.0,
                 **kwargs):
        super().__init__(**kwargs)

        dataset = import_module("pycoral.utils.dataset")
        self.__classify = import_module("pycoral.adapters.classify")
        self.__common = import_module("pycoral.adapters.common")

        validate_not_none(threshold, "threshold is required")
        validate_type(threshold, (float, int), "threshold must be a number")
        validate(0 <= threshold <= 1, "threshold must be between 0 and 1")

        if top_k is not None:
            validate_type(top_k, int, "top_k must be an integer")
            validate(top_k > 0, "top_k must be greater than 0")

        if labels is not None:
            validate_type(labels, dict, "labels must be a dictionary")
            self.__labels = labels
        elif labels_file is not None:
            validate_type(labels_file, str, "labels_file must be a string")
            self.__labels = dataset.read_label_file(labels_file)
        else:
            self.__labels = None

        validate_type(mean, (int, float), "mean must be a number")
        validate_type(std, (int, float), "std must be a number")

        self.interpreter = None
        self.__threshold = threshold
        self.__top_k = top_k
        self.__mean = mean
        self.__std = std
        self.__inference_shape = None

    def run(self, input_data:Any, config:ConfigRegistry=None) -> (List[Any], List[str]):
        """
        Runs the image classification model.

        Arguments:
            input_data (Any): The input data to run the model on.
            config (ConfigRegistry): The configuration for the perceptor.

        Returns:
            (list[Any], list(str)): A tuple containing the detected classes and the labels.
        """
        labels = []
        resized_frame = cv2.resize(input_data, self.__inference_shape)

        params = self.__common.input_details(self.interpreter, "quantization_parameters")
        scales = params["scales"]
        zero_points = params["zero_points"]

        if abs(scales * self.__std - 1) < 1e-5 and abs(self.__mean - zero_points) < 1e-5:
            # Input data does not require preprocessing.
            self.__common.set_input(self.interpreter, resized_frame)
        else:
            normalized_frame = (np.asarray(resized_frame) - self.__mean) / (self.__std * scales) \
                + zero_points
            np.clip(normalized_frame, 0, 255, out=normalized_frame)
            self.__common.set_input(self.interpreter, normalized_frame.astype(np.uint8))

        self.interpreter.invoke()

        detected_classes = self.__classify.get_classes(
            self.interpreter,
            self.__top_k,
            self.__threshold)

        if not self.__labels is None:
            for detected_object in detected_classes:
                labels.append(self.__labels.get(detected_object.id, detected_object.id))

        return detected_classes, labels

    def load(self, accelerator_idx:[int, None]) -> None:
        """
        Loads the image classification model.

        Arguments:
            accelerator_idx (int): The index of the Edge TPU to use.

        Returns:
            None
        """
        CoralPerceptorBase.load(self, accelerator_idx)

        input_shape = self.interpreter.get_input_details()[0]["shape"]
        self.__inference_shape = (input_shape[2], input_shape[1])


# Functions to estimate the PVs of the cerebral cortex and subcortical
# structures. These are wrappers around the core Toblerone functions
# that handle the aggregate various pieces of information into overall results

import warnings
import copy

import numpy as np

from toblerone.utils import STRUCTURES, NP_FLOAT
from toblerone.classes import Surface, Hemisphere

def _cortex(hemispheres, space, struct2ref, supersampler, cores, ones):
    """Estimate the PVs of the cortex.

    Args:
        hemispheres: either a single, or iterable list of, Hemisphere objects.
        space: an ImageSpace within which to operate.
        struct2ref: np.array affine transformation into reference space.
        supersampler: supersampling factor (3-vector) to use for estimation.
        cores: number of processor cores to use.
        ones: debug tool, write ones in all voxels containing triangles.

    Returns:
        4D array, size equal to the reference image, with the PVs arranged
            GM/WM/non-brain in 4th dim
    """

    # Create our own local copy of inputs
    loc_hemispheres = copy.deepcopy(hemispheres)

    if not isinstance(loc_hemispheres, list):
        loc_hemispheres = [loc_hemispheres]

    for h in loc_hemispheres:
        h.inSurf, h.outSurf = [ s.transform(struct2ref) for s in h.surfs ]

    surfs = [ s for h in loc_hemispheres for s in h.surfs ]
    for s in surfs:
        s.index_on(space, cores)
        s.indexed.voxelised = s.voxelise(space, cores)

    # Estimate PV fractions for each surface
    for h in loc_hemispheres:
        if np.any(np.max(np.abs(h.inSurf.points)) >
            np.max(np.abs(h.outSurf.points))):
            warnings.warn("Inner surface vertices appear to be further" +
                " from the origin than the outer vertices. Are the surfaces in" +
                " the correct order?")

        for s, d in zip(h.surfs, ['in', 'out']):
            descriptor = "{} cortex {}".format(h.side, d)
            s._estimate_fractions(supersampler, cores, ones, descriptor)

    # Merge the voxelisation results with PVs
    for h in loc_hemispheres:
        in_pvs = h.inSurf.output_pvs(space).flatten()
        out_pvs = h.outSurf.output_pvs(space).flatten()

        # Combine estimates from each surface into whole hemi PV estimates
        hemiPVs = np.zeros((np.prod(space.size), 3), dtype=NP_FLOAT)
        hemiPVs[:,1] = in_pvs
        hemiPVs[:,0] = np.maximum(0.0, out_pvs - in_pvs)
        hemiPVs[:,2] = 1.0 - (hemiPVs[:,0:2].sum(1))
        h.PVs = hemiPVs

    # Merge the hemispheres, giving priority to GM, then WM, then CSF.
    # Do nothing if just one hemi
    if len(loc_hemispheres) == 1:
        outPVs = loc_hemispheres[0].PVs

    else:
        h1, h2 = loc_hemispheres
        outPVs = np.zeros((np.prod(space.size), 3), dtype=NP_FLOAT)
        outPVs[:,0] = np.minimum(1.0, h1.PVs[:,0] + h2.PVs[:,0])
        outPVs[:,1] = np.minimum(1.0 - outPVs[:,0],
            h1.PVs[:,1] + h2.PVs[:,1])
        outPVs[:,2] = 1.0 - outPVs[:,0:2].sum(1)

    # Sanity checks
    if np.any(outPVs > 1.0):
        raise RuntimeError("PV exceeds 1")

    if np.any(outPVs < 0.0):
        raise RuntimeError("Negative PV returned")

    if not np.all(outPVs.sum(1) == 1.0):
        raise RuntimeError("PVs do not sum to 1")

    return outPVs.reshape((*space.size, 3))

def _structure(surf, space, struct2ref, supersampler, ones, cores):
    """
    Estimate the PVs of a structure denoted by a single surface. Note
    that the results should be interpreted simply as "fraction of each
    voxel lying within the structure", and it is ambiguous as to what tissue
    lies outside the structure

    Args:
        surf: Surface object
        space: ImageSpace to estimate within
        struct2ref: np.array affine transformation into reference space.
        supersampler: supersampling factor (3-vector) to use for estimation
        ones: debug tool, write ones in voxels containing triangles
        cores: number of processor cores to use

    Returns:
        an array of size refSpace.size containing the PVs.
    """

    # Create our own local copy of inputs
    loc_surf = copy.deepcopy(surf)
    loc_surf = loc_surf.transform(struct2ref)
    loc_surf.index_on(space, cores)
    loc_surf.indexed.voxelised = loc_surf.voxelise(space, cores)
    loc_surf._estimate_fractions(supersampler, cores, ones)

    return loc_surf.output_pvs(space)

def stack_images(images):
    """
    Combine the results of estimate_complete() into overall PV maps
    for each tissue. Note that the below logic is entirely specific
    to the surfaces produced by FreeSurfer, FIRST and how they may be
    combined with FAST estimates. If you're going off-piste anywhere else
    then you probably DON'T want to re-use this logic.

    Args:
        dictionary of PV maps, keyed as follows: all FIRST subcortical
        structures named by their FIST convention (eg L_Caud); FAST's estimates
        named as FAST_CSF/WM/GM; cortex estimates as cortex_GM/WM/non_brain

    Returns:
        single 4D array of PVs, arranged GM/WM/non-brain in the 4th dim
    """

    # The logic is as follows:
    # Initialise everything as non-brain
    # Write in PV estimates from the cortical surfaces. This sets the cortical GM
    # and also all subcortical tissue as WM
    # Layer on top FAST's CSF estimates as non-brain. This is because surface methods
    # dom't pick up ventricular or mid-brain CSF, so wherever FAST has found more CSF
    # than Toblerone we will take that higher estimate
    # Layer in subcortical GM from each individual FIRST structure (except brain stem).
    # After adding in each structure's GM, recalculate CSF as either the existing amount,
    # or reduce the CSF estimate if the total (GM + CSF) > 1
    # If there is any remainder unassigned in the voxel, set that as WM

    # To summarise, the tissues are stacked up as:
    # All CSF
    # Cortical GM
    # Then within the subcortex only:
        # All subcortical volume set as WM
        # Subcortical CSF fixed using FAST
        # Add in subcortical GM for each structure, reducing CSF if required
        # Set the remainder 1 - (GM+CSF) as WM in voxels that were updated

    # Copy the dict of images as we are going to make changes and dont want
    # to play with the caller's copy. Pop unwanted images
    reqd_keys = STRUCTURES + [ 'FAST_GM', 'FAST_WM', 'FAST_CSF',
        'cortex_GM', 'cortex_WM', 'cortex_nonbrain' ]
    reqd_keys.remove('BrStem')
    for k in reqd_keys:
        if k not in images.keys():
            raise RuntimeError(f"Did not find '{k}' key in images dict")

    # Pop out FAST's estimates
    csf = images.pop('FAST_CSF').flatten()
    wm = images.pop('FAST_WM').flatten()
    gm = images.pop('FAST_GM')
    shape = (*gm.shape[0:3], 3)

    # Squash small FAST CSF values
    csf[csf < 0.01] = 0

    # Pop the cortex estimates and initialise output as all CSF
    ctxgm = images.pop('cortex_GM').flatten()
    ctxwm = images.pop('cortex_WM').flatten()
    ctxnon = images.pop('cortex_nonbrain').flatten()
    ctx = np.vstack((ctxgm, ctxwm, ctxnon)).T
    out = np.zeros_like(ctx)
    out[:,2] = 1

    # Then write in Toblerone's cortex estimates from all voxels
    # that contain either WM or GM (on the ctx image)
    mask = np.logical_or(ctx[:,0], ctx[:,1])
    out[mask,:] = ctx[mask,:]

    # Layer in FAST's CSF estimates (to get mid-brain and ventricular CSF).
    # Where FAST has suggested a higher CSF estimate than currently exists,
    # and the voxel does not intersect the cortical ribbon, accept FAST's
    # estimate. Then update the WM estimates, reducing where necessary to allow
    # for the greater CSF volume
    GM_threshold = 0.01
    ctxmask = (ctx[:,0] > GM_threshold)
    to_update = np.logical_and(csf > out[:,2], ~ctxmask)
    tmpwm = out[to_update,1]
    out[to_update,2] = np.minimum(csf[to_update], 1 - out[to_update,0])
    out[to_update,1] = np.minimum(tmpwm, 1 - (out[to_update,2] + out[to_update,0]))

    # Sanity checks: total tissue PV in each vox should sum to 1
    assert (out[to_update,0] <= GM_threshold).all(), 'Some update voxels have GM'
    assert (np.abs(out.sum(1) - 1) < 1e-6).all(), 'Voxel PVs do not sum to 1'
    assert (out > -1e-6).all(), 'Negative PV found'

    # For each subcortical structure, create a mask of the voxels which it
    # relates to. The following operations then apply only to those voxels
    # All subcortical structures interpreted as pure GM
    # Update CSF to ensure that GM + CSF in those voxels < 1
    # Finally, set WM as the remainder in those voxels.
    for k,s in images.items():
        smask = (s.flatten() > 0)
        out[smask,0] = np.minimum(1, out[smask,0] + s.flatten()[smask])
        out[smask,2] = np.minimum(out[smask,2], 1 - out[smask,0])
        out[smask,1] = np.maximum(1 - (out[smask,0] + out[smask,2]), 0)
        assert (out > -1e-6).all(), f'Negative found after {k} layer'
        assert (np.abs(out.sum(1) - 1) < 1e-6).all(), f'PVs sum > 1 after {k} layer'

    # Final sanity check, then rescaling so all voxels sum to unity.
    # assert (out > -1e-6).all()
    out[out < 0] = 0
    sums = out.sum(1)
    assert (np.abs(out.sum(1) - 1) < 1e-6).all(), 'Voxel PVs do not sum to 1'
    out = out / sums[:,None]

    return out.reshape(shape)


from __future__ import absolute_import, division, print_function

import os
import sys
import timeit
import numpy as np
from scipy.ndimage.filters import convolve
import math
from .dataset import Dataset

class DataLoader(object):
    """
    Arguments
        db_config: database configuration dictionary
    """
    def __init__(self, db_config):
        print('DataLoader')
        self.patch_size = None
        self.patch_step = None
        self.base_path = None
        self.n_images = 0
        self.n_videos = 0
        self.n_patches = 0
        self.npat_img_list = []
        self.npat_vid_list = []
        self.d_pat_set = []
        self.r_pat_set = []
        self.score_list = None
        self.fps_list = []
        self.n_frms_idx = []
        self.n_ref_frms_idx = []
        self.db_config = db_config
        self.n_frm_trte = 0
        if db_config is not None:
            self.set_config(db_config)

    def set_config(self, db_config):
        # Database
        self.sel_data = db_config.get('sel_data', None)
        self.scenes = db_config.get('scenes', 'all')
        self.dist_types = db_config.get('dist_types', 'all')
        self.select_database(self.sel_data, self.scenes, self.dist_types)

        # Initialize patch size
        self.patch_size = self.init_patch_size(
            db_config.get('patch_size', None))

        # Random crops
        self.random_crops = int(db_config.get('random_crops', 0))

        # If even patch division
        self.patch_step = self.init_patch_step(
            db_config.get('patch_step', None))
        self.patch_mode = db_config.get('patch_mode', None)
        assert self.patch_mode in ['both_side', 'shift_center', None]

        # Pre-processing
        self.color = db_config.get('color', 'gray')
        assert self.color in ['gray', 'rgb', 'ycbcr']
        self.num_ch = 1 if self.color == 'gray' else 3
        self.local_norm = db_config.get('local_norm', False)

        # etc.
        self.horz_ref = db_config.get('horz_ref', False)
        self.n_frm_trte = db_config.get('n_frm_trte', 50) ######new
        self.std_filt_r = db_config.get('std_filt_r', 1.0)
        self.train_size = db_config.get('train_size', 0.8)
        assert self.train_size > 0 and self.train_size < 1, \
            'train_size(%.2f) is not within 0 and 1' % (self.train_size)
        self.shuffle = db_config.get('shuffle', True)
        self.reverse_mos = db_config.get('reverse_mos', False)
        self.toy_example = db_config.get('toy_example', False)

    def init_patch_size(self, patch_size):
        # initialize patch size and step
        if patch_size:
            if isinstance(patch_size, int):
                patch_size_new = (patch_size, patch_size)
            elif isinstance(patch_size, (list, tuple)):
                assert len(patch_size) == 2
                patch_size_new = tuple(patch_size)
            else:
                raise ValueError('Wrong patch_size: {0}'.format(patch_size))
            self.use_original_size = False
        else:
            patch_size_new = None
            self.use_original_size = True

        return patch_size_new

    def init_patch_step(self, patch_step):
        # initialize patch step
        if not self.use_original_size:
            if isinstance(patch_step, int):
                patch_step_new = (patch_step, patch_step)
            elif isinstance(patch_step, (list, tuple)):
                assert len(patch_step) == 2
                patch_step_new = tuple(patch_step)
            elif patch_step is None:
                assert self.patch_size is not None
                patch_step_new = self.patch_size
            else:
                raise ValueError('Wrong patch_step:', patch_step)
        else:
            patch_step_new = (1, 1)
        return patch_step_new

    def select_database(self, sel_data, scenes, dist_types):
        """
        Select database to be loaded, and check scenes and dist_types.
        """
        if sel_data == 'LIVE_VQA':
            from . import LIVE_VQA
            self.DB_module = LIVE_VQA
        elif sel_data == 'CSIQ_VQA':
            from . import CSIQ_VQA
            self.DB_module = CSIQ_VQA
        else:
            raise ValueError('Improper sel_data: {0}'.format(sel_data))

        self.make_image_list_func = self.DB_module.make_image_list
        if scenes == 'all' or scenes is None:
            scenes = self.DB_module.ALL_SCENES
        if dist_types == 'all' or dist_types is None:
            dist_types = self.DB_module.ALL_DIST_TYPES

        self.scenes = scenes
        self.dist_types = dist_types
        self.sel_data = sel_data

        return scenes, dist_types

    def get_setting_dic(self):
        config_dict = {
            'sel_data': self.sel_data,
            'dist_types': self.dist_types,
            'patch_size': self.patch_size,
            'patch_step': self.patch_step,
            'random_crops': self.random_crops,
            'std_filt_r': self.std_filt_r,
            'horz_ref': self.horz_ref,
            'color': self.color,
            'local_norm': self.local_norm,
            'shuffle': self.shuffle,
            'train_size': self.train_size,
            'reverse_mos': self.reverse_mos,
            'n_images': self.n_images,
            'n_patches': self.n_patches,
        }
        return config_dict

    ###########################################################################

    def show_info(self):
        if self.patch_size is not None:
            print(' - Patch Size: (%d, %d)' % (
                self.patch_size[0], self.patch_size[1]))

        if self.random_crops > 0:
            print(' - Number of random crops: %d' % self.random_crops)
        else:
            print(' - Patch Step = (%d, %d)' % (
                self.patch_step[0], self.patch_step[1]), end='')
            print(' / mode: %s' % self.patch_mode)

        print(' - Color: %s' % self.color, end='')
        if self.local_norm:
            print(' (Local norm.)')
        else:
            print('')

        if self.std_filt_r < 1.0 and self.random_crops == 0:
            print(' - Patch sel. ratio (STD) =', self.std_filt_r)

        print(' - Augmentation :', end='')
        if self.horz_ref:
            print(' Horz. flip')
        else:
            print(' None')
        print(' - Reverse subj. score =', self.reverse_mos)

    ###########################################################################
    # Data loader interface
    def load_data_tr_te(self, tr_te_file=None, dataset_obj=False,
                        imagewise=True):
        """
        Load IQA database and divide into training and testing sets.
        """
        if self.toy_example:
            return self.load_toy_data_tr_te()
        print(' (Load training/testing data (shared ref.))')
        print(' - DB = %s' % (self.sel_data))
        train_scenes, test_scenes = self.divide_tr_te_wrt_ref(
            self.scenes, self.train_size, tr_te_file)

        self.show_info()

        # Get train set
        print('\n (Load training data)')
        data_dict = self.make_image_list_func(train_scenes, self.dist_types)
        self.load_ref_dis_images(data_dict)
        if self.horz_ref:
            self.data_augmentation_horz_refl()

        train_dataset = Dataset()
        train_dataset.put_data(
            self.d_pat_set,
            self.r_pat_set,
            self.d_pat_set_diff,
            self.r_pat_set_diff,
            self.dis2ref_idx,
            n_frms_idx=self.n_frms_idx,
            score_data=self.score_list,
            npat_img_list=self.npat_img_list,
            npat_vid_list=self.npat_vid_list,
            filt_idx_list=self.filt_idx_list,
            imagewise=imagewise, shuffle=self.shuffle)
        train_dataset.set_patch_config(
            self.patch_step, self.random_crops)

        # Get test set
        print('\n (Load testing data)')
        dist_type_test = self.dist_types
        data_dict = self.make_image_list_func(test_scenes, dist_type_test[1:])
        self.load_ref_dis_images(data_dict)

        test_dataset = Dataset()
        test_dataset.put_data(
            self.d_pat_set,
            self.r_pat_set,
            self.d_pat_set_diff,
            self.r_pat_set_diff,
            self.dis2ref_idx,
            n_frms_idx=self.n_frms_idx,
            score_data=self.score_list,
            npat_img_list=self.npat_img_list,
            npat_vid_list=self.npat_vid_list,
            filt_idx_list=self.filt_idx_list,
            imagewise=imagewise, shuffle=self.shuffle)
        test_dataset.set_patch_config(
            self.patch_step, self.random_crops)

        return train_dataset, test_dataset

    def load_data_rand(self, number, dataset_obj=False, imagewise=True):
        """
        Load IQA database (random images) and
        divide into training and testing sets.
        """
        print(' (Load training/testing data)', end='')
        print(' %d random videos (shared ref.)' % number)
        print(' - DB = %s' % (self.sel_data))
        scenes = self.DB_module.ALL_SCENES
        n_train_refs = int(np.ceil(number * self.train_size))
        n_test_refs = number - n_train_refs

        rand_seq = np.random.permutation(number)
        scenes_sh = [scenes[elem] for elem in rand_seq]
        train_scenes = sorted(scenes_sh[:n_train_refs])
        test_scenes = sorted(scenes_sh[n_train_refs:])

        print((' - Refs.: training  = %d / testing = %d (Ratio = %.2f)' %
               (n_train_refs, n_test_refs, self.train_size)))

        self.show_info()

        # Get train set
        data_dict = self.make_image_list_func(train_scenes, self.dist_types)
        self.load_ref_dis_images(data_dict)
        if self.horz_ref:
            self.data_augmentation_horz_refl()

        train_dataset = Dataset()
        train_dataset.put_data(
            self.d_pat_set, self.r_pat_set, self.dis2ref_idx,
            score_data=self.score_list,
            npat_img_list=self.npat_img_list,
            filt_idx_list=self.filt_idx_list,
            imagewise=imagewise, shuffle=self.shuffle)
        train_dataset.set_patch_config(
            self.patch_step, self.random_crops)

        # Get test set
        dist_type_test = self.dist_types
        data_dict = self.make_image_list_func(test_scenes, dist_type_test[1:])
        self.load_ref_dis_images(data_dict)

        test_dataset = Dataset()
        test_dataset.put_data(
            self.d_pat_set, self.r_pat_set, self.dis2ref_idx,
            score_data=self.score_list,
            npat_img_list=self.npat_img_list,
            filt_idx_list=self.filt_idx_list,
            imagewise=imagewise, shuffle=self.shuffle)
        test_dataset.set_patch_config(
            self.patch_step, self.random_crops)

        return train_dataset, test_dataset

    def load_data_for_test(self, tr_te_file, dataset_obj=False):
        """
        Load data with MOS and image data.
        There are no overlapping reference images between
        training and testing sets.
        """
        print(' - (Load testing data)')
        train_scenes, test_scenes = self.divide_tr_te_wrt_ref(
            self.scenes, self.train_size, tr_te_file)

        self.show_info()

        # Get test set
        # test set does not contain references for fair comparison
        dist_type_test = self.dist_types

        data_dict = self.make_image_list_func(test_scenes, dist_type_test[1:])
        self.load_ref_dis_images(data_dict)

        test_dataset = Dataset()
        test_dataset.put_data(
            self.d_pat_set, self.r_pat_set, self.dis2ref_idx,
            score_data=self.score_list,
            npat_img_list=self.npat_img_list,
            filt_idx_list=self.filt_idx_list,
            imagewise=True, shuffle=False)
        test_dataset.set_patch_config(
            self.patch_step, self.random_crops)

        return test_dataset

    def divide_tr_te_wrt_ref(self, scenes, train_size=0.8, tr_te_file=None):
        """
        Divdie data with respect to scene
        """
        tr_te_file_loaded = False
        if tr_te_file is not None and os.path.isfile(tr_te_file):
            # Load tr_te_file and divide scenes accordingly
            tr_te_file_loaded = True
            with open(tr_te_file, 'r') as f:
                train_scenes = f.readline().strip().split()
                train_scenes = [int(elem) for elem in train_scenes]
                test_scenes = f.readline().strip().split()
                test_scenes = [int(elem) for elem in test_scenes]

            n_train_refs = len(train_scenes)
            n_test_refs = len(test_scenes)
            train_size = (len(train_scenes) /
                          (len(train_scenes) + len(test_scenes)))
        else:
            # Divide scenes randomly
            # Get the numbers of training and testing scenes
            n_scenes = len(scenes)
            n_train_refs = int(np.ceil(n_scenes * train_size))
            n_test_refs = n_scenes - n_train_refs

            # Randomly divide scenes
            rand_seq = np.random.permutation(n_scenes)
            scenes_sh = [scenes[elem] for elem in rand_seq]
            train_scenes = sorted(scenes_sh[:n_train_refs])
            test_scenes = sorted(scenes_sh[n_train_refs:])

            # Write train-test idx list into file
            if tr_te_file is not None:
                fpath, fname = os.path.split(tr_te_file)
                if not os.path.isdir(fpath):
                    os.makedirs(fpath)
                with open(tr_te_file, 'w') as f:
                    for idx in range(n_train_refs):
                        f.write('%d ' % train_scenes[idx])
                    f.write('\n')
                    for idx in range(n_scenes - n_train_refs):
                        f.write('%d ' % test_scenes[idx])
                    f.write('\n')

        print(' - Refs.: training = %d / testing = %d (Ratio = %.2f)' %
              (n_train_refs, n_test_refs, train_size), end='')
        if tr_te_file_loaded:
            print(' (Loaded %s)' % (tr_te_file))
        else:
            print('')
        return train_scenes, test_scenes

    ###########################################################################

    def load_ref_dis_images(self, data_dict):
        self.score_list = data_dict['score_list']
        if self.reverse_mos and self.score_list is not None:
            self.score_list = 1.0 - self.score_list
        self.n_images = data_dict['n_images']
        self.n_videos = data_dict['n_videos']
        self.fps_list = data_dict['fps_list']
        base_path = data_dict['base_path']
        d_img_list = data_dict['d_img_list']
        r_img_list = data_dict['r_img_list']
        r_idx_list = data_dict['r_idx_list']
        scenes = data_dict['scenes']

        res = self.load_ref_frames(base_path, r_img_list, r_idx_list, scenes)
        ref_img2pat_idx, ref_top_left_set, ref_npat_img_list, inv_scenes, ref_total_frm2pat_idx = res
        self.load_dis_frames(
            base_path, d_img_list, r_idx_list, inv_scenes,
            ref_img2pat_idx, ref_top_left_set, ref_npat_img_list, ref_total_frm2pat_idx)

    def load_ref_frames(self, base_path, r_img_list, r_idx_list, scenes):
        """
        Actual routine of loading reference images.
        """
        self.n_ref_images = len(scenes)
        n_dis_images = len(r_img_list)

        # make a list of reference index
        ref_idx_idx_list = []
        for ref_idx in scenes:
            for idx, this_ref_idx in enumerate(r_idx_list):
                if ref_idx == this_ref_idx:
                    ref_idx_idx_list.append(idx)
                    break

                if idx == n_dis_images - 1:
                    raise ValueError('@ No %d index in r_idx_list' % ref_idx)

        new_r_img_list = []
        for idx in ref_idx_idx_list:
            new_r_img_list.append(r_img_list[idx])

        inv_scenes = np.ones(max(scenes) + 1, dtype='int32') * -1
        for idx, scn in enumerate(scenes):
            inv_scenes[scn] = idx

        patch_size = self.patch_size
        patch_step = self.patch_step
        n_frm_trte = self.n_frm_trte

        n_ref_patches = 0
        ref_npat_img_list = []
        ref_total_frm2pat_idx = []
        ref_top_left_set = []
        r_pat_set = []
        r_pat_set_diff = []
        n_ref_frms_idx = []
        #######################################################################
        # Read images
        start_time = timeit.default_timer()
        pass_list = []

        for vid_idx in range(self.n_ref_images):
            # Show progress
            show_progress(float(vid_idx) / self.n_ref_images)
        #######################################################################
        #video read
        #######################################################################
            ref_img2pat_idx = []
            start_frm = 5 # first 5 frms are ignored
            height = 432
            width = 768
            fwidth = (width + 31) // 32 * 32
            fheight = (height + 15) // 16 * 16

            end_frm = int(self.fps_list[ref_idx_idx_list[vid_idx]])
            n_ref_frms_idx.append(n_frm_trte)  #for same number of frms in video
            slide_frms = int(math.ceil((end_frm)/n_frm_trte)) # sliding window frm size
            vid_1 = []
            vid_2 = []
            stream = open('%s/%s' % (base_path, new_r_img_list[vid_idx]), 'rb')
            for i in range(start_frm, end_frm, slide_frms):
                # frame 1 read
                stream.seek(int(i * width * height * 1.5))
                Y = np.fromfile(stream, dtype=np.uint8, count=fwidth * fheight).reshape((fheight, fwidth))
                U = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                V = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                vid_1.append(Y)
                # frame 2 read
                Y = np.fromfile(stream, dtype=np.uint8, count=fwidth * fheight).reshape((fheight, fwidth))
                U = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                V = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                if end_frm > 300: # over 250frm +1frm
                    Y = np.fromfile(stream, dtype=np.uint8, count=fwidth * fheight).reshape((fheight, fwidth))
                    U = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                    V = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                    vid_2.append(Y)
                else:
                    vid_2.append(Y)
            vid_1 = np.array(vid_1)
            vid_2 = np.array(vid_2)

            for frm_no in range(n_frm_trte):
                r_frm_raw_1 = np.squeeze(vid_1[frm_no])
                r_frm_raw_1_norm = r_frm_raw_1.astype('float32') / 255.
                r_frm_raw_2 = np.squeeze(vid_2[frm_no])
                r_frm_raw_2_norm = r_frm_raw_2.astype('float32') / 255.
                #NOTE: abs is conducted in FR_SENS_VQA
                r_frm_diff = r_frm_raw_1_norm - r_frm_raw_2_norm
                cur_h = r_frm_raw_1.shape[0]
                cur_w = r_frm_raw_1.shape[1]

                if self.use_original_size:
                    patch_size = (cur_h, cur_w)
                # Gray or RGB
                if self.local_norm:
                    if self.color == 'gray':
                        # faster
                        r_img_norm = local_normalize_1ch(r_frm_raw_1)
                    else:
                        r_img_norm = local_normalize(r_frm_raw_1, self.num_ch)
                else:
                    r_img_norm = r_frm_raw_1.astype('float32') / 255.

                if self.color == 'gray':
                    r_img_norm = r_img_norm[:, :, None]
                    r_frm_diff = r_frm_diff[:, :, None]

                # numbers of patches along y and x axes
                ny = (cur_h - patch_size[0]) // patch_step[0] + 1
                nx = (cur_w - patch_size[1]) // patch_step[1] + 1
                npat = int(ny * nx)
                ref_npat_img_list.append((npat, ny, nx))
                ref_img2pat_idx.append(n_ref_patches + np.arange(npat))
                n_ref_patches += npat

                if npat == 0:
                    pass_list.append(vid_idx)
                    ref_top_left_set.append(False)
                    continue

                # get non-covered length along y and x axes
                cov_height = patch_step[0] * (ny - 1) + patch_size[0]
                cov_width = patch_step[1] * (nx - 1) + patch_size[1]
                nc_height = cur_h - cov_height
                nc_width = cur_w - cov_width

                # Shift center
                if self.patch_mode == 'shift_center':
                    shift = [(nc_height + 1) // 2, (nc_width + 1) // 2]
                    if shift[0] % 2 == 1:
                        shift[0] -= 1
                    if shift[1] % 2 == 1:
                        shift[1] -= 1
                    shift = tuple(shift)
                else:
                    shift = (0, 0)

                # generate top_left_set of patches
                top_left_set = np.zeros((nx * ny, 2), dtype=np.int)
                for yidx in range(ny):
                    for xidx in range(nx):
                        top = (yidx * patch_step[0] + shift[0])
                        left = (xidx * patch_step[1] + shift[1])
                        top_left_set[yidx * nx + xidx] = [top, left]
                ref_top_left_set.append(top_left_set)

                # Crop the images to patches
                for idx in range(ny * nx):
                    [top, left] = top_left_set[idx]

                    if top + patch_size[0] > cur_h:
                        print('\n@Error: imidx=%d, pat=%d' % (vid_idx, idx), end='')
                        print(' (%d > %d)' % (top + patch_size[0], cur_h))

                    if left + patch_size[1] > cur_w:
                        print('\n@Error: imidx=%d, pat=%d' % (vid_idx, idx), end='')
                        print(' (%d > %d)' % (left + patch_size[1], cur_w))

                    r_crop_norm = r_img_norm[top:top + patch_size[0],
                                             left:left + patch_size[1]]
                    r_crop_diff = r_frm_diff[top:top + patch_size[0],
                                             left:left + patch_size[1]]

                    r_pat_set.append(r_crop_norm)
                    r_pat_set_diff.append(r_crop_diff) #difference map
            ref_total_frm2pat_idx.append(ref_img2pat_idx) #################new

        # Show 100% progress bar
        show_progress(1.0)
        minutes, seconds = divmod(timeit.default_timer() - start_time, 60)
        print(' - It took {:02.0f}:{:05.2f}'.format(minutes, seconds))
        print(' - Loaded num of ref. patches: {:,}'.format(n_ref_patches))

        if len(pass_list) > 0:
            self.n_images -= len(pass_list)
            print(' - Ignored ref. videos due to small size: %s' %
                  ', '.join(str(i) for i in pass_list))

        self.n_ref_patches = n_ref_patches
        self.ref_npat_img_list = ref_npat_img_list
        self.ref_top_left_set = ref_top_left_set
        self.ref_img2pat_idx = ref_img2pat_idx
        self.n_ref_frms_idx = n_ref_frms_idx
        self.ref_total_frm2pat_idx = ref_total_frm2pat_idx #################new
        self.r_pat_set = r_pat_set
        self.r_pat_set_diff = r_pat_set_diff

        return ref_img2pat_idx, ref_top_left_set, ref_npat_img_list, inv_scenes, ref_total_frm2pat_idx

    def load_dis_frames(self, base_path, d_img_list, r_idx_list, inv_scenes,
                        ref_img2pat_idx, ref_top_left_set, ref_npat_img_list, ref_total_frm2pat_idx):
        """
        Actual routine of loading distorted images.
        """
        self.n_images = len(d_img_list)
        self.n_videos = len(d_img_list)

        n_frm_trte = self.n_frm_trte
        d_img_list = d_img_list

        assert self.n_images > 0, \
            'n_videos(%d) is not positive number' % (self.n_images)

        patch_size = self.patch_size
        n_patches = 0
        npat_vid_list =[]
        d_pat_set = []
        d_pat_set_diff = []
        filt_idx_list = []
        dis2ref_idx = []
        n_frms_idx = []   # number of frames of each videos

        #######################################################################
        # Read images
        start_time = timeit.default_timer()
        pat_idx = 0
        pass_list = []
        for vid_idx in range(self.n_images):

            npat_img_list = []

            ref_idx = inv_scenes[r_idx_list[vid_idx]]
            assert ref_idx >= 0
            if ref_top_left_set[ref_idx] is False:
                continue

            # Show progress
            show_progress(float(vid_idx) / self.n_images)
            #######################################################################
            # video read
            #######################################################################
            start_frm = 5  # first 5 frms are ignored
            height = 432
            width = 768
            fwidth = (width + 31) // 32 * 32
            fheight = (height + 15) // 16 * 16
            end_frm = int(self.fps_list[vid_idx])
            n_frms_idx.append(n_frm_trte)  # for same number of frms in video
            slide_frms = int(math.ceil(end_frm/n_frm_trte))  # sliding window frm size
            vid_1 = []
            vid_2 = []

            stream = open('%s/%s' % (base_path, d_img_list[vid_idx]), 'rb')
            for i in range(start_frm, end_frm, slide_frms):
                # frame 1 read
                stream.seek(int(i * width * height * 1.5))
                Y = np.fromfile(stream, dtype=np.uint8, count=fwidth * fheight).reshape((fheight, fwidth))
                U = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                V = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                vid_1.append(Y)
                # frame 2 read
                Y = np.fromfile(stream, dtype=np.uint8, count=fwidth * fheight).reshape((fheight, fwidth))
                U = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                V = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                if end_frm > 300: # over 250frm +1frm
                    Y = np.fromfile(stream, dtype=np.uint8, count=fwidth * fheight).reshape((fheight, fwidth))
                    U = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                    V = np.fromfile(stream, dtype=np.uint8, count=(fwidth// 2) * (fheight // 2))
                    vid_2.append(Y)
                else:
                    vid_2.append(Y)
            vid_1 = np.array(vid_1)
            vid_2 = np.array(vid_2)

            for frm_no in range(n_frm_trte):
                d_frm_raw_1 = np.squeeze(vid_1[frm_no])
                d_frm_raw_1_norm = d_frm_raw_1.astype('float32')/255.
                d_frm_raw_2 = np.squeeze(vid_2[frm_no])
                d_frm_raw_2_norm = d_frm_raw_2.astype('float32')/255.
                #NOTE: abs is conducted in FR_SENS_VQA
                d_frm_diff = d_frm_raw_1_norm - d_frm_raw_2_norm
                cur_h = d_frm_raw_1.shape[0]
                cur_w = d_frm_raw_1.shape[1]

                if self.use_original_size:
                    patch_size = (cur_h, cur_w)

                if cur_h < patch_size[0] or cur_w < patch_size[1]:
                    pass_list.append(vid_idx)
                    continue

                # Local normalization
                if self.local_norm:
                    if self.color == 'gray':
                        # faster
                        d_img_norm = local_normalize_1ch(d_frm_raw_1)
                    else:
                        d_img_norm = local_normalize(d_frm_raw_1, self.num_ch)
                else:
                    d_img_norm = d_frm_raw_1.astype('float32') / 255.

                if self.color == 'gray':
                    d_img_norm = d_img_norm[:, :, None]
                    d_frm_diff = d_frm_diff[:, :, None]

                top_left_set = ref_top_left_set[ref_idx]
                cur_n_patches = top_left_set.shape[0]

                if self.random_crops > 0: # NONE
                    if self.random_crops < cur_n_patches:
                        n_crops = self.random_crops
                        rand_perm = np.random.permutation(cur_n_patches)
                        sel_patch_idx = sorted(rand_perm[:n_crops])
                        top_left_set = top_left_set[sel_patch_idx].copy()
                    else:
                        n_crops = cur_n_patches
                        sel_patch_idx = np.arange(cur_n_patches)

                    npat_filt = n_crops
                    npat_img_list.append((npat_filt, 1, npat_filt))
                    n_patches += npat_filt

                    idx_set = list(range(npat_filt))
                    filt_idx_list.append(idx_set)

                else:
                    # numbers of patches along y and x axes
                    npat, ny, nx = ref_npat_img_list[ref_idx]
                    npat_filt = int(npat * self.std_filt_r)

                    npat_img_list.append((npat_filt, ny, nx))
                    n_patches += npat_filt

                    if self.std_filt_r < 1.0: # std_filt_r == 1 default
                        std_set = np.zeros((nx * ny))
                        for idx, top_left in enumerate(top_left_set):
                            top, left = top_left
                            std_set[idx] = np.std(
                                d_img_norm[top:top + patch_size[0],
                                      left:left + patch_size[1]])

                    # Filter the patches with low std
                    if self.std_filt_r < 1.0:
                        idx_set = sorted(list(range(len(std_set))),
                                         key=lambda x: std_set[x], reverse=True)
                        idx_set = sorted(idx_set[:npat_filt])
                    else: # std_filt_r == 1 default
                        idx_set = list(range(npat_filt))
                    filt_idx_list.append(idx_set)

                # Crop the images to patches
                for idx in idx_set:
                    [top, left] = top_left_set[idx]

                    if top + patch_size[0] > cur_h:
                        print('\n@Error: imidx=%d, pat=%d' % (vid_idx, idx), end='')
                        print(' (%d > %d)' % (top + patch_size[0], cur_h))

                    if left + patch_size[1] > cur_w:
                        print('\n@Error: imidx=%d, pat=%d' % (vid_idx, idx), end='')
                        print(' (%d > %d)' % (left + patch_size[1], cur_w))

                    d_crop_norm = d_img_norm[top:top + patch_size[0],
                                             left:left + patch_size[1]]
                    d_crop_diff = d_frm_diff[top:top + patch_size[0],
                                             left:left + patch_size[1]]

                    d_pat_set.append(d_crop_norm)
                    d_pat_set_diff.append(d_crop_diff)

                    if self.random_crops > 0:
                        dis2ref_idx.append(
                            ref_total_frm2pat_idx[ref_idx][sel_patch_idx][idx])
                    else:
                        # print(' ref_idx=%d, frm_no=%d, idx=%d \n' % (ref_idx, frm_no, idx), end='')
                        dis2ref_idx.append(ref_total_frm2pat_idx[ref_idx][frm_no][idx])

                    pat_idx += 1
            npat_vid_list.append(npat_img_list)  #################new
        # Show 100 % progress bar
        show_progress(1.0)
        minutes, seconds = divmod(timeit.default_timer() - start_time, 60)
        print(' - It took {:02.0f}:{:05.2f}'.format(minutes, seconds))
        print(' - Loaded num of patches: {:,}'.format(n_patches))

        if len(pass_list) > 0:
            self.n_images -= len(pass_list)
            print(' - Ignored frame list due to small size: %s' %
                  ', '.join(str(i) for i in pass_list))

        self.n_patches = n_patches
        self.npat_img_list = npat_img_list
        self.npat_vid_list = npat_vid_list
        self.n_frms_idx = n_frms_idx
        self.d_pat_set = d_pat_set
        self.d_pat_set_diff = d_pat_set_diff
        self.filt_idx_list = filt_idx_list
        self.dis2ref_idx = dis2ref_idx

    def make_toy_examples(self, patch_size=None, n_images=10):
        if patch_size is None:
            patch_size = self.patch_size
            if self.patch_size is None:
                patch_size = [64, 64]
                self.patch_size = patch_size
        n_ch = 1 if self.color == 'gray' else 3
        ny = 2
        nx = 2
        score_list = np.zeros(n_images, dtype='float32')
        npat_img_list = []
        filt_idx_list = []
        n_patches = 0
        for im_idx in range(n_images):
            npat = ny * nx
            npat_img_list.append((npat, ny, nx))
            n_patches += npat
            filt_idx_list.append(list(range(npat)))

        d_pat_set = []
        pat_shape = (patch_size[0], patch_size[1], n_ch)
        dummy_pat = np.ones(pat_shape, dtype='float32') * 0.5
        for idx in range(n_patches):
            d_pat_set.append(dummy_pat)
        print(' - Generated toy examples: %d x' % n_patches, pat_shape)

        dis2ref_idx = []
        for idx in range(n_patches):
            dis2ref_idx.append(idx)

        self.n_images = n_images
        self.score_list = score_list
        self.n_patches = n_patches
        self.npat_img_list = npat_img_list
        self.d_pat_set = d_pat_set
        self.r_pat_set = d_pat_set
        self.filt_idx_list = filt_idx_list
        self.dis2ref_idx = dis2ref_idx

    def data_augmentation_horz_refl(self):
        # Patches augmentation
        pat_idx = 0
        for pat_idx in range(self.n_patches):
            # Flip horizontal
            self.d_pat_set.append(self.d_pat_set[pat_idx][:, ::-1])
            self.dis2ref_idx.append(
                self.dis2ref_idx[pat_idx] + self.n_ref_patches)

        for pat_idx in range(self.n_ref_patches):
            # Flip horizontal
            self.r_pat_set.append(self.r_pat_set[pat_idx][:, ::-1])

        # Image data augmentation
        if self.score_list is not None:
            self.score_list = np.tile(self.score_list, 2)
        self.npat_img_list += self.npat_img_list
        self.filt_idx_list += self.filt_idx_list

        self.n_patches *= 2
        self.n_images *= 2
        self.n_ref_patches *= 2

        if self.score_list is not None:
            assert self.score_list.shape[0] == self.n_images, (
                'n_score_list: %d != n_videos: %d' %
                (self.score_list.shape[0], self.n_images))
        assert len(self.npat_img_list) == self.n_images, (
            'n_npat_img_list: %d != n_videos: %d' %
            (len(self.npat_img_list), self.n_images))
        assert len(self.filt_idx_list) == self.n_images, (
            'n_filt_idx_list: %d != n_videos: %d' %
            (len(self.filt_idx_list), self.n_images))

        print(' - Augmented patches: {0:,}'.format(self.n_patches), end=' ')
        print('(x2 horz. reflection)')

def show_progress(percent):
    hashes = '#' * int(round(percent * 20))
    spaces = ' ' * (20 - len(hashes))
    sys.stdout.write("\r - Load videos: [{0}] {1}%".format(
        hashes + spaces, int(round(percent * 100))))
    sys.stdout.flush()

def convert_color(img, color):
    """ Convert image into gray or RGB or YCbCr.
    """
    assert len(img.shape) in [2, 3]
    if color == 'gray':
        # if d_img_raw.shape[2] == 1:
        if len(img.shape) == 2:  # if gray
            img_ = img[:, :, None]
        elif len(img.shape) == 3:  # if RGB
            if img.shape[2] > 3:
                img = img[:, :, :3]
            img_ = rgb2gray(img)[:, :, None]
    elif color == 'rgb':
        if len(img.shape) == 2:  # if gray
            img_ = gray2rgb(img)
        elif len(img.shape) == 3:  # if RGB
            if img.shape[2] > 3:
                img = img[:, :, :3]
            img_ = img
    elif color == 'ycbcr':
        if len(img.shape) == 2:  # if gray
            img_ = rgb2ycbcr(gray2rgb(img))
        elif len(img.shape) == 3:  # if RGB
            if img.shape[2] > 3:
                img = img[:, :, :3]
            img_ = rgb2ycbcr(img)
    else:
        raise ValueError("Improper color selection: %s" % color)
    return img_

def convert_color2(img, color):
    """ Convert image into gray or RGB or YCbCr.
    (In case of gray, dimension is not increased for
    the faster local normalization.)
    """
    assert len(img.shape) in [2, 3]
    if color == 'gray':
        # if d_img_raw.shape[2] == 1:
        if len(img.shape) == 3:  # if RGB
            if img.shape[2] > 3:
                img = img[:, :, :3]
            img_ = rgb2gray(img)
    elif color == 'rgb':
        if len(img.shape) == 2:  # if gray
            img_ = gray2rgb(img)
        elif len(img.shape) == 3:  # if RGB
            if img.shape[2] > 3:
                img = img[:, :, :3]
            img_ = img
    elif color == 'ycbcr':
        if len(img.shape) == 2:  # if gray
            img_ = rgb2ycbcr(gray2rgb(img))
        elif len(img.shape) == 3:  # if RGB
            if img.shape[2] > 3:
                img = img[:, :, :3]
            img_ = rgb2ycbcr(img)
    else:
        raise ValueError("Improper color selection: %s" % color)

    return img_

def gray2rgb(im):
    w, h = im.shape
    ret = np.empty((w, h, 3), dtype=np.uint8)
    ret[:, :, :] = im[:, :, np.newaxis]
    return ret

def rgb2gray(rgb):
    assert rgb.shape[2] == 3
    return np.dot(rgb[..., :3], [0.299, 0.587, 0.114])

def rgb2ycbcr(rgb):
    xform = np.array([[.299, .587, .114],
                      [-.1687, -.3313, .5],
                      [.5, -.4187, -.0813]])
    ycbcr = np.dot(rgb[..., :3], xform.T)
    ycbcr[:, :, [1, 2]] += 128
    return ycbcr

def ycbcr2rgb(ycbcr):
    xform = np.array([[1, 0, 1.402],
                      [1, -0.34414, -.71414],
                      [1, 1.772, 0]])
    rgb = ycbcr.astype('float32')
    rgb[:, :, [1, 2]] -= 128
    return rgb.dot(xform.T)

k = np.float32([1, 4, 6, 4, 1])
k = np.outer(k, k)
kern = k / k.sum()

def local_normalize_1ch(img, const=127.0):
    mu = convolve(img, kern, mode='nearest')
    mu_sq = mu * mu
    im_sq = img * img
    tmp = convolve(im_sq, kern, mode='nearest') - mu_sq
    sigma = np.sqrt(np.abs(tmp))
    structdis = (img - mu) / (sigma + const)

    # Rescale within 0 and 1
    # structdis = (structdis + 3) / 6
    structdis = 2. * structdis / 3.
    return structdis

def local_normalize(img, num_ch=1, const=127.0):
    if num_ch == 1:
        mu = convolve(img[:, :, 0], kern, mode='nearest')
        mu_sq = mu * mu
        im_sq = img[:, :, 0] * img[:, :, 0]
        tmp = convolve(im_sq, kern, mode='nearest') - mu_sq
        sigma = np.sqrt(np.abs(tmp))
        structdis = (img[:, :, 0] - mu) / (sigma + const)

        # Rescale within 0 and 1
        # structdis = (structdis + 3) / 6
        structdis = 2. * structdis / 3.
        norm = structdis[:, :, None]
    elif num_ch > 1:
        norm = np.zeros(img.shape, dtype='float32')
        for ch in range(num_ch):
            mu = convolve(img[:, :, ch], kern, mode='nearest')
            mu_sq = mu * mu
            im_sq = img[:, :, ch] * img[:, :, ch]
            tmp = convolve(im_sq, kern, mode='nearest') - mu_sq
            sigma = np.sqrt(np.abs(tmp))
            structdis = (img[:, :, ch] - mu) / (sigma + const)

            # Rescale within 0 and 1
            # structdis = (structdis + 3) / 6
            structdis = 2. * structdis / 3.
            norm[:, :, ch] = structdis

    return norm


from collections import Counter
from random import random
from nltk import word_tokenize
from torch import nn
from torch.autograd import Variable
import pandas as pd
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence,pad_packed_sequence
import re
import numpy as np
import torch
import torch.nn.functional as F
import os
from utils import *
import sys
from reward_function import compute_reward

def train_seq2seq_batch(data_batch, model, optimizer,pos_weight,device):
	document = data_batch['document']
	label = data_batch['labels']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)

	if torch.cuda.is_available():
		document = document.to(device)
		label = label.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)

	out = model(document,input_length,indicators,begin,end,device)

	mask = label.gt(-1).float()
	# loss = F.binary_cross_entropy(out,label,weight = mask,reduction='sum',pos_weight=pos_weight)

	loss = F.binary_cross_entropy_with_logits(out,label,weight = mask,reduction='sum',pos_weight=pos_weight)
	model.zero_grad()
	loss.backward()
	optimizer.step()
	l = loss.data
	del document,label,input_length,indicators,end,begin,loss,out
	torch.cuda.empty_cache()
	return l,total_data

def train_seq2seq_batch_rl(data_batch, model, optimizer,pos_weight,device,lamb = 0.1,gamma=0.7):

	sigmoid = torch.nn.Sigmoid()
	document = data_batch['document']
	label = data_batch['labels']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']
	# summary_representation = data_batch['summary_representation']
	sentence_lengths = data_batch['sentence_lengths']
	sentences_batch = [[' '.join(s) for s in d] for d in data_batch['sent_txt']]
	references_batch = data_batch['refs']
	# rouge_matrix_batch =data_batch['rouge_matrix']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)

	if torch.cuda.is_available():
		document = document.to(device)
		label = label.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)
		# summary_representation = summary_representation.to(device)

	out= model(document,input_length,indicators,begin,end,device)

	output,_ = pad_packed_sequence(document)

	out1 = out.squeeze(-1)
	scores = sigmoid(out1).data
	scores = scores.permute(1,0)

	reward,rl_label = compute_reward(scores,input_length,output,sentences_batch,references_batch,device,sentence_lengths,lamb=lamb)
	mask = label.gt(-1).float()

	# Option 1
	loss_ce = F.binary_cross_entropy_with_logits(out,label,weight = mask,reduction='sum',pos_weight=pos_weight)
	mask = mask*reward
	loss_rl = F.binary_cross_entropy_with_logits(out,rl_label,weight = mask,reduction='sum',pos_weight=pos_weight)
	loss = (1-gamma)*loss_ce+gamma*loss_rl

	model.zero_grad()
	loss.backward()

	optimizer.step()
	l = loss.data
	del document,label,input_length,indicators,end,begin,loss,out
	torch.cuda.empty_cache()
	return l,total_data

def train_seq2seq_batch_neusum(data_batch, model, optimizer,pos_weight,device):
	regression_crit = nn.KLDivLoss(reduction='none')
	document = data_batch['document']
	selections = data_batch['selections']
	scores_gain = data_batch['scores_gain']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']
	# summary_representation = data_batch['summary_representation']
	sentence_lengths = data_batch['sentence_lengths']
	# sentences_batch = data_batch['sentences']
	# references_batch = data_batch['refs']
	# rouge_matrix_batch =data_batch['rouge_matrix']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)

	max_length = max(input_length)
	oracle_length = [s.shape[0] for s in selections]
	max_step = max(oracle_length)

	doc_sent_mask = torch.zeros(len(input_length),max_length)
	for i,l in enumerate(input_length):
		doc_sent_mask[i,:l]=1

	loss_mask = torch.zeros(len(oracle_length),max_step,max_length)
	for i,l in enumerate(oracle_length):
		# batch*step
		loss_mask[i,:l]=doc_sent_mask[i]

	selections = pad_sequence(selections,batch_first=True,padding_value=-1) #batch*longest_oracle
	scores_gain =  make_scores_gain(scores_gain,max_step,max_length)# batch*step*doc_length
	# print(scores_gain.shape)

	if torch.cuda.is_available():
		document = document.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)
		doc_sent_mask = doc_sent_mask.to(device)
		scores_gain = scores_gain.to(device)
		loss_mask = loss_mask.to(device)

		# summary_representation = summary_representation.to(device)

	doc_sent_scores = model(document,input_length,indicators,begin,end,device,max_step,selections,doc_sent_mask)
	loss = regression_loss(doc_sent_scores, scores_gain, loss_mask, regression_crit)

	model.zero_grad()
	loss.backward()

	optimizer.step()
	l = loss.data
	del document,input_length,indicators,end,begin,loss,doc_sent_mask,scores_gain,loss_mask,doc_sent_scores
	torch.cuda.empty_cache()
	return l,total_data

def make_scores_gain(scores_gain,max_step,max_doc_length,temp=200):
	####scores_gain: [[step*doc_length] for batch times]
	####out: [batch*step*doc_length]
	out = torch.zeros(len(scores_gain),max_step,max_doc_length)
	for i,score in enumerate(scores_gain):
		score = F.softmax(temp*score,dim=1)
		out[i,:score.shape[0],:score.shape[1]] = score
	return out

def regression_loss(pred_scores, gold_scores, mask, crit):
	"""
	:param pred_scores: (step, doc_len, batch)
	:param gold_scores: (batch, step, doc_len)
	:param mask: (batch, step,doc_len)
	:param crit:
	:return:
	"""

	pred_scores = pred_scores.permute(2,0,1)  # (batch, step, doc_len)
	if isinstance(crit, nn.KLDivLoss):
		# TODO: we better use log_softmax(), not log() here. log_softmax() is more numerical stable.
		pred_scores = torch.log(pred_scores + 1e-8)
	# gold_scores = gold_scores.view(*pred_scores.size())
	loss = crit(pred_scores, gold_scores)
	loss = loss * mask
	reduce_loss = loss.sum()
	return reduce_loss

def train_seq2seq_batch_newloss(data_batch, model, optimizer,pos_weight,device,beta=0.7):

	sigmoid = torch.nn.Sigmoid()
	document = data_batch['document']
	label = data_batch['labels']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']
	# summary_representation = data_batch['summary_representation']
	sentence_lengths = data_batch['sentence_lengths']
	# sentences_batch = data_batch['sentences']
	# references_batch = data_batch['summary_text']
	# rouge_matrix_batch =data_batch['rouge_matrix']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)

	if torch.cuda.is_available():
		document = document.to(device)
		label = label.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)
		# summary_representation = summary_representation.to(device)

	out= model(document,input_length,indicators,begin,end,device)

	output,_ = pad_packed_sequence(document)

	out1 = out.squeeze(-1)
	scores = sigmoid(out1)
	scores = scores.permute(1,0)

	mask = label.gt(-1).float()

	# Option 1
	loss_ce = F.binary_cross_entropy_with_logits(out,label,weight = mask,reduction='sum',pos_weight=pos_weight)
	o1 = output.permute(1,0,2).unsqueeze(3)
	o1.requires_grad = False
	# sim_mat =  F.cosine_similarity(o1,o1.permute(0,3,2,1),dim=2)
	# try:
	# 	sim_mat =  torch.stack([F.cosine_similarity(o1[i],o1[i].permute(2,1,0),dim=1) for i in range(o1.size()[0])],0)
	# except:

	o1 = o1.cpu()

	mask = (torch.eye(o1.shape[1], o1.shape[1])==1)
	sim_mat =  torch.stack([F.cosine_similarity(o1[i],o1[i].permute(2,1,0),dim=1).masked_fill_(mask, 0) for i in range(o1.size()[0])],0)
	sim_mat = sim_mat.to(device)
	loss_redundancy = torch.sum(torch.bmm(torch.bmm(scores.unsqueeze(1), sim_mat),scores.unsqueeze(2)))

	loss = (1-beta)*loss_ce+beta*loss_redundancy

	model.zero_grad()
	loss.backward()

	optimizer.step()
	l = loss.data
	del document,label,input_length,indicators,end,begin,loss,out, sim_mat,loss_redundancy
	torch.cuda.empty_cache()
	return l,total_data

def eval_seq2seq(val_dataloader,model,hyp_path,ref_path,word_length_limit, sent_length_limit,\
				pos_weight,device,remove_stopwords,stemmer,\
				meteor=False,lcs=False,saveScores=False,save_filename=None,eval_red=False,\
				use_trigramblock=False,use_mmr=False,lambd=0,\
				use_neusum=False):
	model.eval()
	total_loss=0
	total_data=0
	hyp_path_list = []
	ref_path_list = []
	total_correct = 0
	all_selections = []
	all_oracle = []
	all_sections=[]
	all_ids=[]
	# all_attn_weight = []

	sigmoid = torch.nn.Sigmoid()
	for i,data in enumerate(val_dataloader):
		if use_mmr:
			summaryfiles,referencefiles,loss,num_data,select_ids,oracle,sections = eval_seq2seq_batch_rl(sigmoid,data, model,hyp_path,ref_path, \
																						word_length_limit, sent_length_limit,pos_weight,device,\
																						lambd=lambd)
		elif use_neusum:
			summaryfiles,referencefiles,loss,num_data,select_ids,oracle,sections = eval_seq2seq_batch_neusum(data, model,hyp_path,ref_path, \
																						word_length_limit, sent_length_limit,pos_weight,device)
		else:
			summaryfiles,referencefiles,loss,num_data,select_ids,oracle,sections = eval_seq2seq_batch(sigmoid,data, model,hyp_path,ref_path, \
																						word_length_limit, sent_length_limit,pos_weight,device,\
																						use_trigramblock)

		hyp_path_list.extend(summaryfiles)
		ref_path_list.extend(referencefiles)
		# all_attn_weight.extend(attn_weight_batch)
		all_selections.extend(select_ids)
		all_ids.extend(data['id'])
		# all_sections.extend(sections)
		# all_oracle.extend(oracle)
		total_loss+=loss
		total_data+=num_data
		del data
		del loss
		# break

		# if i%200==1:
		# 	print('Batch %d, Loss: %f'%(i,total_loss/float(total_data)))

	rouge2,df = get_rouge(hyp_path_list, ref_path_list, remove_stopwords,stemmer,lcs)
	if eval_red:
		all_unigram_ratio,all_bigram_ratio,all_trigram_ratio,all_redundancy = get_redundancy_scores(hyp_path_list)

	if meteor:
		model_type = type(model).__name__
		get_meteor(hyp_path_list, ref_path_list,model_type)
	if saveScores:
		all_sections.append([0])
		# all_oracle.append([0])
		all_selections.append([0])
		all_ids.append('avg')
		all_unigram_ratio.append(sum(all_unigram_ratio)/len(all_unigram_ratio))
		all_bigram_ratio.append(sum(all_bigram_ratio)/len(all_bigram_ratio))
		all_trigram_ratio.append(sum(all_trigram_ratio)/len(all_trigram_ratio))
		all_redundancy.append(sum(all_redundancy)/len(all_redundancy))
		# all_attn_weight.append([0])
		df['id'] = pd.Series(all_ids,index =df.index)
		df['selections'] = pd.Series(np.array(all_selections),index =df.index)
		df['unigram_ratio'] = pd.Series(np.array(all_unigram_ratio),index =df.index)
		df['bigram_ratio'] = pd.Series(np.array(all_bigram_ratio),index =df.index)
		df['trigram_ratio'] = pd.Series(np.array(all_trigram_ratio),index =df.index)
		df['redundancy_score'] = pd.Series(np.array(all_redundancy),index =df.index)

		# df['oracle'] = pd.Series(np.array(all_oracle),index =df.index)
		# df['sections'] = pd.Series(np.array(all_sections),index =df.index)
		# df['attn_weight'] = pd.Series(np.array(all_attn_weight),index =df.index)
		df.to_csv('%s.csv'%(save_filename))
	return rouge2, total_loss/float(total_data)

def eval_seq2seq_batch_rl(sigmoid,data_batch,model,hyp_path, ref_path, word_length_limit, sent_length_limit, pos_weight,device,lambd=1):
	document = data_batch['document']
	label = data_batch['labels']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']
	# rouge_matrix_batch =data_batch['rouge_matrix']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)
	if torch.cuda.is_available():
		document = document.to(device)
		label = label.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)

	reference = data_batch['refs']

	ids = data_batch['id']

	out= model(document,input_length,indicators,begin,end,device)

	output,_ = pad_packed_sequence(document)
	mask = label.gt(-1).float()
	loss = F.binary_cross_entropy_with_logits(out,label,weight = mask,reduction='sum',pos_weight=pos_weight)
	out = out.squeeze(-1)
	scores = sigmoid(out).data
	scores = scores.permute(1,0)

	summaryfile_batch, reffile_batch,selections= predict_redundancy_max(scores, ids, data_batch['sent_txt'], hyp_path, ref_path,reference, word_length_limit, sent_length_limit,output,device,lambd)
	label = label.squeeze(-1)
	label = label.permute(1,0)

	all_oracle = [list((label[i]==1).nonzero().squeeze(-1).cpu().numpy()) for i in range(label.shape[0])]
	sections = [list(torch.unique(end[i],sorted=True).cpu().numpy()) for i in range(end.shape[0])]
	del document,label,input_length,indicators,end,begin
	return summaryfile_batch,reffile_batch,loss.data,total_data,selections,all_oracle,sections

def eval_seq2seq_batch_neusum(data_batch, model,hyp_path,ref_path, word_length_limit, sent_length_limit,\
								pos_weight,device):
	regression_crit = nn.KLDivLoss(reduction='none')
	document = data_batch['document']
	selections = data_batch['selections']
	scores_gain = data_batch['scores_gain']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']
	sentence_lengths = data_batch['sentence_lengths']
	ids = data_batch['id']

	reference = data_batch['refs']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)

	max_length = max(input_length)
	oracle_length = [s.shape[0] for s in selections]
	max_step = max(oracle_length)

	doc_sent_mask = torch.zeros(len(input_length),max_length)
	for i,l in enumerate(input_length):
		doc_sent_mask[i,:l]=1

	loss_mask = torch.zeros(len(oracle_length),max_step,max_length)
	for i,l in enumerate(oracle_length):
		# batch*step
		loss_mask[i,:l]=doc_sent_mask[i]

	selections = pad_sequence(selections,batch_first=True,padding_value=-1) #batch*longest_oracle
	scores_gain =  make_scores_gain(scores_gain,max_step,max_length)# batch*step*doc_length
	# print(scores_gain.shape)

	if torch.cuda.is_available():
		document = document.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)
		doc_sent_mask = doc_sent_mask.to(device)
		scores_gain = scores_gain.to(device)
		loss_mask = loss_mask.to(device)

		# summary_representation = summary_representation.to(device)

	doc_sent_scores = model(document,input_length,indicators,begin,end,device,max(40,max_step),selections,doc_sent_mask)
	# print(doc_sent_scores.shape)
	loss = regression_loss(doc_sent_scores[:max_step], scores_gain, loss_mask, regression_crit)
	doc_sent_scores = doc_sent_scores.permute(2,0,1)
	# print(doc_sent_scores.shape)
	summaryfile_batch, reffile_batch,selections= predict_neusum(doc_sent_scores, ids, data_batch['sent_txt'], hyp_path, ref_path,reference,word_length_limit, sent_length_limit)

	l = loss.data
	torch.cuda.empty_cache()

	# all_oracle = [list((label[i]==1).nonzero().squeeze(-1).cpu().numpy()) for i in range(label.shape[0])]
	all_oracle=None
	sections = [list(torch.unique(end[i],sorted=True).cpu().numpy()) for i in range(end.shape[0])]
	del document,input_length,indicators,end,begin,loss,doc_sent_mask,scores_gain,loss_mask,doc_sent_scores
	return summaryfile_batch,reffile_batch,l,total_data,selections,all_oracle,sections

def eval_seq2seq_batch(sigmoid,data_batch,model,hyp_path,ref_path, word_length_limit, sent_length_limit,\
						pos_weight,device,use_trigramblock=False):
	document = data_batch['document']
	label = data_batch['labels']
	input_length = data_batch['input_length']
	indicators = data_batch['indicators']
	padded_lengths = data_batch['padded_lengths']

	total_data = torch.sum(input_length)
	end = torch.clamp(torch.cumsum(padded_lengths,1),0,input_length[0])
	begin = torch.cat((torch.zeros((len(input_length),1),dtype=torch.long),end[:,:-1]),1)
	if torch.cuda.is_available():
		document = document.to(device)
		label = label.to(device)
		input_length = input_length.to(device)
		indicators = indicators.to(device)
		end = end.to(device)
		begin= begin.to(device)

	reference = data_batch['refs']

	ids = data_batch['id']

	out= model(document,input_length,indicators,begin,end,device)

	mask = label.gt(-1).float()
	loss = F.binary_cross_entropy_with_logits(out,label,weight = mask,reduction='sum',pos_weight=pos_weight)
	out = out.squeeze(-1)
	scores = sigmoid(out).data
	scores = scores.permute(1,0)

	if use_trigramblock:
		summaryfile_batch, reffile_batch,selections= predict_trigram_block(scores, ids, data_batch['sent_txt'], hyp_path, ref_path,reference,word_length_limit, sent_length_limit)
	else:
		summaryfile_batch, reffile_batch,selections= predict(scores, ids, data_batch['sent_txt'], hyp_path, ref_path,reference,word_length_limit, sent_length_limit)

	label = label.squeeze(-1)
	label = label.permute(1,0)

	all_oracle = [list((label[i]==1).nonzero().squeeze(-1).cpu().numpy()) for i in range(label.shape[0])]
	sections = [list(torch.unique(end[i],sorted=True).cpu().numpy()) for i in range(end.shape[0])]
	del document,label,input_length,indicators,end,begin
	return summaryfile_batch,reffile_batch,loss.data,total_data,selections,all_oracle,sections

def predict(score_batch, ids, src_txt_list, hyp_path, ref_path,tgt_txt, word_length_limit, sent_length_limit):
	#score_batch = [batch,seq_len]
	summaryfile_batch = []
	reffile_batch = []
	# all_ids = []
	selections = []

	for i in range(len(src_txt_list)):
		summary = []
		scores = score_batch[i,:len(src_txt_list[i])]
		sorted_linenum = [x for _,x in sorted(zip(scores,list(range(len((src_txt_list[i]))))),reverse=True)]
		selected_ids = []
		wc = 0
		uc = 0
		for j in sorted_linenum:
			selected_ids.append(j)
			summary.append(' '.join(src_txt_list[i][j]))
			wc+=len(src_txt_list[i][j])
			uc+=1

			if uc>=sent_length_limit:
				break
			if wc>=word_length_limit:
				break

		summary='\n'.join(summary)
		selections.append(selected_ids)

		fname = hyp_path+ids[i]+'.txt'
		of = open(fname,'w')
		of.write(summary)
		of.close()
		summaryfile_batch.append(fname)

		refname = ref_path+ids[i]+'.txt'
		of = open(refname,'w')
		of.write(tgt_txt[i])
		of.close()
		reffile_batch.append(refname)

	return summaryfile_batch, reffile_batch,selections

def predict_neusum(score_batch, ids, src_txt_list, hyp_path, ref_path,tgt_txt, word_length_limit, sent_length_limit):
	#score_batch = [batch,step,doc_length]
	summaryfile_batch = []
	reffile_batch = []
	# all_ids = []
	selections = []
	for i in range(len(src_txt_list)):
		summary = []

		selected_ids = []
		wc = 0
		uc = 0
		step = 0
		while len(selected_ids)<=len(src_txt_list[i]):
			scores = score_batch[i,step,:len(src_txt_list[i])]
			# sorted_linenum = [x for _,x in sorted(zip(scores,list(range(len((src_txt_list[i]))))),reverse=True)]
			j = torch.argmax(scores).item()
			selected_ids.append(j)
			summary.append(' '.join(src_txt_list[i][j]))
			wc+=len(src_txt_list[i][j])
			uc+=1
			step+=1
			# print(step)
			if uc>=sent_length_limit:
				break
			if wc>=word_length_limit:
				break

		summary='\n'.join(summary)
		selections.append(selected_ids)

		fname = hyp_path+ids[i]+'.txt'
		of = open(fname,'w')
		of.write(summary)
		of.close()
		summaryfile_batch.append(fname)

		refname = ref_path+ids[i]+'.txt'
		of = open(refname,'w')
		of.write(tgt_txt[i])
		of.close()
		reffile_batch.append(refname)

	return summaryfile_batch, reffile_batch,selections

def predict_trigram_block(score_batch, ids, src_txt_list, hyp_path, ref_path,tgt_txt, word_length_limit, sent_length_limit,attn_weight = None):
	#score_batch = [batch,seq_len]
	correct_num = 0
	summaryfile_batch = []
	reffile_batch=[]
	selections = []
	for i in range(len(src_txt_list)):
		summary = []
		scores = score_batch[i,:len(src_txt_list[i])]
		sorted_linenum = [x for _,x in sorted(zip(scores,list(range(len((src_txt_list[i]))))),reverse=True)]
		selected_ids = []
		uc = 0
		wc = 0
		current_trigrams = set()
		for j in sorted_linenum:
			check,tmp_current_trigrams = check_trigram(current_trigrams, src_txt_list[i][j], edu=False)
			# no trigram overlap
			if check:
				current_trigrams = tmp_current_trigrams
			else:
				continue
			summary.append(' '.join(src_txt_list[i][j]))
			selected_ids.append(j)
			wc+=len(src_txt_list[i][j])
			uc+=1
			if uc>=sent_length_limit:
				break
			if wc>=word_length_limit:
				break

		summary='\n'.join(summary)
		selections.append(selected_ids)

		fname = hyp_path+ids[i]+'.txt'
		of = open(fname,'w')
		of.write(summary)
		summaryfile_batch.append(fname)

		refname = ref_path+ids[i]+'.txt'
		of = open(refname,'w')
		of.write(tgt_txt[i])
		of.close()
		reffile_batch.append(refname)

	return summaryfile_batch, reffile_batch,selections

def predict_redundancy_max(score_batch, ids, src_txt_list,hyp_path,ref_path, \
						tgt_txt, word_length_limit, sent_length_limit,\
						output,device,lamb,attn_weight = None):
	#score_batch = [batch,seq_len]
	correct_num = 0
	summaryfile_batch = []
	reffile_batch=[]
	selections=[]
	# dim = output.size()[-1]
	# dim=output[0].size()[-1]
	for i in range(len(src_txt_list)):
		summary = []

		scores = score_batch[i,:len(src_txt_list[i])]
		sorted_linenum = [x for _,x in sorted(zip(scores,list(range(len(src_txt_list[i])))),reverse=True)]
		wc = 0
		uc=0
		selected_ids = []

		summary_representation=[]

		###### sent representation
		all_sent = output[:scores.size()[0],i,:].unsqueeze(2)

		while len(selected_ids)<=len(src_txt_list[i]):
			j = sorted_linenum[0]
			summary.append(' '.join(src_txt_list[i][j]))
			selected_ids.append(j)

			###### sent representation
			summary_representation.append(output[j,i,:])
			s = torch.stack(summary_representation,1).unsqueeze(0)

			redundancy_score =torch.max(F.cosine_similarity(all_sent,s,1),1)[0]
			# print(redundancy_score)
			# print(redundancy_score)
			scores[j] = -100
			final_scores = lamb*scores - ((1-lamb)*redundancy_score)
			sorted_linenum = [x for _,x in sorted(zip(final_scores,list(range(len(src_txt_list[i])))),reverse=True)]

			wc+=len(src_txt_list[i][j])
			uc+=1
			if uc>=sent_length_limit:
				break
			if wc>=word_length_limit:
				break

		summary='\n'.join(summary)
		selections.append(selected_ids)

		fname = hyp_path+ids[i]+'.txt'
		of = open(fname,'w')
		of.write(summary)
		summaryfile_batch.append(fname)

		refname = ref_path+ids[i]+'.txt'
		of = open(refname,'w')
		of.write(tgt_txt[i])
		of.close()
		reffile_batch.append(refname)

	return summaryfile_batch, reffile_batch,selections


import sys,os,cv2,time
from keras.preprocessing.image import img_to_array
from keras.models import load_model
import numpy as np
from PyQt5.uic import loadUi
from PyQt5.QtWidgets import QApplication, QAction, qApp, QMainWindow
from PyQt5 import QtCore
from PyQt5 import QtWidgets
from PyQt5 import QtGui
from PyQt5.QtCore import pyqtSignal
from PyQt5 import *

class Window(QMainWindow):
    def __init__(self):
        super(Window, self).__init__()
        loadUi('graph.ui', self)
        self.setWindowTitle('Emotion Recognition')
        width,height=1110,738
        self.setFixedSize(width, height)
        self.pushButton.clicked.connect(self.start)
        self.pushButton_2.clicked.connect(self.closed)
    def closed(self):
        self.close()
    def start(self):
        self.capture = cv2.VideoCapture(0)
        self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.timer = QtCore.QTimer(self)
        self.timer.timeout.connect(self._frame)
        self.timer.start(4)

    def _frame(self):
        ret, self.image = self.capture.read()
        faceCascade = cv2.CascadeClassifier(
            'Cascade/haarcascade-frontalface-default.xml')
        font = cv2.FONT_HERSHEY_SIMPLEX

        #Dictionary for emotion recognition model output and emotions
        emotions = {0: 'Angry', 1: 'Fear', 2: 'Happy',
                    3: 'Sad', 4: 'Surprised', 5: 'Neutral'}

        #Upload images of emojis from emojis folder
        emoji = []
        for index in range(6):
            emotion = emotions[index]
            emoji.append(cv2.imread('emoji/' + emotion + '.png', -1))
        model = load_model('models/emotion.h5', compile=False)
        frame_count = 0

        if ret:

            gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)

            #Face detection takes approx 0.07 seconds
            start_time = time.time()
            faces = faceCascade.detectMultiScale(
                gray,
                scaleFactor=1.1,
                minNeighbors=7,
                minSize=(100, 100),)

            y0 = 15
            for index in range(6):
                emotions[index] + ': ', (5, y0)
                y0 += 15

            try:
                # Flag for showing probability graph of only one face

                if len(faces) > 0:
                    for x, y, width, height in faces:
                        cropped_face = gray[y:y + height, x:x + width] #when detect face after cropped for test our model
                        test_image = cv2.resize(cropped_face, (48, 48)) #resize for test with model(Note model image size 48x48
                        test_image = test_image.reshape([-1, 48, 48, 1]) #reshape and grayscale

                        test_image = np.multiply(test_image, 1.0 / 255.0) #convert img to numpy

                        # Probablities of all classes
                        #Finding class probability takes approx 0.05 seconds
                        start_time = time.time()
                        if frame_count % 5 == 0:
                            probab = model.predict(test_image)[0] * 100 #test

                            label = np.argmax(probab)
                            probab_predicted = int(probab[label])
                            predicted_emotion = emotions[label]

                            frame_count = 0

                        frame_count += 1

                        # Drawing on frame
                        font_size = width / 300
                        filled_rect_ht = int(height / 5)

                        #Resizing emoji according to size of detected face
                        emoji_face = emoji[(label)]
                        emoji_face = cv2.resize(
                            emoji_face, (filled_rect_ht, filled_rect_ht))

                        #Positioning emojis on frame
                        emoji_x1 = x + width - filled_rect_ht
                        emoji_x2 = emoji_x1 + filled_rect_ht
                        emoji_y1 = y + height
                        emoji_y2 = emoji_y1 + filled_rect_ht

                        #Drawing rectangle and showing output values on frame
                        cv2.rectangle(self.image, (x, y), (x + width,
                                                      y + height), (155, 155, 0), 2)
                        cv2.rectangle(self.image, (x-1, y+height), (x+1 + width, y + height+filled_rect_ht),
                                      (155, 155, 0), cv2.FILLED)
                        cv2.putText(self.image, predicted_emotion+' ' + str(probab_predicted)+'%',
                                    (x, y + height + filled_rect_ht-10), font, font_size, (255, 255, 255), 1, cv2.LINE_AA)

                        # Showing emoji on frame
                        for c in range(0, 3):
                            self.image[emoji_y1:emoji_y2, emoji_x1:emoji_x2, c] = emoji_face[:, :, c] * (emoji_face[:, :, 3] / 255.0) + self.image[emoji_y1:emoji_y2, emoji_x1:emoji_x2, c] * (1.0 - emoji_face[:, :, 3] / 255.0)

            except Exception as error:

                pass

        self.display(self.image, 1)

    def display(self, img, windows=1):
        qformat = QtGui.QImage.Format_Indexed8
        if len(img.shape) == 3:
            if img.shape[2] == 4:
                qformat = QtGui.QImage.Format_RGBA8888
            else:
                qformat = QtGui.QImage.Format_RGB888
        outimage = QtGui.QImage(
            img, img.shape[1], img.shape[0], img.strides[0], qformat)
        outimage = outimage.rgbSwapped()
        if windows == 1:
            self.label.setPixmap(QtGui.QPixmap.fromImage(outimage))
            self.label.setScaledContents(True)
if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = Window()

    window.show()
    sys.exit(app.exec_())


from __future__ import absolute_import
from ann_benchmarks.algorithms.base import BaseANN

import subprocess

import struct
import subprocess
import sys
import os
import glob
import numpy as np
import random
import string

class Countrymaam(BaseANN):
    def __init__(self, metric, params):
        self._metric = metric
        self._index = params.get("index", "kd-tree")
        self._n_trees = params.get("n_trees", 8)
        self._leaf_size = params.get("leaf_size", 8)

    def fit(self, X):
        X = X.astype(np.float64)
        suffix = "".join(random.choices(string.ascii_lowercase, k=16))
        index_file_path = f"index_{suffix}_{os.getpid()}.bin"
        p = subprocess.Popen([
            "countrymaam",
            "train",
            "--dim", str(len(X[0])),
            "--index", self._index,
            "--leaf-size", str(self._leaf_size),
            "--tree-num", str(self._n_trees),
       	    "--output", index_file_path
        ], stdin=subprocess.PIPE)

        p.stdin.write(struct.pack(f"={X.size}d", *np.ravel(X)))
        p.communicate()
        p.stdin.close()

        self._pipe = subprocess.Popen([
            "countrymaam",
            "predict",
            "--dim", str(len(X[0])),
            "--index", self._index,
       	    "--input", index_file_path
        ], stdin=subprocess.PIPE, stdout=subprocess.PIPE)

    def set_query_arguments(self, search_k):
        self._search_k = search_k

    def query(self, v, n):
        v = v.astype(np.float64)
        self._pipe.stdin.write(struct.pack(f"=i", self._search_k))
        self._pipe.stdin.write(struct.pack(f"=i", n))
        self._pipe.stdin.write(struct.pack(f"={v.size}d", *v))
        self._pipe.stdin.flush()

        rn = struct.unpack("=i", self._pipe.stdout.read(4))[0]
        ret = [0] * rn
        for i in range(rn):
            ret[i] = struct.unpack("=i", self._pipe.stdout.read(4))[0]
        return np.array(ret)

    def __str__(self):
        return f"Countrymaam(index={self._index}, leaf_size={self._leaf_size} n_trees={self._n_trees}, search_k={self._search_k})"