import numpy as np
import os
from skimage.io import imread, imsave
from skimage.transform import estimate_transform, warp
from time import time
import cv2 as cv2
from .predictor import PosPrediction
from skimage.transform import rescale, resize

from ..utils.estimate_pose import estimate_pose
from ..utils.rotate_vertices import frontalize
from ..utils.render_app import get_visibility, get_uv_mask, get_depth_image
from ..utils.write import write_obj_with_colors, write_obj_with_texture
from ..utils.cv_plot import plot_vertices,plot_pose_box

class PRN:
    ''' Joint 3D Face Reconstruction and Dense Alignment with Position Map Regression Network
    Args:
        is_dlib(bool, optional): If true, dlib is used for detecting faces.
        prefix(str, optional): If run at another folder, the absolute path is needed to load the data.
    '''
    def __init__(self, prefix = '.'):

        # resolution of input and output image size.
        self.resolution_inp = 256
        self.resolution_op = 256

        #---- load PRN
        self.pos_predictor = PosPrediction(self.resolution_inp, self.resolution_op)
        prn_path = os.path.join(prefix, 'net-data/256_256_resfcn256_weight')
        if not os.path.isfile(prn_path + '.data-00000-of-00001'):
            print("please download PRN trained model first.")
            exit()
        self.pos_predictor.restore(prn_path)

        # uv file
        self.uv_kpt_ind = np.loadtxt(prefix + '/uv-data/uv_kpt_ind.txt').astype(np.int32) # 2 x 68 get kpt
        self.face_ind = np.loadtxt(prefix + '/uv-data/face_ind.txt').astype(np.int32) # get valid vertices in the pos map
        self.triangles = np.loadtxt(prefix + '/uv-data/triangles.txt').astype(np.int32) # ntri x 3
        self.canonical_vertices = np.load(prefix + '/uv-data/canonical_vertices.npy')

        self.uv_coords = self.generate_uv_coords()

    def generate_uv_coords(self):
        resolution = self.resolution_op
        uv_coords = np.meshgrid(range(resolution),range(resolution))
        uv_coords = np.transpose(np.array(uv_coords), [1,2,0])
        uv_coords = np.reshape(uv_coords, [resolution**2, -1])
        uv_coords = uv_coords[self.face_ind, :]
        uv_coords = np.hstack((uv_coords[:,:2], np.zeros([uv_coords.shape[0], 1])))
        return uv_coords

    def net_forward(self, image):
        ''' The core of out method: regress the position map of a given image.
        Args:
            image: (256,256,3) array. value range: 0~1
        Returns:
            pos: the 3D position map. (256, 256, 3) array.
        '''
        return self.pos_predictor.predict(image)

    def processImg(self, image, recs):
        pos = self.get_pos(image,recs)
        kpt_3d = self.get_landmarks(pos)
        kpt = kpt_3d[:,:2]
        return kpt

    def process3DFile(self,image,recs,depth=False, pose = False,name='default'):
        img_show ={}
        img_inf = {}

        [h, w, c] = image.shape
        pos,vertices,save_color,img_vertices_show = self.get_3D(image,recs)
        kpt_3d = self.get_landmarks(pos)
        img_show["vertices"] = img_vertices_show
        img_inf["vertices"] = vertices
        img_inf["landmarks_3d"] = kpt_3d
        img_inf["color"] = save_color
        img_inf["triangles"] = self.triangles

        kpt = self.get_landmarks(pos)[:,:2]

        if depth:
            depth_image_show ,depth_inf= self.get_depth(vertices,h,w)
            img_show["depth"] = depth_image_show
            img_inf["depth"] = depth_inf
        if pose:
            camera_matrix, pose, image_pose_show = self.get_pose(image,vertices,kpt)
            img_show["pose"] = image_pose_show
            img_inf["pose"] = pose

        return img_show,img_inf

    def get_3D(self,image,recs):
        pos = self.get_pos(image, recs)
        vertices = self.get_vertices(pos)
        colors = self.get_colors(image, vertices)
        image_show = plot_vertices(image, vertices)

        return pos,vertices,colors,image_show

    def get_depth(self,vertices,h,w):
        depth_inf = get_depth_image(vertices, self.triangles, h, w, True)
        depth_image = np.uint8((depth_inf/np.max(depth_inf))*255)
        return depth_image,depth_inf

    def get_pose(self,image,vertices,kpt):
        camera_matrix, pose = estimate_pose(vertices,self.canonical_vertices)
        image_pose = plot_pose_box(image, camera_matrix, kpt,color=(0,255,0))
        return camera_matrix, pose,image_pose

    def get_pos(self, input, image_info = None):
        ''' process image with crop operation.
        Args:
            input: (h,w,3) array or str(image path). image value range:1~255.
            image_info(optional): the bounding box information of faces. if None, will use dlib to detect face.

        Returns:
            pos: the 3D position map. (256, 256, 3).
        '''
        if isinstance(input, str):
            try:
                image = imread(input)
            except IOError:
                print("error opening file: ", input)
                return None
        else:
            image = input

        if image.ndim < 3:
            image = np.tile(image[:,:,np.newaxis], [1,1,3])

        if image_info is not None:
            if np.max(image_info.shape) > 4: # key points to get bounding box
                kpt = image_info
                if kpt.shape[0] > 3:
                    kpt = kpt.T
                left = np.min(kpt[0, :]); right = np.max(kpt[0, :])
                top = np.min(kpt[1,:]); bottom = np.max(kpt[1,:])
            else:  # bounding box
                bbox = image_info
                left = bbox[0]; right = bbox[2]; top = bbox[1]; bottom = bbox[3]
            old_size = (right - left + bottom - top)/2
            center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0+ old_size*0.14])
            size = int(old_size*1.42)
        else:
            print('warning: no detected face')
            return None

        # crop image
        src_pts = np.array([[center[0]-size/2, center[1]-size/2], [center[0] - size/2, center[1]+size/2], [center[0]+size/2, center[1]-size/2]])
        DST_PTS = np.array([[0,0], [0,self.resolution_inp - 1], [self.resolution_inp - 1, 0]])
        tform = estimate_transform('similarity', src_pts, DST_PTS)

        image = image / 255.
        cropped_image = warp(image, tform.inverse, output_shape=(self.resolution_inp, self.resolution_inp))
        # run our net
        #st = time()
        cropped_pos = self.net_forward(cropped_image)
        #print 'net time:', time() - st

        # restore
        cropped_vertices = np.reshape(cropped_pos, [-1, 3]).T
        z = cropped_vertices[2,:].copy()/tform.params[0,0]
        cropped_vertices[2,:] = 1
        vertices = np.dot(np.linalg.inv(tform.params), cropped_vertices)
        vertices = np.vstack((vertices[:2,:], z))
        pos = np.reshape(vertices.T, [self.resolution_op, self.resolution_op, 3])

        return pos

    def get_landmarks(self, pos):
        '''
        Args:
            pos: the 3D position map. shape = (256, 256, 3).
        Returns:
            kpt: 68 3D landmarks. shape = (68, 3).
        '''
        kpt = pos[self.uv_kpt_ind[1,:], self.uv_kpt_ind[0,:], :]
        return kpt

    def get_vertices(self, pos):
        '''
        Args:
            pos: the 3D position map. shape = (256, 256, 3).
        Returns:
            vertices: the vertices(point cloud). shape = (num of points, 3). n is about 40K here.
        '''
        all_vertices = np.reshape(pos, [self.resolution_op**2, -1])
        vertices = all_vertices[self.face_ind, :]

        return vertices

    def get_colors_from_texture(self, texture):
        '''
        Args:
            texture: the texture map. shape = (256, 256, 3).
        Returns:
            colors: the corresponding colors of vertices. shape = (num of points, 3). n is 45128 here.
        '''
        all_colors = np.reshape(texture, [self.resolution_op**2, -1])
        colors = all_colors[self.face_ind, :]

        return colors

    def get_colors(self, image, vertices):
        '''
        Args:
            pos: the 3D position map. shape = (256, 256, 3).
        Returns:
            colors: the corresponding colors of vertices. shape = (num of points, 3). n is 45128 here.
        '''
        [h, w, _] = image.shape
        vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1)  # x
        vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1)  # y
        ind = np.round(vertices).astype(np.int32)

        colors = image[ind[:,1], ind[:,0], :] # n x 3
        colors = colors[:, [2, 1, 0]]
        return colors


import importlib
import json
import numpy as np
import os
import sys
import time
import uuid
import copy

from hypergan.discriminators import *
from hypergan.distributions import *
from hypergan.generators import *
from hypergan.inputs import *
from hypergan.samplers import *
from hypergan.trainers import *

import hyperchamber as hc
from hyperchamber import Config
from hypergan.ops import TensorflowOps
import hypergan as hg

from hypergan.gan_component import ValidationException, GANComponent
from .base_gan import BaseGAN

from hypergan.distributions.uniform_distribution import UniformDistribution

class AlignedAliGANTest(BaseGAN):
    """
    """
    def __init__(self, *args, **kwargs):
        BaseGAN.__init__(self, *args, **kwargs)

    def required(self):
        """
        `input_encoder` is a discriminator.  It encodes X into Z
        `discriminator` is a standard discriminator.  It measures X, reconstruction of X, and G.
        `generator` produces two samples, input_encoder output and a known random distribution.
        """
        return "generator discriminator ".split()

    def create(self):
        config = self.config
        ops = self.ops

        with tf.device(self.device):
            #x_input = tf.identity(self.inputs.x, name='input')
            xa_input = tf.identity(self.inputs.xa, name='xa_i')
            xb_input = tf.identity(self.inputs.xb, name='xb_i')

            ue = UniformDistribution(self, config.latent)
            ue2 = UniformDistribution(self, config.latent)
            ue3 = UniformDistribution(self, config.latent)
            ue4 = UniformDistribution(self, config.latent)
            if config.same_g:
                gb = self.create_component(config.generator, input=ue3.sample, name='a_generator')
                gb = self.create_component(config.generator, input=ue4.sample, name='a_generator', reuse=True)
            else:
                def _append_xa(gan, config, net):
                    _x = gan.inputs.xa
                    s = [int(x) for x in net.get_shape()]
                    shape = [s[1], s[2]]
                    _x = tf.image.resize_images(_x, shape, 1)
                    return _x
                def _append_xb(gan, config, net):
                    _x = gan.inputs.xb
                    s = [int(x) for x in net.get_shape()]
                    shape = [s[1], s[2]]
                    _x = tf.image.resize_images(_x, shape, 1)
                    return _x

                #config.generator['layer_filter'] = _append_xb
                #gb = self.create_component(config.generator, input=ue3.sample, name='a_generator')
                config.generator['layer_filter'] = _append_xa
                gb = self.create_component(config.generator, input=ue4.sample, name='b_generator')

            za = gb.controls["z"]
            zb = gb.controls["z"]

            self.uniform_sample = gb.sample

            zua = ue.sample
            zub = ue2.sample

            xa = xa_input
            xb = xb_input

            t0 = tf.concat([xb, xb], axis=3)
            t1 = tf.concat([gb.sample, xb], axis=3)
            t0 = xb
            t1 = gb.sample
            f0 = za
            f1 = zb
            f2 = zub
            stack = [t0, t1]
            stacked = ops.concat(stack, axis=0)
            features = ops.concat([f0, f1], axis=0)

            config.discriminator['layer_filter'] = _append_xa
            d = self.create_component(config.discriminator, name='d_ab', input=stacked)

            self.za = za
            self.discriminator = d
            l = self.create_loss(config.loss, d, xb_input, gb.sample, len(stack))
            loss1 = l
            d1_lambda = config.d1_lambda
            d2_lambda = config.d2_lambda
            d_loss1 = d1_lambda * l.d_loss
            g_loss1 = d1_lambda * l.g_loss

            d_vars1 = d.variables()

            d_loss = l.d_loss
            g_loss = l.g_loss

            metrics = {
                    'g_loss': l.g_loss,
                    'd_loss': l.d_loss
                }

            if d2_lambda > -1:
                t0 = xa
                t1 = gb.sample
                f0 = zb
                f1 = za
                stack = [t0, t1]
                stacked = ops.concat(stack, axis=0)
                features = ops.concat([f0, f1], axis=0)

                config.discriminator['layer_filter'] = _append_xa
                d = self.create_component(config.discriminator, name='d2_ab', input=stacked, features=[features])
                self.discriminator2 = d
                l = self.create_loss(config.loss, d, xb_input, gb.sample, len(stack))

                d_vars2 = d.variables()
                metrics["gb_gloss"]=l.g_loss
                metrics["gb_dloss"]=l.d_loss
                loss2=l
                g_loss2 = d2_lambda * loss2.g_loss
                d_loss2 = d2_lambda * loss2.d_loss
            else:
                d_loss2 = 0.0
                g_loss2 = 0.0
                d_vars2 = []

            if(config.alpha):
                t0 = zub
                t1 = zb
                netzd = tf.concat(axis=0, values=[t0,t1])
                z_d = self.create_component(config.z_discriminator, name='z_discriminator', input=netzd)

                loss3 = self.create_component(config.loss, discriminator = z_d, x=xa_input, generator=gb, split=2)
                d_vars1 += z_d.variables()
                metrics["za_gloss"]=loss3.g_loss
                metrics["za_dloss"]=loss3.d_loss
                d_loss1 += loss3.d_loss
                g_loss1 += loss3.g_loss

            if(config.ug):
                t0 = xb
                t1 = ugb
                netzd = tf.concat(axis=0, values=[t0,t1])
                z_d = self.create_component(config.discriminator, name='ug_discriminator', input=netzd)

                loss3 = self.create_component(config.loss, discriminator = z_d, x=xa_input, generator=gb, split=2)
                d_vars1 += z_d.variables()
                metrics["za_gloss"]=loss3.g_loss
                metrics["za_dloss"]=loss3.d_loss
                d_loss1 += loss3.d_loss
                g_loss1 += loss3.g_loss

            if(config.alpha2):
                t0 = zua
                t1 = za
                netzd = tf.concat(axis=0, values=[t0,t1])
                z_d = self.create_component(config.z_discriminator, name='z_discriminator2', input=netzd)

                loss3 = self.create_component(config.loss, discriminator = z_d, x=xa_input, generator=gb, split=2)
                d_vars2 += z_d.variables()
                metrics["za_gloss"]=loss3.g_loss
                metrics["za_dloss"]=loss3.d_loss
                d_loss2 += loss3.d_loss
                g_loss2 += loss3.g_loss

            if config.ug2:
                t0 = xa
                t1 = ugb
                netzd = tf.concat(axis=0, values=[t0,t1])
                z_d = self.create_component(config.discriminator, name='ug_discriminator2', input=netzd)

                loss3 = self.create_component(config.loss, discriminator = z_d, x=xa_input, generator=gb, split=2)
                d_vars2 += z_d.variables()
                metrics["za_gloss"]=loss3.g_loss
                metrics["za_dloss"]=loss3.d_loss
                d_loss2 += loss3.d_loss
                g_loss2 += loss3.g_loss

            loss = hc.Config({
                'd_fake':loss1.d_fake,
                'd_real':loss1.d_real,
                'sample': [tf.add_n([d_loss1, d_loss2]), tf.add_n([g_loss1, g_loss2])]
            })
            self._g_vars = gb.variables() + gb.variables()
            self._d_vars = d_vars1 + d_vars2
            self.loss=loss
            self.generator = gb
            trainer = self.create_component(config.trainer)
            self.initialize_variables()

        self.trainer = trainer
        self.latent = hc.Config({'sample':zb})
        self.generator = gb
        self.encoder = gb # this is the other gan
        self.uniform_distribution = hc.Config({"sample":zb})#uniform_encoder
        self.zb = zb
        self.z_hat = gb.sample
        self.x_input = xa_input

        self.uga = gb.sample
        self.ugb = gb.sample
        self.cyca = gb.sample
        self.cycb = gb.sample
        self.xba = gb.sample
        self.xab = gb.sample

        rgb = tf.cast((self.generator.sample+1)*127.5, tf.int32)
        self.generator_int = tf.bitwise.bitwise_or(rgb, 0xFF000000, name='generator_int')

    def d_vars(self):
        return self._d_vars
    def g_vars(self):
        return self._g_vars

    def create_discriminator(self, _input, reuse=False):
        return self.create_component(self.config.discriminator, name='d_ab', input=_input, features=[tf.zeros_like(self.za)], reuse=reuse)

    def create_loss(self, loss_config, discriminator, x, generator, split):
        loss = self.create_component(loss_config, discriminator = discriminator, x=x, generator=generator, split=split)
        return loss

    def create_encoder(self, x_input, name='input_encoder'):
        config = self.config
        input_encoder = dict(config.input_encoder or config.g_encoder or config.generator)
        encoder = self.create_component(input_encoder, name=name, input=x_input)
        return encoder

    def create_z_discriminator(self, z, z_hat):
        config = self.config
        z_discriminator = dict(config.z_discriminator or config.discriminator)
        z_discriminator['layer_filter']=None
        net = tf.concat(axis=0, values=[z, z_hat])
        encoder_discriminator = self.create_component(z_discriminator, name='z_discriminator', input=net)
        return encoder_discriminator

    def create_cycloss(self, x_input, x_hat):
        config = self.config
        ops = self.ops
        distance = config.distance or ops.lookup('l1_distance')
        pe_layers = self.gan.skip_connections.get_array("progressive_enhancement")
        cycloss_lambda = config.cycloss_lambda
        if cycloss_lambda is None:
            cycloss_lambda = 10

        if(len(pe_layers) > 0):
            mask = self.progressive_growing_mask(len(pe_layers)//2+1)
            cycloss = tf.reduce_mean(distance(mask*x_input,mask*x_hat))

            cycloss *= mask
        else:
            cycloss = tf.reduce_mean(distance(x_input, x_hat))

        cycloss *= cycloss_lambda
        return cycloss

    def create_z_cycloss(self, z, x_hat, encoder, generator):
        config = self.config
        ops = self.ops
        total = None
        distance = config.distance or ops.lookup('l1_distance')
        if config.z_hat_lambda:
            z_hat_cycloss_lambda = config.z_hat_cycloss_lambda
            recode_z_hat = encoder.reuse(x_hat)
            z_hat_cycloss = tf.reduce_mean(distance(z_hat,recode_z_hat))
            z_hat_cycloss *= z_hat_cycloss_lambda
        if config.z_cycloss_lambda:
            recode_z = encoder.reuse(generator.reuse(z))
            z_cycloss = tf.reduce_mean(distance(z,recode_z))
            z_cycloss_lambda = config.z_cycloss_lambda
            if z_cycloss_lambda is None:
                z_cycloss_lambda = 0
            z_cycloss *= z_cycloss_lambda

        if config.z_hat_lambda and config.z_cycloss_lambda:
            total = z_cycloss + z_hat_cycloss
        elif config.z_cycloss_lambda:
            total = z_cycloss
        elif config.z_hat_lambda:
            total = z_hat_cycloss
        return total

    def input_nodes(self):
        "used in hypergan build"
        if hasattr(self.generator, 'mask_generator'):
            extras = [self.mask_generator.sample]
        else:
            extras = []
        return extras + [
                self.x_input
        ]

    def output_nodes(self):
        "used in hypergan build"

        if hasattr(self.generator, 'mask_generator'):
            extras = [
                self.mask_generator.sample,
                self.generator.g1x,
                self.generator.g2x
            ]
        else:
            extras = []
        return extras + [
                self.encoder.sample,
                self.generator.sample,
                self.uniform_sample,
                self.generator_int
        ]


# coding=utf-8

from enum import Enum, IntEnum, auto
import time
import random
import math
import numpy
import cv2
from utils import readimage, writeimage
import matching

class Phases(IntEnum):
    """任务执行的阶段"""
    BEGIN = 0
    RUNNING = 1
    END = 2

class Results(Enum):
    """任务执行的结果"""
    PASS = auto()
    SUCCESS = auto()
    FAIL = auto()

tasks = []

def getTasks():
    """获取任务列表"""
    return tasks

def registerTask(task, name, desc):
    """注册任务(已弃用,请使用装饰器注册任务)"""
    task.name = name
    task.description = desc
    tasks.append(task)
    return

def Task(name, desc):
    """用于自动注册任务的装饰器"""
    def decorator(cls):
        registerTask(cls, name, desc)
        return cls
    return decorator

class TaskBase(object):
    """自动挂机任务的基类"""
    name = ""
    description = ""

    def __init__(self):
        """初始化"""
        return

    def init(self):
        self.image = None
        return

    def begin(self, player, t):
        """开始任务"""
        return Results.FAIL

    def run(self, player, t):
        """执行任务"""
        return Results.FAIL

    def end(self, player, t):
        """结束任务"""
        return Results.FAIL

    def getImageCache(self):
        """获取用于预览的图像,如果不想显示请返回None"""
        return self.image

@Task("自动客潮", "请将界面停留在餐厅")
class TaskKeChao(TaskBase):
    """客潮自动化"""

    def __init__(self):
        super().__init__()
        self.templateButton = readimage("kechao_btn")
        self.templateDish = readimage("kechao_dish_part")
        self.templateTitle = readimage("kechao_title_part")
        return

    def init(self):
        super().init()
        self.lastTime = 0
        # 0:餐厅界面 1:客潮对话框 2:客潮进行中 3:客潮结束结算界面
        self.step = 0
        self.pointCache = []
        return

    def begin(self, player, t):
        """需要玩家位于餐厅界面"""
        self.image = player.screenshot()
        if self.step == 0:
            points = findcircle(self.image, 25)
            for x, y in points:
                if x > (self.image.shape[1] * 0.9) and y < (self.image.shape[0] * 0.2):
                    # 找到客潮按钮
                    player.clickaround(x, y)
                    self.lastTime = t
                    self.step = 1
                    return Results.PASS
            if t - self.lastTime > 3:
                # 没找到客潮按钮且超时
                print("未找到客潮按钮, 请确认您正位于餐厅界面")
                return Results.FAIL
        elif self.step == 1 or self.step == 2:
            if t - self.lastTime > 2:
                points = findtemplate(self.image, self.templateButton)
                for x, y in points:
                    # 找到开始按钮
                    if self.step == 2:
                        print("点击按钮后没有开始, 可能是客潮开启次数不足")
                        return Results.FAIL
                    player.clickaround(x, y)
                    self.lastTime = t
                    self.step = 2
                    return Results.PASS
                # 找不到开始按钮
                if self.step == 2:
                    # 点过开始按钮了, 进入客潮
                    self.lastTime = t
                    print("进入客潮")
                    return Results.SUCCESS
                # 还没点过开始按钮, 回退到第0步
                self.lastTime = t
                self.step = 0
                return Results.PASS
        return Results.PASS

    def run(self, player, t):
        """客潮挂机中"""
        self.image = player.screenshot()
        # 处理点的缓存
        self.pointCache = [(x, y, time - 1) for x, y, time in self.pointCache if time > 1]
        # 识别圆来寻找菜(旧版本用模版匹配, 效果不好)
        points = findcircle(self.image, 25)
        points2 = []
        for x, y in points:
            if x > (self.image.shape[1] * 0.9): continue
            if y > (self.image.shape[0] * 0.8):
                # 客潮结束回到餐厅
                self.lastTime = t
                self.step = 3
                print("客潮结束")
                return Results.SUCCESS
            cv2.circle(self.image, (x, y), 25, (0, 0, 255), 3)
            if not self.containpoint(x, y):
                points2.append((x, y))
        if len(points2) > 0:
            x, y = random.choice(points2)
            player.clickaround(x, y)
            self.pointCache.append((x, y, 10))
            self.lastTime = t
            return Results.PASS
        if t - self.lastTime > 15:
            # 没人点菜, 停止挂机?
            print("超过15秒钟没有客人点菜了, 停止挂机")
            return Results.FAIL
        return Results.PASS

    def end(self, player, t):
        """客潮结束"""
        self.image = player.screenshot()
        if self.step == 3:
            if t - self.lastTime > 2:
                points = findtemplate(self.image, self.templateTitle)
                for x, y in points:
                    # 正位于客潮结算界面
                    filename = "KeChao_" + time.strftime("%Y-%m-%d-%H-%M-%S")
                    writeimage(filename, self.image)
                    print("已将客潮结算界面截图保存至: saved/%s.png" % filename)
                    player.clickaround(x, y)
                    self.lastTime = t
                    return Results.PASS
                # 结算完了
                self.lastTime = t
                return Results.SUCCESS
        return Results.PASS

    def containpoint(self, x, y):
        for cx, cy, time in self.pointCache:
            if math.sqrt(math.pow(int(x) - int(cx), 2) + math.pow(int(y) - int(cy), 2)) < 5:
                return True
        return False

class TaskMiniGame(TaskBase):
    """活动小游戏挂机任务的基类"""

    def __init__(self):
        super().__init__()
        self.templateButton = readimage("minigame_btn")
        return

    def init(self):
        self.lastTime = 0
        # 是否点击过开始按钮了
        self.started = False
        return

    def begin(self, player, t):
        """需要玩家位于小游戏界面"""
        self.image = player.screenshot()
        if not self.started:
            points = findtemplate(self.image, self.templateButton)
            for x, y in points:
                cv2.circle(self.image, (x, y), 40, (0, 0, 255), 2)
                player.click(x, y)
                self.lastTime = t
                self.started = True
                return Results.PASS
            if t - self.lastTime > 3:
                # 没找到开始按钮且超时
                print("未找到开始按钮, 请确认您正位于小游戏界面")
                return Results.FAIL
        elif t - self.lastTime > 1:
            self.lastTime = t
            return Results.SUCCESS
        return Results.PASS

try:
    from constant import *
except:
    # **若想使用请自行修改以下数据**
    # 消除的时间间隔
    TIME_INTERVAL = 0.5
    # 游戏区域距离屏幕左方的距离
    MARGIN_LEFT = 0
    # 游戏区域距离屏幕顶部的距离
    MARGIN_TOP = 0
    # 横向方块数量
    HORIZONTAL_NUM = 10
    # 纵向方块数量
    VERTICAL_NUM = 10
    # 方块宽度
    SQUARE_WIDTH = 100
    # 方块高度
    SQUARE_HEIGHT = 100
    # 切片处理时的左上和右下坐标
    SUB_LT_X = 20
    SUB_LT_Y = 20
    SUB_RB_X = 80
    SUB_RB_Y = 80

@Task("自动小游戏-千人千面", "需自行修改代码进行配置")
class TaskQianRenQianMian(TaskMiniGame):
    """千人千面自动连连看"""

    def init(self):
        super().init()
        self.result = None
        self.pair = None
        return

    def run(self, player, t):
        """小游戏挂机中"""
        self.image = player.screenshot()
        for j in range(VERTICAL_NUM):
            for i in range(HORIZONTAL_NUM):
                x = MARGIN_LEFT + i * SQUARE_WIDTH
                y = MARGIN_TOP + j * SQUARE_HEIGHT
                cv2.rectangle(self.image, (x, y), (x + SQUARE_WIDTH, y + SQUARE_HEIGHT), (0, 255, 0), 1)
        if self.result is None:
            # 图像切片并保存在数组中
            squares = []
            for j in range(VERTICAL_NUM):
                for i in range(HORIZONTAL_NUM):
                    x = MARGIN_LEFT + i * SQUARE_WIDTH
                    y = MARGIN_TOP + j * SQUARE_HEIGHT
                    square = self.image[y : y + SQUARE_HEIGHT, x : x + SQUARE_WIDTH]
                    # 每个方块向内缩小一部分防止边缘不一致造成干扰
                    square = square[SUB_LT_Y : SUB_RB_Y, SUB_LT_X : SUB_RB_X]
                    squares.append(square)
            # 相同的方块作为一种类型放在数组中
            types = []
            for square in squares:
                if self.isbackground(square):
                    continue
                if not self.isimageexist(square, types):
                    types.append(square)
            # 将切片处理后的图片数组转换成相对应的数字矩阵
            self.result = []
            num = 0
            for j in range(VERTICAL_NUM):
                line = []
                for i in range(HORIZONTAL_NUM):
                    if self.isbackground(squares[num]):
                        line.append(0)
                    else:
                        for t in range(len(types)):
                            if isimagesame(squares[num], types[t]):
                                line.append(t + 1)
                                break
                    num += 1
                self.result.append(line)
            return Results.PASS
        # 执行自动消除
        if t - self.lastTime >= TIME_INTERVAL:
            self.lastTime = t
            # 第二次选择
            if self.pair is not None:
                player.click(self.pair[0] + SQUARE_WIDTH / 2, self.pair[1] + SQUARE_HEIGHT / 2)
                self.pair = None
                return Results.PASS
            # 定位第一个选中点
            for i in range(len(self.result)):
                for j in range(len(self.result[0])):
                    if self.result[i][j] != 0:
                        # 定位第二个选中点
                        for m in range(len(self.result)):
                            for n in range(len(self.result[0])):
                                if self.result[m][n] != 0:
                                    if matching.canConnect(i, j, m, n, self.result):
                                        # 执行消除算法并进行第一次选择
                                        self.result[i][j] = 0
                                        self.result[m][n] = 0
                                        x1 = MARGIN_LEFT + j * SQUARE_WIDTH
                                        y1 = MARGIN_TOP + i * SQUARE_HEIGHT
                                        x2 = MARGIN_LEFT + n * SQUARE_WIDTH
                                        y2 = MARGIN_TOP + m * SQUARE_HEIGHT
                                        player.click(x1 + SQUARE_WIDTH / 2, y1 + SQUARE_HEIGHT / 2)
                                        self.pair = (x2, y2)
                                        return Results.PASS
            # TODO 判断一下出现结束画面才算完毕, 否则等待一会后重新规划
            print("自动消除运行完毕")
            return Results.SUCCESS
        return Results.PASS

    def isbackground(self, img):
        # TODO 是否有更好的算法?
        # OpenCV的顺序是BGR不是RGB...
        return abs(img[:, :, 0].mean() - 54) <= 10 and abs(img[:, :, 1].mean() - 70) <= 20 and abs(img[:, :, 2].mean() - 105) <= 15

    def isimageexist(self, img, img_list):
        for existed_img in img_list:
            if isimagesame(img, existed_img):
                return True
        return False

@Task("自动小游戏", "更多自动小游戏敬请期待...")
class TaskMoreMiniGames(TaskBase):
    """算了放弃了, 毁灭吧赶紧的"""

    def begin(self, player, t):
        print("我不想做了, 如果您需要的话可以自行编写挂机任务, 然后提交pr")
        return Results.FAIL

def isimagesame(img1, img2, threshold = 0.5):
    # TODO 是否有更好的算法?
    # b = numpy.subtract(existed_img, img)
    # return not numpy.any(b)
    result = cv2.matchTemplate(img1, img2, cv2.TM_CCOEFF_NORMED)
    location = numpy.where(result >= threshold)
    for pt in zip(*location[::-1]):
        return True
    return False

def findtemplate(image, template, threshold = 0.75, outline = False):
    theight, twidth = template.shape[:2]
    result = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
    # result = cv2.normalize(result, None, 0, 1, cv2.NORM_MINMAX)
    location = numpy.where(result >= threshold)
    lx, ly = 0, 0
    points = []
    for pt in zip(*location[::-1]):
        x, y = pt[0] + int(twidth / 2), pt[1] + int(theight / 2)
        # 去掉重复点
        if x - lx < twidth or y - ly < theight: continue
        points.append((x, y))
        lx, ly = x, y
        if outline:
            cv2.rectangle(image, pt, (pt[0] + twidth, pt[1] + theight), (0, 255, 0), 1)
    return points

def findcircle(image, r, outline = False):
    gimg = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gimg = cv2.medianBlur(gimg, 5)
    result = cv2.HoughCircles(gimg, cv2.HOUGH_GRADIENT, 1, int(r / 2), None, 100, 40, r - 10, r + 10)
    points = []
    # 寻找不到会返回None
    if result is not None:
        result = numpy.uint16(numpy.around(result))
        for p in result[0,:]:
            points.append((p[0], p[1]))
            if outline:
                cv2.circle(image, (p[0], p[1]), p[2], (0, 255, 0), 1)
    return points


from __future__ import division, unicode_literals, print_function
import base64
import copy
import gzip
import math
import os
import operator
import traceback
import random
import signal
import sys
from collections import defaultdict, OrderedDict
from functools import partial
from multiprocessing import Queue, Manager
from queue import Empty
from string import Template

import pandas as pd
from pythomics.proteomics.config import CARBON_NEUTRON, PROTON
from pythomics.proteomics.parsers import GuessIterator
from pythomics.proteomics import config
from scipy.interpolate import UnivariateSpline

from .reader import Reader
from .worker import Worker
from .utils import (
    find_prior_scan,
    get_scan_id_from_rt,
    get_scans_under_peaks,
    naninfmean,
    naninfsum,
    perform_ml,
    get_formatted_mass,
)
from . import peaks
from .cpeaks_wrapper import find_nearest_indices

description = """
PyQuant is a quantification program for mass spectrometry data. It attempts to be a general implementation to quantify
an assortment of datatypes and allows a high degree of customization for how data is to be quantified.
"""

ION_CUTOFF = 2

CRASH_SIGNALS = {
    signal.SIGSEGV,
}

def run_pyquant():
    from . import pyquant_parser

    if len(sys.argv) == 1:
        pyquant_parser.print_help()
        sys.exit()

    # for some reason this is being undefined in windows
    import numpy as np

    args = pyquant_parser.parse_args()
    isotopologue_limit = args.isotopologue_limit
    isotopologue_limit = isotopologue_limit if isotopologue_limit else None
    labels_needed = args.labels_needed
    overlapping_mz = args.overlapping_labels
    threads = args.p
    skip = args.skip
    out = args.out
    html = args.html
    resume = args.resume
    calc_stats = not args.disable_stats
    msn_for_id = args.msn_id
    mass_accuracy_correction = args.no_mass_accuracy_correction
    raw_data_only = not (args.search_file or args.tsv)
    scans_to_select = set(args.scan if args.scan else [])
    msn_for_quant = args.msn_quant_from if args.msn_quant_from else msn_for_id - 1
    if msn_for_quant == 0:
        msn_for_quant = 1
    reporter_mode = args.reporter_ion
    msn_ppm = args.msn_ppm
    if args.msn_rt_window:
        msn_rt_window = [tuple(map(float, i.split("-"))) for i in args.msn_rt_window]
    else:
        msn_rt_window = None
    ref_label = str(args.reference_label) if args.reference_label else None
    quant_method = args.quant_method
    if msn_for_quant == 1 and quant_method is None:
        quant_method = "integrate"
    elif msn_for_quant > 1 and quant_method is None:
        quant_method = "sum"

    if args.isobaric_tags:
        reporter_mode = True
        mass_accuracy_correction = False
        calc_stats = False
        msn_for_quant = 2
    if args.ms3:
        msn_for_quant = 3
        mass_accuracy_correction = False
    if args.gcms:
        msn_for_quant = 1
        msn_for_id = 1
        args.precursor_ppm = args.msn_ppm
        isotopologue_limit = 1
        args.msn_all_scans = True
        args.no_mass_accuracy_correction = True
        args.no_contaminant_detection = True
        args.no_rt_guide = True
        args.disable_stats = True

    mrm_pair_info = (
        pd.read_csv(args.mrm_map, sep="\t") if args.mrm and args.mrm_map else None
    )

    scan_filemap = {}
    found_scans = {}
    raw_files = {}
    mass_labels = {}  # {'Light': {0: set([])}} if not reporter_mode else {}

    name_mapping = {}

    if args.label_scheme:
        label_info = pd.read_csv(
            args.label_scheme.name, sep="\t", header=None, dtype="str"
        )
        try:
            label_info.columns = ["Label", "AA", "Mass", "UserName"]
            name_mapping = dict(
                [(v["Label"], v["UserName"]) for i, v in label_info.iterrows()]
            )
        except ValueError:
            label_info.columns = ["Label", "AA", "Mass"]
            name_mapping = dict(
                [(v["Label"], v["Label"]) for i, v in label_info.iterrows()]
            )
        for group_name, group_info in label_info.groupby("Label"):
            masses = {}
            label_name = name_mapping.get(group_name, group_name)
            for mass, mass_aas in group_info.groupby("Mass"):
                mass_val = float(mass)
                mass_list = mass_aas["AA"].values.tolist()
                try:
                    masses[mass_val].add(mass_list)
                except KeyError:
                    masses[mass_val] = set(mass_list)
            mass_labels.update({label_name: masses})
    if args.label_method:
        mass_labels = config.LABEL_SCHEMES[args.label_method]

    sample = args.sample
    sys.stderr.write("Loading Scans:\n")

    # options determining modes to quantify
    all_msn = False  # we just have a raw file
    ion_search = False  # we have an ion we want to find

    if args.peptide_file:
        peptides = set((i.strip().upper() for i in args.peptide_file))
    elif args.peptide:
        peptides = (
            set(args.peptide) if isinstance(args.peptide, list) else set([args.peptide])
        )
    else:
        peptides = None

    if peptides:
        peptides = set(map(lambda x: x.upper(), peptides))

    input_found = None
    if args.search_file:
        results = GuessIterator(
            args.search_file.name, full=True, store=False, peptide=peptides
        )
        input_found = "ms"
    elif args.tsv:
        results = pd.read_csv(args.tsv, sep="\t")
        input_found = "tsv"

    if args.search_file:
        source_file = args.search_file.name
    elif args.tsv:
        source_file = args.tsv.name
    elif args.scan_file:
        source_file = args.scan_file[0].name
    elif args.scan_file_dir:
        source_file = os.path.split(args.scan_file_dir)[1]
    else:
        pyquant_parser.error(
            "No input files provided. Please provide one via --search-file, --tsv, --scan-file, or --scan-file-dir."
        )

    if args.scan_file:
        nfunc = (
            lambda i: (
                os.path.splitext(os.path.split(i.name)[1])[0],
                os.path.abspath(i.name),
            )
            if hasattr(i, "name")
            else (os.path.splitext(os.path.split(i)[1])[0], os.path.abspath(i))
        )
        scan_filemap = dict([nfunc(i) for i in args.scan_file])
    else:
        if args.scan_file_dir:
            raw_file = args.scan_file_dir
        else:
            raw_file = os.path.abspath(os.path.split(source_file)[0])
        if os.path.isdir(raw_file):
            scan_filemap = dict(
                [
                    (os.path.splitext(i)[0], os.path.abspath(os.path.join(raw_file, i)))
                    for i in os.listdir(raw_file)
                    if i.lower().endswith("mzml")
                ]
            )
        else:
            scan_filemap[
                os.path.splitext(os.path.split(raw_file)[1])[0]
            ] = os.path.abspath(raw_file)

    result_iterator = []
    replicate_file_mapper = {}

    if input_found == "tsv":
        # TODO: Create classes for abstracting this and making it testable
        if args.maxquant:
            peptide_col = "Sequence"
            precursor_col = "Mass"
            rt_col = "Retention time"
            charge_col = "Charge"
            file_col = "Raw file"
            scan_col = "MS/MS Scan Number"
            label_col = "Labeling State"
            if "ms2" in source_file:
                scan_col = "Scan number"
                label_col = None
        else:
            peptide_col = args.peptide_col
            precursor_col = args.mz
            rt_col = args.rt
            charge_col = args.charge
            file_col = args.source
            label_col = args.label
            scan_col = args.scan_col

        result_iterator = results.iterrows()

        def defloat(value):
            # Pandas will coerce columns with missing data to floats, resulting in scan ids like
            # 34284.0
            try:
                return int(value)
            except Exception:
                pass
            return value

        def tsv_formatter(row):
            row_index, row_info = row
            peptide = row_info[peptide_col].strip() if peptide_col in row_info else ""
            specId = None if args.scan_from_rt else str(defloat(row_info[scan_col]))
            fname = row_info[file_col] if file_col in row_info else raw_file
            charge = float(row_info[charge_col]) if charge_col in row_info else 1
            precursor_mass = (
                row_info[precursor_col] if precursor_col in row_info else None
            )
            rt_value = row_info[rt_col] if rt_col in row_info else None
            label = (
                name_mapping.get(row_info[label_col]) if label_col in row_info else None
            )

            if math.isnan(precursor_mass) or not charge:
                return None

            if args.maxquant:
                precursor_mass = precursor_mass / charge + PROTON

            return {
                "file": fname,
                "id": specId,
                "mass": float(precursor_mass),
                "theor_mass": float(precursor_mass),
                "rt": rt_value,
                "charge": charge,
                "peptide": peptide,
                "modifications": None,
                "label": label,
                "key": (fname, specId, peptide, precursor_mass),
            }

        result_formatter = tsv_formatter

    elif input_found == "ms":
        if not (args.label_scheme or args.label_method):
            mass_labels.update(results.getSILACLabels())
        result_iterator = results.getScans(modifications=False, fdr=True)

        def ms_formatter(scan):
            peptide = scan.peptide
            specId = scan.id
            fname = scan.file

            return {
                "file": fname,
                "id": specId,
                "mass": float(scan.mass),
                "rt": scan.rt,
                "charge": scan.charge,
                "peptide": peptide,
                "mod_peptide": scan.modifiedPeptide,
                "theor_mass": scan.getTheorMass(),
                "modifications": scan.getModifications(),
                "accession": getattr(scan, "acc", None),
                "key": (fname, specId, peptide, scan.mass),
            }

        result_formatter = ms_formatter

    for index, scan in enumerate(result_iterator):
        if index % 1000 == 0:
            sys.stderr.write(".")
        if scan is None:
            continue

        result_info = result_formatter(scan)
        if result_info is None:
            continue

        peptide = result_info["peptide"]

        if peptides and peptide.upper() not in peptides:
            continue
        if not peptides and (sample != 1.0 and random.random() > sample):
            continue

        specId = result_info["id"]
        if scans_to_select and str(specId) not in scans_to_select:
            continue

        fname = result_info["file"]
        mass_key = result_info["key"]
        if mass_key in found_scans:
            continue

        d = {
            "file": fname,
            "quant_scan": {},
            "id_scan": {
                "id": specId,
                "theor_mass": result_info.get("theor_mass"),
                "peptide": peptide,
                "mod_peptide": result_info.get("mod_peptide"),
                "rt": result_info["rt"],
                "charge": result_info["charge"],
                "modifications": result_info["modifications"],
                "mass": result_info["mass"],
                "accession": result_info.get("accession"),
            },
        }
        found_scans[mass_key] = d
        fname = os.path.splitext(fname)[0]
        if args.mva:
            # find the most similar name, add in a setting for this
            import difflib

            if fname in replicate_file_mapper:
                fname = replicate_file_mapper[fname]
                if fname is None:
                    continue
            else:
                s = difflib.SequenceMatcher(None, fname)
                seq_matches = []
                for i in scan_filemap:
                    s.set_seq2(i)
                    seq_matches.append((s.ratio(), i))
                seq_matches = [
                    i
                    for i in sorted(
                        seq_matches, key=operator.itemgetter(0), reverse=True
                    )
                    if i[0] != 1
                ]
                replicate_file_mapper[fname] = seq_matches[0][1]
                fname = seq_matches[0][1]
        if fname not in scan_filemap:
            fname = os.path.split(fname)[1]
            if fname not in scan_filemap:
                if skip:
                    continue
                sys.stderr.write(
                    "{0} not found in filemap. Filemap is {1}.".format(
                        fname, scan_filemap
                    )
                )
                return 1
        try:
            raw_files[fname].append(d)
        except KeyError:
            raw_files[fname] = [d]
        del scan

    if scan_filemap and raw_data_only:
        # determine if we want to do ms1 ion detection, ms2 ion detection, all ms2 of each file
        if args.msn_ion or args.msn_peaklist:
            ion_search = True
            ions_of_interest = []
            if args.msn_peaklist:
                for i in args.msn_peaklist:
                    if not i.strip():
                        continue
                    try:
                        ions_of_interest.append(i.strip())
                    except Exception as e:
                        traceback.print_exc()
            if args.msn_ion:
                ions_of_interest += args.msn_ion
            ions_selected = (
                [
                    sorted(list(set(map(float, ion.split(",")))))
                    for ion in ions_of_interest
                ]
                if ions_of_interest
                else []
            )
            d = {"ions": ions_selected, "rt_info": args.msn_ion_rt}
            for i in scan_filemap:
                raw_files[i] = d
        else:
            all_msn = True
            for i in scan_filemap:
                raw_files[i] = [1]
    if not scan_filemap and input_found is None:
        sys.stderr.write(
            "No valid input entered. PyQuant requires at least a raw file or a processed dataset."
        )
        return 1
    sys.stderr.write("\nScans loaded.\n")

    pq_dir = os.path.split(__file__)[0]

    pyquant_html_file = os.path.join(pq_dir, "static", "pyquant_output.html")

    workers = []
    completed = 0
    sys.stderr.write("Beginning quantification.\n")
    scan_count = len(found_scans)

    # before we start, get the information we will be generating
    labels = list(mass_labels.keys())
    if not labels:
        if ion_search:
            if args.require_all_ions:
                labels = ["_".join(map(str, ion_set)) for ion_set in ions_selected]
            else:
                labels = list(
                    map(str, set([ion for ion_set in ions_selected for ion in ion_set]))
                )
        else:
            labels = [""]
    if not scan_filemap and raw_data_only:
        RESULT_ORDER = []
    else:
        RESULT_ORDER = [
            ("peptide", "Peptide"),
            ("modifications", "Modifications"),
            ("accession", "Accession"),
        ]
    RESULT_ORDER.extend(
        [("charge", "Charge"), ("ms1", "MS{} Spectrum ID".format(msn_for_quant)),]
    )
    if msn_for_quant != msn_for_id:
        RESULT_ORDER.extend(
            [("scan", "MS{} Spectrum ID".format(msn_for_id)),]
        )
    RESULT_ORDER.extend(
        [("rt", "Retention Time"),]
    )

    if ion_search and not args.msn_all_scans:
        RESULT_ORDER.extend([("ions_found", "Ions Found")])

    PEAK_REPORTING = []
    if not reporter_mode and args.peaks_n != 1:
        PEAK_REPORTING.extend(
            [
                ("label", "Peak Label"),
                ("auc", "Peak Area"),
                ("amp", "Peak Max"),
                ("peak_mean", "Peak Center"),
                ("peak_width", "RT Width"),
                ("mean_diff", "Mean Offset"),
                ("snr", "SNR"),
                ("sbr", "SBR"),
                ("sdr", "Density"),
                ("coef_det", "R^2"),
                ("residual", "Residual"),
            ]
        )
    iterator = ["_".join(map(str, sorted(labels)))] if args.merge_labels else labels
    for label in iterator:
        label_key = "{} ".format(label) if label != "" else label
        RESULT_ORDER.extend(
            [("{}_precursor".format(label), "{}Precursor".format(label_key)),]
        )
        if msn_for_quant == 1:
            if not args.no_mass_accuracy_correction:
                RESULT_ORDER.extend(
                    [
                        (
                            "{}_calibrated_precursor".format(label),
                            "{}Calibrated Precursor".format(label_key),
                        ),
                    ]
                )
            RESULT_ORDER.extend(
                [("{}_isotopes".format(label), "{}Isotopes Found".format(label_key)),]
            )
        RESULT_ORDER.extend(
            [
                (
                    "{}_intensity".format(label),
                    "{}Intensity".format(label_key),
                    partial(naninfsum, empty=0),
                ),
            ]
        )
        if args.peaks_n == 1:
            RESULT_ORDER.extend(
                [
                    (
                        "{}_peak_width".format(label),
                        "{}RT Width".format(label_key),
                        partial(naninfmean, empty=0),
                    ),
                    (
                        "{}_peak_mean".format(label),
                        "{}Peak Center".format(label_key),
                        partial(naninfmean, empty=np.NaN),
                    ),
                    (
                        "{}_mean_diff".format(label),
                        "{}Mean Offset".format(label_key),
                        partial(naninfmean, empty=np.NaN),
                    ),
                ]
            )

        if msn_for_quant == 1:
            RESULT_ORDER.extend(
                [
                    (
                        "{}_residual".format(label),
                        "{}Residual".format(label_key),
                        partial(naninfsum, empty=np.NaN),
                    ),
                    (
                        "{}_coef_det".format(label),
                        "{}R^2".format(label_key),
                        partial(naninfmean, empty=0),
                    ),
                    (
                        "{}_snr".format(label),
                        "{}SNR".format(label_key),
                        partial(naninfmean, empty=0),
                    ),
                ]
            )
        if not args.merge_labels and not args.no_ratios and label != "":
            for label2 in labels:
                label_key2 = "{}".format(label2) if label2 != "" else label2
                if label_key != label_key2 and (
                    ref_label is None or ref_label.lower() == label_key2.lower()
                ):
                    RESULT_ORDER.extend(
                        [
                            (
                                "{}_{}_ratio".format(label, label2),
                                "{}/{}".format(label_key.strip(), label_key2.strip()),
                            ),
                        ]
                    )
                    if calc_stats:
                        RESULT_ORDER.extend(
                            [
                                (
                                    "{}_{}_confidence".format(label, label2),
                                    "{}/{} Confidence".format(
                                        label_key.strip(), label_key2.strip()
                                    ),
                                ),
                            ]
                        )

    headers = ["Raw File"] + [i[1] for i in RESULT_ORDER]
    if resume and os.path.exists(out):
        if not out:
            sys.stderr.write("You may only resume runs with a file output.\n")
            return -1
        out = open(out, "a+")
        out_path = out.name
    else:
        if out:
            out = open(out, "w+")
            out_path = out.name
        else:
            out = sys.stdout
            out_path = source_file
        out.write("{0}\n".format("\t".join(headers)))

    if html:

        def table_rows(html_list, res=None):
            # each item is a string like a\tb\tc
            if html_list:
                d = html_list.pop(0)
            else:
                return res
            line = d["table"]
            html_extra = d.get("html", {})
            keys = d["keys"]
            if res is None:
                res = "<tr>"
            out = ["<td></td>"]  # for graph controls
            for col_index, (i, v) in enumerate(zip(line.split("\t"), keys)):
                out.append("<td>{0}</td>".format(i))
            res += "\n".join(out) + "</tr>"
            return table_rows(html_list, res=res), html_extra

        if resume:
            html_out = open("{0}.html".format(out_path), "a")
        else:
            html_out = open("{0}.html".format(out_path), "w")
            template = []
            for i in open(pyquant_html_file, "r"):
                if "HTML BREAK" in i:
                    break
                template.append(i)
            html_template = Template("".join(template))
            html_out.write(
                html_template.substitute(
                    {
                        "title": source_file,
                        "table_header": "\n".join(
                            [
                                "<th>{0}</th>".format(i)
                                for i in ["Controls", "Raw File"]
                                + [i[1] for i in RESULT_ORDER]
                            ]
                        ),
                    }
                )
            )

    skip_map = set([])
    if resume:
        resume_name = "{}.tmp".format(out.name)
        if os.path.exists(resume_name):
            with open(resume_name, "r") as temp_file:
                for index, entry in enumerate(temp_file):
                    info = pd.io.json.loads(entry)["res_list"]
                    # key is filename, peptide, charge, target scan id, modifications
                    key = tuple(map(str, (info[0], info[1], info[3], info[5], info[2])))
                    skip_map.add(key)
            temp_file = open(resume_name, "a")
        else:
            temp_file = open(resume_name, "w")
    else:
        temp_file = open("{}.tmp".format(out.name), "w")

    silac_shifts = {}
    for silac_label, silac_masses in mass_labels.items():
        for mass, aas in silac_masses.items():
            try:
                silac_shifts[get_formatted_mass(mass)] |= aas
            except Exception as e:
                silac_shifts[get_formatted_mass(mass)] = aas

    for filename in raw_files.keys():
        raw_scans = raw_files[filename]
        filepath = scan_filemap[filename]
        if not len(raw_scans):
            continue
        in_queue = Queue()
        result_queue = Queue()
        reader_in = Queue()
        reader_outs = {"main": Queue()}
        for i in range(threads):
            reader_outs[i] = Queue()

        msn_map = []
        scan_rt_map = {}
        msn_rt_map = {}
        scan_charge_map = {}

        raw = GuessIterator(filepath, full=False, store=False)
        sys.stderr.write("Processing {}.\n".format(filepath))

        # params in case we are doing ion search or replicate analysis
        ion_tolerance = args.precursor_ppm / 1e6 if args.mva else args.msn_ppm / 1e6
        ion_search_list = []
        replicate_search_list = defaultdict(list)
        scans_to_fetch = []

        # figure out the splines for mass accuracy correction
        calc_spline = (
            not mass_accuracy_correction and not raw_data_only and not args.neucode
        )
        spline_x = []
        spline_y = []
        spline = None

        scan_info_map = defaultdict(dict)

        for index, scan in enumerate(raw):
            if index % 100 == 0:
                sys.stderr.write(".")
            if scan is None:
                continue
            scan_id = scan.id
            msn_map.append((scan.ms_level if not args.mrm else scan.mass, scan_id))
            rt = scan.rt
            if scan.ms_level == msn_for_quant:
                msn_rt_map[scan_id] = int(scan.title) if args.mrm else rt
            scan_rt_map[scan_id] = rt
            scan_info_map[scan_id]["parent"] = scan.parent
            scan_info_map[scan_id]["msn"] = scan.ms_level
            scan_info_map[scan_id]["precursor"] = scan.mass
            scan_charge_map[scan_id] = scan.charge
            if msn_rt_window and not any(
                [(i[0] < float(rt) < i[1]) for i in msn_rt_window]
            ):
                continue
            if scan.parent:
                try:
                    scan_info_map[scan.parent]["children"].add(scan_id)
                except KeyError:
                    scan_info_map[scan.parent]["children"] = set([scan_id])
            if not scans_to_select or str(scan_id) in scans_to_select:
                if ion_search:
                    if scan.ms_level == msn_for_id:
                        scans_to_fetch.append(scan_id)
                elif all_msn:
                    # we are quantifying all msn spectra of a given type
                    if msn_for_id == scan.ms_level:
                        # find the closest scan to this, which will be the parent scan
                        spectra_to_quant = (
                            find_prior_scan(msn_map, scan_id, ms_level=msn_for_quant)
                            if msn_for_quant != msn_for_id
                            else scan_id
                        )
                        d = {
                            "quant_scan": {"id": spectra_to_quant},
                            "id_scan": {
                                "id": scan_id,
                                "rt": scan.rt,
                                "charge": scan.charge,
                                "mass": float(scan.mass),
                                "product_ion": float(scan.product_ion)
                                if args.mrm
                                else None,
                            },
                        }
                        ion_search_list.append((spectra_to_quant, d))
                elif args.mva:
                    if scan.ms_level == msn_for_quant:
                        scans_to_fetch.append(scan_id)
            if not raw_data_only and calc_spline:
                if hasattr(scan, "theor_mass"):
                    theor_mass = scan.getTheorMass()
                    observed_mass = scan.mass
                    mass_error = (theor_mass - observed_mass) / theor_mass * 1e6
                    spline_x.append(observed_mass)
                    spline_y.append(mass_error)
            del scan

        if calc_spline and len(spline_x):
            spline_df = pd.DataFrame(
                zip(spline_x, spline_y), columns=["Observed", "Error"]
            )
            spline_df = spline_df[
                (spline_df["Error"] < 25) & (spline_df["Error"] > -25)
            ].dropna()
            spline_df.sort("Observed", inplace=True)
            spline_df.drop_duplicates("Observed", inplace=True)
            if len(spline_df) > 10:
                spline = UnivariateSpline(
                    spline_df["Observed"].astype(float).values,
                    spline_df["Error"].astype(float).values,
                    s=1e6,
                )
            else:
                spline = None

        reader = Reader(
            reader_in,
            reader_outs,
            raw_file=filepath,
            spline=spline,
            rt_window=msn_rt_window,
        )
        reader.start()
        rep_map = defaultdict(set)
        if ion_search or args.mva:
            ions = (
                [
                    i["id_scan"].get("theor_mass", i["id_scan"]["mass"])
                    for i in raw_scans
                ]
                if args.mva
                else raw_scans["ions"]
            )
            rt_info = raw_scans.get("rt_info") if not args.mva else None
            last_scan_ions = defaultdict(set)
            for scan_id in scans_to_fetch:
                this_scan_ions = defaultdict(set)
                reader_in.put((0, scan_id, None, None))
                scan = reader_outs[0].get()
                if scan is None:
                    continue
                if args.msn_all_scans:
                    # if not args.require_all_ions:
                    #     ions = set([j for i in ions for j in i])
                    for ion_index, ion_set in enumerate(ions):
                        d = {
                            "quant_scan": {"id": scan_id, "scans": scans_to_fetch},
                            "id_scan": {
                                "id": scan_id,
                                "theor_mass": ion_set[0],
                                "rt": rt_info[ion_index] if rt_info else scan["rt"],
                                "charge": 1,
                                "mass": ion_set[0],
                                "ions_found": None,
                                "ion_set": ion_set,
                            },
                            "combine_xics": args.require_all_ions,
                        }
                        ion_search_list.append((scan_id, d))
                    break
                scan_mzs = scan["vals"]
                mz_vals = scan_mzs[scan_mzs[:, 1] > 0][:, 1]
                scan_mzs = scan_mzs[scan_mzs[:, 1] > 0][:, 0]
                if not np.any(scan_mzs):
                    continue
                mass, charge, rt, centroided = (
                    scan["mass"],
                    scan["charge"],
                    scan["rt"],
                    scan["centroid"],
                )
                ions_found = []
                added = set([])
                for ion_set in ions:
                    if isinstance(ion_set, float):
                        ion_set = [ion_set]
                    for ion_index, (ion, nearest_mz_index) in enumerate(
                        zip(
                            ion_set,
                            find_nearest_indices(
                                scan_mzs, np.array(ion_set, dtype=np.float)
                            ),
                        )
                    ):
                        nearest_mz = scan_mzs[nearest_mz_index]
                        found_ions = []
                        if peaks.get_ppm(ion, nearest_mz) < ion_tolerance:
                            if args.mva:
                                d = raw_scans[ion_index]
                                scan_rt = d["id_scan"]["rt"]
                                if (
                                    scan_rt - args.rt_window
                                    < rt
                                    < scan_rt + args.rt_window
                                ):
                                    found_ions.append(
                                        {
                                            "ion": ion,
                                            "nearest_mz": nearest_mz,
                                            "charge": d["id_scan"]["charge"],
                                            "scan_info": d,
                                            "ion_set": ion_set,
                                        }
                                    )
                            else:
                                ion_rt = rt_info[ion_index] if rt_info else None
                                if ion_rt is None or (
                                    rt - args.rt_window < ion_rt < rt + args.rt_window
                                ):
                                    found_ions.append(
                                        {
                                            "ion": ion,
                                            "nearest_mz": nearest_mz,
                                            "rt": ion_rt,
                                            "ion_set": ion_set,
                                        }
                                    )
                        elif args.require_all_ions:
                            # this means an ion in the set was not found, if we require all ions, continue to the next set
                            continue
                        # if we require all ions, only put in a single copy since we'll be finding the same ions in every
                        # scan anyways, otherwise put the entire set in since we may be missing ions
                        ions_found += (
                            found_ions[0] if args.require_all_ions else found_ions
                        )
                if ions_found:
                    # we have two options here. If we are quantifying a preceeding scan or the ion itself per scan
                    isotope_ppm = args.isotope_ppm / 1e6
                    if msn_for_quant == msn_for_id or args.mva:
                        for ion_dict in ions_found:
                            ion, nearest_mz = ion_dict["ion"], ion_dict["nearest_mz"]
                            ion_found = ",".join(map(str, ion_dict["ion_set"]))
                            spectra_to_quant = scan_id
                            # we are quantifying the ion itself
                            if charge == 0 and not args.mva:
                                # see if we can figure out the charge state
                                charge_states = []
                                for i in range(1, 5):
                                    charge_peaks_found = 0
                                    peak_height = 0
                                    for j in range(1, 3):
                                        next_peak = ion + CARBON_NEUTRON / float(
                                            i
                                        ) * float(j)
                                        closest_mz = peaks.find_nearest_index(
                                            scan_mzs, next_peak
                                        )
                                        if (
                                            peaks.get_ppm(
                                                next_peak, scan_mzs[closest_mz]
                                            )
                                            < isotope_ppm * 1.5
                                        ):
                                            charge_peaks_found += 1
                                            peak_height += mz_vals[closest_mz]
                                    charge_states.append(
                                        (charge_peaks_found, i, peak_height)
                                    )
                                charge_states = sorted(
                                    charge_states,
                                    key=operator.itemgetter(0, 2),
                                    reverse=True,
                                )
                                if int(ion_dict["charge"]) not in [
                                    i[0] for i in charge_states
                                ]:
                                    continue
                                elif args.mva:
                                    charge_to_use = ion_dict["charge"]
                                elif charge_states:
                                    if charge_states[0][1] == 1:
                                        charge_to_use = (
                                            charge_states[1][1]
                                            if charge_states[1][2] != 0
                                            else 1
                                        )
                                    else:
                                        charge_to_use = charge_states[0][1]
                                else:
                                    charge_to_use = 1
                            else:
                                charge_to_use = charge
                                if args.mva:
                                    rep_key = (
                                        ion_dict["scan_info"]["id_scan"]["rt"],
                                        ion_dict["scan_info"]["id_scan"]["mass"],
                                        charge_to_use,
                                    )
                                    rep_map[rep_key].add(rt)
                            this_scan_ions[ion].add(charge_to_use)
                            if charge_to_use in last_scan_ions[ion]:
                                continue
                            last_scan_ions[ion].add(charge_to_use)
                            if args.mva:
                                d = copy.deepcopy(ion_dict["scan_info"])
                                d["id_scan"]["id"] = scan_id
                                theo_mass = d["id_scan"].get("theor_mass")
                                if theo_mass:
                                    d["id_scan"]["mass"] = theo_mass
                                spectra_to_quant = find_prior_scan(
                                    msn_map, scan_id, ms_level=msn_for_quant
                                )
                                d["quant_scan"]["id"] = spectra_to_quant
                                d["replicate_scan_rt"] = rt
                            else:
                                d = {
                                    "quant_scan": {"id": scan_id},
                                    "id_scan": {
                                        "id": scan_id,
                                        "theor_mass": ion,
                                        "rt": rt if ion_rt is None else ion_rt,
                                        "charge": charge_to_use,
                                        "mass": float(nearest_mz),
                                        "ions_found": ion_found,
                                        "ion_set": ion_dict["ion_set"],
                                    },
                                    "combine_xics": True,
                                }
                            key = (scan_id, d["id_scan"]["theor_mass"], charge_to_use)
                            if key in added:
                                continue
                            added.add(key)
                            if args.mva:
                                replicate_search_list[
                                    (d["id_scan"]["rt"], ion_dict["ion"])
                                ].append((spectra_to_quant, d))
                            else:
                                ion_search_list.append((spectra_to_quant, d))
                    else:
                        # we are identifying the ion in a particular scan, and quantifying a preceeding scan
                        # find the closest scan to this, which will be the parent scan
                        spectra_to_quant = find_prior_scan(
                            msn_map, scan_id, ms_level=msn_for_quant
                        )
                        d = {
                            "quant_scan": {"id": spectra_to_quant},
                            "id_scan": {
                                "id": scan_id,
                                "rt": rt,
                                "charge": charge,
                                "mass": float(mass),
                                "ions_found": ";".join(
                                    map(
                                        lambda x: "{}({})".format(ion, nearest_mz),
                                        ions_found,
                                    )
                                ),
                            },
                        }
                        ion_search_list.append((spectra_to_quant, d))
                to_remove = []
                for ion, charges in last_scan_ions.items():
                    for charge in charges:
                        if charge not in this_scan_ions[ion]:
                            to_remove.append((ion, charge))
                for ion, charge in to_remove:
                    last_scan_ions[ion].discard(charge)
                del scan

        if args.mva:
            x = []
            y = []
            for i, v in rep_map.items():
                for j in v:
                    x.append(i[0])
                    y.append(j)
            from sklearn.linear_model import LinearRegression

            rep_mapper = LinearRegression()
            try:
                rep_mapper.fit(np.array(x).reshape(len(x), 1), y)
            except ValueError:
                rep_mapper = None

        if ion_search or all_msn:
            raw_scans = [
                i[1] for i in sorted(ion_search_list, key=operator.itemgetter(0))
            ]
        if args.mva:
            raw_scans = []
            for i in replicate_search_list:
                ion_rt, ion = i
                best_scan = sorted(
                    [
                        (
                            np.abs(
                                (
                                    rep_mapper.predict(
                                        np.array(j[1]["replicate_scan_rt"]).reshape(
                                            -1, 1
                                        )
                                    )
                                    if rep_mapper
                                    else j[1]["replicate_scan_rt"]
                                )
                                - ion_rt
                            ),
                            j[1],
                        )
                        for j in replicate_search_list[i]
                    ],
                    key=operator.itemgetter(0),
                )[0][1]
                raw_scans.append(best_scan)

        quant_msn_map = (
            [i for i in msn_map if i[0] == msn_for_quant] if not args.mrm else msn_map
        )
        manager = Manager()
        scan_mask = manager.dict()
        target_scan_rt_series = pd.Series(
            {
                scan_id: scan_rt_map[scan_id]
                for scan_msn, scan_id in msn_map
                if scan_msn == msn_for_id
            }
        )

        for i in range(threads):
            worker = Worker(
                queue=in_queue,
                results=result_queue,
                raw_name=filepath,
                mass_labels=mass_labels,
                debug=args.debug,
                html=html,
                mono=not args.spread,
                precursor_ppm=args.precursor_ppm,
                isotope_ppm=args.isotope_ppm,
                isotope_ppms=None,
                msn_rt_map=msn_rt_map,
                reporter_mode=reporter_mode,
                reader_in=reader_in,
                reader_out=reader_outs[i],
                thread=i,
                quant_method=quant_method,
                spline=spline,
                isotopologue_limit=isotopologue_limit,
                labels_needed=labels_needed,
                quant_msn_map=quant_msn_map,
                overlapping_mz=overlapping_mz,
                min_resolution=args.min_resolution,
                min_scans=args.min_scans,
                mrm_pair_info=mrm_pair_info,
                mrm=args.mrm,
                peak_cutoff=args.peak_cutoff,
                replicate=args.mva,
                ref_label=ref_label,
                max_peaks=args.max_peaks,
                parser_args=args,
                scans_to_skip=scan_mask,
            )
            workers.append(worker)
            worker.start()

        # TODO:
        # combine information from scans (where for instance, we have fragmented both the heavy/light
        # peptides -- we want to use those masses before calculating where it should be). This may not
        # be possible for all types of input though, figure this out.
        quant_map = {}
        scans_to_submit = []

        # this is to fix the header at the end to include peak information if we have multiple peaks
        most_peaks_found = 0
        mrm_added = set([])
        exclusion_masses = (
            mrm_pair_info.loc[
                :,
                [
                    i
                    for i in mrm_pair_info.columns
                    if i.lower() not in ("light", "retention time")
                ],
            ].values.flatten()
            if args.mrm
            else set([])
        )
        for scan_index, raw_scan_info in enumerate(raw_scans):
            target_scan = raw_scan_info["id_scan"]
            quant_scan = raw_scan_info["quant_scan"]
            if args.scan_from_rt and target_scan["id"] is None:
                target_scan["id"] = get_scan_id_from_rt(
                    target_scan_rt_series, target_scan["rt"]
                )

            scanId = target_scan["id"]
            scan_mass = target_scan.get("mass")
            if args.mrm:
                if scan_mass in mrm_added:
                    continue
                mrm_added.add(scan_mass)
                if scan_mass in exclusion_masses:
                    continue

            if quant_scan.get("id") is None:
                scan_to_quant = None
                if msn_for_quant > msn_for_id:
                    # Cases like MS3 hit this, go down until we hit the correct MS level and add
                    # those scans to be quantified
                    children = scan_info_map[scanId].get("children", [scanId])
                    quant_scan["scans"] = []
                    while children:
                        child = children.pop()
                        child_scan_info = scan_info_map[child]
                        scan_to_quant_ms = child_scan_info["msn"]
                        if scan_to_quant_ms < msn_for_quant:
                            children += child_scan_info.get("children", [])
                        if scan_to_quant_ms == msn_for_quant:
                            if scan_to_quant is None:
                                scan_to_quant = child
                            quant_scan["scans"].append(child)
                elif msn_for_quant == msn_for_id:
                    scan_to_quant = scanId
                else:
                    # we will hit this in a normal proteomic run
                    # figure out the ms-1 from the ms level we are at
                    scan_info = scan_info_map[scanId]
                    current_scan = scan_info["parent"]
                    try:
                        while current_scan:
                            scan_info = scan_info_map[scanId]
                            current_scan = scan_info["parent"]
                            scan_to_quant_ms = scan_info[scan_to_quant]["msn"]
                            if scan_to_quant_ms == msn_for_quant:
                                scan_to_quant = current_scan
                                break
                    except KeyError:
                        scan_to_quant = None
                if scan_to_quant is not None:
                    msn_to_quant = scan_to_quant
                else:
                    msn_to_quant = find_prior_scan(
                        msn_map, scanId, ms_level=msn_for_quant
                    )
                quant_scan["id"] = msn_to_quant

            rt = target_scan.get("rt", scan_rt_map.get(scanId))
            if rt is None:
                rt = float(msn_rt_map[msn_to_quant])
                target_scan["rt"] = rt

            if args.mva and rep_mapper is not None:
                target_scan["rt"] = rep_mapper.predict(
                    np.array(float(target_scan["rt"])).reshape(-1, 1)
                )[0]

            mods = target_scan.get("modifications")
            charge = target_scan.get("charge")
            if charge is None or charge == 0:
                charge = int(scan_charge_map.get(scanId, 0))
            if charge == 0:
                continue
            charge = int(charge)

            if msn_for_quant != 1:
                lowest_label = min([j for i, v in mass_labels.items() for j in v])
                target_scan["theor_mass"] = lowest_label
                target_scan["precursor"] = lowest_label
            else:
                mass_shift = 0

                if mods is not None:
                    shift = 0
                    for mod in filter(lambda x: x, mods.split("|")):
                        aa, pos, mass, _ = mod.split(",", 3)
                        mass = get_formatted_mass(mass)
                        if aa in silac_shifts.get(mass, {}):
                            shift += mass
                    mass_shift += float(shift) / float(charge)
                else:
                    # assume we are the light version, include all the labels we are looking for here
                    pass

                target_scan["theor_mass"] = (
                    target_scan.get("theor_mass", target_scan.get("mass")) - mass_shift
                )
                target_scan["precursor"] = (
                    target_scan["mass"] - mass_shift
                    if not args.neucode
                    else target_scan["theor_mass"]
                )
            # key is filename, peptide, charge, target scan id, modifications
            key = (
                filename,
                target_scan.get("peptide", ""),
                target_scan.get("charge"),
                target_scan.get("id"),
                target_scan.get("modifications"),
            )
            if resume:
                if tuple(map(str, key)) in skip_map:
                    completed += 1
                    continue
            params = {"scan_info": raw_scan_info}
            scans_to_submit.append((target_scan["rt"], params))

        # sort by RT so we can minimize our memory footprint by throwing away scans we no longer need
        scans_to_submit.sort(key=operator.itemgetter(0))
        if ion_search or all_msn or args.mva:
            scan_count = len(scans_to_submit)
        for i in scans_to_submit:
            in_queue.put(i[1])

        sys.stderr.write("{0} processed and placed into queue.\n".format(filename))

        # kill the workers
        [in_queue.put(None) for i in range(threads)]
        RESULT_DICT = {i[0]: "NA" for i in RESULT_ORDER}
        scans_to_export = set([])
        export_mapping = defaultdict(set)
        key_map = list(msn_rt_map.keys())
        rt_scan_map = pd.Series(key_map, index=[msn_rt_map[i] for i in key_map])
        rt_scan_map.sort_index(inplace=True)

        result = None
        while workers or result is not None:
            try:
                result = result_queue.get(timeout=0.1)
            except Empty:
                # kill expired workers
                result = None
                to_del = []
                for i, v in enumerate(workers):
                    if not v.is_alive():
                        v.terminate()
                        exit_code = v.exitcode
                        if exit_code in CRASH_SIGNALS:
                            print(
                                "thread has been killed using params {}".format(
                                    v.params
                                )
                            )
                        to_del.append(
                            {
                                "worker_index": i,
                                "thread_id": v.thread,
                                "exitcode": exit_code,
                            }
                        )
                workers_to_add = []
                for worker_dict in sorted(
                    to_del, key=operator.itemgetter("worker_index"), reverse=True
                ):
                    worker_index = worker_dict["worker_index"]
                    if worker_dict["exitcode"] in CRASH_SIGNALS:
                        thread_index = worker_dict["thread_id"]
                        worker = Worker(
                            queue=in_queue,
                            results=result_queue,
                            raw_name=filepath,
                            mass_labels=mass_labels,
                            debug=args.debug,
                            html=html,
                            mono=not args.spread,
                            precursor_ppm=args.precursor_ppm,
                            isotope_ppm=args.isotope_ppm,
                            isotope_ppms=None,
                            msn_rt_map=msn_rt_map,
                            reporter_mode=reporter_mode,
                            reader_in=reader_in,
                            reader_out=reader_outs[thread_index],
                            thread=thread_index,
                            quant_method=quant_method,
                            spline=spline,
                            isotopologue_limit=isotopologue_limit,
                            labels_needed=labels_needed,
                            quant_msn_map=[i for i in msn_map if i[0] == msn_for_quant]
                            if not args.mrm
                            else msn_map,
                            overlapping_mz=overlapping_mz,
                            min_resolution=args.min_resolution,
                            min_scans=args.min_scans,
                            mrm_pair_info=mrm_pair_info,
                            mrm=args.mrm,
                            peak_cutoff=args.peak_cutoff,
                            replicate=args.mva,
                            ref_label=ref_label,
                            max_peaks=args.max_peaks,
                            parser_args=args,
                        )
                        workers_to_add.append(worker)
                        worker.start()
                    del workers[worker_index]
                workers += workers_to_add
            if result is not None:
                completed += 1
                if completed % 10 == 0:
                    sys.stderr.write(
                        "\r{0:2.2f}% Completed".format(completed / scan_count * 100)
                    )
                    sys.stderr.flush()
                res_dict = copy.deepcopy(RESULT_DICT)
                for i in RESULT_ORDER:
                    res_dict[i[0]] = result.get(i[0], "NA")
                peak_report = []
                for label_name in labels:
                    peaks_found = result.get("{}_peaks".format(label_name), {})
                    if args.export_mzml:
                        from . import PER_FILE, PER_ID, PER_PEAK

                        scans = get_scans_under_peaks(rt_scan_map, peaks_found)
                        flattened_scans = set(
                            [
                                _peak_scan_id
                                for _peak_isotope, _peak_isotopes in scans.items()
                                for _xic_peak_index, _xic_peak_scans in _peak_isotopes.items()
                                for _peak_scan_id in _xic_peak_scans
                            ]
                        )
                        scans_to_export |= flattened_scans
                        if args.export_mode == PER_FILE:
                            export_mapping[
                                "{}_{}.mzML".format(out_path, filename)
                            ] |= flattened_scans
                    if len(peaks_found) > most_peaks_found:
                        most_peaks_found = len(peaks_found)

                    xic_peak_summary = OrderedDict()
                    for isotope_index, isotope_peaks in peaks_found.items():
                        if args.export_mzml and args.export_mode == PER_ID:
                            export_mapping[
                                "{out}_{raw}_{ms1}_{precursor}.mzML".format(
                                    **{
                                        "out": out_path,
                                        "raw": filename,
                                        "precursor": result.get(
                                            "{}_precursor".format(label_name)
                                        ),
                                        "ms1": result.get("ms1"),
                                    }
                                )
                            ] |= set(
                                [
                                    _xic_peak_scans
                                    for _xic_peak_index, _xic_peak_scans in scans[
                                        isotope_index
                                    ].items()
                                    for _peak_scan in _xic_peak_scans
                                ]
                            )
                        for xic_peak_index, xic_peak_info in isotope_peaks.items():
                            if args.export_mzml and args.export_mode == PER_PEAK:
                                export_mapping[
                                    "{out}_{raw}_{ms1}_{precursor}_{isotope}_{peak}.mzML".format(
                                        **{
                                            "out": out_path,
                                            "raw": filename,
                                            "peak": xic_peak_index,
                                            "isotope": isotope_index,
                                            "precursor": result.get(
                                                "{}_precursor".format(label_name)
                                            ),
                                            "ms1": result.get("ms1"),
                                        }
                                    )
                                ] |= scans[isotope_index][xic_peak_index]
                            if args.peaks_n == 1:
                                # TODO: fix this terrible loop
                                for i, v in enumerate(RESULT_ORDER):
                                    for j in xic_peak_info:
                                        if "{}_{}".format(label_name, j) == v[0]:
                                            try:
                                                xic_peak_summary[i].append(
                                                    xic_peak_info[j]
                                                )
                                            except KeyError:
                                                xic_peak_summary[i] = [xic_peak_info[j]]
                            else:
                                peak_report.append(
                                    list(
                                        map(
                                            str,
                                            (
                                                xic_peak_info.get(i[0], "NA")
                                                for i in PEAK_REPORTING
                                            ),
                                        )
                                    )
                                )

                    if args.peaks_n == 1:
                        for index, values in xic_peak_summary.items():
                            result_info = RESULT_ORDER[index]
                            if callable(result_info[-1]):
                                res_dict[result_info[0]] = result_info[-1](values)
                            else:
                                res_dict[result_info[0]] = values[0]
                res_dict["filename"] = filename
                # first entry of peak report is label, sort alphabetically
                peak_report.sort(key=operator.itemgetter(0))
                res_dict["peak_report"] = peak_report
                # This is the tsv output we provide
                res_list = (
                    [filename]
                    + [res_dict.get(i[0], "NA") for i in RESULT_ORDER]
                    + ["\t".join(i) for i in peak_report]
                )
                res = "{0}\n".format("\t".join(map(str, res_list)))
                out.write(res)
                out.flush()

                temp_file.write(
                    pd.io.json.dumps(
                        {"res_dict": res_dict, "html": result.get("html", {})}
                    )
                )
                temp_file.write("\n")
                temp_file.flush()

        if scans_to_export:
            raw_file = GuessIterator(open(filepath, "rb"), full=True, store=False)
            for export_filename, scans in export_mapping.items():
                with open(export_filename, "wb") as o:
                    raw_file.parser.writeScans(handle=o, scans=sorted(scans))

        reader_in.put(None)

        del msn_map
        del scan_rt_map
        del raw

    out.flush()
    out.close()
    # fix the header if we need to. We reopen the file because Windows doesn't like it when we read on a file with 'w+'
    if args.peaks_n != 1 and out != sys.stdout:
        tmp_file = "{}_ot".format(out_path)
        with open(out.name, "r") as out:
            with open(tmp_file, "w") as o:
                new_header = out.readline().strip().split("\t")
                for peak_num in range(most_peaks_found):
                    new_header.extend(
                        ["Peak {} {}".format(peak_num, i[1]) for i in PEAK_REPORTING]
                    )
                o.write("{}\n".format("\t".join(new_header)))
                for row in out:
                    o.write(row)
        import shutil

        shutil.move(tmp_file, out_path)

    temp_file.close()
    df_data = []
    html_data = []
    peak_data = []

    for j in open(temp_file.name, "r"):
        result = pd.io.json.loads(j if isinstance(j, str) else j.decode("utf-8"))
        res_dict = result["res_dict"]
        res_list = [res_dict["filename"]] + [
            res_dict.get(i[0], "NA") for i in RESULT_ORDER
        ]
        peak_data.append(res_dict["peak_report"])
        df_data.append(res_list)
        html_data.append(result["html"])
    data = pd.DataFrame.from_records(
        df_data, columns=[i for i in headers if i != "Confidence"]
    )
    header_mapping = []
    order_names = [i[1] for i in RESULT_ORDER]
    for i in data.columns:
        try:
            header_mapping.append(RESULT_ORDER[order_names.index(i)][0])
        except ValueError:
            header_mapping.append(False)
    if calc_stats:
        perform_ml(data, mass_labels)
        data.to_csv("{}_stats".format(out_path), sep=str("\t"), index=None)

    if html:
        html_map = []
        peak_map = {"header": [i[1] for i in PEAK_REPORTING], "data": []}
        for index, (row_index, row) in enumerate(data.iterrows()):
            res = "\t".join(row.astype(str))
            to_write, html_info = table_rows(
                [
                    {
                        "table": res.strip(),
                        "html": html_data[index],
                        "keys": header_mapping,
                    }
                ]
            )
            html_map.append(html_info)
            peak_map["data"].append(peak_data[index])
            html_out.write(to_write)
            html_out.flush()

        template = []
        append = False
        for i in open(pyquant_html_file, "r"):
            if "HTML BREAK" in i:
                append = True
            elif append:
                template.append(i)
        html_template = Template("".join(template))
        peak_map = pd.io.json.dumps(peak_map)
        html_map = pd.io.json.dumps(html_map)
        html_out.write(
            html_template.safe_substitute(
                {
                    "peak_output": base64.b64encode(
                        gzip.zlib.compress(bytes(peak_map, "utf-8"), 9,)
                    ).decode("utf-8"),
                    "html_output": base64.b64encode(
                        gzip.zlib.compress(bytes(html_map, "utf-8"), 9,)
                    ).decode("utf-8"),
                }
            )
        )

    os.remove(temp_file.name)