import sys
sys.path.insert(0, "../data_gen/")
sys.path.insert(0, "../unet/")

import pandas as pd
from dataset import getKpKeys, getKpNum, getFlipMapID, get_kp_index_from_allkeys, generate_input_mask
from kpAnno import KpAnno
from post_process import post_process_heatmap
from keras.models import load_model
import os
from refinenet_mask_v3 import euclidean_loss
import numpy as np
import cv2
from resnet101 import Scale
from utils import load_annotation_from_df
from collections import defaultdict
import copy
from data_process import pad_image_inference

class Evaluation(object):
    def __init__(self, category, modelFile):
        self.category = category
        self.train_img_path = "../../data/train"
        if modelFile is not None:
            self._initialize(modelFile)

    def init_from_model(self, model):
        self._load_anno()
        self.net = model

    def eval(self, multiOut=False, details=False, flip=True):
        xdf = self.annDataFrame
        scores = list()
        xdict = dict()
        xcategoryDict = defaultdict(list)
        for _index, _row in xdf.iterrows():
            imgId = _row['image_id']
            category = _row['image_category']
            imgFile = os.path.join(self.train_img_path, imgId)
            gtKpAnno = self._get_groundtruth_kpAnno(_row)
            if flip:
                predKpAnno = self.predict_kp_with_flip(imgFile, category)
            else:
                predKpAnno = self.predict_kp(imgFile, category, multiOut)
            neScore = Evaluation.calc_ne_score(category, predKpAnno, gtKpAnno)
            scores.extend(neScore)
            if details:
                xcategoryDict[category].extend(neScore)
        if details:
            return sum(scores)/len(scores), xcategoryDict
        else:
            return sum(scores)/len(scores)

    def _initialize(self, modelFile):
        self._load_anno()
        self._initialize_network(modelFile)

    def _initialize_network(self, modelFile):
        self.net = load_model(modelFile, custom_objects={'euclidean_loss': euclidean_loss, 'Scale': Scale})

    def _load_anno(self):
        '''
        Load annotations from train.csv
        '''
        self.annfile = os.path.join("../../data/train/Annotations", "val_split.csv")

        # read into dataframe
        xpd = pd.read_csv(self.annfile)
        xpd = load_annotation_from_df(xpd, self.category)
        self.annDataFrame = xpd

    def _get_groundtruth_kpAnno(self, dfrow):
        mlist = dfrow[getKpKeys(self.category)]
        imgName, kpStr = mlist[0], mlist[1:]
        # read kp annotation from csv file
        kpAnnlst = [KpAnno.readFromStr(_kpstr) for _kpstr in kpStr]
        return kpAnnlst

    def _net_inference_with_mask(self, imgFile, imgCategory):
        import cv2
        from data_process import normalize_image, pad_image_inference
        assert (len(self.net.input_layers) > 1), "input layer need to more than 1"

        # load image and preprocess
        img = cv2.imread(imgFile)

        img, scale = pad_image_inference(img, 512, 512)
        img   = normalize_image(img)
        input_img = img[np.newaxis, :, :, :]

        input_mask = generate_input_mask(imgCategory, (512, 512, getKpNum(self.category)) )
        input_mask = input_mask[np.newaxis, :, :, :]

        # inference
        heatmap = self.net.predict([input_img, input_mask, input_mask])

        return (heatmap, scale)

    def _heatmap_sum(self, heatmaplst):
        outheatmap = np.copy(heatmaplst[0])
        for i in range(1, len(heatmaplst), 1):
            outheatmap += heatmaplst[i]
        return outheatmap

    def predict_kp(self, imgFile, imgCategory, multiOutput=False):

        xnetout, scale = self._net_inference_with_mask(imgFile, imgCategory)

        if multiOutput:
            #todo: fixme, it is tricky that the previous stage has beeter performance than last stage's output.
            #todo: here, we are using multiple stage's output sum.
            heatmap = self._heatmap_sum(xnetout)
        else:
            heatmap = xnetout

        detectedKps = post_process_heatmap(heatmap, kpConfidenceTh=0.2)

        # scale to padded resolution 256X256 -> 512X512
        scaleTo512 = 2.0

        # apply scale to original resolution
        detectedKps = [KpAnno(_kp.x*scaleTo512/scale , _kp.y*scaleTo512/scale, _kp.visibility) for _kp in detectedKps]

        return detectedKps

    def predict_kp_with_flip(self, imgFile, imgCategory):
        #  inference with flip and original image
        heatmap, scale = self._net_inference_flip(imgFile, imgCategory)

        detectedKps = post_process_heatmap(heatmap, kpConfidenceTh=0.2)

        # scale to padded resolution 256X256 -> 512X512
        scaleTo512 = 2.0

        # apply scale to original resolution
        detectedKps = [KpAnno(_kp.x * scaleTo512 / scale, _kp.y * scaleTo512 / scale, _kp.visibility) for _kp in
                       detectedKps]

        return detectedKps

    def _net_inference_flip(self, imgFile, imgCategory):
        import cv2
        from data_process import normalize_image, pad_image_inference
        assert (len(self.net.input_layers) > 1), "input layer need to more than 1"

        batch_size =2

        input_img  = np.zeros(shape=(batch_size, 512, 512, 3), dtype=np.float)
        input_mask = np.zeros(shape=(batch_size, 256, 256, getKpNum(self.category)), dtype=np.float)

        # load image and preprocess
        orgimage = cv2.imread(imgFile)

        padimg, scale = pad_image_inference(orgimage, 512, 512)
        flipimg = cv2.flip(padimg, flipCode=1)

        input_img[0,:,:,:] = normalize_image(padimg)
        input_img[1,:,:,:] = normalize_image(flipimg)

        mask = generate_input_mask(imgCategory, (512, 512, getKpNum(self.category)))
        input_mask[0,:,:,:] = mask
        input_mask[1,:,:,:] = mask

        # inference
        if len(self.net.input_layers) == 2:
            heatmap = self.net.predict([input_img, input_mask])
        elif len(self.net.input_layers) == 3:
            heatmap = self.net.predict([input_img, input_mask, input_mask])
        else:
            assert (0), str(len(self.net.input_layers)) + " should be 2 or 3 "

        # sum heatmap
        avgheatmap = self._heatmap_sum(heatmap)

        orgheatmap = avgheatmap[0,:,:,:]

        # convert to same sequency with original heatmap
        flipheatmap = avgheatmap[1,:,:,:]
        flipheatmap = self._flip_out_heatmap(flipheatmap)

        # average original and flip heatmap
        outheatmap = flipheatmap + orgheatmap
        outheatmap = outheatmap[np.newaxis, :, :, :]

        return (outheatmap, scale)

    def predict_kp_with_rotate(self, imgFile, imgCategory):
        #  inference with rotated image
        rotateheatmap = self._net_inference_rotate(imgFile, imgCategory)
        rotateheatmap = rotateheatmap[np.newaxis, :, :, :]

        # original image and flip image
        orgflipmap, scale = self._net_inference_flip(imgFile, imgCategory)
        mflipmap = cv2.resize(orgflipmap[0,:,:,:], None, fx=2.0/scale, fy=2.0/scale)

        # add mflipmap and rotateheatmap
        avgheatmap = mflipmap[np.newaxis, :, :, :]

        b, h, w , c = rotateheatmap.shape
        avgheatmap[:, 0:h, 0:w,:] += rotateheatmap

        # generate key point locations
        detectedKps = post_process_heatmap(avgheatmap, kpConfidenceTh=0.2)

        return detectedKps

    def _net_inference_rotate(self, imgFile, imgCategory):
        from data_process import normalize_image, pad_image_inference, rotate_image_with_invrmat

        # load image and preprocess
        orgimage = cv2.imread(imgFile)

        anglelst = [-20, -10, 10, 20]

        input_img  = np.zeros(shape=(len(anglelst), 512, 512, 3), dtype=np.float)
        input_mask = np.zeros(shape=(len(anglelst), 256, 256, getKpNum(self.category)), dtype=np.float)

        mlist = list()
        for i, angle in enumerate(anglelst):
            rotateimg, invRotMatrix, orgImgSize = rotate_image_with_invrmat(orgimage, angle)
            padimg, scale = pad_image_inference(rotateimg, 512, 512)
            _img = normalize_image(padimg)
            input_img[i, :, :, :] = _img
            mlist.append((scale, invRotMatrix))

        mask = generate_input_mask(imgCategory, (512, 512, getKpNum(self.category)))
        for i, angle in enumerate(anglelst):
            input_mask[i, :,:,:] = mask

        # inference
        heatmap = self.net.predict([input_img, input_mask, input_mask])
        heatmap = self._heatmap_sum(heatmap)

        # rotate back to original resolution
        sumheatmap =  np.zeros(shape=(orgimage.shape[0], orgimage.shape[1], getKpNum(self.category)), dtype=np.float)
        for i, item in enumerate(mlist):
            _heatmap = heatmap[i, :, :, :]
            _scale, _invRotMatrix = item
            _heatmap = cv2.resize(_heatmap, None, fx=2.0 / _scale, fy=2.0 / _scale)
            _invheatmap = cv2.warpAffine(_heatmap, _invRotMatrix, (orgimage.shape[1], orgimage.shape[0]))
            sumheatmap += _invheatmap

        return sumheatmap

    def _flip_out_heatmap(self, flipout):
        outmap = np.zeros(flipout.shape, dtype=np.float)
        for i in range(flipout.shape[-1]):
            flipid = getFlipMapID(self.category, i)
            mask = np.copy(flipout[:, :, i])
            outmap[:, :, flipid] = cv2.flip(mask, flipCode=1)
        return outmap

    @staticmethod
    def get_normized_distance(category, gtKp):
        '''

        :param category:
        :param gtKp:
        :return: if ground truth's two points do not exist, return a big number 1e6
        '''

        if category in ['skirt' ,'trousers']:
            ##waistband left and right
            waistband_left_index  = get_kp_index_from_allkeys('waistband_left')
            waistband_right_index = get_kp_index_from_allkeys('waistband_right')

            if gtKp[waistband_left_index].visibility != -1 and gtKp[waistband_right_index].visibility != -1:
                distance = KpAnno.calcDistance(gtKp[waistband_left_index], gtKp[waistband_right_index])
            else:
                distance = 1e6
            return distance
        elif category in ['blouse', 'dress', 'outwear']:
            armpit_left_index  = get_kp_index_from_allkeys('armpit_left')
            armpit_right_index = get_kp_index_from_allkeys('armpit_right')
            ##armpit_left armpit_right'
            if gtKp[armpit_left_index].visibility != -1 and gtKp[armpit_right_index].visibility != -1:
                distance = KpAnno.calcDistance(gtKp[armpit_left_index], gtKp[armpit_right_index])
            else:
                distance = 1e6
            return distance
        else:
            assert (0), category + " not implemented in _get_normized_distance"

    @staticmethod
    def calc_ne_score(category, dtKp, gtKp):

        assert (len(dtKp) == len(gtKp)), "predicted keypoint number should be the same as ground truth keypoints" + \
                                         str(dtKp) + " vs " + str(gtKp)

        # calculate normalized error as score
        normalizedDistance = Evaluation.get_normized_distance(category, gtKp)

        mlist = list()
        for i in range(len(gtKp)):
            if gtKp[i].visibility == 1:
                dk = KpAnno.calcDistance(dtKp[i], gtKp[i])
                mlist.append( dk/normalizedDistance)

        return mlist


import pandas as pd
from pyteomics import mzxml, mgf, mass
from collections import defaultdict
import csodiaq_base_functions as cbf
import csodiaq_base_functions_oldMethod as cbf2
from random import randint
import re
from Bio import SeqIO
import bisect
from matplotlib import pyplot
from matplotlib_venn import venn2
import numpy as np
import pickle
import bz2
from numba import njit
from timeit import default_timer as timer
import io
import os
from numba.core import types
from numba.typed import Dict

pd.set_option('display.max_columns', None)

def fdr_calculation(df, returnType=0): #***NOTE***make return type
    # initializing the two return values at 0
    fdrValues = []
    indices = []
    numDecoys = 0
    df.fillna("nan",inplace=True)
    # for every row in the dataframe
    count = 0

    for index, row in df.iterrows():
        # current criteria for 'decoys' is to have 'decoy' in the protein name. This may change in the future.
        if (returnType==0 or returnType==1): decoy = 'DECOY' in row['Name']
        elif returnType==2: decoy = row['decoy']
        if decoy:
            numDecoys += 1

        # calculates the FDR up to this point in the data frame.
        curFDR = numDecoys/(count+1)

        # conditional statement comparing the current FDR to the FDR Cutoff. If larger, function values are returned.
        if curFDR > 0.01:

            # if the number of rows has not yet reached the minimum number that allows for the FDR cutoff, 0 is returned instead.
            if len(fdrValues) < 1/0.01:
                if returnType: return [], 0
                else: return 0, 0
            if returnType==1: return fdrValues, numDecoys-1
            if returnType==2: return indices, numDecoys-1
            else: return len(fdrValues), numDecoys-1
        fdrValues.append(curFDR)
        indices.append(index)
        count += 1

    if returnType: return fdrValues, numDecoys-1
    else: return len(fdrValues), numDecoys-1

def return_frag_mzs(peptide, z):
    mzValues = []
    digPat = r'\+\d+\.\d+'
    digs = re.findall(digPat, peptide)
    pepFrags = re.split(digPat, peptide)
    modValues = {}
    seq = ''
    while len(digs) != 0:
        dig = digs.pop(0)
        frag = pepFrags.pop(0)
        seq += frag
        modValues[len(seq)] = float(dig[1:])/z
    seq += pepFrags[0]
    for i in range(1, len(seq)-1):
        mz = mass.fast_mass(sequence=seq[i:], ion_type='y', charge=z)
        mz += sum([modValues[x] for x in modValues if x > i])
        mzValues.append(mz)

    for i in range(len(seq)-1, 1, -1):
        mz = mass.fast_mass(sequence=seq[:i], ion_type='b', charge=z)
        mz += sum([modValues[x] for x in modValues if x <= i])
        mzValues.append(mz)
    return mzValues

def approx(x, y, ppmTol):
    if x==y: return 1e-7
    ppmDiff = ((x-y)*1000000)/x
    return (ppmDiff if abs(ppmDiff) < ppmTol else 0)

def approx_list(x, l, ppmTol=10):
    for i in range(len(l)):
        if approx(x, l[i], ppmTol): return i
    return -1

#if 'protein' in spec['params']

def clean_mgf_file(file):
    spectra = mgf.read(file)
    fasta = 'C:/Users/ccranney/Desktop/Caleb_Files/data/2019-03-14-td-UP000005640.fasta'
    fDict = {}
    longPep = ''
    for record in SeqIO.parse(open(fasta,'r'),'fasta'):
        fDict[len(longPep)] = record.id
        longPep += str(record.seq) + '.'
    cleaned = []
    count = 0
    pepCount = 0
    for spec in spectra:
        count += 1
        #if count % 40==0: break
        #mzValues = return_frag_mzs(spec['params']['seq'],1)
        #peaks = list(tuple(zip(spec['m/z array'],spec['intensity array'])))
        #for i in range(len(peaks)-1,-1,-1):
        #    if approx_list(peaks[i][0],mzValues)==-1: peaks.pop(i)
        #if len(peaks)==0: continue
        #peaks.sort(key=lambda x:x[0])
        #spec['m/z array'],spec['intensity array'] = map(list,zip(*peaks))
        #'''
        decoy = False
        if 'protein' in spec['params'] and 'DECOY' in spec['params']['protein']: decoy = True
        else:
            seq = re.sub(r'\+\d+\.\d+', '', spec['params']['seq'])
            listOfI = [m.start() for m in re.finditer(seq, longPep)]
            sorted_keys = sorted(fDict.keys())
            proteins = set()
            for i in listOfI:
                insertion_point = bisect.bisect_left(sorted_keys,i)
            # adjust, as bisect returns not exactly what we want
                if insertion_point==len(sorted_keys) or sorted_keys[insertion_point]!=i:
                    insertion_point-=1
                protein = fDict[sorted_keys[insertion_point]]
                proteins.add(fDict[sorted_keys[insertion_point]])
            if len(proteins)==0: proteins.add(spec['params']['seq'])

        if decoy: proteins = ['DECOY_0_'+x for x in proteins]

        protein = str(len(proteins)) + '/' + '/'.join(sorted(proteins))
        spec['params']['protein'] = protein
        if protein != '0/': pepCount += 1
        #'''
        cleaned.append(spec)
        if count % 1000 == 0: print(count); print(pepCount); print(protein)

    cleanedFile = re.sub('(.*).mgf', r'\1_proteinsAdded.mgf', file)
    mgf.write(cleaned, cleanedFile)

print('\n'*30)

'''
scanCount = defaultdict(int)
msplit = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/HeLa-50ppm/MSPLIT_HeLa.txt', sep='\t')
scans = msplit['Scan#']
for scan in scans:
    scanCount[scan]+=1

#hits, decoys = fdr_calculation(msplit)
#print(hits, decoys)
'''
'''
import operator

sorted_tuples = sorted(scanCount.items(), key=operator.itemgetter(1),reverse=True)
temp = []
for x in sorted_tuples[:20]: temp.append(x[0])

#print(msplit[msplit['Scan#']==22501])

#hits, decoys = fdr_calculation(msplit)

#print(hits, decoys)

csodiaq = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/HeLa-50ppm/CsoDIAq-file1_HeLa_160min_DIA_106win_1.csv')
#print(len(csodiaq))
'''
'''
windows = set()
spectralCount = 0
with mzxml.read('C:/Users/ccranney/Desktop/Caleb_Files/data/HeLa_160min_DIA_106win_1.mzXML', use_index=True) as spectra:
#with mzxml.read('C:/Users/ccranney/Desktop/Caleb_Files/data/18302_REP2_500ng_HumanLysate_SWATH_2 (2)..mzXML', use_index=True) as spectra:

    for scan in temp:
        spec = spectra.get_by_id(str(scan))
        print(scan)
        print('PrecursorMz: '+str(spec['precursorMz'][0]['precursorMz']))
        print('Window: '+str(spec['precursorMz'][0]['windowWideness']))

        tempDf = msplit[msplit['Scan#']==scan]
        print('Mz Values:')
        print(tempDf['Mz.1'])
        print('\n')

    for spec in spectra:
        spectralCount += 1
        if spectralCount % 10000 == 0: print(spectralCount)
        if 'precursorMz' in spec:
            windows.add(spec['precursorMz'][0]['windowWideness'])
            if spec['precursorMz'][0]['windowWideness'] == 50.0: print(spec['num'])
for x in windows: print(x)
#'''
'''
file = 'C:/Users/ccranney/Desktop/Caleb_Files/data/18302_REP2_500ng_HumanLysate_SWATH_2 (2)..mzXML'
with open(file, 'rb') as f:
    for _ in range(100): # first 10 lines
        print(f.readline())
#'''
'''
peptide = 'A+42.01057AAAAAGAGPEM+15.9949VR'
#peptide = 'AAAAA'
cbf.return_frag_mzs(peptide, 1)
#'''
'''
file = 'C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy.mgf'
fasta = 'C:/Users/ccranney/Desktop/Caleb_Files/data/2019-03-14-td-UP000005640.fasta'
cbf.clean_mgf_file(file, fasta)
#'''
'''
mgf1 = mgf.read('C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy.mgf')
mgf2 = []
count = 0
print('enter processing')
for spec in mgf1:
    spec['intensity array'] = spec['intensity array'][:10]
    spec['m/z array'] = spec['m/z array'][:10]
    mgf2.append(spec)
    #spec['params']['title'] = count
    #count += 1
    #if count == 10: break
print(mgf2)
print('enter writing')
mgf.write(mgf1, 'C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy3.mgf')
#'''
'''
fasta = 'C:/Users/ccranney/Desktop/Caleb_Files/data/2019-03-14-td-UP000005640.fasta'
fDict = {}
longPep = ''
for record in SeqIO.parse(open(fasta,'r'),'fasta'):
    fDict[len(longPep)] = record.id
    longPep += str(record.seq) + '.'

print(len(fDict))
for i in sorted(fDict)[:3]:
    print(i)
    if i != 0: print(longPep[i-1])
    print(fDict[i])
print(longPep[:200])

pep = 'KMMKRRGINVSDE'
listOfI = [m.start() for m in re.finditer(pep, longPep)]
print(listOfI)
sorted_keys = sorted(fDict.keys())
proteins = set()
for i in listOfI:
    insertion_point = bisect.bisect_left(sorted_keys,i)
# adjust, as bisect returns not exactly what we want
    if insertion_point==len(sorted_keys) or sorted_keys[insertion_point]!=i:
        insertion_point-=1
    proteins.add(fDict[sorted_keys[insertion_point]])
    print(insertion_point)
    print(fDict[sorted_keys[insertion_point]])

proteins = str(len(proteins)) + '/' + '/'.join(sorted(proteins))
print(proteins)
#'''
'''
fasta = 'C:/Users/ccranney/Desktop/Caleb_Files/data/2019-03-14-td-UP000005640.fasta'

record_dict = SeqIO.index(fasta, "fasta")
spectra = mgf.read('C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy_cleaned.mgf')
count = 0
num = 5000
for spec in spectra:
    count += 1
    if count % 1000 == 0: print(count)
    if randint(0,num) % num != 0: continue

    seq = re.sub(r'\+\d+\.\d+', '', spec['params']['seq'])
    proteins = spec['params']['protein'].split('/')[1:]
    for protein in proteins:
        if 'DECOY' in protein: continue
        temp = str(record_dict[protein].seq)
        if seq not in temp:
            print('FAIL ' + 'X'*50)
            print(seq)
            print(temp)
            print('\n')
        else:
            print('SUCCESS')
            print('\n')
#'''
'''
peptides = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/MGF-ID1/CsoDIAq-file1_ID1_corrected_peptideFDR.csv')
peptides = peptides.drop_duplicates(subset='peptide', keep='first').reset_index(drop=True)
print(len(peptides))
proteins = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/MGF-ID1/CsoDIAq-file1_ID1_corrected_proteinFDR.csv')
proteins = proteins[proteins['uniquePeptide']==1]
#proteins = proteins.drop_duplicates(subset='leadingProtein', keep='first').reset_index(drop=True)
prots = set()
for index, row in proteins.iterrows():
    prots.update(row['leadingProtein'].split('/'))
print(len(prots))
#'''
'''
# for determining the intensities of an mzxml file - I'm finding they're either all the same or has a lot of zeroes.
count = 0
#with mzxml.read('C:/Users/ccranney/Desktop/wiffFiles/mzxml-round1/18302_REP2_500ng_HumanLysate_SWATH_2.mzxml', use_index=True) as spectra:
with mzxml.read('C:/Users/ccranney/Desktop/Caleb_Files/data/HeLa_160min_DIA_106win_1.mzxml', use_index=True) as spectra:
    for spec in spectra:
        if 'precursorMz' not in spec: continue
        count += 1
        if count > 10: break
        print(list(spec['intensity array']))
#'''
'''
df = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0.txt',sep='\t')
spectra = mgf.read('C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy_proteinsAdded.mgf')
pepDict = {}
count = 0
for spec in spectra:
    count += 1
    if count % 1000 == 0: print(count)
    pepDict[spec['params']['seq']] = spec['params']['protein']

proteins = []
for index, row in df.iterrows():
    proteins.append(pepDict[row['Peptide']])

df['protein'] = proteins
df.to_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0_proteinsAdded.csv', index=False)
#'''
'''
# Adding MSPLIT proteins to final file
inFile = 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0_proteinsAdded.csv'
specFile = 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0_proteinsAdded_spectralFDR.csv'
pepFile = 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0_proteinsAdded_peptideFDR.csv'
protFile = 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0_proteinsAdded_proteinFDR.csv'
cbf.write_csodiaq_fdr_outputs(inFile, specFile, pepFile, protFile)
#'''
'''
# for comparing msplit with csodiaq
msplit = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_qehfx_09_30_0_proteinsAdded_peptideFDR.csv')
#msplit = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_ID1_02-10-0_temp.csv')
msplit['ID'] = list(zip(msplit['Scan#'].tolist(),msplit['peptide'].tolist(),msplit['z.1'].tolist()))
#msplit = msplit[msplit['uniquePeptide']==1]
#msplit = msplit.drop_duplicates(subset='leadingProtein', keep='first').reset_index(drop=True)
#msplit.set_index('ID',inplace=True)
#msplit = msplit[msplit['protein'].str.contains('DECOY')]

csodiaq = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/WIFF files/CsoDIAq-file1_01_qehfx_lab_SA_R1_corrected_peptideFDR.csv')
#csodiaq = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/WIFF files/CsoDIAq-file1_01_qehfx_lab_SA_R1_corrected_peptideFDR.csv')
csodiaq['ID'] = list(zip(csodiaq['scan'].tolist(),csodiaq['peptide'].tolist(),csodiaq['zLIB'].tolist()))
#csodiaq = csodiaq[csodiaq['uniquePeptide']==1]
#csodiaq = csodiaq.drop_duplicates(subset='leadingProtein', keep='first').reset_index(drop=True)
#csodiaq.set_index('ID',inplace=True)
#csodiaq = csodiaq[csodiaq['protein'].str.contains('DECOY')]

#mValues = set()
#cValues = set()
#for proteinGroup in msplit['leadingProtein']:
#    proteins = re.findall('(DECOY_0_)?(sp\|\w{6}\|)', proteinGroup)
#    mValues.update([p for p in proteins if 'DECOY' not in p])
#for proteinGroup in csodiaq['leadingProtein']:
#    proteins = re.findall('(DECOY_0_)?(sp\|\w{6}\|)', proteinGroup)
#    cValues.update([p for p in proteins if 'DECOY' not in p])

mValues = set(msplit['peptide'])
cValues = set(csodiaq['peptide'])

pyplot.figure(figsize=(12,12))
out = venn2([set(cValues),set(mValues)], set_labels = ["CsoDIAq", "MSPLIT-DIA"])
for text in out.set_labels:
    text.set_fontsize(32)
for text in out.subset_labels:
    text.set_fontsize(24)
pyplot.show()

#inM = set(msplit['ID']) - set(csodiaq['ID'])
#inC = set(csodiaq['ID']) - set(msplit['ID'])
#both = set(csodiaq['ID']).intersection(set(msplit['ID']))

#'''
'''
print(len(msplit['ID']))
print('\n')
print(len(inM))
print('\n')
print(len(csodiaq['ID']))
print('\n')
print(len(inC))
print('\n')
print(len(both))

msplit = msplit[msplit['ID'].isin(both)].sort_values('ID').reset_index(drop=True)
csodiaq = csodiaq[csodiaq['ID'].isin(both)].sort_values('ID').reset_index(drop=True)

for i in range(10):
    print(msplit.loc[i])
    print(csodiaq.loc[i])
    print('\n')
#'''
'''
oldMGF = mgf.read('C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy.mgf')
protMGF = mgf.read('C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy_proteinsAdded.mgf')

count = 0
for spec in oldMGF:
    count += 1
    if count % 5 ==0: break#print(count)
    spec2 = protMGF.get_by_id(spec['params']['title'])
    print(spec2)
    if sorted(spec['m/z array']) != sorted(spec2['m/z array']):
        print('m/z array')
        print(len(spec['m/z array']))
        print(len(spec2['m/z array']))
        print('\n')
    if sorted(spec['intensity array']) != sorted(spec2['intensity array']):
        print('intensity array')
        print(len(spec['intensity array']))
        print(len(spec2['intensity array']))
        print('\n')

    if len(spec['charge']) != len(spec2['charge']):
        print('charge')
        print(len(spec['charge']))
        print(len(spec2['charge']))
        print('\n')
    if len(spec['mask']) != len(spec2['mask']):
        print('mask')
        print(len(spec['mask']))
        print(len(spec2['mask']))
        print('\n')

    if spec['params']['title'] != spec2['params']['title']:
        print(spec['params']['title'])
        print(spec2['params']['title'])
        print('\n')
    if spec['params']['charge'] != spec2['params']['charge']:
        print(spec['params']['charge'])
        print(spec2['params']['charge'])
        print('\n')
    if spec['params']['pepmass'] != spec2['params']['pepmass']:
        print(spec['params']['pepmass'])
        print(spec2['params']['pepmass'])
        print('\n')
    if spec['params']['seq'] != spec2['params']['seq']:
        print(spec['params']['seq'])
        print(spec2['params']['seq'])
        print('\n')
    if spec['params']['scan'] != spec2['params']['scan']:
        print(spec['params']['scan'])
        print(spec2['params']['scan'])
        print('\n')
#'''
'''
# simple script for seeing a few scans in mzxml files

#with mzxml.read('C:/Users/ccranney/Desktop/Caleb_Files/data/HeLa_160min_DIA_106win_1_large.mzXML', use_index=True) as spectra:
with mzxml.read('C:/Users/ccranney/Desktop/wiffFiles/K562_10ug_DDA_Top100_r02-K562_Sample10.mzXML', use_index=True) as spectra:
    count = 1
    windows = defaultdict(list)
    for spec in spectra:
        print(list(spec['intensity array']))
        if 'precursorMz' in spec: count += 1; windows[spec['precursorMz'][0]['windowWideness']].append(spec['num'])
        if count % 10 == 0: break
    #for key, value in windows.items():
        #print(key)
        #print(value)
#'''
'''
# simple script for seeing the first few lines of a csv data file
df = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/docker-shared/consensus.tsv', sep='\t')
#df.sort_values('MaCC_Score', ascending=False, inplace=True)
print(df.head(200))
print(df.columns)
#'''
'''
# checking if there are protein-marked decoys is the pan human library
#df = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/docker-shared/human.tsv',sep='\t')
df = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/lib_tsv.tsv',sep='\t')
#print(df.head(10))
#df = df[df['ProteinName'].str.contains('DECOY')]
df = df[df['decoy']==1]
#df = df.head(100)
#df.to_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/lib_tsv_abbrev.csv')
print(len(df))
#'''
'''
# checking hov FullUniMode differs from PeptideSequence
diffs = set()
df = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/lib_tsv.tsv',sep='\t')
for index, row in df.iterrows():
    if row['PeptideSequence'] != row['FullUniModPeptideName']:
        diffs.add((row['PeptideSequence'],row['FullUniModPeptideName']))

print(len(diffs))
print(sorted(diffs)[:20])
#'''
'''
# Counting peptides in targetted re-analysis files
files = [
    'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CsoDIAq-file1_20210222_DISPA_hela_03_corrected_mostIntenseTargs_-30.0.csv',
    'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CsoDIAq-file1_20210222_DISPA_hela_03_corrected_mostIntenseTargs_-40.0.csv',
    'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CsoDIAq-file1_20210222_DISPA_hela_03_corrected_mostIntenseTargs_-50.0.csv',
    'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CsoDIAq-file1_20210222_DISPA_hela_03_corrected_mostIntenseTargs_-60.0.csv',
    'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CsoDIAq-file1_20210222_DISPA_hela_03_corrected_mostIntenseTargs_-70.0.csv',
    'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CsoDIAq-file1_20210222_DISPA_hela_03_corrected_mostIntenseTargs_-80.0.csv',
]
count = 0
for file in files:
    print(file)
    df = pd.read_csv(file)
    peptides = df['Compound']
    for pep in peptides:
        peps = pep.split('/')
        count += len(peps)
print(count)
#'''
'''
# convert tsv to csv
df = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/docker-shared/qehfx.tsv', sep='\t')
df.to_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/docker-shared/qehfx.csv')
#'''
'''
#counting MGF spectra in a file
spectra = mgf.read('C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy.mgf')
count = 0
for spec in spectra: count += 1
print(count)
'''
'''
#testing table manipulations
with bz2.BZ2File('C:/Users/ccranney/Desktop/Caleb_Files/data/output/compressTest/compressed_processObj_canDeleteAfter/window0', 'rb') as pickleFile: data = pickle.load(pickleFile)
print(data.head(10))

def calc_macc(df):
    lib = df['libIntensity']
    que = df['queIntensity']
    AB = lib.multiply(que).sum()
    A = lib.pow(2).sum()
    B = que.pow(2).sum()
    cosine = cbf.cosine_similarity([AB, A, B])
    return (len(df)**(1/5))*cosine
#test = data.groupby(['libID','queID']).apply(calc_macc)
#print(data.head(10))
tempDict = {}
grouped = data.groupby(['libID','queID'])
print(grouped.ngroups)
data2 = grouped.filter(lambda x: len(x.index) > 2.)
grouped = data2.groupby(['libID','queID'])
tempList = [(21, 42743),
(21, 44234),
(21, 44376),
(21, 44447),
(21, 47003),
(21, 64256),
(21, 64611),
(21, 64966),
(21, 71782),
(21, 71853),
(21, 71924),
(21, 71995),
(21, 72066),
(21, 72208),
(21, 72279),
(21, 72350),
(21, 72421),
(21, 72705),
(21, 72776),
(21, 72847),
(21, 72989),
(21, 73060),
(21, 73131),
(21, 73202),
(21, 73273),
(21, 73344),
(21, 76965),
(21, 77036),
(21, 77107),
(21, 77178),
(21, 77249),
(21, 77320),
(21, 77391),
(21, 77604),
(21, 77817),
(21, 85272),
(41, 85201),
(60, 52825),
(60, 53180),
(60, 65250),
(60, 65321),
(60, 65392),
(60, 69439),
(60, 69510),
(60, 78598),
(60, 78669),
(60, 78811),
(84, 37276),
(84, 37347),
(84, 37489),
(84, 37560),
(86, 36424),
(86, 42388),
(86, 42459),
(86, 42530),
(86, 42601),
(86, 42814),
(86, 42956),
(86, 43382),
(86, 43524),
(86, 49204),
(86, 49346),
(86, 50908),
(86, 59925),
(96, 34862),
(96, 34933),
(96, 35004),
(96, 35075),
(96, 35146),
(96, 37915),
(96, 42246),
(96, 42672),
(96, 42743),
(96, 42814),
(96, 44589),
(96, 44660),
(96, 44731),
(96, 44802),
(96, 44873),
(96, 44944),
(96, 45015),
(96, 45086),
(96, 45228),
(96, 45370),
(96, 45441),
(96, 45583),
(96, 52896),
(96, 56446),
(96, 56517),
(96, 56588),
(96, 56659),
(96, 61203),
(96, 65037),
(96, 78740),
(96, 78953),
(96, 79024),
(96, 79166),
(96, 79237)]
for x in tempList: tempDict[x]=0
#data2['MaCC'] = grouped.apply(lambda x,d: d[x.name])
data2['MaCC'] = data2.groupby(['libID','queID']).apply(lambda x,d: print(d[x.name]),tempDict)
#d[tuple(x.name)]
#count = 0
#for name, group in grouped:
#    print(type(name))
#    count += 1
#    if count==10: break
#print(grouped.ngroups)
#data2['MaCC'] = grouped.apply(calc_macc)
#for name, group in grouped:
#    macc = calc_mass(group)
#print(data2.head(10))
'''
'''
# Writing/testing reduction function
def reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches):
    curLibTag = matchLibTags[0]
    curQueTag = matchQueTags[0]
    count = 1
    #maccScores = []
    AB = matchLibIntensities[0]*matchQueIntensities[0]
    A = matchLibIntensities[0]**2
    B = matchQueIntensities[0]**2
    returnLibTags = []
    returnLibIntensities = []
    returnQueTags = []
    returnQueIntensities = []
    returnPpmMatchs = []
    returnMaccScores = []
    length = len(matchLibTags)
    for i in range(1,length):
        if matchLibTags[i] != curLibTag or matchQueTags[i] != curQueTag:
            if count > 2:
                cosine = cbf.cosine_similarity(AB, A, B)
                macc = (count**(1/5))*cosine
                magnitude = (A**0.5) * (B**0.5) # (sqrt(sum(A^2))*sqrt(sum(B^2)))
                returnLibTags.extend(matchLibTags[i-count:i])
                returnLibIntensities.extend(matchLibIntensities[i-count:i])
                returnQueTags.extend(matchQueTags[i-count:i])
                returnQueIntensities.extend(matchQueIntensities[i-count:i])
                returnPpmMatchs.extend(ppmMatches[i-count:i])
                returnMaccScores.extend([macc]*count)
            count = 1
            curLibTag = matchLibTags[i]
            curQueTag = matchQueTags[i]
            AB = matchLibIntensities[i]*matchQueIntensities[i]
            A = matchLibIntensities[i]**2
            B = matchQueIntensities[i]**2
        else:
            AB += matchLibIntensities[i]*matchQueIntensities[i]
            A += matchLibIntensities[i]**2
            B += matchQueIntensities[i]**2
            count += 1
    if count > 2:
        cosine = cbf.cosine_similarity(AB, A, B)
        macc = (count**(1/5))*cosine
        returnLibTags.extend(matchLibTags[length-count:])
        returnLibIntensities.extend(matchLibIntensities[length-count:])
        returnQueTags.extend(matchQueTags[length-count:])
        returnQueIntensities.extend(matchQueIntensities[length-count:])
        returnPpmMatchs.extend(ppmMatches[length-count:])
        returnMaccScores.extend([macc]*count)
        count = 1

    return returnLibTags, returnLibIntensities, returnQueTags, returnQueIntensities, returnPpmMatchs, returnMaccScores

libTags = np.repeat(np.arange(0,10),3)
libIntensities = np.arange(1.0,31.0)
queTags = np.repeat(np.arange(0,5),6)
queIntensities = np.arange(31.0,61.0)
ppmMatches = np.arange(61.0,91.0)
#print(libTags)
#print(libIntensities)
#print(queTags)
#print(queIntensities)
#print(ppmMatches)

matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, maccScores = reduce_final_df(libTags, libIntensities, queTags, queIntensities, ppmMatches)

#print(len(libTags))
#print(len(matchLibTags))
#print(len(set(maccScores)))
#print(libTags)
#print(matchLibTags)
print(maccScores)
'''
'''
# Testing last part of "compression" attempt
@njit
def match_score_decoys(matchLibTags, matchQueTags, maccScores, decoys):
    curLibTag = matchLibTags[0]
    curQueTag = matchQueTags[0]
    returnMaccs = [maccScores[0]]
    returnDecoys = [decoys[0]]
    length = len(matchLibTags)
    for i in range(1,length):
        if matchLibTags[i] != curLibTag or matchQueTags[i] != curQueTag:
            returnMaccs.append(maccScores[i])
            returnDecoys.append(decoys[i])
            curLibTag = matchLibTags[i]
            curQueTag = matchQueTags[i]
    #return returnLibTags, returnLibIntensities, returnQueTags, returnQueIntensities, returnPpmMatchs, returnMaccScores
    return returnMaccs, returnDecoys

@njit
def fdr_calculation2(maccs, decoys): #***NOTE***make return type
    # initializing the two return values at 0
    numDecoys = 0

    # for every row in the dataframe
    count = 0
    for i in range(len(maccs)):
        #print(str(maccs[i])+':'+str(decoys[i]))
        if decoys[i]: numDecoys += 1
        # calculates the FDR up to this point in the data frame.
        curFDR = numDecoys/(count+1)

        # conditional statement comparing the current FDR to the FDR Cutoff. If larger, function values are returned.
        if curFDR > 0.01:

            # if the number of rows has not yet reached the minimum number that allows for the FDR cutoff, 0 is returned instead.
            if count < 1/0.01: return -1
            return maccs[i-1]
        count += 1

    return maccs[-1]

def collect_ppm_values(ppmMatches, maccScores, maccCutoff):
    ppms = []
    length = len(ppmMatches)
    for i in range(1,length):
        if maccScores[i] >= maccCutoff: ppms.append(ppmMatches[i])
    #return returnLibTags, returnLibIntensities, returnQueTags, returnQueIntensities, returnPpmMatchs, returnMaccScores
    return ppms

#test = np.arange(10,0,-1)
#print(type(test.argsort()))
#print(type(np.where((test>2)*(test<8))[0]))
with bz2.BZ2File('C:/Users/ccranney/Desktop/Caleb_Files/data/output/CompressTest/compressed_processObj_canDeleteAfter/windowallFinals', 'rb') as pickleFile: data = pickle.load(pickleFile)
finalLibTags, finalLibIntensities, finalQueTags, finalQueIntensities, finalPpmMatches, finalMaccScores, finalDecoys = data
print(len(finalDecoys))
print(timer())
#temp1, temp2 = match_score_decoys(np.array([],dtype=int),np.array([],dtype=int),np.array([],dtype=float),np.array([],dtype=int))
#print(timer())
maccs, decoys = match_score_decoys(finalLibTags, finalQueTags, finalMaccScores, finalDecoys)
print(timer())
print(len(maccs))
print(len(set(finalMaccScores)))
print(len(set(maccs)))
maccs = np.array(maccs)
decoys = np.array(decoys)
i1 = (-maccs).argsort()
maccs = maccs[i1]
decoys = decoys[i1]
maccCutoff = fdr_calculation2(maccs, decoys)
print(maccCutoff)
cutoff = 1.4357304070871908
ppms = collect_ppm_values(finalPpmMatches, finalMaccScores, maccCutoff)
offset, tolerance = cbf.find_offset_tol(ppms, 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/CompressTest/CsoDIAq-file1_01_qehfx_lab_SA_R1_histogram.png', stdev=0)
print(offset,tolerance)
lowend = offset-tolerance
highend = offset+tolerance
ppmIndices = np.where((finalPpmMatches>lowend)*(finalPpmMatches<highend))[0]
print(len(ppmIndices))
tempTags = finalLibTags[ppmIndices]
print(len(tempTags))
print(ppmIndices[:10])
print(tempTags[:10])
corLibTags, corLibIntensities, corQueTags, corQueIntensities, corDecoys = [x[ppmIndices] for x in [finalLibTags, finalLibIntensities, finalQueTags, finalQueIntensities, finalDecoys]]
corLibTags, corLibIntensities, corQueTags, corQueIntensities, corDecoys, corMaccScores = cbf.reduce_final_df(corLibTags, corLibIntensities, corQueTags, corQueIntensities, corDecoys)
maccs, decoys = match_score_decoys(corLibTags, corQueTags, corMaccScores, corDecoys)
i1 = (-maccs).argsort()
maccs = maccs[i1]
decoys = decoys[i1]
maccCutoff = fdr_calculation2(maccs, decoys)
'''
'''
#testing numpy compression
t1 = np.arange(10)
t2 = np.arange(10,20)
t3 = np.arange(20,30)
compressed_array = io.BytesIO()
compressed_array2 = io.BytesIO()

#np.savez_compressed(compressed_array, t1=t1, t2=t2, t3=t3)
np.savez_compressed(compressed_array, t1=t1)
np.savez_compressed(compressed_array2, t1=t2)
#np.savez_compressed(compressed_array, t3=t3)

compList = [compressed_array, compressed_array2]
for x in compList:
    x.seek(0)
    decompressed_array = np.load(x)
    print(decompressed_array['t1'])
    #print(decompressed_array['t2'])
#print(decompressed_array['t3'])
#'''
'''
# testing the old method for calculating FDR rates
test = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/100reps_mgf_10iniTol/CsoDIAq-file17_20200719_MAGIC_MCF7_1128repro_17_corrected_full.csv').sort_values('cosine', ascending=False).reset_index(drop=True)

bestHits = 0
bestOutput = [0,0,0,0]
for i in range(3,11):
    temp = test[test['shared'] > i-1].reset_index(drop=True)
    hits, decoys = cbf.fdr_calculation(temp)
    if hits > bestHits: bestOutput = [0, hits, temp.loc[hits-1]['cosine'], 'Naive CsoDIAq, corrected']; bestHits = hits
    if hits != 0: print([0, hits, temp.loc[hits-1]['cosine'], 'Naive CsoDIAq, corrected'])
    else: print([0, hits, 'null', 'Naive CsoDIAq, corrected'])

print(len(test))
print(bestOutput)
#'''
'''
# comparing two data frames (csodiaq outputs) and checking for similarities/differences
df1 = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/output/100reps_tsv_10iniTol_6min/CsoDIAq-file64_20200719_MAGIC_MCF7_1128repro_64_full.csv')
df2 = pd.read_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/csodiaq_lib-human-noloss-400to2000-pt2mz-6peaks_exp-100reps-rep64.csv')

df1 = df1.sort_values('peptide').reset_index(drop=True)
df2 = df2.sort_values('peptide').reset_index(drop=True)

df1 = df1.loc[:, df1.columns.intersection([
                    'scan', 'MzEXP', 'peptide', 'protein', 'MzLIB', 'zLIB',
                    'cosine', 'name', 'Peak(Query)', 'Peaks(Library)', 'shared',
                    'CompensationVoltage', 'totalWindowWidth'])]
df2 = df2.loc[:, df2.columns.intersection([
                    'scan', 'MzEXP', 'peptide', 'protein', 'MzLIB', 'zLIB',
                    'cosine', 'name', 'Peak(Query)', 'Peaks(Library)', 'shared',
                    'CompensationVoltage', 'totalWindowWidth'])]
print(df1.head(10))
print(df2.head(10))

#df1.drop(['fileName','MaCC_Score'],axis=1, inplace=True)
#df2.drop(['fileName','zEXP'],axis=1, inplace=True)

#df1 = df1.sort_values('peptide').reset_index(drop=True)
#df2 = df2.sort_values('peptide').reset_index(drop=True)

#print(df1.loc[0])
#print()
#print(df2.loc[0])

merged = df1.merge(df2, how='outer', indicator=True)
merged.to_csv('C:/Users/ccranney/Desktop/Caleb_Files/data/old_v_new.csv', index=False)
print(merged.head(10))
#'''
'''
# deleting all files but the ones I want from a folder
head = 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/Output/'
allFiles = os.listdir(head)
#for x in allFiles: print(x); print(re.match('C:/Users/ccranney/Desktop/Caleb_Files/data/output/Output/csodiaq_lib-human-noloss-400to2000-pt2mz-31peaks_exp-100reps-rep(\d{2,3})(_corrected)?\.csv', head+x))
files = [head+x for x in list(os.listdir(head)) if not re.match('C:/Users/ccranney/Desktop/Caleb_Files/data/output/Output/csodiaq_lib-human-noloss-400to2000-pt2mz-6peaks_exp-100reps-rep(\d{2,3})(_corrected)?\.csv', head+x)]
for x in files: os.remove(x)
print(len(files))
#'''

fileName = 'C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy_proteinsAdded.mgf'
#fileName = 'C:/Users/ccranney/Desktop/Caleb_Files/data/lib_tsv.tsv'
libPeaks = [10]
expPath = 'C:/Users/ccranney/Desktop/Caleb_Files/data/100mzxml/'
outPath = 'C:/Users/ccranney/Desktop/Caleb_Files/data/output/test_libPeaks_minMatch/'
files = [expPath+x for x in list(os.listdir(expPath))]
#files = ['C:/Users/ccranney/Desktop/Caleb_Files/data/ID1.mzxml']
minMatch = [3]
for lp in libPeaks:
    lib = cbf.library_file_to_dict(fileName, lp)
    for i in range(len(files)):
        outHeader = 'CsoDIAq-file' +str(i+1)+'_'+ '.'.join(files[i].split('/')[-1].split('.')[:-1]) + 'MaCC' + '_libPeak' + str(lp)
        ppmTol = 10
        offset = 0
        queryPooling = np.inf
        corrected = 0
        histFile = 0
        for match in minMatch:
            if match < lp:
                outFile = outPath + outHeader + '_minMatch'+ str(match)+'_corrected.csv'
                cbf.pooled_spectra_analysis(  files[i],
                #cbf2.pooled_spectra_analysis(  files[i],
                                                    outFile,
                                                    lib,
                                                    ppmTol,
                                                    offset,
                                                    queryPooling,
                                                    corrected,
                                                    histFile,
                                                    match)

#java -Xmx2500M -cp C:/Users/ccranney/Desktop/Caleb_Files/MSPLIT-DIAv1.0/MSPLIT-DIAv02102015.jar org.Spectrums.SWATHMSPLITSearch 02 10 0 C:/Users/ccranney/Desktop/Caleb_Files/data/HeLa_160min_DIA_106win_1.mzXML C:/Users/ccranney/Desktop/Caleb_Files/data/human.faims.fixed.decoy.mgf C:/Users/ccranney/Desktop/Caleb_Files/data/output/msplit_HeLa_02-10-0.tsv
#docker run -it -v C:/Users/ccranney/Desktop/Caleb_Files/data/docker-shared:/data openswath/openswath
#TargetedFileConverter -in phl004_consensus_openms24.TraML -out consensus.tsv
#OpenSwathWorkflow -in 01_qehfx_lab_SA_R1.mzXML -tr consensus.tsv -out_features qehfx.featureXML
'''
                        decoyList = [idToDecoyDict[x] for x in libTags]
                        #libMzs = np.array([x[0] for x in pooledLibSpectra])
                        queMzs = np.array([x[0] for x in pooledQueSpectra])
                        tempTag = 'identify'
                        returns = initialize_return_values(tempTag)

                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys = spectra_peak_comparison(libMzs, libIntensities, libTags, queMzs, queIntensities, queTags, decoyList, ppmTol, ppmYOffset)
                        #data = pd.DataFrame({'libID':matchLibTags, 'queID':matchQueTags, 'libIntensity':matchLibIntensities, 'queIntensity':matchQueIntensities, 'ppmDiff':ppmMatches})
                        #data = data.groupby(['libID','queID']).filter(lambda x: len(x.index) > 2.)
                        #print(len(data))
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys = [np.array(x) for x in [matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys]]
                        i1 = matchQueTags.argsort(kind='mergesort')
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys = [x[i1] for x in [matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys]]
                        i2 = matchLibTags.argsort(kind='mergesort')
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys = [x[i2] for x in [matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys]]
                        #reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches)
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys, maccScores = reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys)
                        test = [idToDecoyDict[x] for x in matchLibTags]
def reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, matchDecoys):
    curLibTag = matchLibTags[0]
    curQueTag = matchQueTags[0]
    count = 1
    AB = matchLibIntensities[0]*matchQueIntensities[0]
    A = matchLibIntensities[0]**2
    B = matchQueIntensities[0]**2
    returnLibTags = []
    returnLibIntensities = []
    returnQueTags = []
    returnQueIntensities = []
    returnPpmMatches = []
    returnDecoys = []
    returnMaccScores = []
    length = len(matchLibTags)
    for i in range(1,length):
        if matchLibTags[i] != curLibTag or matchQueTags[i] != curQueTag:
            if count > 2:
                test = 0
                cosine = cosine_similarity(AB, A, B)
                macc = (count**(1/5))*cosine
                magnitude = (A**0.5) * (B**0.5) # (sqrt(sum(A^2))*sqrt(sum(B^2)))
                returnLibTags.extend(matchLibTags[i-count:i])
                returnLibIntensities.extend(matchLibIntensities[i-count:i])
                returnQueTags.extend(matchQueTags[i-count:i])
                returnQueIntensities.extend(matchQueIntensities[i-count:i])
                returnPpmMatches.extend(ppmMatches[i-count:i])
                returnDecoys.extend(matchDecoys[i-count:i])
                returnMaccScores.extend([macc]*count)

            count = 1
            curLibTag = matchLibTags[i]
            curQueTag = matchQueTags[i]
            AB = matchLibIntensities[i]*matchQueIntensities[i]
            A = matchLibIntensities[i]**2
            B = matchQueIntensities[i]**2
        else:
            AB += matchLibIntensities[i]*matchQueIntensities[i]
            A += matchLibIntensities[i]**2
            B += matchQueIntensities[i]**2
            count += 1
    if count > 2:
        test = 0
        cosine = cosine_similarity(AB, A, B)
        macc = (count**(1/5))*cosine
        returnLibTags.extend(matchLibTags[length-count:])
        returnLibIntensities.extend(matchLibIntensities[length-count:])
        returnQueTags.extend(matchQueTags[length-count:])
        returnQueIntensities.extend(matchQueIntensities[length-count:])
        returnPpmMatches.extend(ppmMatches[length-count:])
        returnDecoys.extend(matchDecoys[length-count:])
        returnMaccScores.extend([macc]*count)

    #return returnLibTags, returnLibIntensities, returnQueTags, returnQueIntensities, returnPpmMatchs, returnMaccScores
    return returnLibTags, returnLibIntensities, returnQueTags, returnQueIntensities, returnPpmMatches, returnDecoys, returnMaccScores

####################################
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches = spectra_peak_comparison(libMzs, libIntensities, libTags, queMzs, queIntensities, queTags, ppmTol, ppmYOffset)

                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches = [np.array(x) for x in [matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches]]
                        i1 = matchQueTags.argsort(kind='mergesort')
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches = [x[i1] for x in [matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches]]
                        i2 = matchLibTags.argsort(kind='mergesort')
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches = [list(x[i2]) for x in [matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches]]
                        #reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches)
                        matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, maccScores = reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches)
                        print(len(set(maccScores)))
                        #data = pd.DataFrame({'libID':matchLibTags, 'queID':matchQueTags, 'libIntensity':matchLibIntensities, 'queIntensity':matchQueIntensities, 'ppmDiff':ppmMatches})
                        #data = pd.concat(dfs)
                        #data = data.groupby(['libID','queID']).filter(lambda x: len(x.index) > 2.)
                        #data = data.groupby(['libID','queID']).apply(add_macc_column, maccDict)

                        #data, NAlist = reduce_mem_usage(data)
                        #allDfs.append(data)
                        #with bz2.BZ2File(pickleDir+pickleHeader+str(count), 'w') as pickleFile: pickle.dump(data, pickleFile)

                        pooledQueSpectra.clear()
                        del returns

                count += 1
                if count % printCutoff == 0:
                    time = timer()
                    #print('\nNumber of Pooled Experimental Spectra Analyzed: ' + str(count))
                    #print('Number of Spectra in Current Pooled Spectra: ' + str(len(scans)))
                    #print('Time Since Last Checkpoint: ' + str(round(time-prevtime,2)) + ' Seconds', flush=True)
                    print(round(time-prevtime,2),flush=True)
                    prevtime = time

    # Prints the final number of experimental spectra analyzed.
    print('Total Time (seconds): ' + str(timer()))
    print('Count: '+str(count),flush=True)
    if tag=='identify': return ppmList
#''
#@njit
def reduce_final_df(matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches):
    curLibTag = matchLibTags[-1]
    curQueTag = matchQueTags[-1]
    count = 1
    maccScores = []
    AB = matchLibIntensities[-1]*matchQueIntensities[-1]
    A = matchLibIntensities[-1]**2
    B = matchQueIntensities[-1]**2
    for i in reversed(range(len(matchLibTags)-1)):
        if matchLibTags[i] != curLibTag and matchQueTags[i] != curQueTag:
            if count == 1:
                matchLibTags.pop(i-1)
                matchLibIntensities.pop(i-1)
                matchQueTags.pop(i-1)
                matchQueIntensities.pop(i-1)
                ppmMatches.pop(i-1)
            else:
                cosine = cosine_similarity([AB, A, B])
                macc = (count**(1/5))*cosine
                maccScores.extend([macc]*count)
                count = 1
            curLibTag = matchLibTags[i]
            curQueTag = matchQueTags[i]
            AB = matchLibIntensities[i]*matchQueIntensities[i]
            A = matchLibIntensities[i]**2
            B = matchQueIntensities[i]**2
        else:
            AB += matchLibIntensities[-1]*matchQueIntensities[-1]
            A += matchLibIntensities[-1]**2
            B += matchQueIntensities[-1]**2
            count += 1

    return matchLibTags, matchLibIntensities, matchQueTags, matchQueIntensities, ppmMatches, maccScores[::-1]

def pooled_spectra_analysis(expSpectraFile, outFile, lib, ppmTol, ppmYOffset, queryPooling, spectraKeys=None):

    pickleDir = '/'.join(outFile.split('/')[:-1]) + '/compressed_processObj_canDeleteAfter/'
    if not os.path.exists(pickleDir):
        os.mkdir(pickleDir)
    pickleHeader = 'window'

    # query data file is loaded
    with mzxml.read(expSpectraFile, use_index=True) as spectra:

        # query data is looped over and scans are grouped by mz windows for future pooling.
        #  Additionally, for the second condition, various variables are saved for future reference.
        queScanDict = defaultdict(list)
        queValDict = defaultdict(dict)
        spectralCount = 0

        for spec in spectra:
            if 'precursorMz' not in spec: continue
            queScanDict[spec['precursorMz'][0]['precursorMz'],spec['precursorMz'][0]['windowWideness']].append(spec['num'])
            #if scan=='1199' or scan=='1200': print((spec['precursorMz'][0]['precursorMz'],spec['precursorMz'][0]['windowWideness']))
            peaksCount = spec['peaksCount']
            if 'compensationVoltage' in spec: CV = spec['compensationVoltage']
            else: CV = ''
            queValDict[spec['num']]['peaksCount'] = peaksCount
            queValDict[spec['num']]['CV'] = CV
            spectralCount += 1

        print('Number of Unpooled MS/MS Query Spectra: ' + str(spectralCount))
        print('Number of Pooled MS/MS Query Spectra/Mz Windows: ' + str(len(queScanDict)),flush=True)

        # To enhance the print experience, status prints will be given at intervals tailored to the number of identified windows.
        #  example: if there are 1-99 pooled query spectra, print statements are made after every pooled query spectra analysis is complete.
        #           if there are 100-999, print after every 10 pooled spectra. And so on.
        printCutoff = 100
        while printCutoff < len(queScanDict): printCutoff*=10
        printCutoff /= 100

        # 'lib' dictionary keys are kept as a separate list in this analysis. Note that they are sorted by precursor m/z.
        allLibKeys = lib.keys()
        allLibKeys = sorted(allLibKeys)

        # outfile is opened in advance so results can be written directly to the file as they are produced (mitigating memory use)
        # The second condition of this function returns a list of PPM differences that can be used for correction. Initialized here
        ppmList = []

        # Count variable keeps track of the number of query spectra that have been analyzed for time tracking purposes.
        count = 0

        # Library keys were saved as an integer to save on time and simplify other parts of the algorithm. I believe the purpose is now defunct, but it's harmless, so I'm keeping it.
        idToKeyDict = {}
        for key in allLibKeys: idToKeyDict[lib[key]['ID']] = key

        # tracking time for print statements.
        prevtime = timer()

        print('Enter Pooled Spectra Analysis:')
        print(str(timedelta(seconds=prevtime)), flush=True)

        columns=['libID','queID','libIntensity','queIntensity','ppmDiff']
        allDfs = []
        allData = []
        maccDecoys = []
        decs = []
        # looping through all the windows that the query data corresponds to
        for precMz_win, scans in queScanDict.items():

            # Determining all library spectra that should be pooled for this particular query data window
            top_mz = precMz_win[0] + precMz_win[1] / 2
            bottom_mz = precMz_win[0] - precMz_win[1] / 2
            libKeys = lib_mz_match_query_window( top_mz, bottom_mz, allLibKeys )
            if len(libKeys) == 0: continue
            pooledLibSpectra = pool_lib_spectra(lib, libKeys)

            # begin pooling query spectra
            pooledQueSpectra = []
            returns = initialize_return_values('identify')
            for i in range(len(scans)):

                # adding each scan in the window to the pooled spectra
                scan = scans[i]
                spec = spectra.get_by_id(scan)
                intensity = [x**0.5 for x in spec['intensity array']]
                peakIDs = [int(scan) for x in range(spec['peaksCount'])]
                pooledQueSpectra += list(zip(spec['m/z array'],intensity,peakIDs))

                # to reduce memory use for particularly large files, the user can limit the number of query spectra that are pooled. That's what this conditional statement takes care of.
                if (i % queryPooling == 0 and i!=0) or i == len(scans)-1:
                    pooledQueSpectra.sort()

                    # each peak in the pooled library and query spectra is compared, and the necessary data is extracted from matching peaks (as determined by the ppm tolerance)
                    libMzs = np.array([x[0] for x in pooledLibSpectra])
                    queMzs = np.array([x[0] for x in pooledQueSpectra])

                    pMatch, jMatch, ppmMatch = spectra_peak_comparison(libMzs, queMzs, ppmTol, ppmYOffset)
                    for i in range(len(pMatch)):
                        libPeak = pooledLibSpectra[pMatch[i]]
                        quePeak = pooledQueSpectra[jMatch[i]]

                        update_return_values(returns, pooledLibSpectra[pMatch[i]], pooledQueSpectra[jMatch[i]], pMatch[i], jMatch[i], ppmMatch[i], 'identify')
                        #data.append([libPeak[2], quePeak[2], libPeak[1], quePeak[1], ppmMatch[i]])
                    # output values are saved to the output file

                    pooledQueSpectra.clear()

            cosDict, countDict, ionDict, ppmDict = returns
            data = []
            maccs = []

            for key, value in countDict.items():
                if value > 2:
                    macc = (value**(1/5))*cosine_similarity(cosDict[key])
                    #df = pd.DataFrame(ppmDict[key], columns=columns)
                    #df['MaCC'] = pd.Series(macc, index=df.index)
                    #dfs.append(df)
                    #data += ppmDict[key]
                    #maccDict[key] = macc
                    #data.extend(ppmDict[key])
                    #allData += ppmDict[key]
                    maccs += [macc] * value

                    #libKey = idToKeyDict[key[0]]
                    decoy = 0
                    if 'DECOY' in lib[idToKeyDict[key[0]]]['ProteinName']: decoy = 1
                    #decs += [decoy] * value
                    maccDecoys.append((macc,decoy))

            #data = pd.concat(dfs)
            data = pd.DataFrame(data, columns=['libID','queID','libIntensity','queIntensity','ppmDiff'])
            data['MaCC_Score'] = maccs
            #data = data.groupby(['libID','queID']).filter(lambda x: len(x.index) > 2.)
            #data = data.groupby(['libID','queID']).apply(add_macc_column, maccDict)

            data, NAlist = reduce_mem_usage(data)
            #allDfs.append(data)
            #with bz2.BZ2File(pickleDir+pickleHeader+str(count), 'w') as pickleFile: pickle.dump(data, pickleFile)

            #compressedWindows.append(zlib.compress(pickle.dumps(data)))
            #print('size: '+str(sys.getsizeof(zlib.compress(pickle.dumps(data),zlib.Z_BEST_COMPRESSION))), flush=True)
            #print('size: '+str(sys.getsizeof(data)), flush=True)
            #return []
            #compressedWindows.append(data)
            # print statements for the user to track progress.
            count += 1
            #if count % 1000 == 0:
            if count % printCutoff == 0:
                time = timer()
                #print('\nNumber of Pooled Experimental Spectra Analyzed: ' + str(count))
                #print('Number of Spectra in Current Pooled Spectra: ' + str(len(scans)))
                #print('Time Since Last Checkpoint: ' + str(round(time-prevtime,2)) + ' Seconds', flush=True)
                print(round(time-prevtime,2),flush=True)
                prevtime = time

    pickle.dump(maccDict, open('C:/Users/ccranney/Desktop/Caleb_Files/data/output/CompressTest/maccDict.p', 'wb'))

    #lib = pickle.load(open(args['outDirectory']+'mgf_lib.p', 'rb'))
    #print('size: '+str(sys.getsizeof(compressedWindows)))
    # Prints the final number of experimental spectra analyzed.
    print('Total Time (seconds): ' + str(timer()))
    print('Count: '+str(count),flush=True)
    return ppmList
#'''


"""Module for a multisource sequence-to-sequence model.

The model has either one or two encoders (corresponding to the number of input sources).
The attention mechanism concatenates the context vectors from both encoders, which is then used by the decoder to generate the output.

The model also has adaptations for the low-resource setting: copy mechanism, coverage mechanism, and diagonal attention loss.
These can be switched on or off as hyperparameters to the model.

The OCR post-correction training process uses early stopping on the validation set character error rate.

The model uses beam search for generation.

Copyright (c) 2021, Shruti Rijhwani
All rights reserved.

This source code is licensed under the BSD-style license found in the
LICENSE file in the root directory of this source tree.
"""

import _dynet as dy
from constants import (
    LSTM_NUM_OF_LAYERS,
    EMBEDDING_DIM,
    HIDDEN_DIM,
    ATTENTION_SIZE,
    UNK,
    EOS,
    COV_LOSS_WEIGHT,
    DIAG_LOSS_WEIGHT,
)
from utils import DataReader, Hypothesis
import math
import numpy as np
import random
import logging

class TwoSourceModel:
    def __init__(
        self,
        src1_vocab,
        src2_vocab,
        tgt_vocab,
        single,
        pointer_gen,
        coverage,
        diag_loss,
        load_model,
        model_file,
        beam_size,
        best_val_cer,
    ):
        self.model = dy.ParameterCollection()

        self.src1_vocab = src1_vocab
        self.src2_vocab = src2_vocab
        self.tgt_vocab = tgt_vocab

        self.src1_lookup = self.model.add_lookup_parameters(
            (src1_vocab.length(), EMBEDDING_DIM)
        )
        self.src2_lookup = self.model.add_lookup_parameters(
            (src2_vocab.length(), EMBEDDING_DIM)
        )
        self.tgt_lookup = self.model.add_lookup_parameters(
            (tgt_vocab.length(), EMBEDDING_DIM)
        )

        self.enc1_fwd_lstm = dy.CoupledLSTMBuilder(
            LSTM_NUM_OF_LAYERS, EMBEDDING_DIM, HIDDEN_DIM, self.model
        )
        self.enc1_bwd_lstm = dy.CoupledLSTMBuilder(
            LSTM_NUM_OF_LAYERS, EMBEDDING_DIM, HIDDEN_DIM, self.model
        )
        self.pret1_w = self.model.add_parameters((src1_vocab.length(), HIDDEN_DIM))
        self.pret1_b = self.model.add_parameters((src1_vocab.length()))

        self.enc2_fwd_lstm = dy.CoupledLSTMBuilder(
            LSTM_NUM_OF_LAYERS, EMBEDDING_DIM, HIDDEN_DIM, self.model
        )
        self.enc2_bwd_lstm = dy.CoupledLSTMBuilder(
            LSTM_NUM_OF_LAYERS, EMBEDDING_DIM, HIDDEN_DIM, self.model
        )
        self.pret2_w = self.model.add_parameters((src2_vocab.length(), HIDDEN_DIM))
        self.pret2_b = self.model.add_parameters((src2_vocab.length()))

        self.att1_w1 = self.model.add_parameters((ATTENTION_SIZE, HIDDEN_DIM * 2))
        self.att1_w2 = self.model.add_parameters(
            (ATTENTION_SIZE, HIDDEN_DIM * LSTM_NUM_OF_LAYERS * 2)
        )
        self.att1_v = self.model.add_parameters((1, ATTENTION_SIZE))

        self.att2_w1 = self.model.add_parameters((ATTENTION_SIZE, HIDDEN_DIM * 2))
        self.att2_w2 = self.model.add_parameters(
            (ATTENTION_SIZE, HIDDEN_DIM * LSTM_NUM_OF_LAYERS * 2)
        )
        self.att2_v = self.model.add_parameters((1, ATTENTION_SIZE))

        self.dec_lstm = dy.CoupledLSTMBuilder(
            LSTM_NUM_OF_LAYERS, HIDDEN_DIM * 4 + EMBEDDING_DIM, HIDDEN_DIM, self.model
        )
        self.W_s = self.model.add_parameters((HIDDEN_DIM, HIDDEN_DIM * 4))
        self.b_s = self.model.add_parameters((HIDDEN_DIM))
        self.dec_w = self.model.add_parameters((tgt_vocab.length(), HIDDEN_DIM))
        self.dec_b = self.model.add_parameters((tgt_vocab.length()))

        # Pointer-generator parameters
        self.ptr_w_c = self.model.add_parameters((1, 2 * HIDDEN_DIM))
        self.ptr_w_s = self.model.add_parameters((1, 2 * HIDDEN_DIM))
        self.ptr_w_x = self.model.add_parameters((1, EMBEDDING_DIM + 4 * HIDDEN_DIM))

        # Coverage parameters
        self.w_cov = self.model.add_parameters((ATTENTION_SIZE, 1))

        self.single_source = single
        self.pointer_gen = pointer_gen
        self.coverage = coverage
        self.diag_loss = diag_loss
        self.model_file = model_file

        if load_model:
            self.model.populate(load_model)
            logging.info("Loaded model: {}".format(load_model))

        self.beam_size = beam_size
        self.best_val_cer = best_val_cer

    def save(self):
        self.model.save(self.model_file)

    def run_lstm(self, init_state, input_vecs):
        out_vectors = init_state.transduce(input_vecs)
        return out_vectors

    def embed_idx(self, idx_list, embed_lookup):
        return [embed_lookup[idx] for idx in idx_list]

    def encode(self, embeds, fwd_lstm, bwd_lstm):
        embeds_rev = list(reversed(embeds))
        fwd_vectors = self.run_lstm(fwd_lstm.initial_state(), embeds)
        bwd_vectors = self.run_lstm(bwd_lstm.initial_state(), embeds_rev)
        bwd_vectors = list(reversed(bwd_vectors))
        vectors = [dy.concatenate(list(p)) for p in zip(fwd_vectors, bwd_vectors)]
        return vectors

    def encoder_forward(self, src1, src2):
        embedded_src1 = self.embed_idx(src1, self.src1_lookup)
        if self.single_source:
            embedded_src2 = [dy.vecInput(EMBEDDING_DIM) for idx in src2]
        else:
            embedded_src2 = self.embed_idx(src2, self.src2_lookup)

        encoded_src1 = self.encode(
            embedded_src1, self.enc1_fwd_lstm, self.enc1_bwd_lstm
        )
        encoded_src2 = self.encode(
            embedded_src2, self.enc2_fwd_lstm, self.enc2_bwd_lstm
        )

        src1_mat = dy.concatenate_cols(encoded_src1)
        src1_w1dt = self.att1_w1 * src1_mat
        src2_mat = dy.concatenate_cols(encoded_src2)
        src2_w1dt = self.att2_w1 * src2_mat

        if not self.single_source:
            start = (
                self.W_s * dy.concatenate([encoded_src1[-1], encoded_src2[-1]])
                + self.b_s
            )
        else:
            start = (
                self.W_s
                * dy.concatenate([encoded_src1[-1], dy.vecInput(2 * HIDDEN_DIM)])
                + self.b_s
            )

        last_output_embeddings = self.tgt_lookup[self.tgt_vocab.str2int(EOS)]
        c1_t = dy.vecInput(2 * HIDDEN_DIM)
        c2_t = dy.vecInput(2 * HIDDEN_DIM)
        decoder_state = self.dec_lstm.initial_state([start, dy.tanh(start)]).add_input(
            dy.concatenate([c1_t, c2_t, last_output_embeddings])
        )
        return src1_mat, src2_mat, src1_w1dt, src2_w1dt, decoder_state

    def attend(self, input_mat, state, w1dt, w2, v, coverage):
        w2dt = w2 * dy.concatenate(list(state.s()))
        if coverage:
            w1dt = w1dt + self.w_cov * dy.transpose(coverage)
        a_t = dy.transpose(v * dy.tanh(dy.colwise_add(w1dt, w2dt)))
        a_t = dy.softmax(a_t)
        return a_t, (input_mat * a_t)

    def get_pointergen_probs(self, c_t, state, x_t, a_t, probs, src1):
        if not self.pointer_gen:
            return probs, 1.0
        unk_idx = self.tgt_vocab.str2int(UNK)
        p_gen = dy.logistic(
            self.ptr_w_c * c_t
            + self.ptr_w_s * dy.concatenate(list(state.s()))
            + self.ptr_w_x * x_t
        )
        gen_probs = probs * p_gen
        copy_probs = a_t * (1 - p_gen)
        copy_probs_update = []
        for i in gen_probs:
            copy_probs_update.append([i])
        for char, prob in zip(src1, copy_probs):
            cur_idx = self.tgt_vocab.str2int(self.src1_vocab.int2str(char))
            if cur_idx == unk_idx:
                continue
            if isinstance(cur_idx, int):
                copy_probs_update[cur_idx].append(prob)
            else:
                for idx in cur_idx:
                    copy_probs_update[idx].append(prob / len(cur_idx))
        sum_probs = dy.concatenate([dy.esum(exps) for exps in copy_probs_update])
        return sum_probs, p_gen.scalar_value()

    def get_coverage(self, a_t, prev_coverage, training=True):
        if not self.coverage:
            if not training:
                return None
            return dy.scalarInput(0), None
        coverage = a_t + prev_coverage
        if training:
            return (
                dy.sum_elems(dy.min_dim(dy.concatenate([a_t, coverage], d=1), d=1)),
                coverage,
            )
        return coverage

    def get_diag_loss(self, a_t, t):
        if self.diag_loss < 0:
            return dy.scalarInput(0)
        off_diag_elems = [dy.scalarInput(0)]
        for i, prob in enumerate(a_t):
            if i < (t - self.diag_loss) or i > (t + self.diag_loss):
                off_diag_elems.append(prob)
        return dy.esum(off_diag_elems)

    def decode_loss(self, src1, src2, tgt):
        src1_mat, src2_mat, src1_w1dt, src2_w1dt, decoder_state = self.encoder_forward(
            src1, src2
        )
        _, prev_coverage = self.get_coverage(
            a_t=dy.vecInput(len(src1)), prev_coverage=dy.vecInput(len(src1))
        )

        loss = []
        cov_loss = []
        diag_loss = []

        embedded_tgt = self.embed_idx(tgt, self.tgt_lookup)
        last_output_embeddings = self.tgt_lookup[self.tgt_vocab.str2int(EOS)]

        for t, (char, embedded_char) in enumerate(zip(tgt, embedded_tgt)):
            a_t, c1_t = self.attend(
                src1_mat,
                decoder_state,
                src1_w1dt,
                self.att1_w2,
                self.att1_v,
                prev_coverage,
            )
            if not self.single_source:
                _, c2_t = self.attend(
                    src2_mat, decoder_state, src2_w1dt, self.att2_w2, self.att2_v, None
                )
            else:
                c2_t = dy.vecInput(2 * HIDDEN_DIM)

            x_t = dy.concatenate([c1_t, c2_t, last_output_embeddings])
            decoder_state = decoder_state.add_input(x_t)

            out_vector = self.dec_w * decoder_state.output() + self.dec_b
            probs = dy.softmax(out_vector)
            probs, _ = self.get_pointergen_probs(
                c1_t, decoder_state, x_t, a_t, probs, src1
            )

            loss.append(-dy.log(dy.pick(probs, char)))
            cov_loss_cur, prev_coverage = self.get_coverage(a_t, prev_coverage)
            cov_loss.append(cov_loss_cur)
            diag_loss.append(self.get_diag_loss(a_t, t))

            last_output_embeddings = embedded_char

        loss = dy.esum(loss)
        cov_loss = dy.esum(cov_loss)
        diag_loss = dy.esum(diag_loss)
        return loss + COV_LOSS_WEIGHT * cov_loss + DIAG_LOSS_WEIGHT * diag_loss

    def get_loss(self, src1, src2, tgt):
        return self.decode_loss(src1, src2, tgt)

    def generate_beam(self, src1, src2):
        src1_mat, src2_mat, src1_w1dt, src2_w1dt, decoder_state = self.encoder_forward(
            src1, src2
        )

        hypothesis_list = [
            Hypothesis(
                text_list=[self.tgt_vocab.str2int(EOS)],
                decoder_state=decoder_state,
                c1_t=dy.vecInput(2 * HIDDEN_DIM),
                c2_t=dy.vecInput(2 * HIDDEN_DIM),
                prev_coverage=self.get_coverage(
                    a_t=dy.vecInput(len(src1)),
                    training=False,
                    prev_coverage=dy.vecInput(len(src1)),
                ),
                score=0.0,
                p_gens=[],
            )
        ]
        completed_list = []

        for t in range(int(len(src1) * 1.1)):
            new_hyp_list = []
            new_hyp_scores = []
            for hyp in hypothesis_list:
                last_output_embeddings = self.tgt_lookup[hyp.text_list[-1]]

                a_t, c1_t = self.attend(
                    src1_mat,
                    hyp.decoder_state,
                    src1_w1dt,
                    self.att1_w2,
                    self.att1_v,
                    hyp.prev_coverage,
                )
                if not self.single_source:
                    _, c2_t = self.attend(
                        src2_mat,
                        hyp.decoder_state,
                        src2_w1dt,
                        self.att2_w2,
                        self.att2_v,
                        None,
                    )
                else:
                    c2_t = dy.vecInput(2 * HIDDEN_DIM)

                x_t = dy.concatenate([c1_t, c2_t, last_output_embeddings])
                decoder_state = hyp.decoder_state.add_input(x_t)

                probs = dy.softmax(self.dec_w * decoder_state.output() + self.dec_b)
                probs, cur_p_gen = self.get_pointergen_probs(
                    c1_t, decoder_state, x_t, a_t, probs, src1
                )
                probs = probs.npvalue()

                for ind in range(len(probs)):
                    text_list = hyp.text_list + [ind]
                    p_gens = hyp.p_gens + [cur_p_gen]
                    score = (hyp.score + math.log(probs[ind])) / (len(text_list) ** 0.0)
                    coverage = self.get_coverage(a_t, hyp.prev_coverage, training=False)
                    new_hyp_list.append(
                        Hypothesis(
                            text_list=text_list,
                            decoder_state=decoder_state,
                            c1_t=c1_t,
                            c2_t=c2_t,
                            prev_coverage=coverage,
                            score=score,
                            p_gens=p_gens,
                        )
                    )
                    new_hyp_scores.append(score)

            top_inds = np.argpartition(np.array(new_hyp_scores), -self.beam_size)[
                -self.beam_size :
            ]
            new_hyp_list = np.array(new_hyp_list)[top_inds]

            hypothesis_list = []

            for new_hyp in new_hyp_list:
                if new_hyp.text_list[-1] == self.tgt_vocab.str2int(EOS) and t > 0:
                    completed_list.append(new_hyp)
                else:
                    hypothesis_list.append(new_hyp)

            if len(completed_list) >= self.beam_size:
                break

        if len(completed_list) == 0:
            sorted(hypothesis_list, key=lambda x: x.score, reverse=True)
            completed_list = [hypothesis_list[0]]

        for hyp in completed_list:
            hyp.text_list = [self.tgt_vocab.int2str(i) for i in hyp.text_list]

        top_hyp = sorted(completed_list, key=lambda x: x.score, reverse=True)[0]
        return "".join(top_hyp.text_list).replace(EOS, "").strip(), top_hyp.p_gens[1:-1]


import os
import cv2
import configparser
import numpy as np
import tensorflow as tf

class NSFWModule:
    # 默认必须使用gpu
    def __init__(self, conf: configparser.ConfigParser, gpu_id: int):

        model_path = conf.get("nsfw", "model_path")
        max_memory = int(conf.get("nsfw", "max_memory"))
        gpus = tf.config.list_physical_devices('GPU')
        gpu_used = None
        for item in gpus:
            # print(item)
            if str(gpu_id) in item.name:
                gpu_used = item
        if gpu_used is None:
            raise ValueError("Your machine mayn't have %d gpu, please check your 'gpu_id'!" % (gpu_id))
        tf.config.experimental.set_virtual_device_configuration(
            gpu_used,
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=max_memory)])
        self.gpu = '/device:GPU:' + str(gpu_id)
        with tf.device(self.gpu):
            self.model = self.__load_model(model_path)

    def __load_model(self, model_path: str):
        if model_path is None or not os.path.exists(model_path):
            raise ValueError("saved_model_path must be the valid directory of a saved model to load.")

        model = tf.keras.models.load_model(model_path)  # , custom_objects={'KerasLayer': hub.KerasLayer})
        return model

    def detect(self, batch_image: np.ndarray):
        batch_image = np.asarray([cv2.resize(cv2.cvtColor(item, cv2.COLOR_BGR2RGB), (224, 224))
                                  for item in batch_image], dtype=np.float32) / 255
        result = []
        with tf.device(self.gpu):
            model_preds = self.model.predict(batch_image)
            categories = ['drawings', 'hentai', 'neutral', 'porn', 'sexy']
            for item in model_preds.tolist():
                class_ = item.index(max(item))
                result.append(categories[class_])
        result = [item if item != "neutral" else None for item in result]
        return result


from abc import ABC, abstractmethod
import numpy as np

class Layer(ABC):

    @abstractmethod
    def forward(self,
                x: np.ndarray):
        pass

    @abstractmethod
    def backward(self,
                 prev_error: np.ndarray):
        pass

    @abstractmethod
    def update_parameters(self,
                          alpha: float,
                          beta: float):
        pass