# coding: utf-8
import os
import mxnet as mx
import numpy as np
import math
import cv2
from multiprocessing import Pool
from itertools import repeat
from helper import nms, adjust_input, generate_bbox, detect_first_stage_warpper

class MtcnnDetector(object):
    """
        Joint Face Detection and Alignment using Multi-task Cascaded Convolutional Neural Networks
        see https://github.com/kpzhang93/MTCNN_face_detection_alignment
        this is a mxnet version
    """
    def __init__(self,
                 model_folder='.',
                 minsize = 20,
                 threshold = [0.6, 0.7, 0.8],
                 factor = 0.709,
                 num_worker = 1,
                 accurate_landmark = False,
                 ctx=mx.cpu()):
        """
            Initialize the detector

            Parameters:
            ----------
                model_folder : string
                    path for the models
                minsize : float number
                    minimal face to detect
                threshold : float number
                    detect threshold for 3 stages
                factor: float number
                    scale factor for image pyramid
                num_worker: int number
                    number of processes we use for first stage
                accurate_landmark: bool
                    use accurate landmark localization or not

        """
        self.num_worker = num_worker
        self.accurate_landmark = accurate_landmark

        # load 4 models from folder
        models = ['det1', 'det2', 'det3','det4']
        models = [ os.path.join(model_folder, f) for f in models]

        self.PNets = []
        for i in range(num_worker):
            workner_net = mx.model.FeedForward.load(models[0], 1, ctx=ctx)
            self.PNets.append(workner_net)

        #self.Pool = Pool(num_worker)

        self.RNet = mx.model.FeedForward.load(models[1], 1, ctx=ctx)
        self.ONet = mx.model.FeedForward.load(models[2], 1, ctx=ctx)
        self.LNet = mx.model.FeedForward.load(models[3], 1, ctx=ctx)

        self.minsize   = float(minsize)
        self.factor    = float(factor)
        self.threshold = threshold

    def convert_to_square(self, bbox):
        """
            convert bbox to square

        Parameters:
        ----------
            bbox: numpy array , shape n x 5
                input bbox

        Returns:
        -------
            square bbox
        """
        square_bbox = bbox.copy()

        h = bbox[:, 3] - bbox[:, 1] + 1
        w = bbox[:, 2] - bbox[:, 0] + 1
        max_side = np.maximum(h,w)
        square_bbox[:, 0] = bbox[:, 0] + w*0.5 - max_side*0.5
        square_bbox[:, 1] = bbox[:, 1] + h*0.5 - max_side*0.5
        square_bbox[:, 2] = square_bbox[:, 0] + max_side - 1
        square_bbox[:, 3] = square_bbox[:, 1] + max_side - 1
        return square_bbox

    def calibrate_box(self, bbox, reg):
        """
            calibrate bboxes

        Parameters:
        ----------
            bbox: numpy array, shape n x 5
                input bboxes
            reg:  numpy array, shape n x 4
                bboxex adjustment

        Returns:
        -------
            bboxes after refinement

        """
        w = bbox[:, 2] - bbox[:, 0] + 1
        w = np.expand_dims(w, 1)
        h = bbox[:, 3] - bbox[:, 1] + 1
        h = np.expand_dims(h, 1)
        reg_m = np.hstack([w, h, w, h])
        aug = reg_m * reg
        bbox[:, 0:4] = bbox[:, 0:4] + aug
        return bbox

    def pad(self, bboxes, w, h):
        """
            pad the the bboxes, alse restrict the size of it

        Parameters:
        ----------
            bboxes: numpy array, n x 5
                input bboxes
            w: float number
                width of the input image
            h: float number
                height of the input image
        Returns :
        ------s
            dy, dx : numpy array, n x 1
                start point of the bbox in target image
            edy, edx : numpy array, n x 1
                end point of the bbox in target image
            y, x : numpy array, n x 1
                start point of the bbox in original image
            ex, ex : numpy array, n x 1
                end point of the bbox in original image
            tmph, tmpw: numpy array, n x 1
                height and width of the bbox

        """
        tmpw, tmph = bboxes[:, 2] - bboxes[:, 0] + 1,  bboxes[:, 3] - bboxes[:, 1] + 1
        num_box = bboxes.shape[0]

        dx , dy= np.zeros((num_box, )), np.zeros((num_box, ))
        edx, edy  = tmpw.copy()-1, tmph.copy()-1

        x, y, ex, ey = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]

        tmp_index = np.where(ex > w-1)
        edx[tmp_index] = tmpw[tmp_index] + w - 2 - ex[tmp_index]
        ex[tmp_index] = w - 1

        tmp_index = np.where(ey > h-1)
        edy[tmp_index] = tmph[tmp_index] + h - 2 - ey[tmp_index]
        ey[tmp_index] = h - 1

        tmp_index = np.where(x < 0)
        dx[tmp_index] = 0 - x[tmp_index]
        x[tmp_index] = 0

        tmp_index = np.where(y < 0)
        dy[tmp_index] = 0 - y[tmp_index]
        y[tmp_index] = 0

        return_list = [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph]
        return_list = [item.astype(np.int32) for item in return_list]

        return  return_list

    def slice_index(self, number):
        """
            slice the index into (n,n,m), m < n
        Parameters:
        ----------
            number: int number
                number
        """
        def chunks(l, n):
            """Yield successive n-sized chunks from l."""
            for i in range(0, len(l), n):
                yield l[i:i + n]
        num_list = range(number)
        return list(chunks(num_list, self.num_worker))

    def detect_face_limited(self, img, det_type=2):
        height, width, _ = img.shape
        if det_type>=2:
          total_boxes = np.array( [ [0.0, 0.0, img.shape[1], img.shape[0], 0.9] ] ,dtype=np.float32)
          num_box = total_boxes.shape[0]

          # pad the bbox
          [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = self.pad(total_boxes, width, height)
          # (3, 24, 24) is the input shape for RNet
          input_buf = np.zeros((num_box, 3, 24, 24), dtype=np.float32)

          for i in range(num_box):
              tmp = np.zeros((tmph[i], tmpw[i], 3), dtype=np.uint8)
              tmp[dy[i]:edy[i]+1, dx[i]:edx[i]+1, :] = img[y[i]:ey[i]+1, x[i]:ex[i]+1, :]
              input_buf[i, :, :, :] = adjust_input(cv2.resize(tmp, (24, 24)))

          output = self.RNet.predict(input_buf)

          # filter the total_boxes with threshold
          passed = np.where(output[1][:, 1] > self.threshold[1])
          total_boxes = total_boxes[passed]

          if total_boxes.size == 0:
              return None

          total_boxes[:, 4] = output[1][passed, 1].reshape((-1,))
          reg = output[0][passed]

          # nms
          pick = nms(total_boxes, 0.7, 'Union')
          total_boxes = total_boxes[pick]
          total_boxes = self.calibrate_box(total_boxes, reg[pick])
          total_boxes = self.convert_to_square(total_boxes)
          total_boxes[:, 0:4] = np.round(total_boxes[:, 0:4])
        else:
          total_boxes = np.array( [ [0.0, 0.0, img.shape[1], img.shape[0], 0.9] ] ,dtype=np.float32)
        num_box = total_boxes.shape[0]
        [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = self.pad(total_boxes, width, height)
        # (3, 48, 48) is the input shape for ONet
        input_buf = np.zeros((num_box, 3, 48, 48), dtype=np.float32)

        for i in range(num_box):
            tmp = np.zeros((tmph[i], tmpw[i], 3), dtype=np.float32)
            tmp[dy[i]:edy[i]+1, dx[i]:edx[i]+1, :] = img[y[i]:ey[i]+1, x[i]:ex[i]+1, :]
            input_buf[i, :, :, :] = adjust_input(cv2.resize(tmp, (48, 48)))

        output = self.ONet.predict(input_buf)
        #print(output[2])

        # filter the total_boxes with threshold
        passed = np.where(output[2][:, 1] > self.threshold[2])
        total_boxes = total_boxes[passed]

        if total_boxes.size == 0:
            return None

        total_boxes[:, 4] = output[2][passed, 1].reshape((-1,))
        reg = output[1][passed]
        points = output[0][passed]

        # compute landmark points
        bbw = total_boxes[:, 2] - total_boxes[:, 0] + 1
        bbh = total_boxes[:, 3] - total_boxes[:, 1] + 1
        points[:, 0:5] = np.expand_dims(total_boxes[:, 0], 1) + np.expand_dims(bbw, 1) * points[:, 0:5]
        points[:, 5:10] = np.expand_dims(total_boxes[:, 1], 1) + np.expand_dims(bbh, 1) * points[:, 5:10]

        # nms
        total_boxes = self.calibrate_box(total_boxes, reg)
        pick = nms(total_boxes, 0.7, 'Min')
        total_boxes = total_boxes[pick]
        points = points[pick]

        if not self.accurate_landmark:
            return total_boxes, points

        #############################################
        # extended stage
        #############################################
        num_box = total_boxes.shape[0]
        patchw = np.maximum(total_boxes[:, 2]-total_boxes[:, 0]+1, total_boxes[:, 3]-total_boxes[:, 1]+1)
        patchw = np.round(patchw*0.25)

        # make it even
        patchw[np.where(np.mod(patchw,2) == 1)] += 1

        input_buf = np.zeros((num_box, 15, 24, 24), dtype=np.float32)
        for i in range(5):
            x, y = points[:, i], points[:, i+5]
            x, y = np.round(x-0.5*patchw), np.round(y-0.5*patchw)
            [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = self.pad(np.vstack([x, y, x+patchw-1, y+patchw-1]).T,
                                                                    width,
                                                                    height)
            for j in range(num_box):
                tmpim = np.zeros((tmpw[j], tmpw[j], 3), dtype=np.float32)
                tmpim[dy[j]:edy[j]+1, dx[j]:edx[j]+1, :] = img[y[j]:ey[j]+1, x[j]:ex[j]+1, :]
                input_buf[j, i*3:i*3+3, :, :] = adjust_input(cv2.resize(tmpim, (24, 24)))

        output = self.LNet.predict(input_buf)

        pointx = np.zeros((num_box, 5))
        pointy = np.zeros((num_box, 5))

        for k in range(5):
            # do not make a large movement
            tmp_index = np.where(np.abs(output[k]-0.5) > 0.35)
            output[k][tmp_index[0]] = 0.5

            pointx[:, k] = np.round(points[:, k] - 0.5*patchw) + output[k][:, 0]*patchw
            pointy[:, k] = np.round(points[:, k+5] - 0.5*patchw) + output[k][:, 1]*patchw

        points = np.hstack([pointx, pointy])
        points = points.astype(np.int32)

        return total_boxes, points

    def detect_face(self, img, det_type=0):
        """
            detect face over img
        Parameters:
        ----------
            img: numpy array, bgr order of shape (1, 3, n, m)
                input image
        Retures:
        -------
            bboxes: numpy array, n x 5 (x1,y2,x2,y2,score)
                bboxes
            points: numpy array, n x 10 (x1, x2 ... x5, y1, y2 ..y5)
                landmarks
        """

        # check input
        height, width, _ = img.shape
        if det_type==0:
            MIN_DET_SIZE = 12

            if img is None:
                return None

            # only works for color image
            if len(img.shape) != 3:
                return None

            # detected boxes
            total_boxes = []

            minl = min( height, width)

            # get all the valid scales
            scales = []
            m = MIN_DET_SIZE/self.minsize
            minl *= m
            factor_count = 0
            while minl > MIN_DET_SIZE:
                scales.append(m*self.factor**factor_count)
                minl *= self.factor
                factor_count += 1

            #############################################
            # first stage
            #############################################
            #for scale in scales:
            #    return_boxes = self.detect_first_stage(img, scale, 0)
            #    if return_boxes is not None:
            #        total_boxes.append(return_boxes)

            sliced_index = self.slice_index(len(scales))
            total_boxes = []
            for batch in sliced_index:
                #local_boxes = self.Pool.map( detect_first_stage_warpper, \
                #        zip(repeat(img), self.PNets[:len(batch)], [scales[i] for i in batch], repeat(self.threshold[0])) )
                local_boxes = map( detect_first_stage_warpper, \
                        zip(repeat(img), self.PNets[:len(batch)], [scales[i] for i in batch], repeat(self.threshold[0])) )
                total_boxes.extend(local_boxes)

            # remove the Nones
            total_boxes = [ i for i in total_boxes if i is not None]

            if len(total_boxes) == 0:
                return None

            total_boxes = np.vstack(total_boxes)

            if total_boxes.size == 0:
                return None

            # merge the detection from first stage
            pick = nms(total_boxes[:, 0:5], 0.7, 'Union')
            total_boxes = total_boxes[pick]

            bbw = total_boxes[:, 2] - total_boxes[:, 0] + 1
            bbh = total_boxes[:, 3] - total_boxes[:, 1] + 1

            # refine the bboxes
            total_boxes = np.vstack([total_boxes[:, 0]+total_boxes[:, 5] * bbw,
                                     total_boxes[:, 1]+total_boxes[:, 6] * bbh,
                                     total_boxes[:, 2]+total_boxes[:, 7] * bbw,
                                     total_boxes[:, 3]+total_boxes[:, 8] * bbh,
                                     total_boxes[:, 4]
                                     ])

            total_boxes = total_boxes.T
            total_boxes = self.convert_to_square(total_boxes)
            total_boxes[:, 0:4] = np.round(total_boxes[:, 0:4])
        else:
            total_boxes = np.array( [ [0.0, 0.0, img.shape[1], img.shape[0], 0.9] ] ,dtype=np.float32)

        #############################################
        # second stage
        #############################################
        num_box = total_boxes.shape[0]

        # pad the bbox
        [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = self.pad(total_boxes, width, height)
        # (3, 24, 24) is the input shape for RNet
        input_buf = np.zeros((num_box, 3, 24, 24), dtype=np.float32)

        for i in range(num_box):
            tmp = np.zeros((tmph[i], tmpw[i], 3), dtype=np.uint8)
            tmp[dy[i]:edy[i]+1, dx[i]:edx[i]+1, :] = img[y[i]:ey[i]+1, x[i]:ex[i]+1, :]
            input_buf[i, :, :, :] = adjust_input(cv2.resize(tmp, (24, 24)))

        output = self.RNet.predict(input_buf)

        # filter the total_boxes with threshold
        passed = np.where(output[1][:, 1] > self.threshold[1])
        total_boxes = total_boxes[passed]

        if total_boxes.size == 0:
            return None

        total_boxes[:, 4] = output[1][passed, 1].reshape((-1,))
        reg = output[0][passed]

        # nms
        pick = nms(total_boxes, 0.7, 'Union')
        total_boxes = total_boxes[pick]
        total_boxes = self.calibrate_box(total_boxes, reg[pick])
        total_boxes = self.convert_to_square(total_boxes)
        total_boxes[:, 0:4] = np.round(total_boxes[:, 0:4])

        #############################################
        # third stage
        #############################################
        num_box = total_boxes.shape[0]

        # pad the bbox
        [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = self.pad(total_boxes, width, height)
        # (3, 48, 48) is the input shape for ONet
        input_buf = np.zeros((num_box, 3, 48, 48), dtype=np.float32)

        for i in range(num_box):
            tmp = np.zeros((tmph[i], tmpw[i], 3), dtype=np.float32)
            tmp[dy[i]:edy[i]+1, dx[i]:edx[i]+1, :] = img[y[i]:ey[i]+1, x[i]:ex[i]+1, :]
            input_buf[i, :, :, :] = adjust_input(cv2.resize(tmp, (48, 48)))

        output = self.ONet.predict(input_buf)

        # filter the total_boxes with threshold
        passed = np.where(output[2][:, 1] > self.threshold[2])
        total_boxes = total_boxes[passed]

        if total_boxes.size == 0:
            return None

        total_boxes[:, 4] = output[2][passed, 1].reshape((-1,))
        reg = output[1][passed]
        points = output[0][passed]

        # compute landmark points
        bbw = total_boxes[:, 2] - total_boxes[:, 0] + 1
        bbh = total_boxes[:, 3] - total_boxes[:, 1] + 1
        points[:, 0:5] = np.expand_dims(total_boxes[:, 0], 1) + np.expand_dims(bbw, 1) * points[:, 0:5]
        points[:, 5:10] = np.expand_dims(total_boxes[:, 1], 1) + np.expand_dims(bbh, 1) * points[:, 5:10]

        # nms
        total_boxes = self.calibrate_box(total_boxes, reg)
        pick = nms(total_boxes, 0.7, 'Min')
        total_boxes = total_boxes[pick]
        points = points[pick]

        if not self.accurate_landmark:
            return total_boxes, points

        #############################################
        # extended stage
        #############################################
        num_box = total_boxes.shape[0]
        patchw = np.maximum(total_boxes[:, 2]-total_boxes[:, 0]+1, total_boxes[:, 3]-total_boxes[:, 1]+1)
        patchw = np.round(patchw*0.25)

        # make it even
        patchw[np.where(np.mod(patchw,2) == 1)] += 1

        input_buf = np.zeros((num_box, 15, 24, 24), dtype=np.float32)
        for i in range(5):
            x, y = points[:, i], points[:, i+5]
            x, y = np.round(x-0.5*patchw), np.round(y-0.5*patchw)
            [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = self.pad(np.vstack([x, y, x+patchw-1, y+patchw-1]).T,
                                                                    width,
                                                                    height)
            for j in range(num_box):
                tmpim = np.zeros((tmpw[j], tmpw[j], 3), dtype=np.float32)
                tmpim[dy[j]:edy[j]+1, dx[j]:edx[j]+1, :] = img[y[j]:ey[j]+1, x[j]:ex[j]+1, :]
                input_buf[j, i*3:i*3+3, :, :] = adjust_input(cv2.resize(tmpim, (24, 24)))

        output = self.LNet.predict(input_buf)

        pointx = np.zeros((num_box, 5))
        pointy = np.zeros((num_box, 5))

        for k in range(5):
            # do not make a large movement
            tmp_index = np.where(np.abs(output[k]-0.5) > 0.35)
            output[k][tmp_index[0]] = 0.5

            pointx[:, k] = np.round(points[:, k] - 0.5*patchw) + output[k][:, 0]*patchw
            pointy[:, k] = np.round(points[:, k+5] - 0.5*patchw) + output[k][:, 1]*patchw

        points = np.hstack([pointx, pointy])
        points = points.astype(np.int32)

        return total_boxes, points

    def list2colmatrix(self, pts_list):
        """
            convert list to column matrix
        Parameters:
        ----------
            pts_list:
                input list
        Retures:
        -------
            colMat:

        """
        assert len(pts_list) > 0
        colMat = []
        for i in range(len(pts_list)):
            colMat.append(pts_list[i][0])
            colMat.append(pts_list[i][1])
        colMat = np.matrix(colMat).transpose()
        return colMat

    def find_tfrom_between_shapes(self, from_shape, to_shape):
        """
            find transform between shapes
        Parameters:
        ----------
            from_shape:
            to_shape:
        Retures:
        -------
            tran_m:
            tran_b:
        """
        assert from_shape.shape[0] == to_shape.shape[0] and from_shape.shape[0] % 2 == 0

        sigma_from = 0.0
        sigma_to = 0.0
        cov = np.matrix([[0.0, 0.0], [0.0, 0.0]])

        # compute the mean and cov
        from_shape_points = from_shape.reshape(from_shape.shape[0]/2, 2)
        to_shape_points = to_shape.reshape(to_shape.shape[0]/2, 2)
        mean_from = from_shape_points.mean(axis=0)
        mean_to = to_shape_points.mean(axis=0)

        for i in range(from_shape_points.shape[0]):
            temp_dis = np.linalg.norm(from_shape_points[i] - mean_from)
            sigma_from += temp_dis * temp_dis
            temp_dis = np.linalg.norm(to_shape_points[i] - mean_to)
            sigma_to += temp_dis * temp_dis
            cov += (to_shape_points[i].transpose() - mean_to.transpose()) * (from_shape_points[i] - mean_from)

        sigma_from = sigma_from / to_shape_points.shape[0]
        sigma_to = sigma_to / to_shape_points.shape[0]
        cov = cov / to_shape_points.shape[0]

        # compute the affine matrix
        s = np.matrix([[1.0, 0.0], [0.0, 1.0]])
        u, d, vt = np.linalg.svd(cov)

        if np.linalg.det(cov) < 0:
            if d[1] < d[0]:
                s[1, 1] = -1
            else:
                s[0, 0] = -1
        r = u * s * vt
        c = 1.0
        if sigma_from != 0:
            c = 1.0 / sigma_from * np.trace(np.diag(d) * s)

        tran_b = mean_to.transpose() - c * r * mean_from.transpose()
        tran_m = c * r

        return tran_m, tran_b

    def extract_image_chips(self, img, points, desired_size=256, padding=0):
        """
            crop and align face
        Parameters:
        ----------
            img: numpy array, bgr order of shape (1, 3, n, m)
                input image
            points: numpy array, n x 10 (x1, x2 ... x5, y1, y2 ..y5)
            desired_size: default 256
            padding: default 0
        Retures:
        -------
            crop_imgs: list, n
                cropped and aligned faces
        """
        crop_imgs = []
        for p in points:
            shape  =[]
            for k in range(len(p)/2):
                shape.append(p[k])
                shape.append(p[k+5])

            if padding > 0:
                padding = padding
            else:
                padding = 0
            # average positions of face points
            mean_face_shape_x = [0.224152, 0.75610125, 0.490127, 0.254149, 0.726104]
            mean_face_shape_y = [0.2119465, 0.2119465, 0.628106, 0.780233, 0.780233]

            from_points = []
            to_points = []

            for i in range(len(shape)/2):
                x = (padding + mean_face_shape_x[i]) / (2 * padding + 1) * desired_size
                y = (padding + mean_face_shape_y[i]) / (2 * padding + 1) * desired_size
                to_points.append([x, y])
                from_points.append([shape[2*i], shape[2*i+1]])

            # convert the points to Mat
            from_mat = self.list2colmatrix(from_points)
            to_mat = self.list2colmatrix(to_points)

            # compute the similar transfrom
            tran_m, tran_b = self.find_tfrom_between_shapes(from_mat, to_mat)

            probe_vec = np.matrix([1.0, 0.0]).transpose()
            probe_vec = tran_m * probe_vec

            scale = np.linalg.norm(probe_vec)
            angle = 180.0 / math.pi * math.atan2(probe_vec[1, 0], probe_vec[0, 0])

            from_center = [(shape[0]+shape[2])/2.0, (shape[1]+shape[3])/2.0]
            to_center = [0, 0]
            to_center[1] = desired_size * 0.4
            to_center[0] = desired_size * 0.5

            ex = to_center[0] - from_center[0]
            ey = to_center[1] - from_center[1]

            rot_mat = cv2.getRotationMatrix2D((from_center[0], from_center[1]), -1*angle, scale)
            rot_mat[0][2] += ex
            rot_mat[1][2] += ey

            chips = cv2.warpAffine(img, rot_mat, (desired_size, desired_size))
            crop_imgs.append(chips)

        return crop_imgs


"""BDD100K tracking evaluation with CLEAR MOT metrics."""
import time
from functools import partial
from multiprocessing import Pool
from typing import Callable, Dict, List, Tuple, Union, overload

import motmetrics as mm
import numpy as np
import pandas as pd
from scalabel.label.transforms import box2d_to_bbox
from scalabel.label.typing import Frame, Label

from ..common.logger import logger

EvalResults = Dict[str, Dict[str, float]]
Frames = List[Frame]
Files = List[str]
FramesList = List[Frames]
FilesList = List[Files]
FramesFunc = Callable[[Frames, Frames, float, float], List[mm.MOTAccumulator]]
FilesFunc = Callable[[Files, Files, float, float], List[mm.MOTAccumulator]]

METRIC_MAPS = {
    "idf1": "IDF1",
    "mota": "MOTA",
    "motp": "MOTP",
    "num_false_positives": "FP",
    "num_misses": "FN",
    "num_switches": "IDSw",
    "mostly_tracked": "MT",
    "partially_tracked": "PT",
    "mostly_lost": "ML",
    "num_fragmentations": "FM",
}

SUPER_CLASSES = {
    "HUMAN": ["pedestrian", "rider"],
    "VEHICLE": ["car", "truck", "bus", "train"],
    "BIKE": ["motorcycle", "bicycle"],
}
CLASSES = [c for cs in SUPER_CLASSES.values() for c in cs]
IGNORE_CLASSES = ["trailer", "other person", "other vehicle"]

def parse_objects(objects: List[Label]) -> List[np.ndarray]:
    """Parse objects under Scalable formats."""
    bboxes, labels, ids, ignore_bboxes = [], [], [], []
    for obj in objects:
        box2d = obj.box2d
        if box2d is None:
            continue
        bbox = box2d_to_bbox(box2d)
        category = obj.category
        if category in CLASSES:
            if obj.attributes is not None and bool(
                obj.attributes.get("crowd", False)
            ):
                ignore_bboxes.append(bbox)
            else:
                bboxes.append(bbox)
                labels.append(CLASSES.index(category))
                ids.append(obj.id)
        elif category in IGNORE_CLASSES:
            ignore_bboxes.append(bbox)
        else:
            raise KeyError("Unknown category.")
    return list(map(np.array, [bboxes, labels, ids, ignore_bboxes]))

def intersection_over_area(preds: np.ndarray, gts: np.ndarray) -> np.ndarray:
    """Returns the intersection over the area of the predicted box."""
    out = np.zeros((len(preds), len(gts)))
    for i, p in enumerate(preds):
        for j, g in enumerate(gts):
            w = min(p[0] + p[2], g[0] + g[2]) - max(p[0], g[0])
            h = min(p[1] + p[3], g[1] + g[3]) - max(p[1], g[1])
            out[i][j] = max(w, 0) * max(h, 0) / float(p[2] * p[3])
    return out

def acc_single_video_mot(
    gts: List[Frame],
    results: List[Frame],
    iou_thr: float = 0.5,
    ignore_iof_thr: float = 0.5,
) -> List[mm.MOTAccumulator]:
    """Accumulate results for one video."""
    num_classes = len(CLASSES)
    assert len(gts) == len(results)
    get_frame_index = (
        lambda x: x.frame_index if x.frame_index is not None else 0
    )
    gts = sorted(gts, key=get_frame_index)
    results = sorted(results, key=get_frame_index)
    accs = [mm.MOTAccumulator(auto_id=True) for i in range(num_classes)]
    for gt, result in zip(gts, results):
        assert gt.frame_index == result.frame_index
        gt_bboxes, gt_labels, gt_ids, gt_ignores = parse_objects(
            gt.labels if gt.labels is not None else []
        )
        pred_bboxes, pred_labels, pred_ids, _ = parse_objects(
            result.labels if result.labels is not None else []
        )
        for i in range(num_classes):
            gt_inds, pred_inds = gt_labels == i, pred_labels == i
            gt_bboxes_c, gt_ids_c = gt_bboxes[gt_inds], gt_ids[gt_inds]
            pred_bboxes_c, pred_ids_c = (
                pred_bboxes[pred_inds],
                pred_ids[pred_inds],
            )
            if gt_bboxes_c.shape[0] == 0 and pred_bboxes_c.shape[0] != 0:
                distances = np.full((0, pred_bboxes_c.shape[0]), np.nan)
            elif gt_bboxes_c.shape[0] != 0 and pred_bboxes_c.shape[0] == 0:
                distances = np.full((gt_bboxes_c.shape[0], 0), np.nan)
            else:
                distances = mm.distances.iou_matrix(
                    gt_bboxes_c, pred_bboxes_c, max_iou=1 - iou_thr
                )
            if gt_ignores.shape[0] > 0:
                # 1. assign gt and preds
                fps = np.ones(pred_bboxes_c.shape[0]).astype(np.bool8)
                le, ri = mm.lap.linear_sum_assignment(distances)
                for m, n in zip(le, ri):
                    if np.isfinite(distances[m, n]):
                        fps[n] = False
                # 2. ignore by iof
                iofs = intersection_over_area(pred_bboxes_c, gt_ignores)
                ignores = (iofs > ignore_iof_thr).any(axis=1)
                # 3. filter preds
                valid_inds = ~(fps & ignores)
                pred_ids_c = pred_ids_c[valid_inds]
                distances = distances[:, valid_inds]
            if distances.shape != (0, 0):
                accs[i].update(gt_ids_c, pred_ids_c, distances)
    return accs

def aggregate_accs(
    accumulators: List[List[mm.MOTAccumulator]],
) -> Tuple[List[List[str]], List[List[mm.MOTAccumulator]], List[str]]:
    """Aggregate the results of the entire dataset."""
    # accs for each class
    items = CLASSES.copy()
    names: List[List[str]] = [[] for c in CLASSES]
    accs: List[List[str]] = [[] for c in CLASSES]
    for video_ind, _accs in enumerate(accumulators):
        for cls_ind, acc in enumerate(_accs):
            if (
                len(acc._events["Type"])  # pylint: disable=protected-access
                == 0
            ):
                continue
            name = f"{CLASSES[cls_ind]}_{video_ind}"
            names[cls_ind].append(name)
            accs[cls_ind].append(acc)

    # super categories
    for super_cls, classes in SUPER_CLASSES.items():
        items.append(super_cls)
        names.append([n for c in classes for n in names[CLASSES.index(c)]])
        accs.append([a for c in classes for a in accs[CLASSES.index(c)]])

    # overall
    items.append("OVERALL")
    names.append([n for name in names[: len(CLASSES)] for n in name])
    accs.append([a for acc in accs[: len(CLASSES)] for a in acc])

    return names, accs, items

def evaluate_single_class(
    names: List[str], accs: List[mm.MOTAccumulator]
) -> List[float]:
    """Evaluate results for one class."""
    mh = mm.metrics.create()
    summary = mh.compute_many(
        accs, names=names, metrics=METRIC_MAPS.keys(), generate_overall=True
    )
    results = [v["OVERALL"] for k, v in summary.to_dict().items()]
    motp_ind = list(METRIC_MAPS).index("motp")
    if np.isnan(results[motp_ind]):
        num_dets = mh.compute_many(
            accs,
            names=names,
            metrics=["num_detections"],
            generate_overall=True,
        )
        sum_motp = (summary["motp"] * num_dets["num_detections"]).sum()
        motp = mm.math_util.quiet_divide(
            sum_motp, num_dets["num_detections"]["OVERALL"]
        )
        results[motp_ind] = float(1 - motp)
    else:
        results[motp_ind] = 1 - results[motp_ind]
    return results

def render_results(
    summaries: List[List[float]],
    items: List[str],
    metrics: List[str],
) -> EvalResults:
    """Render the evaluation results."""
    eval_results = pd.DataFrame(columns=metrics)
    # category, super-category and overall results
    for i, item in enumerate(items):
        eval_results.loc[item] = summaries[i]
    dtypes = {m: type(d) for m, d in zip(metrics, summaries[0])}
    # average results
    avg_results: List[Union[int, float]] = []
    for i, m in enumerate(metrics):
        v = np.array([s[i] for s in summaries[: len(CLASSES)]])
        v = np.nan_to_num(v, nan=0)
        if dtypes[m] == int:
            avg_results.append(int(v.sum()))
        elif dtypes[m] == float:
            avg_results.append(float(v.mean()))
        else:
            raise TypeError()
    eval_results.loc["AVERAGE"] = avg_results
    eval_results = eval_results.astype(dtypes)

    metric_host = mm.metrics.create()
    metric_host.register(mm.metrics.motp, formatter="{:.1%}".format)
    strsummary = mm.io.render_summary(
        eval_results,
        formatters=metric_host.formatters,
        namemap=METRIC_MAPS,
    )
    strsummary = strsummary.split("\n")
    assert len(strsummary) == len(CLASSES) + len(SUPER_CLASSES) + 3
    split_line = "-" * len(strsummary[0])
    strsummary.insert(1, split_line)
    strsummary.insert(2 + len(CLASSES), split_line)
    strsummary.insert(3 + len(CLASSES) + len(SUPER_CLASSES), split_line)
    strsummary = "".join([f"{s}\n" for s in strsummary])
    strsummary = "\n" + strsummary
    logger.info(strsummary)

    outputs: EvalResults = dict()
    for i, item in enumerate(items):
        outputs[item] = dict()
        for j, metric in enumerate(METRIC_MAPS.values()):
            outputs[item][metric] = summaries[i][j]
    outputs["OVERALL"]["mIDF1"] = eval_results.loc["AVERAGE"]["idf1"]
    outputs["OVERALL"]["mMOTA"] = eval_results.loc["AVERAGE"]["mota"]
    outputs["OVERALL"]["mMOTP"] = eval_results.loc["AVERAGE"]["motp"]

    return outputs

@overload
def evaluate_track(
    acc_single_video: FramesFunc,
    gts: FramesList,
    results: FramesList,
    iou_thr: float = 0.5,
    ignore_iof_thr: float = 0.5,
    nproc: int = 4,
) -> EvalResults:
    ...

@overload
def evaluate_track(
    acc_single_video: FilesFunc,
    gts: FilesList,
    results: FilesList,
    iou_thr: float = 0.5,
    ignore_iof_thr: float = 0.5,
    nproc: int = 4,
) -> EvalResults:
    ...

def evaluate_track(
    acc_single_video: Union[FramesFunc, FilesFunc],
    gts: Union[FramesList, FilesList],
    results: Union[FramesList, FilesList],
    iou_thr: float = 0.5,
    ignore_iof_thr: float = 0.5,
    nproc: int = 4,
) -> EvalResults:
    """Evaluate CLEAR MOT metrics for BDD100K."""
    logger.info("BDD100K tracking evaluation with CLEAR MOT metrics.")
    t = time.time()
    assert len(gts) == len(results)
    metrics = list(METRIC_MAPS.keys())

    logger.info("accumulating...")
    with Pool(nproc) as pool:
        accs = pool.starmap(
            partial(
                acc_single_video,
                iou_thr=iou_thr,
                ignore_iof_thr=ignore_iof_thr,
            ),
            zip(gts, results),
        )
    names, accs, items = aggregate_accs(accs)

    logger.info("evaluating...")
    with Pool(nproc) as pool:
        summaries = pool.starmap(evaluate_single_class, zip(names, accs))

    logger.info("rendering...")
    eval_results = render_results(summaries, items, metrics)
    t = time.time() - t
    logger.info("evaluation finishes with %.1f s.", t)
    return eval_results


#!/usr/bin/env python3

import sys
import os
import argparse
import cv2
from skimage.metrics import structural_similarity
from skimage import img_as_ubyte
from datetime import datetime, timedelta, timezone
import time
import numpy as np
import subprocess as sp
from threading import Thread, Event
from queue import Queue
import curses
from keras.models import load_model

FULL_WIDTH = 1280
FULL_HEIGHT = 720
SMALL_SIZE = (426, 240)

utc = timezone.utc

# https://answers.opencv.org/question/92450/processing-camera-stream-in-opencv-pushing-it-over-rtmp-nginx-rtmp-module-using-ffmpeg/
FFMPEG_COMMAND = [
    'ffmpeg',
    '-y',
    '-f', 'rawvideo',
    '-vcodec', 'rawvideo',
    '-s', f'{FULL_WIDTH}x{FULL_HEIGHT}',
    '-pix_fmt', 'bgr24',
    '-i', '-',
    '-codec', 'copy',
    '-f', 'v4l2',
    '/dev/video3'
]

class ThreadWithSettingsAndMessages(Thread):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.settings = {}
        self.message_queue = Queue()

    def set_setting(self, setting, value):
        self.settings[setting] = value

    def get_setting(self, setting):
        return self.settings[setting]

    def toggle_setting(self, setting):
        self.settings[setting] = not self.settings[setting]
        return self.settings[setting]

    def has_message(self) -> bool:
        return self.message_queue.qsize() > 0

    def get_message(self):
        # don't block
        if self.has_message:
            return self.message_queue.get()
        return None

class ImageCompareThread(ThreadWithSettingsAndMessages):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.work_queue = Queue()
        self.result = None
        self.stoprequest = Event()
        self.settings = {
            'use_hull': False,
            'use_and': False,
            'use_model': True,
        }
        self.base = cv2.imread('base.png', -1)
        self.base_small = cv2.resize(self.base, SMALL_SIZE)

    def run(self):
        # get model here
        # https://github.com/anilsathyan7/Portrait-Segmentation
        model = load_model('deconv_bnoptimized_munet.h5', compile=False)
        while not self.stoprequest.isSet():
            live = self.work_queue.get()
            mask = None
            if self.get_setting('use_model'):
                mask = self.infer_from_model(live, model)
            else:
                mask = self.compare_to_background(live)

            # https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_contours/py_contour_features/py_contour_features.html#
            contours, hierarchy = cv2.findContours(
                mask,
                cv2.RETR_EXTERNAL,
                cv2.CHAIN_APPROX_SIMPLE
            )
            largest_contour = None
            area = 0
            for cnt in contours:
                if cv2.contourArea(cnt) > area:
                    largest_contour = cnt
                    area = cv2.contourArea(cnt)

            contour = None
            if self.settings['use_hull'] is True:
                contour = cv2.convexHull(largest_contour)
            else:
                contour = largest_contour

            contour_mask = np.zeros(self.base_small.shape, np.uint8)

            # https://stackoverflow.com/a/50022122
            contour_mask = cv2.drawContours(contour_mask, [contour], -1, (255, 255, 255), -1)
            contour_mask_big = cv2.resize(contour_mask, (FULL_WIDTH, FULL_HEIGHT))
            contour_blur = cv2.medianBlur(contour_mask_big, 21)

            self.result = contour_blur

    def infer_from_model(self, live, model):
        # based on
        # https://news.ycombinator.com/item?id=22842823
        try:
            start = datetime.now(utc)

            frame = cv2.cvtColor(live, cv2.COLOR_BGR2RGB)
            simg = cv2.resize(frame, (128, 128), interpolation=cv2.INTER_AREA)
            simg = simg.reshape((1, 128, 128, 3)) / 255.0

            preprocessing = datetime.now(utc)

            out = model.predict(simg)
            predicting = datetime.now(utc)

            mask = out.reshape((128, 128, 1))
            mask = img_as_ubyte(mask)
            mask = cv2.resize(mask, SMALL_SIZE)

            postprocessing = datetime.now(utc)

            preprocessing_diff = preprocessing - start
            predicting_diff = predicting - preprocessing
            postprocessing_diff = postprocessing - predicting
            total_diff = postprocessing - start

            message = (
                f"Preprocessing took {preprocessing_diff.total_seconds()*1000:.0f} ms, "
                f"Predicting took {predicting_diff.total_seconds()*1000:.0f} ms, "
                f"Postprocessing took {postprocessing_diff.total_seconds()*1000:.0f} ms, "
                f"total: {total_diff.total_seconds()*1000:.0f} ms."
            )
            self.message_queue.put(message)
            return mask
        except Exception as e:
            message = str(e)
            self.message_queue.put(f"Error on model: {message}")
            return None

    def compare_to_background(self, live):
        # compare images
        start = datetime.now(utc)
        live_small = cv2.resize(live, SMALL_SIZE)
        preprocessing = datetime.now(utc)
        # https://www.pyimagesearch.com/2017/06/19/image-difference-with-opencv-and-python/
        # TODO: this takes to the most of the processing time, make more efficient
        (score, diff) = structural_similarity(self.base_small, live_small, full=True, multichannel=True)
        diff = (diff * 255).astype("uint8")
        ssim_time = datetime.now(utc)

        b, g, r = cv2.split(diff)

        color_mask_b = cv2.threshold(b, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
        color_mask_g = cv2.threshold(g, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
        color_mask_r = cv2.threshold(r, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]

        mask = None
        if self.settings['use_and']:
            mask = cv2.bitwise_and(color_mask_b, color_mask_g, color_mask_r)
        else:
            mask = cv2.bitwise_or(color_mask_b, color_mask_g, color_mask_r)

        now = datetime.now(utc)
        preprocessing_diff = preprocessing - start
        ssim = ssim_time - start
        postprocessing_diff = now - ssim_time
        total_diff = now - start
        message = (
            f"Preprocessing took {preprocessing_diff.total_seconds()*1000:.0f} ms, "
            f"SSIM took {ssim.total_seconds()*1000:.0f} ms, "
            f"Postprocessing took {postprocessing_diff.total_seconds()*1000:.0f} ms, "
            f"total: {total_diff.total_seconds()*1000:.0f} ms."
        )
        self.message_queue.put(message)
        return mask

    def put_work(self, message):
        self.work_queue.put(message)

    def join(self, timeout=None):
        self.stoprequest.set()
        super().join(timeout)

class MaskingThread(ThreadWithSettingsAndMessages):
    def __init__(self, worker_thread: Thread, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.stoprequest = Event()
        self.worker_thread = worker_thread
        self.cap = cv2.VideoCapture()
        self.settings = {
            'use_green': True,
        }
        self.frame_count = 0

    def run(self):
        contour_mask = None

        green = np.zeros((FULL_HEIGHT, FULL_WIDTH, 3), np.uint8)
        # Fill image with green color for greenscreen
        green[:] = (0, 255, 0)
        base_blur = cv2.imread('base_blur.png', -1)

        # loosely based on
        # https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_gui/py_video_display/py_video_display.html

        cap = self.cap
        # try to have higher perf by using V4L2 here
        # https://stackoverflow.com/a/55779890
        cap.open(0, apiPreference=cv2.CAP_V4L2)
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, FULL_WIDTH)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, FULL_HEIGHT)
        cap.set(cv2.CAP_PROP_FPS, 30)

        proc = sp.Popen(FFMPEG_COMMAND, stdin=sp.PIPE, stderr=sp.PIPE, shell=False)

        ret, live = cap.read()
        self.worker_thread.put_work(live)

        while not self.stoprequest.isSet():
            start = datetime.now(utc)
            ret, live = cap.read()
            reading = datetime.now(utc)

            if self.worker_thread.result is not None:
                contour_mask = self.worker_thread.result
                self.worker_thread.result = None
                self.worker_thread.put_work(live)
            # TODO: better wait for first result from other thread
            if contour_mask is not None:
                # based on
                # https://www.learnopencv.com/alpha-blending-using-opencv-cpp-python/
                # Convert uint8 to float
                foreground = live.astype(float)
                background = None
                if self.settings['use_green']:
                    background = green.astype(float)
                else:
                    background = base_blur.astype(float)

                # Normalize the alpha mask to keep intensity between 0 and 1
                alpha = contour_mask.astype(float) / 255

                # Multiply the foreground with the alpha matte
                foreground = cv2.multiply(alpha, foreground)

                # Multiply the background with ( 1 - alpha )
                background = cv2.multiply(1.0 - alpha, background)

                # Add the masked foreground and background.
                outImage = cv2.add(foreground, background).astype('uint8')

                masking = datetime.now(utc)

                self.frame_count += 1
                proc.stdin.write(outImage.tostring())
                writing = datetime.now(utc)

                reading_diff = reading - start
                masking_diff = masking - reading
                writing_diff = writing - masking
                total_diff = writing - start
                message = (
                    f"Reading took {reading_diff.total_seconds()*1000:.0f} ms, "
                    f"masking took {masking_diff.total_seconds()*1000:.0f} ms, "
                    f"writing took {writing_diff.total_seconds()*1000:.0f} ms, "
                    f"total: {total_diff.total_seconds()*1000:.0f} ms."
                )
                self.message_queue.put(message)
        cap.release()

    def join(self, timeout=None):
        self.stoprequest.set()
        super().join(timeout)

    def get_framecount(self):
        frame_count = self.frame_count
        self.frame_count = 0
        return frame_count

def handle_input(
    stdscr, image_compare_thread: ThreadWithSettingsAndMessages, masking_thread: ThreadWithSettingsAndMessages
):
    curses.start_color()
    curses.init_pair(1, curses.COLOR_GREEN, curses.COLOR_BLUE)
    curses.init_pair(2, curses.COLOR_YELLOW, curses.COLOR_BLACK)
    stdscr.bkgd(curses.color_pair(1))
    stdscr.nodelay(True)
    stdscr.refresh()

    # Hull info
    win1 = curses.newwin(5, 20, 0, 0)
    win1.bkgd(curses.color_pair(2))
    win1.box()
    win1.addstr(2, 2, "(H)ull inactive")
    win1.refresh()

    # And/Or info
    win2 = curses.newwin(5, 20, 0, 20)
    win2.bkgd(curses.color_pair(2))
    win2.box()
    win2.addstr(2, 2, "(A)nd inactive")
    win2.refresh()

    green_win = curses.newwin(5, 20, 0, 40)
    green_win.bkgd(curses.color_pair(2))
    green_win.box()
    green_win.addstr(2, 2, "(G)reen active")
    green_win.refresh()

    model_win = curses.newwin(5, 20, 0, 60)
    model_win.bkgd(curses.color_pair(2))
    model_win.box()
    model_win.addstr(2, 2, "(M)odel used")
    model_win.refresh()

    win3 = curses.newwin(5, 20, 0, 100)
    win3.bkgd(curses.color_pair(2))
    win3.box()
    win3.addstr(2, 2, "(Q)uit")
    win3.refresh()

    last_fps_update = datetime.now(utc)

    fps_win = curses.newwin(5, 20, 5, 00)
    fps_win.bkgd(curses.color_pair(2))
    fps_win.box()
    fps_win.refresh()

    # message windows
    masking_message_win = curses.newwin(10, 120, 10, 0)
    masking_message_win.bkgd(curses.color_pair(2))
    masking_message_win.box()
    masking_message_win.refresh()
    masking_message_queue = Queue()

    compare_message_win = curses.newwin(10, 120, 20, 0)
    compare_message_win.bkgd(curses.color_pair(2))
    compare_message_win.box()
    compare_message_win.refresh()
    compare_message_queue = Queue()

    while True:
        key = ''
        try:
            key = str(stdscr.getkey())
        except curses.error:
            pass
        if key == 'h':
            win1.clear()
            win1.box()
            if image_compare_thread.toggle_setting('use_hull'):
                win1.addstr(2, 2, "(H)ull active")
            else:
                win1.addstr(2, 2, "(H)ull inactive")
            win1.refresh()
        elif key == 'a':
            win2.clear()
            win2.box()
            if image_compare_thread.toggle_setting('use_and'):
                win2.addstr(2, 2, "(A)nd active")
            else:
                win2.addstr(2, 2, "(A)nd inactive")
            win2.refresh()
        elif key == 'g':
            green_win.clear()
            green_win.box()
            if masking_thread.toggle_setting('use_green'):
                green_win.addstr(2, 2, "(G)reen active")
            else:
                green_win.addstr(2, 2, "(G)reen inactive")
            green_win.refresh()
        elif key == 'm':
            model_win.clear()
            model_win.box()
            if image_compare_thread.toggle_setting('use_model'):
                model_win.addstr(2, 2, "(M)odel active")
            else:
                model_win.addstr(2, 2, "(M)odel inactive")
            model_win.refresh()
        elif key == 'q':
            break
        elif key == os.linesep:
            break

        if masking_thread.has_message():
            masking_message_win.clear()
            masking_message_win.box()
            masking_message_queue.put(masking_thread.get_message())
            messages = []
            for i in range(0, min(8, masking_message_queue.qsize())):
                message = masking_message_queue.get()
                masking_message_win.addstr(i + 1, 2, message)
                messages.append(message)
            masking_message_win.refresh()
            for message in messages:
                masking_message_queue.put(message)

        if image_compare_thread.has_message():
            compare_message_win.clear()
            compare_message_win.box()
            compare_message_queue.put(image_compare_thread.get_message())
            messages = []
            for i in range(0, min(8, compare_message_queue.qsize())):
                message = compare_message_queue.get()
                compare_message_win.addstr(i + 1, 2, message)
                messages.append(message)
            compare_message_win.refresh()
            for message in messages:
                compare_message_queue.put(message)

        # update fps counter once per second
        if datetime.now(utc) - last_fps_update > timedelta(seconds=1):
            frame_count = masking_thread.get_framecount()
            fps_win.clear()
            fps_win.box()
            fps_win.addstr(2, 2, f"{frame_count} fps")
            fps_win.refresh()
            last_fps_update = datetime.now(utc)

def main(argv=None):
    image_compare_thread = ImageCompareThread()
    image_compare_thread.start()

    masking_thread = MaskingThread(image_compare_thread)
    masking_thread.start()

    try:
        curses.wrapper(handle_input, image_compare_thread, masking_thread)
    except KeyboardInterrupt:
        print("Exiting...")
    image_compare_thread.join()
    masking_thread.join()

if __name__ == '__main__':
    sys.exit(main())


#!/usr/bin/env python3
#  Copyright 2016 Google Inc. All Rights Reserved.
#
# 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.

import random
import os

import fruit_source_generator
import boost_di_source_generator
import no_di_library_source_generator
from makefile_generator import generate_makefile
import argparse
import networkx as nx

def generate_injection_graph(num_components_with_no_deps,
                             num_components_with_deps,
                             num_deps):
    injection_graph = nx.DiGraph()

    num_used_ids = 0
    is_toplevel = [True for i in range(0, num_components_with_no_deps + num_components_with_deps)]
    toplevel_components = set()
    for i in range(0, num_components_with_no_deps):
        id = num_used_ids
        num_used_ids += 1
        toplevel_components.add(id)

    # Then the rest have num_deps deps, chosen (pseudo-)randomly from the previous components.
    # The last few components depend more components with >1 deps, so that the last component transitively depends on
    # everything.
    for i in range(0, num_components_with_deps):
        deps = set()

        if len(toplevel_components) > (num_components_with_deps - 1 - i) * (num_deps - 1):
            # We need at least 1 dep with deps, otherwise the last few components will not be enough
            # to tie together all components.
            num_deps_with_deps = len(toplevel_components) - (num_components_with_deps - 1 - i) * (num_deps - 1)
            deps |= set(random.sample(toplevel_components, num_deps_with_deps))

        # Add other deps to get to the desired num_deps.
        deps |= set(random.sample(range(0, num_components_with_no_deps + i), num_deps - len(deps)))

        toplevel_components -= deps
        for dep in deps:
            is_toplevel[dep] = False

        component_id = num_used_ids
        toplevel_components |= {component_id}
        num_used_ids += 1
        deps_list = list(deps)
        random.shuffle(deps_list)
        for dep in deps_list:
            injection_graph.add_edge(component_id, dep)

    assert len(toplevel_components) == 1, toplevel_components
    toplevel_component = num_used_ids - 1
    assert is_toplevel[toplevel_component]

    return injection_graph

def generate_benchmark(
        di_library,
        compiler,
        cxx_std,
        output_dir,
        num_components_with_no_deps,
        num_components_with_deps,
        num_deps,
        generate_runtime_bench_code,
        use_exceptions=True,
        use_rtti=True,
        fruit_build_dir=None,
        fruit_sources_dir=None,
        boost_di_sources_dir=None,
        generate_debuginfo=False,
        use_new_delete=False,
        use_interfaces=False,
        use_normalized_component=False):
    """Generates a sample codebase using the specified DI library, meant for benchmarking.

    :param boost_di_sources_dir: this is only used if di_library=='boost_di', it can be None otherwise.
    """

    if num_components_with_no_deps < num_deps:
        raise Exception(
            "Too few components with no deps. num_components_with_no_deps=%s but num_deps=%s." % (num_components_with_no_deps, num_deps))
    if num_deps < 2:
        raise Exception("num_deps should be at least 2.")

    # This is a constant so that we always generate the same file (=> benchmark more repeatable).
    random.seed(42)

    injection_graph = generate_injection_graph(num_components_with_no_deps=num_components_with_no_deps,
                                               num_components_with_deps=num_components_with_deps,
                                               num_deps=num_deps)

    if di_library == 'fruit':
        file_content_by_name = fruit_source_generator.generate_files(injection_graph, generate_runtime_bench_code)
        include_dirs = [fruit_build_dir + '/include', fruit_sources_dir + '/include']
        library_dirs = [fruit_build_dir + '/src']
        link_libraries = ['fruit']
    elif di_library == 'boost_di':
        file_content_by_name = boost_di_source_generator.generate_files(injection_graph, generate_runtime_bench_code)
        include_dirs = [boost_di_sources_dir + '/include', boost_di_sources_dir + '/extension/include']
        library_dirs = []
        link_libraries = []
    elif di_library == 'none':
        file_content_by_name = no_di_library_source_generator.generate_files(injection_graph, use_new_delete, use_interfaces, generate_runtime_bench_code)
        include_dirs = []
        library_dirs = []
        link_libraries = []
    else:
        raise Exception('Unrecognized di_library: %s' % di_library)

    include_flags = ' '.join(['-I%s' % include_dir for include_dir in include_dirs])
    library_dirs_flags = ' '.join(['-L%s' % library_dir for library_dir in library_dirs])
    rpath_flags = ' '.join(['-Wl,-rpath,%s' % library_dir for library_dir in library_dirs])
    link_libraries_flags = ' '.join(['-l%s' % library for library in link_libraries])
    other_compile_flags = []
    if generate_debuginfo:
        other_compile_flags.append('-g')
    if not use_exceptions:
        other_compile_flags.append('-fno-exceptions')
    if not use_rtti:
        other_compile_flags.append('-fno-rtti')
    compile_command = '%s -std=%s -MMD -MP -O2 -W -Wall -Werror -DNDEBUG -ftemplate-depth=10000 %s %s' % (compiler, cxx_std, include_flags, ' '.join(other_compile_flags))
    link_command = '%s -std=%s -O2 -W -Wall -Werror %s %s' % (compiler, cxx_std, rpath_flags, library_dirs_flags)
    # GCC requires passing the -lfruit flag *after* all object files to be linked for some reason.
    link_command_suffix = link_libraries_flags

    cpp_files = [file_name
                 for file_name in file_content_by_name.keys()
                 if file_name.endswith('.cpp')]

    file_content_by_name['Makefile'] = generate_makefile(cpp_files, 'main', compile_command, link_command, link_command_suffix)

    os.makedirs(output_dir, exist_ok=True)
    for file_name, file_content in file_content_by_name.items():
        with open('%s/%s' % (output_dir, file_name), 'w') as file:
            file.write(file_content)

    return file_content_by_name.keys()

def main():
    parser = argparse.ArgumentParser(description='Generates source files and a build script for benchmarks.')
    parser.add_argument('--di-library', default='fruit', help='DI library to use. One of {fruit, boost_di, none}. (default: fruit)')
    parser.add_argument('--compiler', help='Compiler to use')
    parser.add_argument('--fruit-sources-dir', help='Path to the fruit sources (only used when di_library==\'fruit\')')
    parser.add_argument('--fruit-build-dir', help='Path to the fruit build dir (only used with --di_library=\'fruit\')')
    parser.add_argument('--boost-di-sources-dir', help='Path to the Boost.DI sources (only used with --di-library==\'boost_di\')')
    parser.add_argument('--num-components-with-no-deps', default=10, help='Number of components with no deps that will be generated')
    parser.add_argument('--num-components-with-deps', default=90, help='Number of components with deps that will be generated')
    parser.add_argument('--num-deps', default=10, help='Number of deps in each component with deps that will be generated')
    parser.add_argument('--output-dir', help='Output directory for generated files')
    parser.add_argument('--cxx-std', default='c++11',
                        help='Version of the C++ standard to use. Typically one of \'c++11\' and \'c++14\'. (default: \'c++11\')')
    parser.add_argument('--use-new-delete', default='false', help='Set this to \'true\' to use new/delete. Only relevant when --di_library=none.')
    parser.add_argument('--use-interfaces', default='false', help='Set this to \'true\' to use interfaces. Only relevant when --di_library=none.')
    parser.add_argument('--use-normalized-component', default='false', help='Set this to \'true\' to create a NormalizedComponent and create the injector from that. Only relevant when --di_library=fruit and --generate-runtime-bench-code=false.')
    parser.add_argument('--generate-runtime-bench-code', default='true', help='Set this to \'false\' for compile benchmarks.')
    parser.add_argument('--generate-debuginfo', default='false', help='Set this to \'true\' to generate debugging information (-g).')
    parser.add_argument('--use-exceptions', default='true', help='Set this to \'false\' to disable exceptions.')
    parser.add_argument('--use-rtti', default='true', help='Set this to \'false\' to disable RTTI.')

    args = parser.parse_args()

    if args.compiler is None:
        raise Exception('--compiler is required.')

    if args.di_library == 'fruit':
        if args.fruit_sources_dir is None:
            raise Exception('--fruit-sources-dir is required with --di-library=\'fruit\'.')
        if args.fruit_build_dir is None:
            raise Exception('--fruit-build-dir is required with --di-library=\'fruit\'.')
    elif args.di_library == 'boost_di':
        if args.boost_di_sources_dir is None:
            raise Exception('--boost-di-sources-dir is required with --di-library=\'boost_di\'.')
    elif args.di_library == 'none':
        pass
    else:
        raise Exception('Unrecognized --di-library: \'%s\'. Allowed values are %s' % (args.di_library, {'fruit', 'boost_di', 'none'}))

    num_components_with_deps = int(args.num_components_with_deps)
    num_components_with_no_deps = int(args.num_components_with_no_deps)
    num_deps = int(args.num_deps)

    if args.output_dir is None:
        raise Exception("output_dir must be specified.")

    generate_benchmark(
        di_library=args.di_library,
        fruit_sources_dir=args.fruit_sources_dir,
        boost_di_sources_dir=args.boost_di_sources_dir,
        output_dir=args.output_dir,
        compiler=args.compiler,
        cxx_std=args.cxx_std,
        num_components_with_deps=num_components_with_deps,
        num_components_with_no_deps=num_components_with_no_deps,
        fruit_build_dir=args.fruit_build_dir,
        num_deps=num_deps,
        generate_debuginfo=(args.generate_debuginfo == 'true'),
        use_new_delete=(args.use_new_delete == 'true'),
        use_interfaces=(args.use_interfaces == 'true'),
        use_normalized_component=(args.use_normalized_component == 'true'),
        generate_runtime_bench_code=(args.generate_runtime_bench_code == 'true'),
        use_exceptions=(args.use_exceptions == 'true'),
        use_rtti=(args.use_rtti == 'true'))

if __name__ == "__main__":
    main()


# python detect_blinks.py --shape-predictor shape_predictor_68_face_landmarks.dat
from pynput.keyboard import Key, Controller
from scipy.spatial import distance as dist
from imutils.video import FileVideoStream
from imutils.video import VideoStream
from imutils import face_utils
import numpy as np
import argparse
import imutils
import time
import dlib
import cv2

def eye_aspect_ratio(eye):
	A = dist.euclidean(eye[1], eye[5])
	B = dist.euclidean(eye[2], eye[4])
	C = dist.euclidean(eye[0], eye[3])
	ear = (A + B) / (2.0 * C)
	return ear

# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-p", "--shape-predictor", required=True,
	help="path to facial landmark predictor")
ap.add_argument("-v", "--video", type=str, default="",
	help="path to input video file")
args = vars(ap.parse_args())

EYE_AR_THRESH = 0.18
EYE_AR_CONSEC_FRAMES = 3
COUNTER = 0
TOTAL = 0

print("[INFO] loading facial landmark...")
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor(args["shape_predictor"])

# grab the indexes of the facial landmarks for the left and
# right eye, respectively
(lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
(rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]

# start the video stream thread
print("[INFO] starting video stream thread...")
# vs = FileVideoStream(args["video"]).start()
vs = VideoStream(src=0).start()
time.sleep(2.0)

time.sleep(1.0)
keyboard = Controller()
# loop over frames from the video stream
while True:
	frame = vs.read()
	frame = imutils.resize(frame, width=450)
	gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

	# detect faces in the grayscale frame
	rects = detector(gray, 0)

	# loop over the face detections
	for rect in rects:
		shape = predictor(gray, rect)
		shape = face_utils.shape_to_np(shape)

		leftEye = shape[lStart:lEnd]
		rightEye = shape[rStart:rEnd]
		leftEAR = eye_aspect_ratio(leftEye)
		rightEAR = eye_aspect_ratio(rightEye)

		# average the eye aspect ratio together for both eyes
		ear = (leftEAR + rightEAR) / 2.0

		# compute the convex hull for the left and right eye, then
		# visualize each of the eyes
		leftEyeHull = cv2.convexHull(leftEye)
		rightEyeHull = cv2.convexHull(rightEye)
		cv2.drawContours(frame, [leftEyeHull], -1, (0, 255, 0), 1)
		cv2.drawContours(frame, [rightEyeHull], -1, (0, 255, 0), 1)

		# check to see if the eye aspect ratio is below the blink
		# threshold, and if so, increment the blink frame counter
		if ear < EYE_AR_THRESH:
			COUNTER += 1
			cv2.putText(frame, "Jump!".format(TOTAL), (10, 30),cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
			keyboard.press(Key.space)
			keyboard.release(Key.space)

		# otherwise, the eye aspect ratio is not below the blink
		# threshold
		else:
			# if the eyes were closed for a sufficient number of
			# then increment the total number of blinks
			if COUNTER >= EYE_AR_CONSEC_FRAMES:
				TOTAL += 1

			# reset the eye frame counter
			COUNTER = 0

	# show the frame
	cv2.imshow("Frame", frame)
	key = cv2.waitKey(1) & 0xFF

	# if the `q` key was pressed, break from the loop
	if key == ord("q"):
		break

# do a bit of cleanup
cv2.destroyAllWindows()
vs.stop()


# -*- coding: utf-8 -*-
# @Time    : 2017/7/12 上午11:16
# @Author  : play4fun
# @File    : 11.ipython.py
# @Software: PyCharm

"""
11.ipython.py:
"""

import cv2
import numpy as np
'''

In [10]: x = 5
In [11]: %timeit y=x**2
10000000 loops, best of 3: 73 ns per loop
In [12]: %timeit y=x*x
10000000 loops, best of 3: 58.3 ns per loop
In [15]: z = np.uint8([5])
In [17]: %timeit y=z*z
1000000 loops, best of 3: 1.25 us per loop
In [19]: %timeit y=np.square(z)
1000000 loops, best of 3: 1.16 us per loop
#Python 的标量运算比 Nump 的标量运算 快。
#对于仅包含一两个 元素的操作 Python 标量比 Numpy 的数组 快。
#但是当数组稍微大一点时 Numpy 就会胜出了
'''

'''
In [35]: %timeit z = cv2.countNonZero(img)
100000 loops, best of 3: 15.8 us per loop
In [36]: %timeit z = np.count_nonzero(img)
1000 loops, best of 3: 370 us per loop
#OpenCV 的函数是 Numpy 函数的 25 倍

#一般情况下 OpenCV 的函数 比 Numpy 函数快。所以对于相同的操
作最好使用 OpenCV 的函数。当然也有例外 尤其是当使用 Numpy 对 图  而 复制   操作时
'''

'''
有些技术和编程方法可以 我们最大的发挥 Python 和 Numpy 的威力。
我们  仅仅提一下相关的,你可以超链接查找更多 细信息。
我们 要说的最重要的一点是：
 先用简单的方式实现你的算法(结果正确最重要)
 当结果正确后 再使用上面的提到的方法找到程序的瓶 来优化它。
1. 尽量避免使用循环 尤其双层三层循环 它们天生就是 常慢的。
2. 算法中尽量使用向量操作 因为 Numpy 和 OpenCV  对向量操作了优化。
3. 利用高速缓存一致性。
4. 没有必要的话就不要复制数组。使用视图来代替复制。数组复制是非常浪费资源的。
'''


#!/usr/bin/env python3
import argparse
import serial
import select
import struct
import sys
import time
import math
import cv2
import numpy as np

# Set the serial port here (this is the name of the USB to serial adapter that shows up on your computer)
PORT = 'COM5'

ser = serial.Serial(port=PORT, baudrate=19200,timeout=1)

DEPTH_BUFFER = 10

STATE_OUT_OF_SYNC   = 0
STATE_SYNC_START    = 1
STATE_SYNC_1        = 2
STATE_SYNC_2        = 3
STATE_SYNC_OK       = 4

# Actual Switch DPAD Values
A_DPAD_CENTER    = 0x08
A_DPAD_U         = 0x00
A_DPAD_U_R       = 0x01
A_DPAD_R         = 0x02
A_DPAD_D_R       = 0x03
A_DPAD_D         = 0x04
A_DPAD_D_L       = 0x05
A_DPAD_L         = 0x06
A_DPAD_U_L       = 0x07

# Enum DIR Values
DIR_CENTER    = 0x00
DIR_U         = 0x01
DIR_R         = 0x02
DIR_D         = 0x04
DIR_L         = 0x08
DIR_U_R       = DIR_U + DIR_R
DIR_D_R       = DIR_D + DIR_R
DIR_U_L       = DIR_U + DIR_L
DIR_D_L       = DIR_D + DIR_L

BTN_NONE         = 0x0000000000000000
BTN_Y            = 0x0000000000000001
BTN_B            = 0x0000000000000002
BTN_A            = 0x0000000000000004
BTN_X            = 0x0000000000000008
BTN_L            = 0x0000000000000010
BTN_R            = 0x0000000000000020
BTN_ZL           = 0x0000000000000040
BTN_ZR           = 0x0000000000000080
BTN_MINUS        = 0x0000000000000100
BTN_PLUS         = 0x0000000000000200
BTN_LCLICK       = 0x0000000000000400
BTN_RCLICK       = 0x0000000000000800
BTN_HOME         = 0x0000000000001000
BTN_CAPTURE      = 0x0000000000002000

DPAD_CENTER      = 0x0000000000000000
DPAD_U           = 0x0000000000010000
DPAD_R           = 0x0000000000020000
DPAD_D           = 0x0000000000040000
DPAD_L           = 0x0000000000080000
DPAD_U_R         = DPAD_U + DPAD_R
DPAD_D_R         = DPAD_D + DPAD_R
DPAD_U_L         = DPAD_U + DPAD_L
DPAD_D_L         = DPAD_D + DPAD_L

LSTICK_CENTER    = 0x0000000000000000
LSTICK_R         = 0x00000000FF000000 #   0 (000)
LSTICK_U_R       = 0x0000002DFF000000 #  45 (02D)
LSTICK_U         = 0x0000005AFF000000 #  90 (05A)
LSTICK_U_L       = 0x00000087FF000000 # 135 (087)
LSTICK_L         = 0x000000B4FF000000 # 180 (0B4)
LSTICK_D_L       = 0x000000E1FF000000 # 225 (0E1)
LSTICK_D         = 0x0000010EFF000000 # 270 (10E)
LSTICK_D_R       = 0x0000013BFF000000 # 315 (13B)

RSTICK_CENTER    = 0x0000000000000000
RSTICK_R         = 0x000FF00000000000 #   0 (000)
RSTICK_U_R       = 0x02DFF00000000000 #  45 (02D)
RSTICK_U         = 0x05AFF00000000000 #  90 (05A)
RSTICK_U_L       = 0x087FF00000000000 # 135 (087)
RSTICK_L         = 0x0B4FF00000000000 # 180 (0B4)
RSTICK_D_L       = 0x0E1FF00000000000 # 225 (0E1)
RSTICK_D         = 0x10EFF00000000000 # 270 (10E)
RSTICK_D_R       = 0x13BFF00000000000 # 315 (13B)

NO_INPUT       = BTN_NONE + DPAD_CENTER + LSTICK_CENTER + RSTICK_CENTER

# Commands to send to MCU
COMMAND_NOP        = 0x00
COMMAND_SYNC_1     = 0x33
COMMAND_SYNC_2     = 0xCC
COMMAND_SYNC_START = 0xFF

# Responses from MCU
RESP_USB_ACK       = 0x90
RESP_UPDATE_ACK    = 0x91
RESP_UPDATE_NACK   = 0x92
RESP_SYNC_START    = 0xFF
RESP_SYNC_1        = 0xCC
RESP_SYNC_OK       = 0x33

# Enums for environments
tetris_enum = [DPAD_U, DPAD_R, DPAD_D, DPAD_L, BTN_A, BTN_X, BTN_R, BTN_NONE]

# Compute x and y based on angle and intensity
def angle(angle, intensity):
    # y is negative because on the Y input, UP = 0 and DOWN = 255
    x =  int((math.cos(math.radians(angle)) * 0x7F) * intensity / 0xFF) + 0x80
    y = -int((math.sin(math.radians(angle)) * 0x7F) * intensity / 0xFF) + 0x80
    return x, y

def lstick_angle(angle, intensity):
    return (intensity + (angle << 8)) << 24

def rstick_angle(angle, intensity):
    return (intensity + (angle << 8)) << 44

# Precision wait
def p_wait(waitTime):
    t0 = time.perf_counter()
    t1 = t0
    while (t1 - t0 < waitTime):
        t1 = time.perf_counter()

# Wait for data to be available on the serial port
def wait_for_data(timeout = 1.0, sleepTime = 0.1):
    t0 = time.perf_counter()
    t1 = t0
    inWaiting = ser.in_waiting
    while ((t1 - t0 < sleepTime) or (inWaiting == 0)):
        time.sleep(sleepTime)
        inWaiting = ser.in_waiting
        t1 = time.perf_counter()

# Read X bytes from the serial port (returns list)
def read_bytes(size):
    bytes_in = ser.read(size)
    return list(bytes_in)

# Read 1 byte from the serial port (returns int)
def read_byte():
    bytes_in = read_bytes(1)
    if len(bytes_in) != 0:
        byte_in = bytes_in[0]
    else:
        byte_in = 0
    return byte_in

# Discard all incoming bytes and read the last (latest) (returns int)
def read_byte_latest():
    inWaiting = ser.in_waiting
    if inWaiting == 0:
        inWaiting = 1
    bytes_in = read_bytes(inWaiting)
    if len(bytes_in) != 0:
        byte_in = bytes_in[0]
    else:
        byte_in = 0
    return byte_in

# Write bytes to the serial port
def write_bytes(bytes_out):
    ser.write(bytearray(bytes_out))
    return

# Write byte to the serial port
def write_byte(byte_out):
    write_bytes([byte_out])
    return

# Compute CRC8
# https://www.microchip.com/webdoc/AVRLibcReferenceManual/group__util__crc_1gab27eaaef6d7fd096bd7d57bf3f9ba083.html
def crc8_ccitt(old_crc, new_data):
    data = old_crc ^ new_data

    for i in range(8):
        if (data & 0x80) != 0:
            data = data << 1
            data = data ^ 0x07
        else:
            data = data << 1
        data = data & 0xff
    return data

# Send a raw packet and wait for a response (CRC will be added automatically)
def send_packet(packet=[0x00,0x00,0x08,0x80,0x80,0x80,0x80,0x00], debug=False):
    if not debug:
        bytes_out = []
        bytes_out.extend(packet)

        # Compute CRC
        crc = 0
        for d in packet:
            crc = crc8_ccitt(crc, d)
        bytes_out.append(crc)
        write_bytes(bytes_out)
        # print(bytes_out)

        # Wait for USB ACK or UPDATE NACK
        byte_in = read_byte()
        commandSuccess = (byte_in == RESP_USB_ACK)
    else:
        commandSuccess = True
    return commandSuccess

# Convert DPAD value to actual DPAD value used by Switch
def decrypt_dpad(dpad):
    if dpad == DIR_U:
        dpadDecrypt = A_DPAD_U
    elif dpad == DIR_R:
        dpadDecrypt = A_DPAD_R
    elif dpad == DIR_D:
        dpadDecrypt = A_DPAD_D
    elif dpad == DIR_L:
        dpadDecrypt = A_DPAD_L
    elif dpad == DIR_U_R:
        dpadDecrypt = A_DPAD_U_R
    elif dpad == DIR_U_L:
        dpadDecrypt = A_DPAD_U_L
    elif dpad == DIR_D_R:
        dpadDecrypt = A_DPAD_D_R
    elif dpad == DIR_D_L:
        dpadDecrypt = A_DPAD_D_L
    else:
        dpadDecrypt = A_DPAD_CENTER
    return dpadDecrypt

# Convert CMD to a packet
def cmd_to_packet(command):
    cmdCopy = command
    low              =  (cmdCopy & 0xFF)  ; cmdCopy = cmdCopy >>  8
    high             =  (cmdCopy & 0xFF)  ; cmdCopy = cmdCopy >>  8
    dpad             =  (cmdCopy & 0xFF)  ; cmdCopy = cmdCopy >>  8
    lstick_intensity =  (cmdCopy & 0xFF)  ; cmdCopy = cmdCopy >>  8
    lstick_angle     =  (cmdCopy & 0xFFF) ; cmdCopy = cmdCopy >> 12
    rstick_intensity =  (cmdCopy & 0xFF)  ; cmdCopy = cmdCopy >>  8
    rstick_angle     =  (cmdCopy & 0xFFF)
    dpad = decrypt_dpad(dpad)
    left_x, left_y   = angle(lstick_angle, lstick_intensity)
    right_x, right_y = angle(rstick_angle, rstick_intensity)

    packet = [high, low, dpad, left_x, left_y, right_x, right_y, 0x00]
    # print (hex(command), packet, lstick_angle, lstick_intensity, rstick_angle, rstick_intensity)
    return packet

# Send a formatted controller command to the MCU
def send_cmd(command=NO_INPUT):
    commandSuccess = send_packet(cmd_to_packet(command))
    return commandSuccess

#Test all buttons except for home and capture
def testbench_btn():
    send_cmd(BTN_A) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_B) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_X) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_Y) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_PLUS) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_MINUS) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_LCLICK) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(BTN_RCLICK) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)

# Test DPAD U / R / D / L
def testbench_dpad():
    send_cmd(DPAD_U) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(DPAD_R) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(DPAD_D) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(DPAD_L) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)

# Test DPAD Diagonals - Does not register on switch due to dpad buttons
def testbench_dpad_diag():
    send_cmd(DPAD_U_R) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(DPAD_D_R) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(DPAD_D_L) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(DPAD_U_L) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)

# Test Left Analog Stick
def testbench_lstick():
    #Test U/R/D/L
    send_cmd(BTN_LCLICK) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(LSTICK_U) ; p_wait(0.5)
    send_cmd(LSTICK_R) ; p_wait(0.5)
    send_cmd(LSTICK_D) ; p_wait(0.5)
    send_cmd(LSTICK_L) ; p_wait(0.5)
    send_cmd(LSTICK_U) ; p_wait(0.5)
    send_cmd(LSTICK_CENTER) ; p_wait(0.5)

    # 360 Circle @ Full Intensity
    for i in range(0,721):
        cmd = lstick_angle(i + 90, 0xFF)
        send_cmd(cmd)
        p_wait(0.001)
    send_cmd(LSTICK_CENTER) ; p_wait(0.5)

    # 360 Circle @ Partial Intensity
    for i in range(0,721):
        cmd = lstick_angle(i + 90, 0x80)
        send_cmd(cmd)
        p_wait(0.001)
    send_cmd(LSTICK_CENTER) ; p_wait(0.5)

# Test Right Analog Stick
def testbench_rstick():
    #Test U/R/D/L
    send_cmd(BTN_RCLICK) ; p_wait(0.5) ; send_cmd() ; p_wait(0.001)
    send_cmd(RSTICK_U) ; p_wait(0.5)
    send_cmd(RSTICK_R) ; p_wait(0.5)
    send_cmd(RSTICK_D) ; p_wait(0.5)
    send_cmd(RSTICK_L) ; p_wait(0.5)
    send_cmd(RSTICK_U) ; p_wait(0.5)
    send_cmd(RSTICK_CENTER) ; p_wait(0.5)

    # 360 Circle @ Full Intensity
    for i in range(0,721):
        cmd = rstick_angle(i + 90, 0xFF)
        send_cmd(cmd)
        p_wait(0.001)
    send_cmd(RSTICK_CENTER) ; p_wait(0.5)

    # 360 Circle @ Partial Intensity
    for i in range(0,721):
        cmd = rstick_angle(i + 90, 0x80)
        send_cmd(cmd)
        p_wait(0.001)
    send_cmd(RSTICK_CENTER) ; p_wait(0.5)

# Test Packet Speed
def testbench_packet_speed(count=100, debug=False):
    sum = 0
    min = 999
    max = 0
    avg = 0
    err = 0

    for i in range(0, count + 1):

        # Send packet and check time
        t0 = time.perf_counter()
        status = send_packet()
        t1 = time.perf_counter()

        # Count errors
        if not status:
            err += 1
            print('Packet Error!')

        # Compute times
        delta = t1 - t0
        if delta < min:
            min = delta
        if delta > max:
            max = delta
        sum = sum + (t1 - t0)

    avg = sum / i
    print('Min =', '{:.3f}'.format(min), 'Max =', '{:.3}'.format(max), 'Avg =', '{:.3f}'.format(avg), 'Errors =', err)

def testbench():
    testbench_btn()
    testbench_dpad()
    testbench_lstick()
    testbench_rstick()
    testbench_packet_speed()
    return

# Force MCU to sync
def force_sync():
    # Send 9x 0xFF's to fully flush out buffer on device
    # Device will send back 0xFF (RESP_SYNC_START) when it is ready to sync
    write_bytes([0xFF,0xFF,0xFF,0xFF,0xFF,0xFF,0xFF,0xFF,0xFF])

    # Wait for serial data and read the last byte sent
    wait_for_data()
    byte_in = read_byte_latest()

    # Begin sync...
    inSync = False
    if byte_in == RESP_SYNC_START:
        write_byte(COMMAND_SYNC_1)
        byte_in = read_byte()
        if byte_in == RESP_SYNC_1:
            write_byte(COMMAND_SYNC_2)
            byte_in = read_byte()
            if byte_in == RESP_SYNC_OK:
                inSync = True
    return inSync

# Start MCU syncing process
def sync():
    inSync = False

    # Try sending a packet
    inSync = send_packet()
    if not inSync:
        # Not in sync: force resync and send a packet
        inSync = force_sync()
        if inSync:
            inSync = send_packet()
    return inSync


import numpy
global ELEV#=[[0 for x in range(17)]for y in range(79)]
ELEV=[[0 for x in range(17)]for y in range(79)]
global NSDEG#[17]
global AA#[17]
global BB#[17]
global SCR,SCR1
#COMMON/MIXC/
global PRSH#(6,3,17,17)
global ESH#(6,3,17)
global AUG#(6,3,17,17,17)
global RAD#[6,3,17,17]
global PRSHBT#(6,3,17)
global IZ#[6,3]
global INIOCC#(6,3,17)
global ISHLMX#(6,3)
global AMZ#[6,3]
#COMMON/UPD/
global NOCC#(6,3,17)
global AUGR#(6,3,17,17,17)
global RADR#(6,3,17,17)
#COMMON/CALCAS2/
global IONSUM0#(10)
global IFLSUM0#(10)
global ESTORE0#(10,28)
global EPHOTON0#(10,28)
global DRXE0#(10,28)
global DRYE0#(10,28)
global DRZE0#(10,28)
global DRX0#(10,28)
global DRY0#(10,28)
global DRZ0#(10,28)
#COMMON/CALCAS3/
global IONSUM#(10)
IONSUM=[0 for x in range(11)]
global IFLSUM#(10)
IFLSUM=[0 for x in range(11)]
global ESTORE#(10,28)
global EPHOTON#(10,28)
global DRXE#(10,28)
global DRYE#(10,28)
global DRZE#(10,28)
global DRX#(10,28)
global DRY#(10,28)
global DRZ#[10,28]

def CALC3(NVAC,KGAS,LGAS,ELECEN,ISHELL,L1):
	# IMPLICIT #real*8(A-H,O-Z)
	# IMPLICIT #integer*8(I-N)
	# COMMON/GENCAS/ELEV(17,79),NSDEG(17),AA(17),BB(17),SCR,SCR1
 #      COMMON/MIXC/PRSH(6,3,17,17),ESH(6,3,17),AUG(6,3,17,17,17),
 #     /RAD(6,3,17,17),PRSHBT(6,3,17),IZ(6,3),INIOCC(6,3,17),ISHLMX(6,3),
 #     /AMZ(6,3)
 #      COMMON/UPD/NOCC(6,3,17),AUGR(6,3,17,17,17),RADR(6,3,17,17)
 #      COMMON/CALCAS2/IONSUM0(10),IFLSUM0(10),ESTORE0(10,28),
 #     /EPHOTON0(10,28),DRXE0(10,28),DRYE0(10,28),DRZE0(10,28),
 #     /DRX0(10,28),DRY0(10,28),DRZ0(10,28)
 #      COMMON/CALCAS3/IONSUM(10),IFLSUM(10),ESTORE(10,28),EPHOTON(10,28),
 #     /DRXE(10,28),DRYE(10,28),DRZE(10,28),DRX(10,28),DRY(10,28),
 #     /DRZ(10,28)
 #      DIMENSION TEMP(17),TEMP1(289)
	#CHARACTER*6
	# SCR="",
	# SCR1=""
	#COMMON/GENCAS/
	global ELEV#[17,79]
	global NSDEG#[17]
	global AA#[17]
	global BB#[17]
	global SCR,SCR1
	#COMMON/MIXC/
	global PRSH#(6,3,17,17)
	global ESH#(6,3,17)
	global AUG#(6,3,17,17,17)
	global RAD#[6,3,17,17]
	global PRSHBT#(6,3,17)
	global IZ#[6,3]
	global INIOCC#(6,3,17)
	global ISHLMX#(6,3)
	global AMZ#[6,3]
	#COMMON/UPD/
	global NOCC#(6,3,17)
	global AUGR#(6,3,17,17,17)
	global RADR#(6,3,17,17)
	#COMMON/CALCAS2/
	global IONSUM0#(10)
	global IFLSUM0#(10)
	global ESTORE0#(10,28)
	global EPHOTON0#(10,28)
	global DRXE0#(10,28)
	global DRYE0#(10,28)
	global DRZE0#(10,28)
	global DRX0#(10,28)
	global DRY0#(10,28)
	global DRZ0#(10,28)
	#COMMON/CALCAS3/
	global IONSUM#(10)
	global IFLSUM#(10)
	global ESTORE#(10,28)
	global EPHOTON#(10,28)
	global DRXE#(10,28)
	global DRYE#(10,28)
	global DRZE#(10,28)
	global DRX#(10,28)
	global DRY#(10,28)
	global DRZ#[10,28]
	TEMP=[0 for x in range(18)]
	TEMP1=[0 for x in range(289)]
	#
	# CALCULATE CASCADE IN GAS KGAS AND MOLECULAR COMPONENT LGAS
	# WITH INTIAL ENERGY DEPOSIT ELECEN AND SHELL VACANCY CREATED AT ISHELL
	#
	ISTART=IONSUM[NVAC]
	ISTARTF=IFLSUM[NVAC]
	API=numpy.arccos(-1.00)
	TWOPI=2.00*API
	def GOTO100():
		ELEFT=ELECEN
		INIT=1
		# SET STARTING ARRAY NOCC EQUAL TO INIOCC
		for I in range(1,17):
			NOCC[KGAS][LGAS][I]=INIOCC[KGAS][LGAS][I]
		IONSUM[NVAC]=ISTART+1
		IFLSUM[NVAC]=ISTARTF
		# STORE PHOTOELECTRON ENERGY AND ANGLE
		ESTORE[NVAC][IONSUM[NVAC]]=ELECEN-ELEV[ISHELL][IZ[KGAS][LGAS]]
		ELECN=ESTORE[NVAC][IONSUM[NVAC]]
		ELEFT=ELEFT-ELECN
		NOCC[KGAS][LGAS][ISHELL]=NOCC[KGAS][LGAS][ISHELL]-1
		# USE PHOTOELECTRON ANGULAR DISTRIBUTION
		APE=AA[ISHELL]
		BPE=BB[ISHELL]
		ANGGEN(APE,BPE,THET)
		if(THET < 0.0):
			THET=THET+API
		R3=DRAND48(RDUM)
		PHI=TWOPI*R3
		DRCOS(DRX0[NVAC][L1],DRY0[NVAC][L1],DRZ0[NVAC][L1],THET,PHI,DRXX,DRYY,DRZZ)
		DRXE[NVAC][IONSUM[NVAC]]=DRXX
		DRYE[NVAC][IONSUM[NVAC]]=DRYY
		DRZE[NVAC][IONSUM[NVAC]]=DRZZ
		# LOOP AROUND CASCADE
		def GOTO4():
			# CHECK FOR ELECTRON SHAKEOFF
			IDUM=1
			if(INIT > 1):
				ELECN=ESTORE[NVAC][IONSUM[NVAC]]
			INSUM=IONSUM[NVAC]
			SHAKE(ISHELL,ELECN,KGAS,LGAS,ESHK,IDUM,INSUM,JVAC)
			#  CALCULATE ENERGY OF ELECTRON
			if(JVAC == 0):
				GOTO2()
			#  ELECTRON + SHAKEOFF
			ELECN=ELECN-ESHK-ELEV[JVAC,IZ[KGAS,LGAS]]
			ESTORE[NVAC][IONSUM[NVAC]]=ELECN
			IONSUM[NVAC]=IONSUM[NVAC]+1
			# MAXIMUM ION CHARGE STATE =28
			if(IONSUM[NVAC]> 28) :
				print(' 3RD GEN ION CHARGE LIMITED TO 28  IONSUM=',IONSUM[NVAC])
				sys.exit()
			# endif
			ESTORE[NVAC][IONSUM[NVAC]]=ESHK
			ELEFT=ELEFT-ESHK-ELEV[JVAC][IZ[KGAS][LGAS]]
			if(ELEFT < 0.0):
				GOTO100()
			# RANDOM EMISSION ANGLE
			R3=DRAND48(RDUM)
			THET=numpy.arccos(1.0-2.0*R3)
			R4=DRAND48(RDUM)
			PHI=TWOPI*R4
			DRXE[NVAC][IONSUM[NVAC]]=numpy.sin(THET)*numpy.cos(PHI)
			DRYE[NVAC][IONSUM[NVAC]]=numpy.sin(THET)*numpy.sin(PHI)
			DRZE[NVAC][IONSUM[NVAC]]=numpy.cos(THET)
			def GOTO2():
				UPDATE(KGAS,LGAS,ISHELL)
				INIT=2
				# CHOOSE FLUORESCENCE OR AUGER TRANSITION
				TSUM=0.0
				for I in range(1,17):
					TSUM=TSUM+RADR[KGAS][LGAS][ISHELL][I]
					for J in range(1,17):
						TSUM=TSUM+AUGR[KGAS][LGAS][ISHELL][I][J]
				# NO MORE TRANSITIONS POSSIBLE
				if(TSUM == 0.0):
					return
				# NORMALISE TO 1.0
				for I in range(1,17):
					RADR[KGAS][LGAS][ISHELL][I]=RADR[KGAS][LGAS][ISHELL][I]/TSUM
					for J in range(1,17):
						AUGR[KGAS][LGAS][ISHELL][I][J]=AUGR[KGAS][LGAS][ISHELL][I][J]/TSUM
				# CREATE CUMULATIVE SUM ARRAY
				TEMP[1]=RADR(KGAS,LGAS,ISHELL,1)
				for I in range(2,17):
					TEMP[I]=RADR[KGAS][LGAS][ISHELL][I]+TEMP[I-1]
				TEMP1[1]=AUGR[KGAS][LGAS][ISHELL][1][1]
				for I in range(2,17):
					TEMP1[I]=AUGR[KGAS][LGAS][ISHELL][I][1]+TEMP1[I-1]
				for J in range(1,16):
					for I in range(1,17):
						TEMP1[I+(J*17)]=AUGR[KGAS][LGAS][ISHELL][I][(J+1)]+TEMP1[I+(J*17)-1]
				# FIND FLUORESCENCE OR AUGER TRANSITION
				R1=DRAND48(RDUM)
				for I in range(1,17):
					if(R1 < TEMP[I]) :
						# STORE PHOTON ENERGY AND UPDATE NOCC
						IFLSUM[NVAC]=IFLSUM[NVAC]+1
						EPHOTON[NVAC][IFLSUM[NVAC]]=ELEV[ISHELL,IZ[KGAS,LGAS]]-ELEV[I,IZ[KGAS,LGAS]]
						ELEFT=ELEFT-EPHOTON[NVAC][IFLSUM[NVAC]]
						if(ELEFT < 0.0):
							GOTO100()
						# RANDOM EMISSION DIRECTION
						R3=DRAND48(RDUM)
						THET=numpy.arccos(1.0-2.0*R3)
						R4=DRAND48(RDUM)
						PHI=TWOPI*R4
						DRX[NVAC][IFLSUM[NVAC]]=numpy.sin(THET)*numpy.cos(PHI)
						DRY[NVAC][IFLSUM[NVAC]]=numpy.sin(THET)*numpy.sin(PHI)
						DRZ[NVAC][IFLSUM[NVAC]]=numpy.cos(THET)
						NOCC[KGAS][LGAS][ISHELL]=NOCC[KGAS][LGAS][ISHELL]+1
						NOCC[KGAS][LGAS][I]=NOCC[KGAS][LGAS][I]-1
						# FIND LOWEST VACANCY
						VACANCY(KGAS,LGAS,ISHELL,ILAST)
						if(ILAST == 1):
							# NO MORE TRANSITIONS POSSIBLE
							return
						# endif
						GOTO2()
					# endif
				CONTINUE
				counter1=1
				counter2=1
				while(counter1):
					counter1=0
					R2=R1-TEMP[17]
					for J in range(1,17):
						if(counter1):
							break
						for I in range(1,17):
							if(R2 < TEMP1[I+((J-1)*17)]) :
								# AUGER OR COSTER KRONIG
								# STORE EJECTED ELECTRON AND UPDATE NOCC
								ETEMP=ELEV[ISHELL,IZ[KGAS,LGAS]]-(ELEV[I,IZ[KGAS,LGAS]]+ELEV[I,IZ[KGAS,LGAS]+1])*0.5-(ELEV[J,IZ[KGAS,LGAS]]+ELEV[J,IZ[KGAS,LGAS]+1])*0.5
								if(ETEMP < 0.0):
									# DO NOT ALLOW NEGATIVE ENERGY TRANSITIONS
									while(counter2):
										R1=DRAND48(RDUM)
										if(R1 < TEMP[17]):
											counter2=1
									counter1=1
									break
									# endif
								IONSUM[NVAC]=IONSUM[NVAC]+1
								if(IONSUM[NVAC]> 28) :
									print(' 3RD GEN ION CHARGE LIMITED TO 28  IONSUM=', IONSUM[NVAC])
									sys.exit()
								# endif
								ESTORE[NVAC][IONSUM[NVAC]]=ETEMP
								ELEFT=ELEFT-ETEMP
								if(ELEFT < 0.0):
									GOTO100()
								# RANDOM EMISSION DIRECTION
								R3=DRAND48(RDUM)
								THET=numpy.arccos(1.0-2.0*R3)
								R4=DRAND48(RDUM)
								PHI=TWOPI*R4
								DRXE[NVAC][IONSUM[NVAC]]=numpy.sin(THET)*numpy.cos(PHI)
								DRYE[NVAC][IONSUM[NVAC]]=numpy.sin(THET)*numpy.sin(PHI)
								DRZE[NVAC][IONSUM[NVAC]]=numpy.cos(THET)
								NOCC[KGAS][LGAS][ISHELL]=NOCC[KGAS][LGAS][ISHELL]+1
								NOCC[KGAS][LGAS][I]=NOCC[KGAS][LGAS][I]-1
								NOCC[KGAS][LGAS][J]=NOCC[KGAS][LGAS][J]-1
								# FIND LOWEST VACANCY
								VACANCY(KGAS,LGAS,ISHELL,ILAST)
								if(ILAST == 1):
									# NO MORE TRANSITIONS POSSIBLE
									returnSCR="",
	# SCR1=""
								# endif
							GOTO4()
							# endif
	print(' ERROR IN CASCADE 3')
	sys.exit()
# end

CALC3(1,1,1,1,1,1)


import glob,os
from operator import itemgetter
from numpy import argmax
from collections import defaultdict
from LSA import LSA

class Cluster_Analysis(LSA):

    def __init__(self,inputpath,outputpath,assembly_suffix='_velvet/',assembly_norm=2500000):
        super(Cluster_Analysis,self).__init__(inputpath,outputpath)
        self.assembly_suffix = assembly_suffix
        self.assembly_norm = assembly_norm
        # DUMB:
        self.read_length = 100

    def process_cluster(self,cluster_prefix):
        total_reads = 0
        singletons = []
        pair_channels = [[] for _ in range(self.num_libs)]
        for sample,lib in self.sample_library.items():
            try:
                sample_reads = self.sort_read_pairs(cluster_prefix+'.'+sample)
            except:
                sample_reads = 0
            if sample_reads > 0:
                singletons.append(self.input_path+'.'.join([cluster_prefix,sample,'singleton','fastq']))
                pair_channels[lib].append(self.input_path+'.'.join([cluster_prefix,sample,'pairs','fastq']))
                total_reads += sample_reads
        pair_channels = [pc for pc in pair_channels if len(pc)>0]
        ec = total_reads*self.read_length/self.assembly_norm
        self.assembly_loc = self.output_path+cluster_prefix+self.assembly_suffix
        self.assemble_cluster(cluster_prefix,singletons,pair_channels,exp_cov=max(ec,10),cov_cutoff=max(ec/15,2))
        h = []
        for i in range(29,61,10):
            try:
                L = open(self.assembly_loc+'_'+str(i)+'/Log').readlines()
                h.append(int(L[-1].split(',')[0].split()[-1]))
            except:
                h.append(0)
        amax = argmax(h)
        os.system('cp '+self.assembly_loc+'_'+str(range(29,61,10)[amax])+'/contigs.fa '+self.assembly_loc+'contigs.fa')
        self.align_assembly(cluster_prefix)
        self.phyler_assembly(cluster_prefix)

    def assemble_cluster(self,cluster_prefix,sg,pc,exp_cov=50,cov_cutoff=15,min_contig=500):
        os.system('mkdir '+self.assembly_loc)
        i = 1
        s = []
        for c in pc:
            s.append('-shortPaired')
            if i > 1:
                s[-1] += str(i)
            s += c
            i += 1
        os.system('/home/unix/bcleary/src/velvet/velveth '+self.assembly_loc+' 29,61,10 -fastq -short '+' '.join(sg)+' '+' '.join(s))
        for i in range(29,61,10):
            os.system('/home/unix/bcleary/src/velvet/velvetg '+self.assembly_loc+'_'+str(i)+'/ -exp_cov '+str(exp_cov)+' -cov_cutoff '+str(cov_cutoff)+' -min_contig_lgth '+str(min_contig))

    def assemble_cluster_idba(self,cluster_prefix,min_contig=500):
        read_loc = self.input_path+cluster_prefix
        os.system('/bli/bin/Linux/x86_64/idba_ud-1.0.9/bin/fq2fa '+read_loc+'.pairs.fastq '+read_loc+'.pairs.fasta')
        os.system('mkdir '+self.assembly_loc)
        os.system('/bli/bin/Linux/x86_64/idba_ud-1.0.9/bin/idba_ud -r '+read_loc+'.pairs.fasta -o '+self.assembly_loc+' --min_contig '+str(min_contig)+' --pre_correction')

    def align_assembly(self,cluster_prefix):
        os.system('/bli/bin/Linux/x86_64/ncbi-blast-2.2.27/ncbi-blast-2.2.27+/bin/blastn -query '+self.assembly_loc+'contigs.fa -db /import/pool2/projects/mega/search_dbs/prod/BLIS_SEQUENCES/blast/Bacterial_DNA_All_Prokaryotic -task megablast -outfmt "7 qseqid qstart qend sseqid score" -out '+self.output_path+cluster_prefix+'.alignments.txt')
        os.system('/bli/bin/Linux/x86_64/ncbi-blast-2.2.27/ncbi-blast-2.2.27+/bin/blastn -query '+self.assembly_loc+'contigs.fa -db /import/pool2/projects/mega/search_dbs/prod/BLIS_SEQUENCES/blast/Viral_DNA_All_GenBank_Virus_Sequences -task megablast -outfmt "7 qseqid qstart qend sseqid score" -out '+self.output_path+cluster_prefix+'.viral.alignments.txt')

    def phyler_assembly(self,cluster_prefix):
        os.system('/home/unix/bcleary/src/MetaPhylerV1.25/runMetaphyler.pl '+self.assembly_loc+'contigs.fa blastn '+self.output_path+cluster_prefix+'.phyler.blastn 4')

    def assembly_stats(self,cluster_prefix,file_suffix='contigs.fa'):
        AS = {}
        L = []
        f = open(self.assembly_loc+file_suffix)
        line = f.readline()
        while line != '':
            if line[0] == '>':
                L.append(int(line.replace('_',' ').split()[3]))
            line = f.readline()
        f.close()
        if L:
            L.sort(reverse=True)
            t = 0
            for l in L:
                t += l
                if t > sum(L)/2:
                    break
            AS = {'total bp': sum(L),'num scaffolds': len(L),'N50': l,'largest contig': L[0]}
        return AS

    def alignment_stats(self,cluster_prefix):
        f = open(self.output_path+cluster_prefix+'.alignments.txt')
        L = f.readlines()
        f.close()
        A = defaultdict(int)
        current_alignments = {}
        for l in L:
            lst = l.strip().split('\t')
            if len(lst) != 5:
                for k,v in current_alignments.items():
                    A[k[0]] += v
                current_alignments = {}
            else:
                alignid_querypos = (lst[3],(lst[1],lst[2]))
                if alignid_querypos not in current_alignments:
                    current_alignments[alignid_querypos] = int(lst[4])
        return A

    def phyler_stats(self,cluster_prefix,level='genus',type='blastn'):
        L = open(self.output_path+cluster_prefix+'.'.join(['.phyler',type,level,'taxprof'])).readlines()
        A = []
        for l in L[1:]:
            x = l.strip().split()
            A.append((x[0],float(x[1]),int(x[2])))
        return A

    def sort_read_pairs(self,cluster_prefix):
        try:
            f = open(self.input_path+cluster_prefix+'.fastq','r')
        except:
            return 0
        type = self.id_type(f)
        # this is kind of a bummer since files are *mostly* sorted already
        if type == 1:
            sorted_reads = sorted(self.read_generator(f,raw_reads=True,max_reads=10**7),key=lambda (d): d[:d.index(' ')+2])
            def get_id(r):
                return r[:r.index(' ')+2]
        elif type == 2:
            sorted_reads = sorted(self.read_generator(f,raw_reads=True,max_reads=10**7),key=lambda (d): d.split()[0])
            def get_id(r):
                return r.split()[0]
        else:
            sorted_reads = sorted(self.read_generator(f,raw_reads=True,max_reads=10**7),key=lambda (d): d.split()[0])
            def get_id(r):
                return r.split()[0]+'*'
        if len(sorted_reads) > 0:
            pair_file1 = open(self.input_path+cluster_prefix+'.mate1.fastq','w')
            pair_file2 = open(self.input_path+cluster_prefix+'.mate2.fastq','w')
            singleton_file = open(self.input_path+cluster_prefix+'.singleton.fastq','w')
            total_reads = 0
            while len(sorted_reads) > 0:
                current_id = ''
                last = ''
                for r in sorted_reads:
                    r_id = get_id(r)
                    if (r_id[:-1] == current_id[:-1]) and (r_id != current_id):
                        pair_file1.write(last)
                        pair_file2.write(r)
                        total_reads += 2
                        last = ''
                    else:
                        singleton_file.write(last)
                        total_reads += 1
                        last = r
                    current_id = r_id
                if type == 1:
                    sorted_reads = sorted(self.read_generator(f,raw_reads=True,max_reads=10**7),key=lambda (d): d[:d.index(' ')+2])
                elif type == 2:
                    sorted_reads = sorted(self.read_generator(f,raw_reads=True,max_reads=10**7),key=lambda (d): d.split()[0])
                else:
                    sorted_reads = sorted(self.read_generator(f,raw_reads=True,max_reads=10**7),key=lambda (d): d.split()[0])
            pair_file1.close()
            pair_file2.close()
            singleton_file.close()
            f.close()
            os.system('rm '+self.input_path+cluster_prefix+'.fastq')
            os.system('touch '+self.input_path+cluster_prefix+'.fastq')
            return total_reads
        else:
            f.close()
            try:
                r = sum(self.read_counts(cluster_prefix).values())
            except:
                r = 0
            return r

    def read_counts(self,cluster_prefix):
        RC = {}
        RC['paired'] = self.read_count(self.input_path+cluster_prefix+'.mate1.fastq',startchar='@')*2
        RC['singleton'] = self.read_count(self.input_path+cluster_prefix+'.singleton.fastq',startchar='@')
        return RC

    def read_count(self,fp,startchar='>'):
        f = open(fp,'r')
        r = 0
        line = f.readline()
        while line != '':
            if line[0] == startchar:
                r += 1
            line = f.readline()
        f.close()
        return r