#!/usr/bin/env python
# encoding: utf-8
"""
Sample for Multi Pixel Gibbs MCMC.
"""

from collections import OrderedDict
# import pprint
import json
import numpy as np

import fsps

from astropy.utils.console import ProgressBar
from astropy.table import Table, hstack, Column

from sedbot.photconv import micro_jy_to_luminosity
from sedbot.chain import MultiPixelChain

MODEL = None

class MultiPixelGibbsBgSampler(object):
    """MH-in-Gibbs sampler for three level hierarchical models of
    pixel parameters, global parameters, and a linearly sampled background
    parameter.

    Parameters
    ----------
    model_initializer : function
        Function that, when called, creates a model instance.
        This is done to enable multi-processing.
    n_cpu : int
        Number of CPUs to run on.
    """
    def __init__(self, model_initializer,
                 n_cpu=1):
        super(MultiPixelGibbsBgSampler, self).__init__()
        self._n_cpu = n_cpu
        self._model_initializer = model_initializer

        # Set up a local model instance
        self._model = self._model_initializer()

        # Set up a pixel compute pool
        self._map = self._init_pixel_compute_pool()

    def _init_pixel_compute_pool(self):
        """Initialize an ipython.parallel pool for processing pixels."""
        # Put the model in the global namespace for parallel processing
        global MODEL
        MODEL = self._model_initializer()

        # Setup up the pool
        if self._n_cpu == 1:
            map_fcn = map  # single processing
        else:
            map_fcn = None  # TODO
        return map_fcn

    def sample(self, n_iter,
               theta_prop=None,
               phi_prop=None,
               theta0=None,
               phi0=None,
               B0=None,
               chain=None,
               use_bkg_sample_errors=False):
        """Sample for `n_iter` Gibbs steps.

        Parameters
        ----------
        n_iter : int
            Number of Gibbs steps (number of full iterations over all
            parameters)
        theta_prop : ndarray
            Standard deviations of Gaussian proposal distribtions for
            theta (pixel) parameters. If any individual proposal sigma is
            zero then that parameter will not be sampled (but instead kept
            constant in the chain).
        phi_prop : ndarray
            Standard deviations of Gaussian proposal distributions for
            phi (global) parameters. Leave as ``None`` to use the model's
            defaults. If any individual proposal sigma is
            zero then that parameter will not be sampled (but instead kept
            constant in the chain).
        theta0 : ndarray
            Initial values for the theta parameters, a ``(n_pix, n_theta)``
            array.
        phi0 : ndarray
            Initial values for the phi parameters, a ``(n_phi)`` array.
        B0 : ndarray
            Initial values for the background, a ``(n_band)`` array.
            Background is in units of flux per arcsec^2.
        chain : :class:`sedbot.chain.MultiPixelChain`
            A previously-built chain, whose last sample will be used
            as the starting points for this sampling (in liue of setting
            `theta0`, `phi0` and `B0` manually).
        use_bkg_sample_errors : bool
            Use sample standard deviation when updating the background.
        """
        global MODEL

        # Keep these for storage in the HDF5 file later
        self._theta_prop = theta_prop
        self._phi_prop = phi_prop

        if chain is not None:
            theta0 = np.empty((self._model.n_pix, self._model.n_theta),
                              dtype=np.float)
            for i, n in enumerate(self._model.theta_params):
                theta0[:, i] = chain[n][-1]

            phi0 = np.empty(self._model.n_phi, dtype=np.float)
            for i, n in enumerate(self._model.phi_params):
                phi0[i] = chain[n][-1]

            B0 = np.empty(self._model.n_bands, dtype=np.float)
            for i, (n, instr) in enumerate(zip(self._model.observed_bands,
                                               self._model.instruments)):
                name = "B__{0}__{1}".format(instr, n)
                B0[i] = chain[name][-1]
        else:
            assert theta0 is not None
            assert phi0 is not None
            assert B0 is not None

        # initialize memory
        i0 = self._init_chains(n_iter, theta0, phi0, B0)

        with ProgressBar(n_iter) as bar:
            for i in xrange(i0, i0 + n_iter):
                # Sample for all pixels
                args = []
                for ipix in xrange(self._model.n_pix):
                    _post0 = self.pix_lnpost[i - 1, ipix]
                    _theta0 = self.theta[i - 1, ipix, :]
                    _phi0 = self.phi[i - 1, :]
                    _B0 = self.B[i - 1, :]
                    args.append((ipix, _post0, _theta0, _phi0, _B0,
                                 theta_prop))
                results = self._map(pixel_mh_sampler, args)
                for ipix, result in enumerate(results):
                    lnpost, theta, blob, theta_accept = result
                    self.pix_lnpost[i, ipix] = lnpost
                    self.theta[i, ipix, :] = theta
                    if blob is not None:
                        for k, v in blob.iteritems():
                            self.blobs[i][k][ipix] = v
                    else:
                        # repeat previous blob values
                        for k in self.blobs.dtype.fields:
                            self.blobs[i][k][ipix] = self.blobs[i - 1][k][ipix]
                    self._theta_n_accept[ipix, :] += theta_accept

                # Update the background
                model_seds = self.blobs[i]['model_sed']
                B_new, global_lnp, pixel_lnp, pixel_blobs \
                    = self._model.update_background(
                        self.theta[i, :, :],
                        self.phi[i - 1, :],
                        model_seds,
                        sample_errors=use_bkg_sample_errors)
                self.B[i, :] = B_new

                # Sample the global parameters
                phi_new, global_lnp, pixel_lnps, pixel_blobs, n_accept \
                    = self._global_mh(i,
                                      global_lnp,
                                      self.phi[i - 1, :],
                                      self.theta[i, :, :],
                                      self.B[i, :],
                                      phi_prop)
                # fill in new values to chain.
                self.phi[i, :] = phi_new
                self.lnpost[i] = global_lnp
                self._phi_n_accept += n_accept
                # Update ln post and blob data for all pixels too
                # it was done for the pixel-only step, but these values have
                # changed given the new background, etc.
                if pixel_lnps is not None:
                    for ipix, lnp in enumerate(pixel_lnps):
                        self.pix_lnpost[i, ipix] = lnp
                if pixel_blobs is not None:
                    for ipix, blobs in enumerate(pixel_blobs):
                        for k, v in blobs.iteritems():
                            self.blobs[i][k][ipix] = v

                bar.update()

    def _global_mh(self, i, lnpost0, phi, theta, B, phi_prop):
        """Perform MH-in-Gibbs at the global level."""
        phi = np.copy(phi)
        lnpost = lnpost0
        lnpost_pixel = None
        blobs = None
        n_accept = np.zeros(phi.shape[0], dtype=np.int)
        for j in xrange(phi.shape[0]):
            # Gaussian proposal for parameter j, only
            # Skip parameter advancement if proposal sigma is zero.
            if phi_prop[j] > 0.:
                phi_new = np.copy(phi)
                phi_new[j] += phi_prop[j] * np.random.randn()
                global_lnp, pixel_lnp, pixel_blobs \
                    = self._model.sample_global(theta,
                                                phi_new,
                                                B)
                ln_r = global_lnp - lnpost0
                reject = True
                if ~np.isfinite(ln_r):
                    pass
                elif ln_r >= 0.:
                    reject = False
                else:
                    x = np.random.random_sample()
                    if x < np.exp(ln_r):
                        reject = False
                if not reject:
                    # adopt new point
                    phi = phi_new
                    lnpost = global_lnp
                    lnpost_pixel = pixel_lnp
                    blobs = pixel_blobs
                    n_accept[j] = 1
        return phi, lnpost, lnpost_pixel, blobs, n_accept

    def _init_chains(self, n_iter, theta0, phi0, B0):
        """Initialize memory

        Note this could be adapted to extend existing chains.
        """
        # Make parameter chains
        self.theta = np.empty((n_iter + 1,
                               self._model.n_pix,
                               self._model.n_theta),
                              dtype=np.float)
        self.theta.fill(np.nan)
        self.theta[0, :, :] = theta0

        self.phi = np.empty((n_iter + 1,
                             self._model.n_phi),
                            dtype=np.float)
        self.phi.fill(np.nan)
        self.phi[0, :] = phi0

        self.B = np.empty((n_iter + 1,
                           self._model.n_bands),
                          dtype=np.float)
        self.B.fill(np.nan)
        self.B[0, :] = B0

        self.lnpost = np.empty(n_iter + 1, dtype=np.float)
        self.lnpost.fill(np.nan)

        self.pix_lnpost = np.empty((n_iter + 1, self._model.n_pix),
                                   dtype=np.float)
        self.pix_lnpost.fill(np.nan)

        # Make a structured array for the complex blob data
        self.blobs = np.empty(n_iter + 1, dtype=self._model.blob_dtype)
        self.blobs.fill(np.nan)

        # Initialize the starting point of the chains
        global_lnp, pixel_lnp, pixel_blobs \
            = self._model.sample_global(theta0, phi0, B0)
        self.lnpost[0] = global_lnp
        for ipix, lnp in enumerate(pixel_lnp):
            self.pix_lnpost[0, ipix] = lnp
        for ipix, blobs in enumerate(pixel_blobs):
            for k, v in blobs.iteritems():
                self.blobs[0][k][ipix] = v

        # Track acceptance of individual parameters for each pixel
        self._theta_n_accept = np.zeros((self._model.n_pix,
                                         self._model.n_theta),
                                        dtype=np.int)
        self._phi_n_accept = np.zeros(self._model.n_phi, dtype=np.int)

        return 1

    @property
    def median_theta_faccept(self):
        """Median acceptance fraction across all pixels for each parameter"""
        n_samples = float(self.theta.shape[0])
        faccept = np.median(self._theta_n_accept / n_samples, axis=0)
        return faccept

    @property
    def phi_faccept(self):
        """Acceptance fraction of phi parameters."""
        n_samples = float(self.theta.shape[0])
        return self._phi_n_accept / n_samples

    def print_faccept(self):
        """Printed report of the acceptance fraction statistics"""
        faccept = self.median_theta_faccept
        d = OrderedDict(zip(self._model.theta_params, faccept))
        # pp = pprint.PrettyPrinter(indent=4)
        # pp.pprint(d)
        print("theta")
        print(json.dumps(d, indent=4))
        print("phi")
        d = OrderedDict(zip(self._model.phi_params, self.phi_faccept))
        print(json.dumps(d, indent=4))

    @property
    def table(self):
        """An :class:`astropy.table.Table` with the chain."""
        msuns = np.array([fsps.get_filter(n).msun_ab
                          for n in self._model.computed_bands])
        theta_f_accept = json.dumps(dict(zip(self._model.theta_params,
                                             self.median_theta_faccept)))
        phi_f_accept = json.dumps(dict(zip(self._model.phi_params,
                                           self.phi_faccept)))
        meta = OrderedDict((
            ('theta_f_accept', theta_f_accept),
            ('phi_f_accept', phi_f_accept),
            ('observed_bands', self._model.observed_bands),
            ('instruments', self._model.instruments),
            ('computed_bands', self._model.computed_bands),
            ('msun_ab', msuns),
            ('band_indices', self._model.band_indices),
            ('theta_params', self._model.theta_params),
            ('phi_params', self._model.phi_params),
            ('theta_proposal_sigma', self._theta_prop),
            ('phi_proposal_sigma', self._theta_prop),
            ('sed', self._model._seds),
            ('sed_err', self._model._errs),
            ('pixels', self._model.pixel_metadata),
            ('area', self._model._areas)))
        # Make tables for individual chains; stack later
        # FIXME should axis order be changed for theta throughout the sampler?
        # or I can just continue to swarp aces here
        theta_table = Table(np.swapaxes(self.theta, 1, 2),
                            names=self._model.theta_params,
                            meta=meta)
        phi_table = Table(self.phi, names=self._model.phi_params)
        background_names = ["B__{0}__{1}".format(n, b)
                            for n, b in zip(self._model.instruments,
                                            self._model.observed_bands)]
        B_table = Table(self.B, names=background_names)
        blob_table = Table(self.blobs)
        tbl = MultiPixelChain(hstack((theta_table,
                                      phi_table,
                                      B_table,
                                      blob_table)))

        # Add M/L computations for each computed band.
        for i, (band_name, msun) in enumerate(zip(self._model.computed_bands,
                                                  msuns)):
            logLsol = micro_jy_to_luminosity(tbl['model_sed'][:, :, i],
                                             msun,
                                             np.atleast_2d(tbl['d']).T)
            ml = tbl['logMstar'] - logLsol
            colname = "logML_{0}".format(band_name)
            tbl.add_column(Column(name=colname, data=ml))

        return tbl

def pixel_mh_sampler(args):
    """Execute MH-in-Gibbs sampling for a single pixel.

    This function is design to be called by a ``map`` function, hence all
    arguments are wrapped in a single tuple, ``args``.

    Parameters
    ----------
    ipix : int
        Pixel index
    post0 : float
        Ln Posterior probability *for the pixel* given initial values of
        theta0 and phi0.
    theta0 : ndarray
        Initial values of theta parameters.
    phi0 : ndarray
        Initial values of the phi parameters.
    B0 : ndarray
        Initial values of the B, background, parameters.
    theta_prop : ndarray
        Array of Gaussian proposal standard deviations. If an individual
        proposal sigma is zero then that value will not be sampled,
        and kept constant.

    Returns
    -------
    post : float
        Final posterior probability *for this pixel*.
    theta : ndarray
        Theta parameters from this Gibbs step.
    blob : dict
        Stellar population metadata for this pixel given the updated `theta`
        parameters.
    n_accept : ndarray
        Array of length ``n_theta`` whose values are 1 for parameters
        that were updated, and zeros for values that remain constant from
        the previous iteration.
    """
    global MODEL
    ipix, post0, theta0, phi0, B0, theta_prop = args
    theta = np.copy(theta0)
    post = post0
    # MH on each parameter
    blob = None
    n_accept = np.zeros(theta.shape[0], dtype=np.int)
    for i in xrange(theta0.shape[0]):
        # Gaussian proposal for parameter i, only
        if theta_prop[i] > 0.:
            theta_new = theta.copy()
            theta_new[i] += theta_prop[i] * np.random.randn()
            lnpost_new, blob_new = MODEL.sample_pixel(theta_new,
                                                      phi0, B0, ipix)
            ln_r = lnpost_new - post
            # print "pix {0:d} param {1:d} ln_r {2:.2f}".format(ipix, i, ln_r),
            reject = True
            if ~np.isfinite(ln_r):
                # print " reject ln_r not finite"
                pass
            elif ln_r >= 0.:
                # print " accepted".format(ipix, i)
                reject = False
            else:
                x = np.random.random_sample()
                if x < np.exp(ln_r):
                    # print " accept {0:.2f} < r".format(x)
                    reject = False
                # else:
                    # print " reject {0:.2f} > r".format(x)
            if not reject:
                # adopt new point
                theta = theta_new
                blob = blob_new
                post = lnpost_new
                n_accept[i] = 1
    return post, theta, blob, n_accept


"""
    Brief: Demo for render labelled car 3d poses to the image
    Author: wangpeng54@baidu.com
    Date: 2018/6/10
"""
import sys
sys.path.append("../")

import argparse
import cv2
import car_models
import numpy as np
import json
import pickle as pkl

import utils.data as data
import utils.utils as uts
import utils.eval_utils as eval_uts
import renderer.render_egl as render
import logging

from collections import OrderedDict

logger = logging.getLogger()
logger.setLevel(logging.INFO)

class CarPoseVisualizer(object):
    def __init__(self, args=None, scale=0.4, linewidth=0.):
        """Initializer
        Input:
            scale: whether resize the image in case image is too large
            linewidth: 0 indicates a binary mask, while > 0 indicates
                       using a frame.
        """
        self.dataset = data.ApolloScape(args)
        self._data_config = self.dataset.get_3d_car_config()

        self.MAX_DEPTH = 1e4
        self.MAX_INST_NUM = 100
        h, w = self._data_config['image_size']

        # must round prop to 4 due to renderer requirements
        # this will change the original size a bit, we usually need rescale
        # due to large image size
        self.image_size = np.uint32(uts.round_prop_to(
            np.float32([h * scale, w * scale])))

        self.scale = scale
        self.linewidth = linewidth
        self.colors = np.random.random((self.MAX_INST_NUM, 3)) * 255

    def load_car_models(self):
        """Load all the car models
        """
        self.car_models = OrderedDict([])
        logging.info('loading %d car models' % len(car_models.models))
        for model in car_models.models:
            car_model = '%s/%s.json' % (self._data_config['car_model_dir'],
                                       model.name)
            with open(car_model, 'rb') as f:
                self.car_models[model.name] = json.load(f)
                self.car_models[model.name]['vertices'] = np.array(
                    self.car_models[model.name]['vertices'])
                # fix the inconsistency between obj and json
                self.car_models[model.name]['vertices'][:, [0, 1]] *= -1

    def render_car_cv2(self,  pose, car_name, image):
            car = self.car_models[car_name]
            pose = np.array(pose)
            # project 3D points to 2d image plane
            rmat = uts.euler_angles_to_rotation_matrix(pose[:3])
            rvect, _ = cv2.Rodrigues(rmat)
            imgpts, jac = cv2.projectPoints(np.float32(car['vertices']), rvect, pose[3:], self.intrinsic, distCoeffs=None)
            mask = np.zeros(image.shape)
            for face in np.array(car['faces']) - 1:
                pts = np.array([[imgpts[idx, 0, 0], imgpts[idx, 0, 1]] for idx in face], np.int32)
                pts = pts.reshape((-1, 1, 2))
                cv2.polylines(mask, [pts], True, (0, 255, 0), thickness=1)

            return mask

    def render_car(self, pose, car_name):
        """Render a car instance given pose and car_name
        """
        car = self.car_models[car_name]
        scale = np.ones((3, ))
        pose = np.array(pose)
        vert = uts.project(pose, scale, car['vertices'])
        K = self.intrinsic
        intrinsic = np.float64([K[0, 0], K[1, 1], K[0, 2], K[1, 2]])
        depth, mask = render.renderMesh_py(np.float64(vert),
                                           np.float64(car['faces']),
                                           intrinsic,
                                           self.image_size[0],
                                           self.image_size[1],
                                           np.float64(self.linewidth))
        return depth, mask

    def compute_reproj_sim(self, car_names, out_file=None):
        """Compute the similarity matrix between every pair of cars.
        """
        if out_file is None:
            out_file = './sim_mat.txt'

        sim_mat = np.eye(len(self.car_model))
        for i in range(len(car_names)):
            for j in range(i, len(car_names)):
                name1 = car_names[i][0]
                name2 = car_names[j][0]
                ind_i = self.car_model.keys().index(name1)
                ind_j = self.car_model.keys().index(name2)
                sim_mat[ind_i, ind_j] = self.compute_reproj(name1, name2)
                sim_mat[ind_j, ind_i] = sim_mat[ind_i, ind_j]

        np.savetxt(out_file, sim_mat, fmt='%1.6f')

    def compute_reproj(self, car_name1, car_name2):
        """Compute reprojection error between two cars
        """
        sims = np.zeros(10)
        for i, rot in enumerate(np.linspace(0, np.pi, num=10)):
            pose = np.array([0, rot, 0, 0, 0,5.5])
            depth1, mask1 = self.render_car(pose, car_name1)
            depth2, mask2 = self.render_car(pose, car_name2)
            sims[i] = eval_uts.IOU(mask1, mask2)

        return np.mean(sims)

    def merge_inst(self,
                   depth_in,
                   inst_id,
                   total_mask,
                   total_depth):
        """Merge the prediction of each car instance to a full image
        """

        render_depth = depth_in.copy()
        render_depth[render_depth <= 0] = np.inf
        depth_arr = np.concatenate([render_depth[None, :, :],
                                    total_depth[None, :, :]], axis=0)
        idx = np.argmin(depth_arr, axis=0)

        total_depth = np.amin(depth_arr, axis=0)
        total_mask[idx == 0] = inst_id

        return total_mask, total_depth

    def rescale(self, image, intrinsic):
        """resize the image and intrinsic given a relative scale
        """

        intrinsic_out = uts.intrinsic_vec_to_mat(intrinsic,
                                                 self.image_size)
        hs, ws = self.image_size
        image_out = cv2.resize(image.copy(), (ws, hs))

        return image_out, intrinsic_out

    def showAnn(self, image_name):
        """Show the annotation of a pose file in an image
        Input:
            image_name: the name of image
        Output:
            depth: a rendered depth map of each car
            masks: an instance mask of the label
            image_vis: an image show the overlap of car model and image
        """

        car_pose_file = '%s/%s.json' % (self._data_config['pose_dir'], image_name)
        with open(car_pose_file) as f:
            car_poses = json.load(f)
        image_file = '%s/%s.jpg' % (self._data_config['image_dir'], image_name)
        image = cv2.imread(image_file, cv2.IMREAD_UNCHANGED)[:, :, ::-1]

        intrinsic = self.dataset.get_intrinsic(image_name)
        image, self.intrinsic = self.rescale(image, intrinsic)

        merged_image = image.copy()
        mask_all = np.zeros(image.shape)
        for i, car_pose in enumerate(car_poses):
            car_name = car_models.car_id2name[car_pose['car_id']].name
            mask = self.render_car_cv2(car_pose['pose'], car_name, image)
            mask_all += mask

        mask_all = mask_all * 255 / mask_all.max()

        alpha = 0.5
        cv2.addWeighted(image.astype(np.uint8), 1, mask_all.astype(np.uint8), 0.2, 0, merged_image)

        cv2.imshow('', merged_image.astype(np.uint8))
        cv2.waitKey(0)

        # return image, mask, depth

    def showAnn_bak(self, image_name):
        """Show the annotation of a pose file in an image
        Input:
            image_name: the name of image
        Output:
            depth: a rendered depth map of each car
            masks: an instance mask of the label
            image_vis: an image show the overlap of car model and image
        """
        car_pose_file = '%s/%s.json' % (
            self._data_config['pose_dir'], image_name)
        with open(car_pose_file) as f:
            car_poses = json.load(f)
        image_file = '%s/%s.jpg' % (self._data_config['image_dir'], image_name)
        image = cv2.imread(image_file, cv2.IMREAD_UNCHANGED)[:, :, ::-1]

        # intrinsic are all used by Camera 5
        intrinsic = self.dataset.get_intrinsic(image_name, 'Camera_5')
        image, self.intrinsic = self.rescale(image, intrinsic)

        self.depth = self.MAX_DEPTH * np.ones(self.image_size)
        self.mask = np.zeros(self.depth.shape)

        for i, car_pose in enumerate(car_poses):
            car_name = car_models.car_id2name[car_pose['car_id']].name
            depth, mask = self.render_car(car_pose['pose'], car_name)
            self.mask, self.depth = self.merge_inst(
                depth, i + 1, self.mask, self.depth)

        self.depth[self.depth == self.MAX_DEPTH] = -1.0
        image = 0.5 * image
        for i in range(len(car_poses)):
            frame = np.float32(self.mask == i + 1)
            frame = np.tile(frame[:, :, None], (1, 1, 3))
            image = image + frame * 0.5 * self.colors[i, :]

        uts.plot_images({'image_vis': np.uint8(image),
                         'depth': self.depth, 'mask': self.mask},
                         layout=[1, 3])

        return image, self.mask, self.depth

class LabelResaver(object):
    """ Resave the raw labeled file to the required json format for evaluation
    """
    #(TODO Peng) Figure out why running pdb it is correct, but segment fault when
    # running
    def __init__(self, args):
        self.visualizer = CarPoseVisualizer(args, scale=0.5)
        self.visualizer.load_car_models()

    def strs_to_mat(self, strs):
        """convert str to numpy matrix
        """
        assert len(strs) == 4
        mat = np.zeros((4, 4))
        for i in range(4):
            mat[i, :] = np.array([np.float32(str_f) for str_f in strs[i].split(' ')])

        return mat

    def read_car_pose(self, file_name):
        """load the labelled car pose
        """
        cars = []
        lines = [line.strip() for line in open(file_name)]
        i = 0
        while i < len(lines):
            car = OrderedDict([])
            line = lines[i].strip()
            if 'Model Name :' in line:
                car_name = line[len('Model Name : '):]
                car['car_id'] = car_models.car_name2id[car_name].id
                pose = self.strs_to_mat(lines[i + 2: i + 6])
                pose[:3, 3] = pose[:3, 3] / 100.0 # convert cm to meter
                rot = uts.rotation_matrix_to_euler_angles(
                        pose[:3, :3], check=False)
                trans = pose[:3, 3].flatten()
                pose = np.hstack([rot, trans])
                car['pose'] = pose
                i += 6
                cars.append(car)
            else:
                i += 1

        return cars

    def convert(self, pose_file_in, pose_file_out):
        """ Convert the raw labelled file to required json format
        Input:
            file_name: str filename
        """
        car_poses = self.read_car_pose(pose_file_in)
        car_num = len(car_poses)
        MAX_DEPTH = self.visualizer.MAX_DEPTH
        image_size = self.visualizer.image_size
        intrinsic = self.visualizer.dataset.get_intrinsic(
                pose_file_in, 'Camera_5')
        self.visualizer.intrinsic = uts.intrinsic_vec_to_mat(intrinsic,
                image_size)
        self.depth = MAX_DEPTH * np.ones(image_size)
        self.mask = np.zeros(self.depth.shape)
        vis_rate = np.zeros(car_num)

        for i, car_pose in enumerate(car_poses):
            car_name = car_models.car_id2name[car_pose['car_id']].name
            depth, mask = self.visualizer.render_car(car_pose['pose'], car_name)
            self.mask, self.depth = self.visualizer.merge_inst(
                depth, i + 1, self.mask, self.depth)
            vis_rate[i] = np.float32(np.sum(mask == (i + 1))) / (
                    np.float32(np.sum(mask)) + np.spacing(1))

        keep_idx = []
        for i, car_pose in enumerate(car_poses):
            area = np.round(np.float32(np.sum(
                self.mask == (i + 1))) / (self.visualizer.scale ** 2))
            if area > 1:
                keep_idx.append(i)

            car_pose['pose'] = car_pose['pose'].tolist()
            car_pose['area'] = int(area)
            car_pose['visible_rate'] = float(vis_rate[i])
            keep_idx.append(i)

        car_poses = [car_poses[idx] for idx in keep_idx]
        with open(pose_file_out, 'w') as f:
            json.dump(car_poses, f, sort_keys=True, indent=4,
                    ensure_ascii=False)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Render car instance and convert car labelled files.')
    parser.add_argument('--image_name', default='180116_053947113_Camera_5',
                        help='the dir of ground truth')
    parser.add_argument('--data_dir', default='../apolloscape/3d_car_instance_sample/',
                        help='the dir of ground truth')
    parser.add_argument('--split', default='/sample_data', help='split for visualization')
    args = parser.parse_args()
    assert args.image_name

    if False:
        # (TODO) Debugging correct but segment fault when running. Useless for visualization
        print('Test converter')
        pose_file_in = './test_files/%s.poses' % args.image_name
        pose_file_out = './test_files/%s.json' % args.image_name
        label_resaver = LabelResaver(args)
        label_resaver.convert(pose_file_in, pose_file_out)

    print('Test visualizer')
    visualizer = CarPoseVisualizer(args)
    visualizer.load_car_models()
    visualizer.showAnn(args.image_name)


"""
(c) RIKEN 2015. All rights reserved.
Author: Keitaro Yamashita

This software is released under the new BSD License; see LICENSE.
"""
import numpy
import itertools
from yamtbx.dataproc.xds import get_xdsinp_keyword
from yamtbx.util import safe_float

class XPARM:
    def __init__(self, xparm_file=None):
        if xparm_file is not None:
            self.parse_xparm_file(xparm_file)
        else:
            self.starting_frame = 1
            self.starting_angle = 0.
            self.osc_range = -1.
            self.rotation_axis = numpy.array((1.,0.,0.))
            self.wavelength = -1.
            self.incident_beam = numpy.array((0.,0.,1.))
            self.nx, self.ny = 0, 0
            self.qx, self.qy = 0., 0.
            self.distance = 0.
            self.origin = numpy.array((0., 0.))
            self.X_axis = numpy.array((1., 0., 0.))
            self.Y_axis = numpy.array((0., 1., 0.))
            self.Z_axis = numpy.array((0., 0., 1.))
            self.spacegroup = 1
            self.unit_cell = numpy.array((100., 100., 100., 90., 90., 90.))
            self.a_axis = numpy.array((100., 0., 0.))
            self.b_axis = numpy.array((0., 100., 0.))
            self.c_axis = numpy.array((0., 0., 100.))
    # __init__()

    def space_group_info(self):
        from cctbx import sgtbx
        g = sgtbx.space_group_info(self.spacegroup).group()
        return g.info()

    def parse_xparm_file(self, xparm_file):
        lines = open(xparm_file).readlines()

        is_new_format = "XPARM.XDS" in lines[0]

        if not is_new_format:
            starting_frame, starting_angle, osc_range, rotx, roty, rotz = lines[0].split()
            wavelength, ibeamx, ibeamy, ibeamz = lines[1].split()
            nx, ny, qx, qy = lines[2].split()
            distance, orgx, orgy = lines[3].split()
            Xx, Xy, Xz = lines[4].split()
            Yx, Yy, Yz = lines[5].split()
            Zx, Zy, Zz = lines[6].split()
            spacegroup, a, b, c, alpha, beta, gamma = lines[7].split()
            ax, ay, az = lines[8].split()
            bx, by, bz = lines[9].split()
            cx, cy, cz = lines[10].split()
        else:
            starting_frame, starting_angle, osc_range, rotx, roty, rotz = lines[1].split()
            wavelength, ibeamx, ibeamy, ibeamz = lines[2].split()
            spacegroup, a, b, c, alpha, beta, gamma = lines[3].split()
            ax, ay, az = lines[4][:15], lines[4][15:30], lines[4][30:45]
            bx, by, bz = lines[5][:15], lines[5][15:30], lines[5][30:45]
            cx, cy, cz = lines[6][:15], lines[6][15:30], lines[6][30:45]
            nseg, nx, ny, qx, qy = lines[7].split()
            orgx, orgy, distance = lines[8].split()
            Xx, Xy, Xz = lines[9].split()
            Yx, Yy, Yz = lines[10].split()
            Zx, Zy, Zz = lines[11].split()

        self.starting_frame = int(starting_frame)
        self.starting_angle = float(starting_angle)
        self.osc_range = float(osc_range)
        self.rotation_axis = numpy.array((float(rotx), float(roty), float(rotz)))
        self.wavelength = float(wavelength)
        self.incident_beam = numpy.array((float(ibeamx), float(ibeamy), float(ibeamz)))
        self.nx = float(nx)
        self.ny = float(ny)
        self.qx = float(qx)
        self.qy = float(qy)
        self.distance = float(distance)
        self.origin = numpy.array((float(orgx), float(orgy)))
        self.X_axis = numpy.array((float(Xx), float(Xy), float(Xz)))
        self.Y_axis = numpy.array((float(Yx), float(Yy), float(Yz)))
        self.Z_axis = numpy.array((float(Zx), float(Zy), float(Zz)))
        self.spacegroup = int(spacegroup)
        self.unit_cell = numpy.array((float(a), float(b), float(c), float(alpha), float(beta), float(gamma)))
        self.a_axis = numpy.array(map(safe_float, (ax, ay, az)))
        self.b_axis = numpy.array(map(safe_float, (bx, by, bz)))
        self.c_axis = numpy.array(map(safe_float, (cx, cy, cz)))
    # parse_xparm_file()

    def set_info_from_xdsinp(self, xdsinp):
        # XXX x, y, z axes

        table = [("STARTING_FRAME", "starting_frame", lambda x: int(x)),
                 ("STARTING_ANGLE", "starting_angle", lambda x: float(x)),
                 ("OSCILLATION_RANGE", "osc_range", lambda x: float(x)),
                 ("ROTATION_AXIS", "rotation_axis", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("X-RAY_WAVELENGTH", "wavelength", lambda x: float(x)),
                 ("INCIDENT_BEAM_DIRECTION", "incident_beam", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("NX", "nx", lambda x: int(x)),
                 ("NY", "ny", lambda x: int(x)),
                 ("QX", "qx", lambda x: float(x)),
                 ("QY", "qy", lambda x: float(x)),
                 ("DETECTOR_DISTANCE", "distance", lambda x: float(x)),
                 ("SPACE_GROUP_NUMBER", "spacegroup", lambda x: int(x)),
                 ("UNIT_CELL_CONSTANTS", "unit_cell", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("UNIT_CELL_A-AXIS", "a_axis", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("UNIT_CELL_B-AXIS", "b_axis", lambda x: numpy.array(map(lambda y:float(y), x.split()))),
                 ("UNIT_CELL_C-AXIS", "c_axis", lambda x: numpy.array(map(lambda y:float(y), x.split())))
                 ]
        inp = dict(get_xdsinp_keyword(xdsinp)) # I believe dict() removes duplicated parameters and keeps last.

        for k, at, f in table:
            if k in inp and inp[k].strip() != "":
                setattr(self, at, f(inp[k]))
        if "ORGX" in inp:
            self.origin[0] = float(inp["ORGX"])
        if "ORGY" in inp:
            self.origin[1] = float(inp["ORGY"])
    # set_info_from_xdsinp()

    def xparm_str(self, old_format=False):
        assert not old_format # Currently, only new format is supported!
        xparm_str = """ XPARM.XDS    VERSION March 30, 2013
%6d%14.4f%10.4f%10.6f%10.6f%10.6f
%15.6f%15.6f%15.6f%15.6f
%6d%12.6f%12.6f%12.6f%8.3f%8.3f%8.3f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%10d%10d%10d%12.6f%12.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%15.6f%15.6f%15.6f
%10d%10d%10d%10d%10d
%8.2f%8.2f%8.2f%9.5f%9.5f%9.5f%9.5f%9.5f%9.5f
""" % (self.starting_frame, self.starting_angle, self.osc_range, self.rotation_axis[0], self.rotation_axis[1], self.rotation_axis[2],
       self.wavelength, self.incident_beam[0], self.incident_beam[1], self.incident_beam[2],
       self.spacegroup, self.unit_cell[0], self.unit_cell[1], self.unit_cell[2], self.unit_cell[3], self.unit_cell[4], self.unit_cell[5],
       self.a_axis[0], self.a_axis[1], self.a_axis[2],
       self.b_axis[0], self.b_axis[1], self.b_axis[2],
       self.c_axis[0], self.c_axis[1], self.c_axis[2],
       1, self.nx, self.ny, self.qx, self.qy,
       self.origin[0], self.origin[1], self.distance,
       self.X_axis[0], self.X_axis[1], self.X_axis[2],
       self.Y_axis[0], self.Y_axis[1], self.Y_axis[2],
       self.Z_axis[0], self.Z_axis[1], self.Z_axis[2],
       1, 1, self.nx, 1, self.ny,
       0., 0., 0., 1., 0., 0., 0., 1., 0.,
       )
        return xparm_str
    # xparm_str()

    def crystal_symmetry(self):
        from cctbx import crystal
        return crystal.symmetry(tuple(self.unit_cell),
                                self.spacegroup)
    # crystal_symmetry()

    def update_cell_based_on_axes(self):
        from yamtbx.util.maths import vectors_angle

        a, b, c = map(lambda x: numpy.linalg.norm(x), (self.a_axis, self.b_axis, self.c_axis))
        al = vectors_angle(self.b_axis, self.c_axis)
        be = vectors_angle(self.c_axis, self.a_axis)
        ga = vectors_angle(self.a_axis, self.b_axis)
        al, be, ga = map(lambda x: numpy.rad2deg(x), (al, be, ga))

        self.unit_cell = numpy.array((a,b,c,al,be,ga))
        print
# class XPARM

def get_xparm_from_integrate_lp(lpfile, frame):
    assert 0 < frame

    keys = {"beam direction": "DIRECT BEAM COORDINATES (REC. ANGSTROEM)",
            "beam center": "DETECTOR ORIGIN (PIXELS) AT",
            "distance": "CRYSTAL TO DETECTOR DISTANCE (mm)",
            "rotation axis": "LAB COORDINATES OF ROTATION AXIS",
            "a axis": "COORDINATES OF UNIT CELL A-AXIS",
            "b axis": "COORDINATES OF UNIT CELL B-AXIS",
            "c axis": "COORDINATES OF UNIT CELL C-AXIS",
            "cell": "UNIT CELL PARAMETERS",
            "spacegroup": "SPACE GROUP NUMBER"
            }
    data = {}

    flag_read = False
    for l in open(lpfile):
        if "PROCESSING OF IMAGES" in l:
            flag_read = False
            l = l.strip()
            first, last = map(lambda x:int(x.strip()), l[l.index("PROCESSING OF IMAGES")+len("PROCESSING OF IMAGES"):].split("..."))
            if first <= frame <= last:
                flag_read = True

        if flag_read:
            for key, s in keys.items():
                if s in l:
                    l = l.strip()
                    val = map(lambda x:float(x.strip()), l[l.index(s)+len(s):].split())
                    data[key] = val

    beam = data["beam direction"]
    rotaxis = data["rotation axis"]
    distance = data["distance"][0]
    orgx, orgy = data["beam center"]
    spacegroup = data["spacegroup"][0]
    a, b, c, alpha, beta, gamma = data["cell"]
    aaxis = data["a axis"]
    baxis = data["b axis"]
    caxis = data["c axis"]

    xp = XPARM("XPARM.XDS")

    xp.rotation_axis = numpy.array(rotaxis)
    xp.incident_beam = numpy.array(beam)
    xp.spacegroup = spacegroup
    xp.unit_cell = numpy.array((a, b, c, alpha, beta, gamma))
    xp.a_axis = numpy.array(aaxis)
    xp.b_axis = numpy.array(baxis)
    xp.c_axis = numpy.array(caxis)
    xp.origin = numpy.array((orgx, orgy))
    xp.distance = distance

    return xp.xparm_str()
# get_xparm_from_integrate_lp()

def find_best_ucaxes_transform(xparm1, xparm2):
    """
    Reference: matrix iteration method in whirligig in CrystFEL-0.6.0 src/whirligig.c: gatinator()

    calculates the matrix m, which gives the best result; axes1 and m * axes2
    returns m, axes1, axes2, rms
    """

    a1, b1, c1 = xparm1.a_axis, xparm1.b_axis, xparm1.c_axis
    a2, b2, c2 = xparm2.a_axis, xparm2.b_axis, xparm2.c_axis

    P = numpy.vstack([a2, b2, c2])
    Q = numpy.vstack([a1, b1, c1])

    min_score = numpy.inf
    best_mat = None

    for x in itertools.product(xrange(-1,2), repeat=9):
        #if x != (1,0,0,0,1,0,0,0,1): continue
        m = numpy.array(x).reshape(3,3)
        if numpy.linalg.det(m) != 1: continue
        #print "Trying"
        #print m

        score = numpy.sum((numpy.dot(m, Q) - P)**2)
        #if x == (1,0,0,0,1,0,0,0,1): print score
        if score < min_score:
            min_score = score
            best_mat = m

    return best_mat, Q, P, numpy.sqrt(min_score/3.)
# find_best_ucaxes_transform()


import pickle
import os
import numpy as np
import tensorflow as tf

import api.config as config

def calculate_score(input_processed_array):
    '''
    The function consumes the models and calculate the score.
    '''

    xgb_model_loaded = pickle.load(open(os.path.dirname(os.getcwd()) + str(config.XGBOOST_MODEL_PATH), "rb"))
    xgb_predict = list(xgb_model_loaded.predict(np.array(input_processed_array).reshape((1,-1))))[0]

    keras_model_loaded =  tf.keras.models.load_model(os.path.dirname(os.getcwd()) + str(config.KERAS_MODEL_PATH))
    keras_predict = list(keras_model_loaded.predict(np.array(input_processed_array).reshape((1,-1))))[0]

    return ((xgb_predict + keras_predict) / 2)[0]


from simglucose.simulation.env import T1DSimEnv as _T1DSimEnv
from simglucose.simulation.env import T1DSimEnvExtendedObs as _T1DSimEnvExtendedObs
from simglucose.patient.t1dpatient import T1DPatient
from simglucose.sensor.cgm import CGMSensor
from simglucose.actuator.pump import InsulinPump
from simglucose.simulation.scenario_gen import RandomScenario
from simglucose.simulation.scenario_gen import RandomBalancedScenario
from simglucose.controller.base import Action
from simglucose.analysis.risk import risk_index
import numpy as np
import pandas as pd
import pkg_resources
import gym
from gym import spaces
from gym.utils import seeding
from datetime import datetime
from typing import Dict
from collections import OrderedDict
PATIENT_PARA_FILE = pkg_resources.resource_filename(
    'simglucose', 'params/vpatient_params.csv')

class T1DSimEnv(gym.Env):
    '''
    A wrapper of simglucose.simulation.env.T1DSimEnv to support gym API
    '''
    metadata = {'render.modes': ['human']}

    SENSOR_HARDWARE = 'Dexcom'
    INSULIN_PUMP_HARDWARE = 'Insulet'

    def __init__(self, patient_name=None, reward_fun=None, seed=None, normalize=False, discrete=False, max_basal=None, state=None, noise=True,  sequence=-1, append_time=False, harrison_benedict = False):
        '''
        patient_name must be 'adolescent#001' to 'adolescent#010',
        or 'adult#001' to 'adult#010', or 'child#001' to 'child#010'
        normalize: normalize observations and actions to [-1,1]
        discrete: use discrete insulin dosis. TODO: should pass the ranges, now it is harcoded in set_spaces and step
        state: pass a vector with the seed an hour values to reproduce a particular env
        noise: if False, the real BG instead of the CGM is used in the observations
        sequence: the number of steps to run with zero insulin before a new observation is returned (-1 if every step returns an observation)
        '''
        # have to hard code the patient_name, gym has some interesting
        # error when choosing the patient
        if patient_name is None:
            patient_name = 'adolescent#001'
        self.patient_name = patient_name
        self.reward_fun = reward_fun
        self.noise=noise
        self.max_basal=max_basal
        self.sequence=sequence
        self.harrison_benedict=harrison_benedict
        self.np_random, myseed = seeding.np_random(seed=seed)
        self.saved_state=None
        self.env, _, _,_,_ = self._create_env_from_random_state()
        self._normalize = normalize
        self._discrete=discrete
        self.append_time=append_time
	#Set the spaces
        self._set_spaces()
        self.episode_steps=0
        self.save_callback = None
        self.saved_state=state
    def set_callback(self,callback):
        self.save_callback=callback
    def step(self, action):
        # This gym only controls basal insulin
        if self._normalize:
          _action=self.rescale_action(action)
          if self._discrete:
            _action=action/5
        else:
          _action=action
        act = Action(basal=_action, bolus=0)
        if self.reward_fun is None:
          _obs, reward, done , info =  self.env.step(act)
        else:
          _obs, reward, done , info = self.env.step(act, reward_fun=self.reward_fun)
        if not done and self.sequence>0:
           for i in range(1, self.sequence-1):
               #Remember: we pass unnormalized actions to the simglucose env
               act=Action(basal=0,bolus=0)
               _obs, r, d, info = self.env.step(act,reward_fun=self.reward_fun)
               #Make sure that one done=True is not replaced by a later done=False
               reward += r
               if d==True:
                   done=d
                   break
        if self.append_time:
            if self._normalize:
               obs=self.normalize_obs(np.array([_obs.CGM,self.tomin()]))
            else:
               obs=np.array([_obs.CGM,self.tomin()])
        else:
            if self._normalize:
               obs=self.normalize_obs(_obs.CGM)
            else:
               obs=_obs.CGM

        if self.save_callback is not None:
           go_on=self.save_callback.step(self.episode_steps)
           if not go_on:
              info['Time-limit truncated at callback'] = not done
              done=True
        self.episode_steps +=1
        return (np.array([obs]), reward, done, info)

    def tomin(self):
       return (self.env.time.hour*60 + self.env.time.minute)
    def reset(self):

        self.env, seed2,seed3,seed4, hour = self._create_env_from_random_state()
        if self.save_callback is not None:
           self.save_callback.reset([seed2,seed3,seed4,hour])
        obs, _, _, _ = self.env.reset()
        if self.append_time:
            if self._normalize:
               obs=self.normalize_obs(np.array([obs.CGM,self.tomin()]))
            else:
               obs=np.array([obs.CGM,self.tomin()])
        else:
            if self._normalize:
               obs=self.normalize_obs(obs.CGM)
            else:
               obs=obs.CGM
        if self._normalize:
           obs=self.normalize_obs(obs)
        self.episode_steps =0
        return np.array([obs])

    def seed(self, seed=None):
        self.np_random, seed1 = seeding.np_random(seed=seed)
        self.env, seed2, seed3 = self._create_env_from_random_state()
        return [seed1, seed2, seed3]

    def _create_env_from_random_state(self):
        # Derive a random seed. This gets passed as a uint, but gets
        # checked as an int elsewhere, so we need to keep it below
        # 2**31.
        seed2 = seeding.hash_seed(self.np_random.randint(0, 1000)) % 2**31
        seed3 = seeding.hash_seed(seed2 + 1) % 2**31
        seed4 = seeding.hash_seed(seed3 + 1) % 2**31
        hour = self.np_random.randint(low=0.0, high=24.0)
        if self.saved_state is not None:
           seed2=self.saved_state[0]
           seed3=self.saved_state[1]
           seed4=self.saved_state[2]
           hour=self.saved_state[3]
           print('Using state', seed2, seed3, seed4, hour)
        start_time = datetime(2021, 1, 1, hour, 0, 0)
        patient = T1DPatient.withName(self.patient_name, random_init_bg=True, seed=seed4)
        sensor = CGMSensor.withName(self.SENSOR_HARDWARE, seed=seed2)
        if self.harrison_benedict:
            vpatient_params = pd.read_csv(PATIENT_PARA_FILE)
            self.kind = self.patient_name.split('#')[0]
            self.bw = vpatient_params.query('Name=="{}"'.format(self.patient_name))['BW'].item()
            scenario=RandomBalancedScenario(start_time=start_time, seed=seed3, harrison_benedict=self.harrison_benedict, bw=self.bw, kind=self.kind)
        else:
            scenario = RandomScenario(start_time=start_time, seed=seed3)
        pump = InsulinPump.withName(self.INSULIN_PUMP_HARDWARE)
        env = _T1DSimEnv(patient, sensor, pump, scenario,noise=self.noise)
        return env, seed2, seed3, seed4, hour

    def render(self, mode='human', close=False):
        self.env.render(close=close)
    @property
    def time(self):
        return self.env.time
    @property
    def scenario(self):
        return self.env.scenario
    def _set_spaces(self):
        self.min_action=0
        if self.max_basal is None:
            self.max_action= self.env.pump._params['max_basal']
        else:
            self.max_action= self.max_basal
        self.min_obs=0
        #Limit to apply normalization
        self.max_obs=500
        if self._normalize:
           if self._discrete:
              self.action_space=spaces.Discrete(150)
           else:
              self.action_space=spaces.Box(low=-1, high=1, shape=(1,))
           if self.append_time:
              self.observation_space= spaces.Box(low=-1, high=np.array([1,1]))
           else:
              self.observation_space= spaces.Box(low=-1, high=1, shape=(1,))
        else:
           if self._discrete:
              self.action_space=spaces.Discrete(300)
           else:
              self.action_space=spaces.Box(low=self.min_action, high=self.max_action, shape=(1,))
           if self.append_time:
              self.observation_space= spaces.Box(low=0, high=np.array([np.inf, 1439]))
           else:
              self.observation_space= spaces.Box(low=0, high=np.inf, shape=(1,))

    def normalize_action(self,action):
        return ((2*action - self.max_action -self.min_action)/(self.max_action-self.min_action))

    def rescale_action(self,action):
        return (0.5*(action*(self.max_action-self.min_action) + self.max_action +self.min_action))

    def normalize_obs(self,obs):
        if self.append_time:
           o=(2*obs[0] - self.max_obs -self.min_obs)/(self.max_obs-self.min_obs)
           t=(2*obs[1] - 1439)/1439
           return np.array([o,t])
        else:
           return ((2*obs - self.max_obs -self.min_obs)/(self.max_obs-self.min_obs))

    def rescale_obs(self,obs):
        return (0.5*(obs*(self.max_obs-self.min_obs) + self.max_obs +self.min_obs))


# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Functions related to sampling from a model."""

import jax
import jax.nn
import jax.numpy as jnp
from absl import logging
from flax.training import common_utils
from jax import lax, random

LARGE_NEGATIVE = -1e12

def multinomial(rng, logits):
    """Draws samples from a multinomial distribution.

    Args:
      rng: A jax.random.PRNGKey.
      logits: An array of shape (..., num_categories) containing unnormalized
        log-probabilities.

    Returns:
      An array of shape (...) containing sampled category indices.
    """
    probs = jax.nn.softmax(logits)
    probs = jnp.cumsum(probs, axis=-1)
    a = jax.random.uniform(rng, logits.shape[:-1] + (1,))
    out = jnp.argmin(a > probs, axis=-1)
    return out

def apply_repetition_penalty(
    sequences,
    logits,
    i,
    repetition_penalty,
    repetition_window,
    repetition_penalty_normalize,
):
    """Apply repetition penalty.

    Repetition penalty is introduced in Section 4.1. of the
    [CTRL paper](https://einstein.ai/presentations/ctrl.pdf).

    It involves reducing the logits corresponding to previously generated tokens.
    Given a list of generated tokens g, the next token probability is given by:

      p_{i + 1} propto exp(logit_i / (T * I_i))

    where I_i = repetition_penalty, if i in g,
                1, otherwise.

    Args:
      sequences: An array of shape (batch_size, max_seq_len) with the sequences.
      logits: An array of shape (batch_size, vocab_size) with next token logits.
      i: An array of shape () with the current sequence positions.
      repetition_penalty: A float indicating the repetition penalty.
      repetition_window: An int indicating the window size for repetition penalty.
        We reduce the probability of tokens that have been generated within
        repetition_window tokens prior to the token to be generated.
      repetition_penalty_normalize: A bool indicating whether to normalize
        the logits (log_softmax) before applying the repetition penalty. This
        ensures that all logits have the same sign.

    Returns:
      An array of shape (batch_size, vocab_size) with updated next token logits.
    """
    max_i = i  # We are currently generating a token for position i + 1.
    min_i = i - repetition_window + 1
    batch_size, vocab_size = logits.shape
    positions = jnp.arange(sequences.shape[1])
    positions = jnp.tile(
        positions[jnp.newaxis, :, jnp.newaxis], [batch_size, 1, vocab_size]
    )
    sequences_onehot = jnp.eye(vocab_size)[sequences]
    sequences_onehot = jnp.where(
        (positions >= min_i) * (positions <= max_i),
        sequences_onehot,
        jnp.zeros_like(sequences_onehot),
    )
    # Compute the indicator that a token appeared at least once in the
    # repetition window. Output shape: (batch_size, vocab_size).
    indicator = jnp.max(sequences_onehot, axis=1)
    # Compute a penalty tensor. The logits are divided by the penalty tensor.
    penalty_tensor = jnp.where(
        indicator, jnp.ones_like(logits) * repetition_penalty, jnp.ones_like(logits)
    )
    if repetition_penalty_normalize:
        logits = jax.nn.log_softmax(logits)
    # Dividing a negative logit by the penalty tensor actually increases the
    # resulting probability. Take the inverse for negative logits.
    penalty_tensor = jnp.where(logits > 0, penalty_tensor, 1 / penalty_tensor)

    logits = logits / penalty_tensor
    return logits

def temperature_sample(
    prompt,
    init_cache,
    tokens_to_logits,
    temperature=1.0,
    top_k=None,
    repetition_penalty=1,
    repetition_window=4,
    repetition_penalty_normalize=False,
    max_decode_len=512,
    rng=None,
    eos_token=None,
    pad_token=None,
    masked_tokens=None,
    use_lax_while_loop=True,
):
    """Temperature sampling for sequence models.

    Args:
      prompt: An array of shape (batch_size, prompt_length) containing the input
        prompt (the model consumes these tokens and starts generation after). For
        generic sampling, the prompt must be a single BOS token.
      init_cache: A flax.nn.attention.Cache object.
      tokens_to_logits: A fast autoregressive decoder function taking single token
        slices and the cache and returning next-token logits and updated cache.
      temperature: A float with the sampling temperature factor. As it approaches
        zero this sampling procedure becomes equivalent to greedy sampling.
      top_k: An int with the number of high probability tokens to keep for
        sampling. If None, keep all.
      repetition_penalty: A float with the repetition penalty. Values smaller
        (greater) than 1 encourage (discourage) repetition; 1 disables penalties.
      repetition_window: An int indicating the window size for repetition penalty.
        We reduce the probability of tokens that have been generated within
        repetition_window tokens prior to the token to be generated.
      repetition_penalty_normalize: A bool indicating whether to normalize the
        logits (log_softmax) before applying the repetition penalty. This ensures
        that all logits have the same sign.
      max_decode_len: An int indicating the maximum sequence length.
      rng: A jax.random.PRNGKey.
      eos_token: An int token id. If not None, we stop decoding a sequence after
        the first instance.
      pad_token: An int token used to pad sequences after the eos token. If none,
        we set pad_token to eos_token.
      masked_tokens: A list of int token id. If not None, we mask these token
        before sampling.
      use_lax_while_loop: A bool; whether to use a lax while loop or python loop.

    Returns:
      An array of shape (batch_size, max_decode_len) containing sequences. If
        variable-length, the sequences are right-padded with the EOS token.
    """
    batch_size = prompt.shape[0]
    eos_token = eos_token if eos_token is not None else -1
    if pad_token is None:
        logging.warn("Pad token is not provided. Using the EOS token.")
    pad_token = pad_token if pad_token is not None else eos_token
    end_marker = jnp.array(eos_token)
    temperature = jnp.array(temperature)

    # Initialize sampling loop state.
    rng0 = rng if rng is not None else random.PRNGKey(0)
    i0 = jnp.array(0)
    ended0 = jnp.zeros((batch_size, 1)).astype(jnp.bool_)

    # Initialize sequences with the prompt followed by the out_of_prompt_marker to
    # indicate when we can start generation.
    out_of_prompt_marker = jnp.array(-2)
    pad = (
        jnp.ones((batch_size, max_decode_len - prompt.shape[1]), dtype=jnp.int32)
        * out_of_prompt_marker
    )
    sequences0 = jnp.concatenate([prompt, pad], axis=1)
    token0 = sequences0[:, 0:1]

    # Sampling loop state is stored in a simple tuple.
    tokens_to_logits_state = None

    sampling_loop_init_state = (
        i0,
        sequences0,
        init_cache,
        token0,
        ended0,
        rng0,
        tokens_to_logits_state,
    )

    def sampling_loop_cond_fn(state):
        """Sampling loop termination condition."""
        (i, _, _, _, ended, _, _) = state
        # Have we reached max decoding length?
        not_at_end = i <= max_decode_len
        # Have all sampled sequences reached an end marker?
        all_sequences_ended = jnp.all(ended)
        return not_at_end & (~all_sequences_ended)

    def sampling_loop_body_fn(state):
        """Sampling loop state update."""
        i, sequences, cache, cur_token, ended, rng, tokens_to_logits_state = state

        # Split RNG for sampling.
        rng1, rng2 = random.split(rng)

        # Call fast-decoder model on current tokens to get raw next-position logits.
        logits, new_cache, new_tokens_to_logits_state = tokens_to_logits(
            cur_token, cache, internal_state=tokens_to_logits_state
        )
        logits = logits / temperature

        # Mask out the BOS token.
        if masked_tokens is not None:
            mask = common_utils.onehot(
                jnp.array(masked_tokens),
                num_classes=logits.shape[-1],
                on_value=LARGE_NEGATIVE,
            )
            mask = jnp.sum(mask, axis=0)[None, :]  # Combine multiple masks together
            logits = logits + mask

        # Apply the repetition penalty.
        if repetition_penalty != 1:
            logits = apply_repetition_penalty(
                sequences,
                logits,
                i,
                repetition_penalty=repetition_penalty,
                repetition_window=repetition_window,
                repetition_penalty_normalize=repetition_penalty_normalize,
            )

        # Mask out everything but the top-k entries.
        if top_k is not None:
            # Compute top_k_index and top_k_threshold with shapes (batch_size, 1).
            top_k_index = jnp.argsort(logits, axis=-1)[:, ::-1][:, top_k - 1 : top_k]
            top_k_threshold = jnp.take_along_axis(logits, top_k_index, axis=-1)
            logits = jnp.where(
                logits < top_k_threshold, jnp.full_like(logits, LARGE_NEGATIVE), logits
            )
        # Sample next token from logits.
        sample = multinomial(rng1, logits)
        next_token = sample.astype(jnp.int32)
        # Only use sampled tokens if we have past the out_of_prompt_marker.
        out_of_prompt = sequences[:, i + 1] == out_of_prompt_marker
        next_token = next_token * out_of_prompt + sequences[:, i + 1] * ~out_of_prompt
        # If end-marker reached for batch item, only emit padding tokens.
        next_token = next_token[:, None]
        next_token_or_endpad = jnp.where(
            ended, jnp.full_like(next_token, pad_token), next_token
        )
        ended |= next_token_or_endpad == end_marker
        # Add current sampled tokens to recorded sequences.
        new_sequences = lax.dynamic_update_slice(
            sequences, next_token_or_endpad, (0, i + 1)
        )
        return (
            i + 1,
            new_sequences,
            new_cache,
            next_token_or_endpad,
            ended,
            rng2,
            new_tokens_to_logits_state,
        )

    # Run sampling loop and collect final state.
    if use_lax_while_loop:
        final_state = lax.while_loop(
            sampling_loop_cond_fn, sampling_loop_body_fn, sampling_loop_init_state
        )
    else:
        final_state = sampling_loop_init_state
        while sampling_loop_cond_fn(final_state):
            final_state = sampling_loop_body_fn(final_state)

    final_sequences = final_state[1]
    # If generation ended early (all sampled sequences reached an end marker)
    # replace all remaining out_of_prompt_marker instances with the pad_token.
    final_sequences = jnp.where(
        final_sequences == out_of_prompt_marker,
        jnp.full_like(final_sequences, pad_token),
        final_sequences,
    )
    return final_sequences


import numpy as np
from pyannote.audio.keras_utils import load_model
from pyannote.audio.signal import Binarize, Peak
from pyannote.audio.features import Precomputed
import my_cluster
from pyannote.core import Annotation
from pyannote.audio.embedding.utils import l2_normalize
from pyannote.database import get_annotated

class SpeakerDiarizationPre(object):
    '''Speaker diarization with affinity propagation'''

    def __init__(self, feature_extraction, sad__pre, scd__pre, emb__pre,
                 sad__onset=0.7, sad__offset=0.7, sad__dimension=1,
                 scd__alpha=0.5, scd__min_duration=1., scd__dimension=1,
                 emb__internal=False,
                 cls__damping=0.8, cls__preference=-20,
                 cls__metric='cosine'):

        super(SpeakerDiarizationPre, self).__init__()

        self.feature_extraction = feature_extraction

        # speech activity detection hyper-parameters
        self.sad__onset = sad__onset
        self.sad__offset = sad__offset
        self.sad__dimension = sad__dimension

        # speaker change detection hyper-parameters
        self.scd__alpha = scd__alpha
        self.scd__min_duration = scd__min_duration
        self.scd__dimension = scd__dimension

        # embedding hyper-parameters
        self.emb__internal = emb__internal

        # clustering hyper-parameters
        self.cls__damping = cls__damping
        self.cls__preference = cls__preference
        self.cls__metric = cls__metric

        step = self.feature_extraction.sliding_window().step

        # initialize speech activity detection module
        self.sad_ = Precomputed(sad__pre)
        self.sad_binarize_ = Binarize(onset=self.sad__onset,
                                      offset=self.sad__offset)

        # initialize speaker change detection module
        self.scd_ = Precomputed(scd__pre)
        self.scd_peak_ = Peak(alpha=self.scd__alpha,
                              min_duration=self.scd__min_duration,
                              percentile=False)

        # initialize speech turn embedding module
        self.emb_ = Precomputed(emb__pre)

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringAP(metric=self.cls__metric,
                                        damping=self.cls__damping,
                                        preference=self.cls__preference)

    def __call__(self, current_file, annotated=False):

        # speech activity detection
        soft_sad = self.sad_(current_file)
        hard_sad = self.sad_binarize_.apply(
            soft_sad, dimension=self.sad__dimension)

        # speaker change detection
        soft_scd = self.scd_(current_file)
        hard_scd = self.scd_peak_.apply(
            soft_scd, dimension=self.scd__dimension)

        # speech turns
        speech_turns = hard_scd.crop(hard_sad)

        if annotated:
            speech_turns = speech_turns.crop(
                get_annotated(current_file))

        # remove small speech turns
        emb = self.emb_(current_file)
        speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]

        # speech turns embedding
        to_stack = [
            np.sum(emb.crop(speech_turn, mode='loose'), axis=0)
            for speech_turn in speech_turns]
        if len(to_stack) < 1:
            return None
        fX = l2_normalize(np.vstack(to_stack))

        # speech turn clustering
        cluster_labels = self.cls_.apply(fX)

        # build hypothesis from clustering results
        hypothesis = Annotation(uri=current_file['uri'])
        for speech_turn, label in zip(speech_turns, cluster_labels):
            hypothesis[speech_turn] = label
        return hypothesis

class SpeakerDiarizationOracleSegAP(object):
    '''Speaker diarization with oracle segmentation and affinity propagation'''

    def __init__(self, feature_extraction, emb__pre,
                 emb__internal=False,
                 cls__damping=0.8, cls__preference=-20,
                 cls__metric='cosine'):

        super(SpeakerDiarizationOracleSegAP, self).__init__()

        self.feature_extraction = feature_extraction

        # embedding hyper-parameters
        self.emb__internal = emb__internal

        # clustering hyper-parameters
        self.cls__damping = cls__damping
        self.cls__preference = cls__preference
        self.cls__metric = cls__metric

        step = self.feature_extraction.sliding_window().step

        # initialize speech turn embedding module
        self.emb_ = Precomputed(emb__pre)

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringAP(metric=self.cls__metric,
                                        damping=self.cls__damping,
                                        preference=self.cls__preference)

    def __call__(self, current_file, annotated=False):

        # speech turns
        speech_turns = current_file['annotation'].get_timeline()

        if annotated:
            speech_turns = speech_turns.crop(
                get_annotated(current_file))

        # remove small speech turns
        emb = self.emb_(current_file)
        speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]

        # speech turns embedding
        to_stack = [
            np.sum(emb.crop(speech_turn, mode='loose'), axis=0)
            for speech_turn in speech_turns]
        if len(to_stack) < 1:
            return None
        fX = l2_normalize(np.vstack(to_stack))

        # speech turn clustering
        cluster_labels = self.cls_.apply(fX)

        # build hypothesis from clustering results
        hypothesis = Annotation(uri=current_file['uri'])
        for speech_turn, label in zip(speech_turns, cluster_labels):
            hypothesis[speech_turn] = label
        return hypothesis

class SpeakerDiarizationHACPre(object):
    '''Speaker diarization with hierarchical agglomerative clustering'''

    def __init__(self, feature_extraction, sad__pre, scd__pre, emb__pre,
                 sad__onset=0.7, sad__offset=0.7, sad__dimension=1,
                 scd__alpha=0.5, scd__min_duration=1., scd__dimension=1,
                 emb__internal=False,
                 cls__method='average', cls__threshold=5,
                 cls__metric='cosine'):

        super(SpeakerDiarizationHACPre, self).__init__()

        self.feature_extraction = feature_extraction

        # speech activity detection hyper-parameters
        self.sad__onset = sad__onset
        self.sad__offset = sad__offset
        self.sad__dimension = sad__dimension

        # speaker change detection hyper-parameters
        self.scd__alpha = scd__alpha
        self.scd__min_duration = scd__min_duration
        self.scd__dimension = scd__dimension

        # embedding hyper-parameters
        self.emb__internal = emb__internal

        # clustering hyper-parameters
        self.cls__method = cls__method
        self.cls__threshold = cls__threshold
        self.cls__metric = cls__metric

        step = self.feature_extraction.sliding_window().step

        # initialize speech activity detection module
        self.sad_ = Precomputed(sad__pre)
        self.sad_binarize_ = Binarize(onset=self.sad__onset,
                                      offset=self.sad__offset)

        # initialize speaker change detection module
        self.scd_ = Precomputed(scd__pre)
        self.scd_peak_ = Peak(alpha=self.scd__alpha,
                              min_duration=self.scd__min_duration,
                              percentile=False)

        # initialize speech turn embedding module
        self.emb_ = Precomputed(emb__pre)

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringHAC(metric=self.cls__metric,
                                        method=self.cls__method,
                                        threshold=self.cls__threshold)

    def __call__(self, current_file, annotated=False):

        # speech activity detection
        soft_sad = self.sad_(current_file)
        hard_sad = self.sad_binarize_.apply(
            soft_sad, dimension=self.sad__dimension)

        # speaker change detection
        soft_scd = self.scd_(current_file)
        hard_scd = self.scd_peak_.apply(
            soft_scd, dimension=self.scd__dimension)

        # speech turns
        speech_turns = hard_scd.crop(hard_sad)

        if annotated:
            speech_turns = speech_turns.crop(
                get_annotated(current_file))

        # remove small speech turns
        emb = self.emb_(current_file)
        speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]

        # speech turns embedding
        to_stack = [
            np.sum(emb.crop(speech_turn, mode='loose'), axis=0)
            for speech_turn in speech_turns]
        if len(to_stack) < 1:
            return None
        fX = l2_normalize(np.vstack(to_stack))

        # speech turn clustering
        cluster_labels = self.cls_.apply(fX)

        # build hypothesis from clustering results
        hypothesis = Annotation(uri=current_file['uri'])
        for speech_turn, label in zip(speech_turns, cluster_labels):
            hypothesis[speech_turn] = label
        return hypothesis

class SpeakerDiarizationPreStages(object):

    def __init__(self, feature_extraction, sad__pre, scd__pre, emb__pre,
                 sad__onset=0.7, sad__offset=0.7, sad__dimension=1,
                 scd__alpha=0.5, scd__min_duration=1., scd__dimension=1,
                 emb__internal=False,
                 cls__damping=0.8, cls__preference=-20,
                 cls__metric='cosine'):

        super(SpeakerDiarizationPreStages, self).__init__()

        self.feature_extraction = feature_extraction

        # speech activity detection hyper-parameters
        self.sad__onset = sad__onset
        self.sad__offset = sad__offset
        self.sad__dimension = sad__dimension

        # speaker change detection hyper-parameters
        self.scd__alpha = scd__alpha
        self.scd__min_duration = scd__min_duration
        self.scd__dimension = scd__dimension

        # embedding hyper-parameters
        self.emb__internal = emb__internal

        # clustering hyper-parameters
        self.cls__damping = cls__damping
        self.cls__preference = cls__preference
        self.cls__metric = cls__metric

        step = self.feature_extraction.sliding_window().step

        # initialize speech activity detection module
        self.sad_ = Precomputed(sad__pre)
        self.sad_binarize_ = Binarize(onset=self.sad__onset,
                                      offset=self.sad__offset)

        # initialize speaker change detection module
        self.scd_ = Precomputed(scd__pre)
        self.scd_peak_ = Peak(alpha=self.scd__alpha,
                              min_duration=self.scd__min_duration,
                              percentile=False)

        # initialize speech turn embedding module
        self.emb_ = Precomputed(emb__pre)

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringAP(metric=self.cls__metric,
                                        damping=self.cls__damping,
                                        preference=self.cls__preference)

    def __call__(self, current_file, annotated=False):

        # speech activity detection
        soft_sad = self.sad_(current_file)
        hard_sad = self.sad_binarize_.apply(
            soft_sad, dimension=self.sad__dimension)

        sad_output = hard_sad.to_annotation()

        # speaker change detection
        soft_scd = self.scd_(current_file)
        hard_scd = self.scd_peak_.apply(
            soft_scd, dimension=self.scd__dimension)

        # speech turns
        speech_turns = hard_scd.crop(hard_sad)

        scd_output = speech_turns.to_annotation()

        if annotated:
            speech_turns = speech_turns.crop(
                get_annotated(current_file))

        # remove small speech turns
        emb = self.emb_(current_file)
        speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]

        # speech turns embedding
        to_stack = [
            np.sum(emb.crop(speech_turn, mode='loose'), axis=0)
            for speech_turn in speech_turns]
        if len(to_stack) < 1:
            return None
        fX = l2_normalize(np.vstack(to_stack))

        # speech turn clustering
        cluster_labels = self.cls_.apply(fX)

        # build hypothesis from clustering results
        hypothesis = Annotation(uri=current_file['uri'])
        for speech_turn, label in zip(speech_turns, cluster_labels):
            hypothesis[speech_turn] = label
        return hypothesis, sad_output, scd_output

class SpeakerDiarizationWeighted(object):

    def __init__(self, feature_extraction, sad__pre, scd__pre, weight__pre, emb__pre,
                 sad__onset=0.7, sad__offset=0.7, sad__dimension=1,
                 scd__alpha=0.5, scd__min_duration=1., scd__dimension=1,
                 emb__internal=False,
                 cls__damping=0.8, cls__preference=-20,
                 cls__metric='cosine'):

        super(SpeakerDiarizationWeighted, self).__init__()

        self.feature_extraction = feature_extraction

        # speech activity detection hyper-parameters
        self.sad__onset = sad__onset
        self.sad__offset = sad__offset
        self.sad__dimension = sad__dimension

        # speaker change detection hyper-parameters
        self.scd__alpha = scd__alpha
        self.scd__min_duration = scd__min_duration
        self.scd__dimension = scd__dimension

        # embedding hyper-parameters
        self.emb__internal = emb__internal

        # clustering hyper-parameters
        self.cls__damping = cls__damping
        self.cls__preference = cls__preference
        self.cls__metric = cls__metric

        step = self.feature_extraction.sliding_window().step

        # initialize speech activity detection module
        self.sad_ = Precomputed(sad__pre)
        self.sad_binarize_ = Binarize(onset=self.sad__onset,
                                      offset=self.sad__offset)

        # initialize speaker change detection module
        self.scd_ = Precomputed(scd__pre)
        self.scd_peak_ = Peak(alpha=self.scd__alpha,
                              min_duration=self.scd__min_duration,
                              percentile=False)

        # initialize weights
        self.weight_ = Precomputed(weight__pre)

        # initialize speech turn embedding module
        self.emb_ = Precomputed(emb__pre)

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringAP(metric=self.cls__metric,
                                        damping=self.cls__damping,
                                        preference=self.cls__preference)

    def __call__(self, current_file, annotated=False):

        # speech activity detection
        soft_sad = self.sad_(current_file)
        hard_sad = self.sad_binarize_.apply(
            soft_sad, dimension=self.sad__dimension)

        # speaker change detection
        soft_scd = self.scd_(current_file)
        hard_scd = self.scd_peak_.apply(
            soft_scd, dimension=self.scd__dimension)

        # speech turns
        speech_turns = hard_scd.crop(hard_sad)

        if annotated:
            speech_turns = speech_turns.crop(
                get_annotated(current_file))

        # remove small speech turns
        emb = self.emb_(current_file)
        speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]

        # weights
        weight = self.weight_(current_file)

        # speech turns embedding
        to_stack = [
            np.mean(emb.crop(speech_turn, mode='loose')*(1-weight.crop(speech_turn, mode='loose')), axis=0)
            for speech_turn in speech_turns]
        if len(to_stack) < 1:
            return None
        fX = l2_normalize(np.vstack(to_stack))

        # speech turn clustering
        cluster_labels = self.cls_.apply(fX)

        # build hypothesis from clustering results
        hypothesis = Annotation(uri=current_file['uri'])
        for speech_turn, label in zip(speech_turns, cluster_labels):
            hypothesis[speech_turn] = label
        return hypothesis

class SpeakerDiarizationOnSceneHAC(object):

    def __init__(self, emb__pre,
                 cls__method='average', cls__threshold=5,
                 cls__metric='cosine'):

        super(SpeakerDiarizationOnSceneHAC, self).__init__()

        # clustering hyper-parameters
        self.cls__method = cls__method
        self.cls__threshold = cls__threshold
        self.cls__metric = cls__metric

        # initialize speech turn embedding module
        self.emb_ = Precomputed(emb__pre)

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringHAC(metric=self.cls__metric,
                                        method=self.cls__method,
                                        threshold=self.cls__threshold)

    def __call__(self, current_file):

        # speech turns
        hypothesis = Annotation(uri=current_file['uri'])
        sencences = current_file['speech_timeline']
        scenes = current_file['scenes']
        # remove small speech turns
        emb = self.emb_(current_file)
        #speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]
        for scene in scenes:
            speech_turns = sencences.crop(scene)
            if len(speech_turns) == 0:
                continue
            if len(speech_turns) == 1:
                hypothesis[speech_turns[0]] = 1
                continue
            # speech turns embedding
            to_stack = [
                np.sum(emb.crop(speech_turn, mode='loose'), axis=0)
                for speech_turn in speech_turns]
            fX = l2_normalize(np.vstack(to_stack))

            # speech turn clustering
            cluster_labels = self.cls_.apply(fX)

            # build hypothesis from clustering results

            for speech_turn, label in zip(speech_turns, cluster_labels):
                hypothesis[speech_turn] = label
        return hypothesis

class SpeakerDiarizationOnEnrollHAC(object):

    def __init__(self,
                 cls__method='average', cls__threshold=5,
                 cls__metric='cosine'):

        super(SpeakerDiarizationOnEnrollHAC, self).__init__()

        # clustering hyper-parameters
        self.cls__method = cls__method
        self.cls__threshold = cls__threshold
        self.cls__metric = cls__metric

        # initialize clustering module
        self.cls_ = my_cluster.ClusteringHAC(metric=self.cls__metric,
                                        method=self.cls__method,
                                        threshold=self.cls__threshold)

    def __call__(self, embedding, speech_turns):
        hypothesis = Annotation()
        #speech_turns = [speech_turn for speech_turn in speech_turns if len(emb.crop(speech_turn, mode='loose')) > 0]

        if len(speech_turns) == 0:
            return hypothesis
        if len(speech_turns) == 1:
            hypothesis[speech_turns[0]] = 1
            return hypothesis
        # speech turns embedding
        to_stack = [
            np.sum(embedding.crop(speech_turn, mode='loose'), axis=0)
            for speech_turn in speech_turns]
        fX = l2_normalize(np.vstack(to_stack))

        # speech turn clustering
        cluster_labels = self.cls_.apply(fX)

        # build hypothesis from clustering results

        for speech_turn, label in zip(speech_turns, cluster_labels):
            hypothesis[speech_turn] = label
        return hypothesis


#!/usr/bin/env python3
import os
import sys
import collections
import argparse
import tables
import itertools
import numpy as np
import pandas as pd
import scipy.stats as stats
import scipy.sparse as sp_sparse

from multiprocessing import Pool
from collections import defaultdict
from scipy import sparse, io
from scipy.sparse import csr_matrix
from multiprocessing import Pool

from _preprocessing import *
from _util import *

np.random.seed(0)

print("Start Analyzing.", file = sys.stderr)

def get_matrix_from_h5(filename, genome):
    with tables.open_file(filename, 'r') as f:
        try:
            dsets = {}
            for node in f.walk_nodes('/' + genome, 'Array'):
                dsets[node.name] = node.read()
            matrix = sp_sparse.csc_matrix((dsets['data'],
                                           dsets['indices'],
                                           dsets['indptr']),
                                           shape=dsets['shape'])
            return GeneBCMatrix(dsets['genes'],
                                dsets['gene_names'],
                                dsets['barcodes'],
                                matrix)
        except tables.NoSuchNodeError:
            raise Exception("Genome %s does not exist in this file." % genome)
        except KeyError:
            raise Exception("File is missing one or more required datasets.")

def save_matrix_to_h5(gbm, filename, genome):
    flt = tables.Filters(complevel=1)
    with tables.open_file(filename, 'w', filters=flt) as f:
        try:
            group = f.create_group(f.root, genome)
            f.create_carray(group, 'genes', obj=gbm.gene_ids)
            f.create_carray(group, 'gene_names', obj=gbm.gene_names)
            f.create_carray(group, 'barcodes', obj=gbm.barcodes)
            f.create_carray(group, 'data', obj=gbm.matrix.data)
            f.create_carray(group, 'indices', obj=gbm.matrix.indices)
            f.create_carray(group, 'indptr', obj=gbm.matrix.indptr)
            f.create_carray(group, 'shape', obj=gbm.matrix.shape)
        except:
            raise Exception("Failed to write H5 file.")

def subsample_matrix(gbm, barcode_indices):
    return GeneBCMatrix(gbm.gene_ids,
                        gbm.gene_names,
                        gbm.barcodes[barcode_indices],
                        gbm.matrix[:, barcode_indices])

def get_expression(gbm, gene_name):
    gene_indices = np.where(gbm.gene_names == gene_name)[0]
    if len(gene_indices) == 0:
        raise Exception("%s was not found in list of gene names." % gene_name)
    return gbm.matrix[gene_indices[0], :].toarray().squeeze()

#Create an N-dimentional nested dictionary
def nested_dict(n, type):
    if n == 1:
        return defaultdict(type)
    else:
        return defaultdict(lambda: nested_dict(n-1, type))

#Function to define the turning point of a cumulative curve by using the minimum derivative method
def turn_point(sgRNA_name, df):
    sgRNA_count  = df.T.filter(items=[sgRNA_name]).sum(axis=1).sort_values(ascending=False)
    sgRNA_cumsum = sgRNA_count.cumsum()

    #get the total cell number of this sgRNA
    cell_num = np.argwhere(sgRNA_count > 0).size

    #calculate the turning point by using the max derivative
    turning_point = sgRNA_cumsum.loc[((sgRNA_cumsum.diff()) / sgRNA_count.sum() > (1/cell_num))].shape

    return(sgRNA_count.iloc[turning_point])

#load sgRNA-barcode data
def load_data(input_file):
    input_fh  = open(input_file, 'r')
    #define parameters
    cell_bc_list   = []
    num_sgRNA_list = np.array([])
    sgRNAs         = []
    umis           = []
    #initiate a 2D dictionary
    data_dict = nested_dict(2, list)
    for line in input_fh:
        cell_bc    = line.strip().split('\t')[0]
        num_sgRNA  = line.strip().split('\t')[2]
        sgRNA_list = line.strip().split('\t')[3].split(';')
        umi_list   = line.strip().split('\t')[5].split(';')
        for i in zip(sgRNA_list, umi_list):
            data_dict[cell_bc][i[0]] = i[1]
    #read the 2D dictionary into the pandas DataFrame
    df = pd.DataFrame(data_dict).fillna(0).astype(int)
    return df

#filter the sgRNA UMI count based on the cutoff values
def filter_umi (df):
    #calculate the cutoff value for each sgRNA in the dataset
    sgRNA_cutoff = [[turn_point(i, df)] for i in list(df.index)]
    for i in range(0, len(sgRNA_cutoff)):
        df.iloc[i, :].loc[df.iloc[i, :] <= sgRNA_cutoff[i]] = 0
    return df

#fine the indecies of cells which contain the input sgRNA
def find_sgrna_cells(sgRNA_list, df, gbm):
    cell_bc = df.loc[:, (df.loc[sgRNA_list].sum() > 0)].T.index
    cell_index = []
    for i in cell_bc:
        current_idx = np.argwhere(gbm.barcodes == i.encode('utf-8'))
        if  current_idx.size > 0:
            cell_index.append(current_idx.item())
    return [x for x in set(cell_index)]

def find_non_zero_cells(df, gbm):
    cell_bc = df.T.loc[df.sum() != 0].index
    cell_index = []
    for i in cell_bc:
        current_idx = np.argwhere(gbm.barcodes == i.encode('utf-8'))
        if  current_idx.size > 0:
            cell_index.append(current_idx.item())
    return [x for x in set(cell_index)]

def hypergeo_test(non_zero_array, sgrna_idx, i):
    #find indecies of cells in which expression of given gene is
    #equal or less than the median of this gene in the whole population
    median_cell_idx  = np.argwhere(non_zero_array <= np.median(non_zero_array))

    #find the same cells subset in the cells with a given sgRNA
    overlap_cell_idx = np.intersect1d(median_cell_idx, sgrna_idx)

    #calculate the median fold change
    other_idx = np.setxor1d(sgrna_idx, range(len(non_zero_array)))

    fc = (np.mean(non_zero_array[sgrna_idx]) + 0.01) / (np.mean(non_zero_array[other_idx]) + 0.01)

    #perform hypergeometric test, get the upper tail
    k = len(overlap_cell_idx)
    M = len(non_zero_array)
    n = len(median_cell_idx)
    N = len(sgrna_idx)
    try:
        pval_up = stats.hypergeom.logcdf(k, M, n, N).item()
    except:
        pval_up = float('nan')

    try:
        pval_down = stats.hypergeom.logsf(k, M, n, N).item()
    except:
        pval_down = float('nan')

    return pval_down, pval_up, fc

#num_sgrna_cell, pval_list = perform_DE(sgrna_idx, input_array, idx, num_processing)
def perform_DE(sgrna_idx, input_array, idx, num_processes, pval_list_down, pval_list_up, fc_list):
    nonzero_pval_list_up = []
    nonzero_pval_list_down = []
    nonzero_fc_list = []
    with Pool(processes=num_processes) as p:
        for pval_down, pval_up, fc in p.starmap(hypergeo_test, zip(
                input_array,
                itertools.repeat(sgrna_idx),
                idx)
        ):
            nonzero_pval_list_down.append(pval_down)
            nonzero_pval_list_up.append(pval_up)
            nonzero_fc_list.append(fc)
    for i in idx:
        pval_list_up[i] = nonzero_pval_list_up.pop(0)
        pval_list_down[i] = nonzero_pval_list_down.pop(0)
        fc_list[i] = nonzero_fc_list.pop(0)
    return len(sgrna_idx), pval_list_up, pval_list_down, fc_list

def FDR(x):
    """
    Assumes a list or numpy array x which contains p-values for multiple tests
    Copied from p.adjust function from R
    """
    o = [i[0] for i in sorted(enumerate(x), key=lambda v:v[1],reverse=True)]
    ro = [i[0] for i in sorted(enumerate(o), key=lambda v:v[1])]
    q = sum([1.0/i for i in range(1,len(x)+1)])
    l = [q*len(x)/i*x[j] for i,j in zip(reversed(range(1,len(x)+1)),o)]
    l = [l[k] if l[k] < 1.0 else 1.0 for k in ro]
    return np.asarray(l)

def create_combo(sgrna_list):
    combo_list = []
    for i in range(1, 11):
        my_array = list(itertools.combinations(enumerate(sgrna_list), i))
        for j in range(len(my_array)):
            idx = list(zip(*my_array[j]))[0]
            combo_list.append(idx)
    #return a list of tuple, which idicate the indexing of each combination
    return combo_list;

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-f', '--hdf5', dest='input_hdf5', required=True,
        type=str,
        help='specify the hdf5 file (output from 10X pipeline).'
    )
    parser.add_argument(
        '-s', '--sgrna', dest='input_sgrna', required=True,
        type=str,
        help='specify the sgrna summary file.'
    )
    parser.add_argument(
        '-g', '--genome', dest='input_genome', required=True,
        type=str,
        help='specify the genome of the hdf5 file.'
    )
    parser.add_argument(
        '-l', '--sgrna_list', dest='sgrna_list', required=True,
        type=str,
        help='specify the sgRNAs need to be tested.'
    )
    parser.add_argument(
        '-o', '--output_dir', dest='output_dir', required=False,
        default = '.',
        help='specify an output directory, default is current dir.'
    )
    parser.add_argument(
        '-t', '--threads', dest = 'threads', required=False,
        type=int,
        default=1,
        help='set number of barcode comparison threads. \
                        The default is 1'
    )
    parser.add_argument(
        '-n', '--norm', dest = 'norm_method', required=False,
        type=str,
        default='cpm',
        help='choose normalization methods: CPM only or \
                        normalize to the Meta-cell.'
    )

    args = parser.parse_args()

    num_processing = args.threads
    norm = args.norm_method

    #check the normalization method
    if (norm != 'cpm') and (norm != 'metacell'):
        print("Incorrect normalization method. Has to be either 'cpm' or 'metacell'.", file = sys.stderr)
        sys.exit(0)

    #read the 10X hdf5 file
    print("Reading HDF5 File...", file = sys.stderr)
    filtered_matrix_h5 = args.input_hdf5
    genome = args.input_genome
    gene_bc_matrix = get_matrix_from_h5(filtered_matrix_h5, genome)

    #load the sgRNA file
    print("Loading sgRNA Annotations...", file = sys.stderr)
    sgrna_file = args.input_sgrna
    sgrna_df_adj = pd.read_pickle(sgrna_file)

    #check the positive control sgRNA in FADS1 region, which has 287 cells
    #perform hypergeometric test for every single gene in the dataframe
    sgrna_dict  = {}
    sgrnas_file = args.sgrna_list
    with open(sgrnas_file) as f:
        for line in f:
            region_id, sgrna_string = line.strip().split("\t")
            sgrnas = sgrna_string.split(";")
            sgrna_dict.update({region_id : sgrnas})

    output_dir = args.output_dir

    #get rid of the cells without any detectable sgRNA
    print("Removing Cells without sgRNA...", file = sys.stderr)
    non_zero_cell_idx = find_non_zero_cells (sgrna_df_adj, gene_bc_matrix)
    nonzero_gbm = subsample_matrix(gene_bc_matrix, non_zero_cell_idx)

    '''
    Read parameters for figure plotting
    '''
    print("Loading Plotting Parameters...", file = sys.stderr)
    cmap = matplotlib.cm.get_cmap('Set2')

    #create cumsum list of the chromsomal position
    length_list = [
        0, 248956422,491149951,689445510,879660065,1061198324,
        1232004303,1391350276,1536488912,1674883629,1808681051,
        1943767673,2077042982,2191407310,2298451028,2400442217,
        2490780562,2574038003,2654411288,2713028904,2777473071,
        2824183054,2875001522
    ]

    chr_order = [
        'chr1', 'chr2', 'chr3', 'chr4', 'chr5', 'chr6', 'chr7',
        'chr8', 'chr9', 'chr10', 'chr11', 'chr12', 'chr13', 'chr14',
        'chr15', 'chr16', 'chr17', 'chr18', 'chr19', 'chr20',
        'chr21', 'chr22', 'chrX'
    ]

    #read the gene annotation file for plotting, which re-ordered the genes based on their
    #relative positions on the genome
    annot_df = pd.read_csv(
        './generate_annotations/plot_annotation.txt',
        header = None,
        sep='\t',
        names = ["idx", "gene_names", "chromosome", "pos", "strand", "color_idx", "chr_idx"]
    )

    #Normalization.
    [g, c] = nonzero_gbm.matrix.shape
    cpm_matrix = np.zeros((g, c))
    uniq_id = set()

    for x in nonzero_gbm.barcodes:
        uniq_id.add(x.decode()[-1])

    if (norm == 'cpm'):
        cpm_matrix = np.array(nonzero_gbm.matrix / nonzero_gbm.matrix.sum(axis = 0) * 1e6)
    elif (norm == 'metacell'):
        for lib in sorted(uniq_id):
            print("Normalizing Batch " + lib, file = sys.stderr)
            index = [i for i,e in enumerate(nonzero_gbm.barcodes) if e.decode()[-1] == lib]
            one_cell_cpm = np.array((nonzero_gbm.matrix[:,index].sum(axis = 1) + 1)
                                    / np.sum(nonzero_gbm.matrix[:,index]) * 1e6).flatten()
            for i in index:
                cell_cpm = nonzero_gbm.matrix[:,i].toarray().flatten()\
                           / np.sum(nonzero_gbm.matrix[:,i]) * 1e6
                cpm_matrix[:,i] = cell_cpm / one_cell_cpm

    #prepare the input array
    nonzero_idx = np.where(np.sum(nonzero_gbm.matrix > 0, axis = 1) > 1)[0]
    idx = list(set(nonzero_idx) & set(annot_df.idx))

    #create input ndarray
    input_array = cpm_matrix[idx]

    for k in sgrna_dict:

        print("Region in processing: " + k[0:], file = sys.stderr)

        #idx index of cells containing the given sgRNA
        sgrna_idx = find_sgrna_cells(sgrna_dict[k], sgrna_df_adj, nonzero_gbm)

        #force the up-tail p-vals of all zero expressed genes to be zero. (actual p-val is 1)
        pval_list_down = np.zeros(len(nonzero_gbm.gene_names))
        pval_list_up = np.zeros(len(nonzero_gbm.gene_names))

        fc_list = np.ones(len(nonzero_gbm.gene_names))

        #perform the differential gene analysis by using Virtual FACS
        num_sgrna_cell, pval_list_up, pval_list_down, fc_list = perform_DE(
            sgrna_idx,
            input_array,
            idx,
            num_processing,
            pval_list_down,
            pval_list_up,
            fc_list
        )

        #save all the output
        io.savemat(
            '%s/%s-up_log-pval'%(output_dir, k[0:]),
            {'matrix':pval_list_up}
        )

        io.savemat(
            '%s/%s-down_log-pval'%(output_dir, k[0:]),
            {'matrix':pval_list_down}
        )

if __name__ == '__main__':
    main()