#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Aug 30 09:19:54 2019

@author: bressler
"""

import SBCcode as sbc
import numpy as np
from collections import Counter
from LED_blocking_scintillation import isBlocked
import runlistscatalogue as rlc
from multibubblescintillationcheck import check_multibub_scintillation
import os
import gc
import traceback
import csv

def merge_normal(runsToMerge, name, pmtgain, PMTwindow):

    allxyzfname = "/pnfs/coupp/persistent/grid_output/SBC-17/output/SimpleXYZ_all.bin"
    xyzf = sbc.DataHandling.ReadBinary.ReadBlock(allxyzfname)
    totevents = 0
    blockedEvents = 0
    totbub = 0
    nosigNotBlockedEvents = 0
    with open("/nashome/b/bressler/sbcoutput/%s_merged.txt"%name,"w+") as fout,\
        open("/nashome/b/bressler/sbcoutput/%s_nosig_merged.txt"%name,"w+") as fnosig:
        fout.write("run event nbub x y z AP1 AP2 at0_0 at0_1 PMTt0 PMTt0_ind PMTphe lag which_acoustic_channel isBlocked\n")

        for run in runsToMerge:
            events = [evnt for evnt in os.listdir('/bluearc/storage/SBC-17-data/'+run) if not os.path.isfile(os.path.join('/bluearc/storage/SBC-17-data/'+run, evnt))]
            Nevents = len(events)
            ch0file=open("/nashome/b/bressler/sbcoutput/%s_PMTmatching_ch0.txt"%run,"r")
            ch1file=open("/nashome/b/bressler/sbcoutput/%s_PMTmatching_ch1.txt"%run,"r")

            APfile = open("/nashome/b/bressler/sbcoutput/XeBCAP_%s.txt"%run,"r")

            indices = [i for i,x in enumerate(xyzf["runid"]) if str(x[0])+"_"+str(x[1]) == run]
            runposreco = {"ev":[xyzf["ev"][indices]],"x":[xyzf["bubX"][indices]],
                          "y":[xyzf["bubY"][indices]],"z":[xyzf["bubZ"][indices]]}
            runreconpath = "/pnfs/coupp/persistent/grid_output/SBC-17/output/%s/"%run
            acousticfilename = runreconpath+"AcousticAnalysis_%s.bin"%run
            a = sbc.DataHandling.ReadBinary.ReadBlock(acousticfilename)
            bubt0 = a["bubble_t0"]

            getbubfile = "/coupp/data/home/coupp/HumanGetBub_output_SBC-17/HumanGetBub_%s.bin"%run
            c = sbc.DataHandling.ReadBinary.ReadBlock(getbubfile)
            eventn = c["ev"]
            count = Counter(eventn)

            AP1s = []
            AP2s = []

            ev_with_scintillation_0 = []
            at0_0s = []
            indices_0 = []
            pmtt0s_0 = []
            lags_0 = []
            phes_0 = []

            ev_with_scintillation_1 = []
            at0_1s = []
            indices_1 = []
            pmtt0s_1 = []
            lags_1 = []
            phes_1 = []

            ch0 = enumerate(ch0file)
            ch1 = enumerate(ch1file)
            for i,l in ch0:
                if i>=1:
                    d=l.split()
                    run_num = d[0]
                    event_num = d[1]
                    ev_with_scintillation_0.append(int(event_num))
                    ind = d[2]
                    indices_0.append(ind)
                    lag = d[3]
                    lags_0.append(lag)
                    pmtt0 = d[4]
                    pmtt0s_0.append(pmtt0)
                    at0 = d[5]
                    at0_0s.append(at0)
                    phe = d[6]
                    phes_0.append(phe)
                    z = d[7]

            for i,l in ch1:
                if i>=1:
                    d=l.split()
                    run_num = d[0]
                    event_num = d[1]
                    ev_with_scintillation_1.append(int(event_num))
                    ind = d[2]
                    indices_1.append(ind)
                    lag = d[3]
                    lags_1.append(lag)
                    pmtt0 = d[4]
                    pmtt0s_1.append(pmtt0)
                    at0 = d[5]
                    at0_1s.append(at0)
                    phe = d[6]
                    phes_1.append(phe)
                    z = d[7]

            ch0file.close()
            ch1file.close()

            ap = enumerate(APfile)
            for i,l in ap:
                if i>=1:
                    d=l.split()
                    AP1s.append(float(d[1])) # 2 for raw
                    AP2s.append(float(d[4])) # 5 for raw

            #for evnt in runposreco["ev"][0]:
                # this was actually where I'm selecting single bubble events by only looping through events with position reconstruction
            for evnt in range(Nevents):
                try:
                    N = count[evnt]
                    if N%2 == 0:
                        n = N/2

                    else:
                        n= 0
                    totevents += 1
                    AP1 = AP1s[evnt]
                    AP2 = AP2s[evnt]
                    if evnt in runposreco['ev'][0]:
                        position_reconstruction_event_index = list(runposreco['ev'][0]).index(evnt)
                        x=runposreco["x"][0][position_reconstruction_event_index]
                        y=runposreco["y"][0][position_reconstruction_event_index]
                        z=runposreco["z"][0][position_reconstruction_event_index]
                    else:
                        x = np.nan
                        y = np.nan
                        z = np.nan
                    at0_0=bubt0[evnt,0]
                    at0_1=bubt0[evnt,1]
                    in0=False
                    in1=False
                    PMTt0 = np.nan
                    sig = np.nan
                    pmtt0ind = np.nan
                    lag = np.nan
                    chan = np.nan

                    ind0 = np.nan
                    ind1 = np.nan
                    if evnt in ev_with_scintillation_0:
                        in0=True
                        ind0 = ev_with_scintillation_0.index(evnt)

                    if evnt in ev_with_scintillation_1:
                        in1=True
                        ind1 = ev_with_scintillation_1.index(evnt)

                    if (in0 and in1) or (in0 and not in1):
                        PMTt0 = pmtt0s_0[ind0]
                        sig = phes_0[ind0]
                        pmtt0ind = indices_0[ind0]
                        lag = lags_0[ind0]
                        chan = 0.

                    elif (in1 and not in0):
                        PMTt0 = pmtt0s_1[ind1]
                        sig = phes_1[ind1]
                        pmtt0ind = indices_1[ind1]
                        lag = lags_1[ind1]
                        chan = 1.
                    if n>1 and not np.isnan(sig):
                        print("multibubble with light")
                    bl = 0

                    if not np.isnan(float(x)):
                        totbub+= 1
                    #check only events with no signal:
                    """
                    if math.isnan(float(sig)) and not math.isnan(float(x)) and n==1:
                        bl0=isBlocked(run,int(evnt),float(at0_0),float(z))
                        bl1=isBlocked(run,int(evnt),float(at0_1),float(z))
                        bl = bl0# or bl1

                        if not bl:
                            fnosig.write("%s %d\n"%(str(run),int(evnt)))
                            nosigNotBlockedEvents += 1
                            #print('not blocked and zero light: %s-%d\n'%(str(run),int(evnt)))
                            #print('phes_0[ind0]: %s'%phes_0[ind0])
                            #print('phes_1[ind1]: %s'%phes_1[ind1])

                        else:
                            blockedEvents +=1
                    elif math.isnan(float(sig)) and  n>1:
                        multibub_check = check_multibub_scintillation(run, evnt, float(at0_0), pmtgain)
                        #bl1 = check_multibub_scintillation(run, evnt, float(at0_1))
                        bl = multibub_check[0]
                        sig = multibub_check[1]
                    """

                    #check all events:
                    if n == 1:
                        bl = isBlocked(run, int(evnt), float(at0_0), float(z), PMTwindow)
                        if not bl:
                            nosigNotBlockedEvents += 1
                        else:
                            blockedEvents += 1
                    elif n>1:
                        multibub_check = check_multibub_scintillation(run, evnt, float(at0_0), pmtgain, PMTwindow)
                        #bl1 = check_multibub_scintillation(run, evnt, float(at0_1))
                        bl = multibub_check[0]
                        sig = multibub_check[1]

                    fout.write("%s %d %d %f %f %f %f %f %f %f %f %.0f %f %f %.0f %0f\n"%(str(run),int(evnt), int(n),
                        float(x),float(y),float(z),AP1,AP2,float(at0_0),float(at0_1),float(PMTt0),
                                                                      float(pmtt0ind),float(sig),
                                                                      float(lag),float(chan),float(bl)))
                    print("wrote %s-%s"%(str(run),str(evnt)))
                    #gc.collect()
                except Exception as x:
                    print(x)
                    print(traceback.format_exc())
    print("fraction of events blocked by LED: %f"%(blockedEvents/totevents))
    print("fraction of bubbles blocked by LED: %f"%(blockedEvents/totbub))
    print("fraction of events with no scintillation and not blocked: %f"%(nosigNotBlockedEvents/totevents))

def merge_handscanned(runsToMerge, name, pmtgain, PMTwindow):

    allxyzfname = "/pnfs/coupp/persistent/grid_output/SBC-17/output/SimpleXYZ_all.bin"
    xyzf = sbc.DataHandling.ReadBinary.ReadBlock(allxyzfname)
    totevents = 0
    blockedEvents = 0
    totbub = 0
    nosigNotBlockedEvents = 0
    with open("/nashome/b/bressler/sbcoutput/%s_merged.txt"%name,"w+") as fout,\
        open("/nashome/b/bressler/sbcoutput/%s_nosig_merged.txt"%name,"w+") as fnosig:
        fout.write("run event nbub x y z AP1 AP2 at0_0 at0_1 PMTt0 PMTt0_ind PMTphe lag which_acoustic_channel isBlocked\n")

        for run in runsToMerge:
            events = [evnt for evnt in os.listdir('/bluearc/storage/SBC-17-data/'+run) if not os.path.isfile(os.path.join('/bluearc/storage/SBC-17-data/'+run, evnt))]
            Nevents = len(events)
            ch1file=open("/nashome/b/bressler/sbcoutput/%s_PMTmatching_ch1.txt"%run,"r")

            APfile = open("/nashome/b/bressler/sbcoutput/XeBCAP_%s.txt"%run,"r")

            indices = [i for i,x in enumerate(xyzf["runid"]) if str(x[0])+"_"+str(x[1]) == run]
            runposreco = {"ev":[xyzf["ev"][indices]],"x":[xyzf["bubX"][indices]],
                          "y":[xyzf["bubY"][indices]],"z":[xyzf["bubZ"][indices]]}

            bubt0_dict = {}
            with open("/nashome/b/bressler/sbcoutput/handscans_acoustic_t0/acoustict0handscan_%s.csv"%name, 'r') as csvfile:
                r = csv.reader(csvfile)
                i=0
                for row in r:
                    if i>0:
                        bubt0_dict[row[0]+'-'+row[1]] = float(row[2])
                    i += 1

            getbubfile = "/coupp/data/home/coupp/HumanGetBub_output_SBC-17/HumanGetBub_%s.bin"%run
            c = sbc.DataHandling.ReadBinary.ReadBlock(getbubfile)
            eventn = c["ev"]
            count = Counter(eventn)

            AP2s = []

            ev_with_scintillation_1 = []
            at0_1s = []
            indices_1 = []
            pmtt0s_1 = []
            lags_1 = []
            phes_1 = []

            ch1 = enumerate(ch1file)

            for i,l in ch1:
                if i>=1:
                    d=l.split()
                    run_num = d[0]
                    event_num = d[1]
                    ev_with_scintillation_1.append(int(event_num))
                    ind = d[2]
                    indices_1.append(ind)
                    lag = d[3]
                    lags_1.append(lag)
                    pmtt0 = d[4]
                    pmtt0s_1.append(pmtt0)
                    at0 = d[5]
                    at0_1s.append(at0)
                    phe = d[6]
                    phes_1.append(phe)
                    z = d[7]

            ch1file.close()

            ap = enumerate(APfile)
            for i,l in ap:
                if i>=1:
                    d=l.split()
                    AP2s.append(float(d[4])) # 5 for raw

            #for evnt in runposreco["ev"][0]:
                # this was actually where I'm selecting single bubble events by only looping through events with position reconstruction
            for evnt in range(Nevents):
                try:
                    N = count[evnt]
                    if N%2 == 0:
                        n = N/2

                    else:
                        n= 0
                    totevents += 1
                    AP2 = AP2s[evnt]
                    if evnt in runposreco['ev'][0]:
                        position_reconstruction_event_index = list(runposreco['ev'][0]).index(evnt)
                        x=runposreco["x"][0][position_reconstruction_event_index]
                        y=runposreco["y"][0][position_reconstruction_event_index]
                        z=runposreco["z"][0][position_reconstruction_event_index]
                    else:
                        x = np.nan
                        y = np.nan
                        z = np.nan
                    at0_1=bubt0_dict[run+'-'+str(evnt)]
                    in1=False
                    PMTt0 = np.nan
                    sig = np.nan
                    pmtt0ind = np.nan
                    lag = np.nan
                    chan = np.nan

                    ind1 = np.nan

                    if evnt in ev_with_scintillation_1:
                        in1=True
                        ind1 = ev_with_scintillation_1.index(evnt)

                    if in1:
                        PMTt0 = pmtt0s_1[ind1]
                        sig = phes_1[ind1]
                        pmtt0ind = indices_1[ind1]
                        lag = lags_1[ind1]
                        chan = 1.
                    if n>1 and not np.isnan(sig):
                        print("multibubble with light")
                    bl = 0

                    if not np.isnan(float(x)):
                        totbub+= 1
                    #check only events with no signal:
                    """
                    if math.isnan(float(sig)) and not math.isnan(float(x)) and n==1:
                        bl0=isBlocked(run,int(evnt),float(at0_0),float(z))
                        bl1=isBlocked(run,int(evnt),float(at0_1),float(z))
                        bl = bl0# or bl1

                        if not bl:
                            fnosig.write("%s %d\n"%(str(run),int(evnt)))
                            nosigNotBlockedEvents += 1
                            #print('not blocked and zero light: %s-%d\n'%(str(run),int(evnt)))
                            #print('phes_0[ind0]: %s'%phes_0[ind0])
                            #print('phes_1[ind1]: %s'%phes_1[ind1])

                        else:
                            blockedEvents +=1
                    elif math.isnan(float(sig)) and  n>1:
                        multibub_check = check_multibub_scintillation(run, evnt, float(at0_0), pmtgain)
                        #bl1 = check_multibub_scintillation(run, evnt, float(at0_1))
                        bl = multibub_check[0]
                        sig = multibub_check[1]
                    """

                    #check all events:
                    if n == 1:
                        bl = isBlocked(run, int(evnt), float(at0_1), float(z), PMTwindow)
                        if not bl:
                            nosigNotBlockedEvents += 1
                        else:
                            blockedEvents += 1
                    elif n>1:
                        multibub_check = check_multibub_scintillation(run, evnt, float(at0_1), pmtgain, PMTwindow)
                        #bl1 = check_multibub_scintillation(run, evnt, float(at0_1))
                        bl = multibub_check[0]
                        sig = multibub_check[1]

                    AP1 = np.nan
                    at0_0 = np.nan
                    fout.write("%s %d %d %f %f %f %f %f %f %f %f %.0f %f %f %.0f %0f\n"%(str(run),int(evnt), int(n),
                        float(x),float(y),float(z),AP1,AP2,float(at0_0),float(at0_1),float(PMTt0),
                                                                      float(pmtt0ind),float(sig),
                                                                      float(lag),float(chan),float(bl)))
                    print("wrote %s-%s"%(str(run),str(evnt)))
                    #gc.collect()
                except Exception as x:
                    print(x)
                    print(traceback.format_exc())
    print("fraction of events blocked by LED: %f"%(blockedEvents/totevents))
    print("fraction of bubbles blocked by LED: %f"%(blockedEvents/totbub))
    print("fraction of events with no scintillation and not blocked: %f"%(nosigNotBlockedEvents/totevents))

def merge(runs, isHandscannedt0, datasetname, PMTgain, PMTwindow):
    if isHandscannedt0:
        merge_handscanned(runs, datasetname, PMTgain, PMTwindow)
    else:
        merge_normal(runs, datasetname, PMTgain, PMTwindow)
def main():
    merge(rlc.cfJuneminus50CCombined, False, 'cfJuneminus50CCombined', 1e7, [-600e-6, 600e-6])
if __name__=="__main__":
    main()


from __future__ import division

import math
import os
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFont
from torch.autograd import Variable
from torchvision.ops import nms

class DecodeBox(nn.Module):
    def __init__(self, anchors, num_classes, img_size):
        super(DecodeBox, self).__init__()
        #-----------------------------------------------------------#
        #   13x13的特征层对应的anchor是[116,90],[156,198],[373,326]
        #   26x26的特征层对应的anchor是[30,61],[62,45],[59,119]
        #   52x52的特征层对应的anchor是[10,13],[16,30],[33,23]
        #-----------------------------------------------------------#
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.bbox_attrs = 5 + num_classes
        self.img_size = img_size

    def forward(self, input):
        #-----------------------------------------------#
        #   输入的input一共有三个，他们的shape分别是
        #   batch_size, 255, 13, 13
        #   batch_size, 255, 26, 26
        #   batch_size, 255, 52, 52
        #-----------------------------------------------#
        batch_size = input.size(0)
        input_height = input.size(2)
        input_width = input.size(3)

        #-----------------------------------------------#
        #   输入为416x416时
        #   stride_h = stride_w = 32、16、8
        #-----------------------------------------------#
        stride_h = self.img_size[1] / input_height
        stride_w = self.img_size[0] / input_width
        #-------------------------------------------------#
        #   此时获得的scaled_anchors大小是相对于特征层的
        #-------------------------------------------------#
        scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors]

        #-----------------------------------------------#
        #   输入的input一共有三个，他们的shape分别是
        #   batch_size, 3, 13, 13, 85
        #   batch_size, 3, 26, 26, 85
        #   batch_size, 3, 52, 52, 85
        #-----------------------------------------------#
        prediction = input.view(batch_size, self.num_anchors,
                                self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()

        # 先验框的中心位置的调整参数
        x = torch.sigmoid(prediction[..., 0])
        y = torch.sigmoid(prediction[..., 1])
        # 先验框的宽高调整参数
        w = prediction[..., 2]
        h = prediction[..., 3]
        # 获得置信度，是否有物体
        conf = torch.sigmoid(prediction[..., 4])
        # 种类置信度
        pred_cls = torch.sigmoid(prediction[..., 5:])

        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor

        #----------------------------------------------------------#
        #   生成网格，先验框中心，网格左上角
        #   batch_size,3,13,13
        #----------------------------------------------------------#
        grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
            batch_size * self.num_anchors, 1, 1).view(x.shape).type(FloatTensor)
        grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(
            batch_size * self.num_anchors, 1, 1).view(y.shape).type(FloatTensor)

        #----------------------------------------------------------#
        #   按照网格格式生成先验框的宽高
        #   batch_size,3,13,13
        #----------------------------------------------------------#
        anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
        anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
        anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
        anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)

        #----------------------------------------------------------#
        #   利用预测结果对先验框进行调整
        #   首先调整先验框的中心，从先验框中心向右下角偏移
        #   再调整先验框的宽高。
        #----------------------------------------------------------#
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + grid_x
        pred_boxes[..., 1] = y.data + grid_y
        pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
        pred_boxes[..., 3] = torch.exp(h.data) * anchor_h

        #----------------------------------------------------------#
        #   将输出结果调整成相对于输入图像大小
        #----------------------------------------------------------#
        _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
        output = torch.cat((pred_boxes.view(batch_size, -1, 4) * _scale,
                            conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
        return output.data

def letterbox_image(image, size):
    iw, ih = image.size
    w, h = size
    scale = min(w/iw, h/ih)
    nw = int(iw*scale)
    nh = int(ih*scale)

    image = image.resize((nw,nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128,128,128))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))
    return new_image

def yolo_correct_boxes(top, left, bottom, right, input_shape, image_shape):
    new_shape = image_shape*np.min(input_shape/image_shape)

    offset = (input_shape-new_shape)/2./input_shape
    scale = input_shape/new_shape

    box_yx = np.concatenate(((top+bottom)/2,(left+right)/2),axis=-1)/input_shape
    box_hw = np.concatenate((bottom-top,right-left),axis=-1)/input_shape

    box_yx = (box_yx - offset) * scale
    box_hw *= scale

    box_mins = box_yx - (box_hw / 2.)
    box_maxes = box_yx + (box_hw / 2.)
    boxes =  np.concatenate([
        box_mins[:, 0:1],
        box_mins[:, 1:2],
        box_maxes[:, 0:1],
        box_maxes[:, 1:2]
    ],axis=-1)
    boxes *= np.concatenate([image_shape, image_shape],axis=-1)
    return boxes

def bbox_iou(box1, box2, x1y1x2y2=True):
    """
        计算IOU
    """
    if not x1y1x2y2:
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
    else:
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)

    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * \
                 torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)

    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    return iou

def non_max_suppression(prediction, num_classes, conf_thres=0.5, nms_thres=0.4):
    #----------------------------------------------------------#
    #   将预测结果的格式转换成左上角右下角的格式。
    #   prediction  [batch_size, num_anchors, 85]
    #----------------------------------------------------------#
    box_corner = prediction.new(prediction.shape)
    box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
    box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
    box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
    box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
    prediction[:, :, :4] = box_corner[:, :, :4]

    output = [None for _ in range(len(prediction))]
    for image_i, image_pred in enumerate(prediction):
        #----------------------------------------------------------#
        #   对种类预测部分取max。
        #   class_conf  [num_anchors, 1]    种类置信度
        #   class_pred  [num_anchors, 1]    种类
        #----------------------------------------------------------#
        class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True)

        #----------------------------------------------------------#
        #   利用置信度进行第一轮筛选
        #----------------------------------------------------------#
        conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()

        #----------------------------------------------------------#
        #   根据置信度进行预测结果的筛选
        #----------------------------------------------------------#
        image_pred = image_pred[conf_mask]
        class_conf = class_conf[conf_mask]
        class_pred = class_pred[conf_mask]
        if not image_pred.size(0):
            continue
        #-------------------------------------------------------------------------#
        #   detections  [num_anchors, 7]
        #   7的内容为：x1, y1, x2, y2, obj_conf, class_conf, class_pred
        #-------------------------------------------------------------------------#
        detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)

        #------------------------------------------#
        #   获得预测结果中包含的所有种类
        #------------------------------------------#
        unique_labels = detections[:, -1].cpu().unique()

        if prediction.is_cuda:
            unique_labels = unique_labels.cuda()
            detections = detections.cuda()

        for c in unique_labels:
            #------------------------------------------#
            #   获得某一类得分筛选后全部的预测结果
            #------------------------------------------#
            detections_class = detections[detections[:, -1] == c]

            #------------------------------------------#
            #   使用官方自带的非极大抑制会速度更快一些！
            #------------------------------------------#
            keep = nms(
                detections_class[:, :4],
                detections_class[:, 4] * detections_class[:, 5],
                nms_thres
            )
            max_detections = detections_class[keep]

            # # 按照存在物体的置信度排序
            # _, conf_sort_index = torch.sort(detections_class[:, 4]*detections_class[:, 5], descending=True)
            # detections_class = detections_class[conf_sort_index]
            # # 进行非极大抑制
            # max_detections = []
            # while detections_class.size(0):
            #     # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
            #     max_detections.append(detections_class[0].unsqueeze(0))
            #     if len(detections_class) == 1:
            #         break
            #     ious = bbox_iou(max_detections[-1], detections_class[1:])
            #     detections_class = detections_class[1:][ious < nms_thres]
            # # 堆叠
            # max_detections = torch.cat(max_detections).data

            # Add max detections to outputs
            output[image_i] = max_detections if output[image_i] is None else torch.cat(
                (output[image_i], max_detections))

    return output