#!/usr/bin/env python

#nknguyen at soe ucsc edu
#Feb 07 2012
#TCR repertoire simulation

'''
Simulate CDR3 TCR repertoires to assess the overlapping of clones/sequences between two independent samplings of certain size.
(Purpose: use to determine what is the sufficient sampling size to get all the important, clonally expanded CDR3)

There are 4 main steps involved:
1/ Generate a set of m unique CDR3 nucleotide sequences, call this set R (it is estimated that each individual has about 3x10^6 uniq seqs)
   This steps take into account the CDR3-length distribution and the Amino-acid usage at each position of the CDR3.
   Input: a/ len2count.txt with format: <aa-cdr3Length>\t<count>
          b/ aaUsage.txt with format: <aa-cdr3Length>\t<Position>\t<aa-Letter>\t<count>
          c/ m = size of the repertoire to be created
   Generate: m nucleotide sequences such that when translated to amino acid sequences, follows the input length distribution and aaUsage.

2/ Generate the repertoire (say P) of CDR3 of an individual from set R in (1):
   Let M be the number of cells (number of total sequences) an individual has. It is estimated that there are about 6x10^7 CD8+ T cells per individual.
   Populate repertoire P with the n sequences in R until P has M sequences. This step can assume a lognormal (previously used uniform) distribution of the uniq sequences or model clonal expansion.
   The program model clonal expansion by in input pseudo repertoire with format <topFreq>,<secondTopFreq>,...,<ith-topFreq>. A random sequence will be drawn from n sequences in (a), and will populate P with topFreq*M counts. Then a second random seq will be drawn, etc until there are M sequences or the frequency list is done. If freq list is done and there are < M sequences in P, just randomly fill in P with the rest of the m sequences
   Generate: M CDR3 nucleotide sequences

3/ Sampling:
   Repeat step (1) and (2) to generate another TCR repertoire for the second individual.
   For each individual, randomly choose S sequences from the corresponding repertoire.
   Calculate how much overlap the two samplings have.
   Repeat this step with difference sampling size

4/ a. Repeat step (3) -numSamplingPerSim times. Calculate mean & std
(Or/And)
   b. Repeat steps (1,2,3) N times (N simulations). Calculate mean statistics

(5/ Plots)
'''

import os, sys, re, time, random, copy, gzip
import cPickle as pickle
from optparse import OptionParser
import xml.etree.ElementTree as ET

from jobTree.scriptTree.target import Target
from jobTree.scriptTree.stack import Stack

from sonLib.bioio import logger
from sonLib.bioio import system
from sonLib.bioio import getTempDirectory
from sonLib.bioio import setLogLevel
import numpy as np
from scipy.stats import poisson, lognorm

#from immunoseq.lib import *

################## SEQUENCE OBJECTS ######################
class Seq:
    def __init__(self, nuc, aa):
        self.nuc = nuc
        self.aa = aa
        self.count = 0

    def __cmp__(self, other):
        return cmp(self.nuc, other.nuc)

    def setCount(self, count):
        self.count = count

    def updateCount(self, add):
        self.count += add

class Aaseq:
    def __init__(self, aa):
        self.aa = aa
        self.count = 0

    def __cmp__(self, other):
        return cmp(self.aa, other.aa)

    def setCount(self, count):
        self.count = count

    def updateCount(self, add):
        self.count += add

    def setFreq(self, total):
        if total == 0:
            self.freq = 0
        else:
            self.freq = 100.0*self.count/total

class Seqs(list):
    def add(self, seq):
        leftIndex = 0
        rightIndex = len(self) -1

        if len(self) == 0 or seq > self[rightIndex]:
            self.append(seq)
        elif seq < self[leftIndex]:
            self.insert(0, seq)
        elif seq == self[leftIndex] or seq == self[rightIndex]:
            return
        else:
            while True:
                middleIndex = int( (leftIndex + rightIndex)/2 )
                compare = cmp( self[middleIndex], seq )
                if compare == 0: #sequence is already in the list, do not add, just return
                    break
                elif leftIndex == middleIndex: #add seq to the right of leftIndex
                    self.insert(rightIndex, seq)
                    break
                elif compare > 0: #seq lies somewhere btw [middleIndex, rightIndex]
                    rightIndex = middleIndex
                elif compare < 0:
                    leftIndex = middleIndex

    def search(self, seq):
        leftIndex = 0
        rightIndex = len(self) -1

        if len(self) == 0 or seq > self[rightIndex] or seq < self[leftIndex]: #not in list
            return -1
        elif seq == self[leftIndex]:
            return leftIndex
        elif seq == self[rightIndex]:
            return rightIndex
        else:
            while True:
                middleIndex = int( (leftIndex + rightIndex)/2 )
                compare = cmp( self[middleIndex], seq )
                if compare == 0: #sequence is already in the list, do not add, just return
                    return middleIndex
                elif leftIndex == middleIndex: #Not found
                    return -1
                elif compare > 0: #seq lies somewhere btw [middleIndex, rightIndex]
                    rightIndex = middleIndex
                elif compare < 0:
                    leftIndex = middleIndex

######################### MAIN PIPELINE ################################

#============= SETTING UP SIMULATIONS ==============
class Setup( Target ):
    """Setting up simulations
    """
    def __init__(self, options, doneSims):
        Target.__init__(self, time = 0.00025)
        self.options = options
        self.doneSims = doneSims

    def run(self):
        setLogLevel("DEBUG")
        numSim = self.options.numSim
        maxSimPerRun = self.options.maxSimPerRun
        sims = min( [maxSimPerRun, numSim - self.doneSims] )

        outdir = os.path.join( self.options.outdir, "sims", str(self.doneSims) ) #outdir/sims/batchId
        system("mkdir -p %s" %outdir)
        self.addChildTarget( SimulationBatch(self.options, outdir, sims, self.doneSims) )

        doneSims = self.doneSims + sims
        if doneSims < numSim:
            self.setFollowOnTarget( Setup(self.options, doneSims) ) #recusion call for next batch of simulations
        else:
            summarydir = os.path.join(self.options.outdir, "summary") #outdir/summary
            statsdir = os.path.join(self.options.outdir, "stats") #outdir/stats
            system("mkdir -p %s" % statsdir)
            readPickle = True
            writePickle = False
            writeSummary = True
            self.setFollowOnTarget( Summary(summarydir, self.options.samSize, self.options.numSamples, statsdir, writePickle, writeSummary, readPickle) )

class SimulationBatch( Target ):
    def __init__(self, options, outdir, numSim, batchid):
        Target.__init__(self, time=0.00025)
        self.options = options
        self.outdir = outdir #outdir/sims/batchId
        self.numSim = numSim
        self.batchid = batchid

    def run(self):
        for i in xrange(self.numSim):
            outdir = os.path.join(self.outdir, "sim-%d" %i) #outdir/sims/batchId/sim-Id
            system("mkdir -p %s" %outdir)
            self.addChildTarget( Simulation(self.options, outdir, self.batchid, i) )

        #Combine the results
        samplingsDir = os.path.join(self.options.outdir, "samplings", str(self.batchid)) #outdir/samplings/batchid
        sumdir = os.path.join(self.options.outdir, "summary", str(self.batchid)) #outdir/summary/batchid
        system("mkdir -p %s" %sumdir)
        readPickle = False
        writePickle = True
        writeSummary = False
        #self.setFollowOnTarget( Summary(self.outdir, self.options.samSize, self.options.numSamples, sumdir, writePickle, writeSummary, readPickle) )
        self.setFollowOnTarget( Summary(samplingsDir, self.options.samSize, self.options.numSamples, sumdir, writePickle, writeSummary, readPickle) )

#------------- SIMULATION ---------------------
class Simulation( Target ):
    """Generate s simulations (corresponding to s samples (s individuals))
       Creating TCR repertoire for each individuals
       After got repertoires for All individuals, call following job to sampling from these repertoires
    """
    def __init__(self, options, outdir, batchid, simid):
        Target.__init__(self, time=0.00025)
        self.options = options
        self.outdir = outdir #outdir/sims/batchId/sim-Id
        self.batchid = batchid
        self.simid = simid

    def run(self):
        globalTempDir = self.getGlobalTempDir()
        for i in xrange(self.options.numSamples):
            outdir = os.path.join(globalTempDir, str(i))
            system("mkdir -p %s" %outdir) #simTempDir/sampleID/
            self.addChildTarget( SimulationSingle(self.options, outdir) )

        #Sampling from these repertoire:
        samplingsOutdir = os.path.join(self.options.outdir, "samplings", str(self.batchid), "sim-%d" %self.simid)
        system("mkdir -p %s" %samplingsOutdir) #outdir/samplings/batchId/sim-id
        self.setFollowOnTarget( Samplings(globalTempDir, self.outdir, self.options.numSamples, self.options.samSize, self.options.cutoffs, self.options.numSamplingPerSim, samplingsOutdir) )
        #Calculate pair-wise overlap between the samples
        #self.setFollowOnTarget( Overlap(globalTempDir, self.outdir, self.options.numSamples, self.options.samSize, self.options.cutoffs) )

class SimulationSingle( Target ):
    """Generate repertoire for one sample/individual
       Starts by generating the uniq sequences of certain CDR3 length
       Then calls follow on jobs to use these sequences to create the repertoire
    """
    def __init__(self, options, outdir):
        Target.__init__(self, time = 0.00025)
        self.outdir = outdir #simTempDir/sampleID
        self.options = options

    def run(self):
        globalTempDir = self.getGlobalTempDir()
        #Generate repertoire of uniq sequences for each length:
        for l in self.options.len2freq:#each length
            size = self.options.totalClones*self.options.len2freq[l]
            if l not in self.options.len2aaUsage or size == 0:
                continue
            aaUsage = self.options.len2aaUsage[l]
            aa2codons = self.options.aa2codons
            self.addChildTarget( UniqSeqs(globalTempDir, size, aaUsage, aa2codons) )

        self.setFollowOnTarget( Repertoire(globalTempDir, self.outdir, self.options.totalSeqs, self.options.topFreqs, self.options.samSize) )

#+++++++++++ REPERTOIRE OF UNIQ SEQS/ CLONES ++++++++++++++
class UniqSeqs( Target ):
    """Generate repertoire of uniq sequences of a certain length
    """
    def __init__(self, outdir, size, aaUsage, aa2codons):
        Target.__init__(self)
        self.outdir = outdir
        self.size = size
        self.aaUsage = aaUsage
        self.aa2codons = aa2codons

    def run(self):
        #Generate the sequences
        seqs = getUniqSeqRep(self.size, self.aaUsage, self.aa2codons)
        #Pickling the sequences to outdir
        pickleFile = os.path.join( self.outdir, "uSeqs-l%d.pickle" %len(self.aaUsage) )
        pickle.dump(seqs, gzip.open(pickleFile, "wb"))

#+++++++++++++ PSEUDO REPERTOIRE ++++++++++++++++++
class Repertoire( Target ):
    """Combine the len-specific uniqSeqs repertoires
       From uniqSeq repertoire, generate a pseudo repertoire, using topFreqs if available
       Add jobs to run samplings with different sizes
    """
    def __init__(self, indir, outdir, size, topFreqs, samSizes):
        Target.__init__(self)
        self.indir = indir #simSampleTempDir
        self.outdir = outdir #simTempDir/sampleID
        self.size = size
        self.topFreqs = topFreqs
        self.samSizes = samSizes

    def run(self):
        #Load pickles to get uniqSeqs
        pickleFiles = os.listdir( self.indir )
        seqs = []
        for file in pickleFiles:
            if re.search("pickle", file):
                seqs.extend( pickle.load( gzip.open( os.path.join(self.indir, file) , "rb") ) )
                system( "rm -f %s" %(os.path.join(self.indir, file)) ) #Remove the uniq-sequences pickle file of the specific length after done loading

        getRep( seqs, self.size, self.topFreqs )

        #Print repertoire to server so that it can be read by following jobs:
        pickleFile = os.path.join(self.outdir, "rep.pickle") #simTempDir/sampleID/rep.pickle
        pickle.dump( seqs, gzip.open(pickleFile, "wb") )

#++++++++++++++ SAMPLING from pseudoRepertoire ++++++++++++++++
class Samplings( Target ):
    """(For each repertoier simulation,) Set up samplings for all pairs, all sampling sizes, all number of samplings.
    """
    def __init__(self, indir, outdir, numSamples, samSizes, cutoffs, numSamplings, samplingsOutdir):
        Target.__init__(self, time=0.00025)
        self.indir = indir #simTempDir
        self.outdir = outdir #outdir/sims/batchId/sim-Id
        self.numSamples = numSamples
        self.samSizes = samSizes
        self.cutoffs = cutoffs
        self.numSamplings = numSamplings
        self.samplingsOutdir = samplingsOutdir #outdir/samplings/batchId/sim-id

    def run(self):
        for i in xrange(self.numSamplings):#Each sampling
            for samsize in self.samSizes: #Each sampling size
                outdir = os.path.join(self.outdir, "sampling-%d" %i) #outdir/sims/batchId/sim-Id/sampling-Id
                for j in xrange(self.numSamples - 1):
                    for k in xrange(self.numSamples): #Each pair
                        pair = [j, k]
                        self.addChildTarget( SamplingPair(self.indir, outdir, samsize, pair, self.cutoffs) )
        #Summary the overlapping stats:
        readPickle = False
        writePickle = False
        writeSummary = True
        self.setFollowOnTarget( Summary(self.outdir, self.samSizes, self.numSamples, self.samplingsOutdir, writePickle, writeSummary, readPickle) )

class SamplingPair( Target ):
    """
    """
    def __init__(self, indir, outdir, size, pair, cutoffs):
        Target.__init__(self, time=0.00025)
        self.indir = indir #simTempDir (which has simTempDir/sampleId/rep.pickle)
        self.outdir = outdir #outdir/sims/batchId/sim-Id/sampling-Id
        self.size = size
        self.pair = pair
        self.cutoffs = cutoffs

    def run(self):
        globalTempDir = self.getGlobalTempDir()
        seqfiles = []
        for i, p in enumerate(self.pair):
            repertoireFile = os.path.join(self.indir, str(p), "rep.pickle") #simTempDir/sampleId/rep.pickle
            outfile = os.path.join(globalTempDir, "%d.pickle" %i) #samplingPairTempDir/sampleId.pickle
            seqfiles.append(outfile)
            self.addChildTarget( Sampling(outfile, repertoireFile, self.size) )

        overlapOutdir = os.path.join( self.outdir, "%d-%d" %(self.pair[0], self.pair[1]) ) #outdir/sims/batchId/sim-Id/sampling-Id/sam1-sam2
        system("mkdir -p %s" %overlapOutdir)
        overlapOutfile = os.path.join( overlapOutdir, "%d.txt" %self.size )
        self.setFollowOnTarget( OverlapPairwise(seqfiles, overlapOutfile, self.cutoffs) )

class Sampling( Target ):
    """Sample from the input repertoire for "size" sequences
    """
    def __init__(self, outfile, repertoireFile, samsize):
        Target.__init__(self)
        self.outfile = outfile
        self.repfile = repertoireFile
        self.size = samsize

    def run(self):
        seqs = pickle.load( gzip.open(self.repfile, "rb") )
        samseqs = sampling(seqs, self.size)

        #Pickling seqs:
        pickle.dump( samseqs, gzip.open(self.outfile, "wb") )

#----------- OVERLAP -----------
class OverlapPairwise( Target ):
    """Output file: %cutoff\tclones1\tclones2\toverlap\t%1overlap2\t%2overlap1\t%reads1overlap2\t%reads2overlap1
    """
    def __init__(self, seqsFiles, outfile, cutoffs):
        Target.__init__(self, time=0.001)
        self.outfile = outfile
        self.seqsFile1 = seqsFiles[0]
        self.seqsFile2 = seqsFiles[1]
        self.cutoffs = cutoffs

    def run(self):
        seqs1 = pickle.load( gzip.open(self.seqsFile1, "rb") )
        seqs2 = pickle.load( gzip.open(self.seqsFile2, "rb") )
        reads1, reads2, clones1, clones2, stats1, stats2 = getOverlap(seqs1, seqs2, self.cutoffs)

        #Print stats
        f = open(self.outfile, "w")
        f.write("#%Cutoff\tClones1\tClones2\tOverlap1\tOverlap2\t%1overlap2\t%2overlap1\t%reads1overlap2\t%reads2overlap1\n")
        for i,c in enumerate(self.cutoffs):
            oc1 = stats1["oclones"][i]
            or1 = stats1["oreads"][i]
            t1 = clones1[i]
            r1 = reads1[i]

            oc2 = stats2["oclones"][i]
            or2 = stats2["oreads"][i]
            t2 = clones2[i]
            r2 = reads2[i]

            f.write("%.3f\t%d\t%d\t%d\t%d\t%.2f\t%.2f\t%.2f\t%.2f\n" %( c, t1, t2, oc1, oc2, getPc(oc1, t1), getPc(oc2, t2), getPc(or1, r1), getPc(or2, r2) ))
        f.close()

#================ AVERAGE ACROSS THE SIMULATIONS ====================
class Summary( Target ):
    def __init__(self, indir, samSizes, numSamples, outdir, writePickle, writeSummary, readPickle):
        Target.__init__(self, time=0.001)
        self.indir = indir #outdir/sims/batchId/sim-Id; #outdir/samplings/batchId; #outdir/summary
        self.samSizes = samSizes
        self.numSamples = numSamples
        self.outdir = outdir #outdir/samplings/batchId/sim-Id; #outdir/summary/batchid; #outdir/stats
        self.writePickle = writePickle
        self.writeSummary = writeSummary
        self.readPickle = readPickle

    def run(self):
        #Get the pairs:
        for i in xrange(self.numSamples - 1):
            for j in xrange(self.numSamples):
                pair = "%d-%d" %( i, j )
                outdir = os.path.join(self.outdir, pair) #outdir/samplings/batchId/sim-Id/pair; #outdir/summary/batchid/pair
                system("mkdir -p %s" %outdir)
                for samsize in self.samSizes:
                    self.addChildTarget( SummarySamsize(self.indir, samsize, pair, outdir, self.writePickle, self.writeSummary, self.readPickle) )
        if not self.writePickle:
            self.setFollowOnTarget( CleanupSummary(self.indir) )

class SummarySamsize( Target ):
    def __init__(self, indir, samsize, pair, outdir, writePickle, writeSummary, readPickle):
        Target.__init__(self, time=0.001)
        self.indir = indir #outdir/sims/batchId/sim-Id; #outdir/samplings/batchId; #outdir/summary
        self.samsize = samsize
        self.pair = pair
        self.outdir = outdir #outdir/samplings/batchId/sim-Id/pair; #outdir/summary/batchid/pair; #outdir/stats/pair
        self.writePickle = writePickle
        self.writeSummary = writeSummary
        self.readPickle = readPickle

    def run(self):
        #Get the stats:
        cutoff2stats = {} #key = cutoff, val = [clones1, clones2, overlap1, overlap2, %1oc2, %2oc1, %1or2, %2or1]
        indirs = os.listdir(self.indir)
        for indir in indirs:
            if not os.path.isdir( os.path.join(self.indir,indir) ):
                continue
            file = os.path.join(os.path.join(self.indir,indir), self.pair, "%d.txt" %self.samsize)
            if self.readPickle:
                file = os.path.join(os.path.join(self.indir,indir), self.pair, "%d.pickle" %self.samsize)

            if not os.path.exists(file):
                raise NonExistFileError("File %s does not exist\n" %file)

            if not self.readPickle:
                currCutoff2stats = readStats(file)
            else:
                currCutoff2stats = pickle.load( gzip.open(file, "rb") )
            for c, stats in currCutoff2stats.iteritems():
                if c not in cutoff2stats:
                    cutoff2stats[c] = stats
                else:
                    for i, s in enumerate(stats):
                        cutoff2stats[c][i].extend(s)

        #Calculate average and standard deviation:
        if self.writePickle:
            outfile = os.path.join(self.outdir, "%d.pickle" %self.samsize)
            pickle.dump(cutoff2stats, gzip.open(outfile, "wb"))

        if self.writeSummary:
            for c, stats in cutoff2stats.iteritems():
                for i, s in enumerate(stats):
                    cutoff2stats[c][i] = [np.mean( s ), np.std( s )]
                    #cutoff2stats[c][i] = float(s)/self.denom

            #Print to output file:
            outfile = os.path.join(self.outdir, "%d.txt" %self.samsize)
            f = open(outfile, 'w')
            f.write("#%Cutoff\tClones1\tClones2\toverlap1\toverlap2\t%1overlap2\t%2overlap1\t%reads1overlap2\t%reads2overlap1\t")
            f.write("StdClones1\tStdClones2\tStdOverlap1\tStdOverlap2\tStd%1overlap2\tStd%2overlap1\tStd%reads1overlap2\tStd%reads2overlap1\n")
            for c in sorted( cutoff2stats.keys() ):
                stats = cutoff2stats[c]
                meanStats = [ "%.4f" % s[0] for s in stats ]
                stdStats = [ "%.4f" % s[1] for s in stats ]
                f.write( "%.3f\t%s\t%s\n" %(c, "\t".join(meanStats), "\t".join(stdStats)) )
            f.close()

#=============== CLEANUP =================
class CleanupSummary ( Target ):
    def __init__(self, indir):
        Target.__init__(self, time = 0.00025)
        self.indir = indir

    def run(self):
        system("rm -fR %s" %self.indir)

####################### UTILITIES FUNCTIONS ##########################

#============== calculate overlap =================
def getOverlap( seqs1, seqs2, cutoffs ):
    #Initialize stats:
    stats1 = {"oclones":[], "oreads":[]}
    stats2 = {"oclones":[], "oreads":[]}
    for i in xrange( len(cutoffs) ):
        for k in stats1:
            stats1[k].append(0)
            stats2[k].append(0)

    #get number of clones that pass the cutoff:
    reads1, clones1, total1 = getNumClones(seqs1, cutoffs)
    reads2, clones2, total2 = getNumClones(seqs2, cutoffs)

    if total1 == 0 or total2 == 0:
        return
    #get overlap:
    for s1 in seqs1:
        i2 = seqs2.search(s1)
        if i2 == -1:#not found in repertoire 2
            continue
        s2 = seqs2[i2]
        for i, cutoff in enumerate(cutoffs):
            #if s1.freq >= cutoff and s2.freq >= cutoff:
            if s1.freq >= cutoff:
                stats1["oclones"][i] += 1
                stats1["oreads"][i] += s1.count
            if s2.freq >= cutoff:
                stats2["oclones"][i] += 1
                stats2["oreads"][i] += s2.count
    return reads1, reads2, clones1, clones2, stats1, stats2

def getNumClones(seqs, cutoffs):
    reads = [ 0 for c in cutoffs ]
    clones = [ 0 for c in cutoffs ]
    total = sum([ s.count for s in seqs ])
    for s in seqs:
        s.setFreq(total)
        for i,c in enumerate(cutoffs):
            if s.freq >= c:
                clones[i] += 1
                reads[i] += s.count
    return reads, clones, total

#================ GENERATE PSEUDO REPERTOIRE =========
def sampling( seqs, size ):
    samseqs = Seqs()
    indexList = []
    for i, s in enumerate(seqs):
        indexList.extend( [ i for j in xrange( int(s.count)) ] )

    for j in xrange(size):
        i = random.randint(0, len(indexList) -1) #randomly pick one sequence from the repertoire
        s = Aaseq( seqs[indexList[i]].aa )
        #s = copy.copy( seqs[indexList[i]] )
        sindex = samseqs.search(s) #Search for this sequence in the current set of picked sequences

        #print [ seq.aa for seq in samseqs]
        #print "%s" %s.aa
        #print sindex

        if sindex >=0:
            samseqs[sindex].updateCount(1)
        else:#sequence hasn't picked yet
            s.setCount(1)
            samseqs.add(s)
    return samseqs

#get the repertoire using topFreqs if available, otherwise assume lognormal distribution
#IMCOMPLETE
def getRepLognorm( seqs, size, mean, std ):
    #mean = #mean of log(clonesize)
    #std = #standard deviation of log(clonesize)

    #Generate repertoire:
    topindices = []
    sizeToFill = size

    if topFreqs:
        for i, f in enumerate(topFreqs):
            index = random.randint(0, len(seqs) -1)
            while (index in topindices):
                index = random.randint(0, len(seqs) -1)
            topindices.append(index)
            count = int(f*size)
            seqs[index].setCount( count )
        sizeToFill = size - int( sum(topFreqs)*size )

    for i in xrange(sizeToFill):
        index = random.randint(0, len(seqs) -1)
        while (index in topindices):
            index = random.randint(0, len(seqs) -1)

#get the repertoire using topFreqs if available, otherwise assume uniform distribution
def getRep( seqs, size, topFreqs ):
    #Generate repertoire:
    topindices = []
    sizeToFill = size
    if topFreqs:
        for i, f in enumerate(topFreqs):
            index = random.randint(0, len(seqs) -1)
            while (index in topindices):
                index = random.randint(0, len(seqs) -1)
            topindices.append(index)
            count = int(f*size)
            seqs[index].setCount( count )
        sizeToFill = size - int( sum(topFreqs)*size )

    for i in xrange(sizeToFill):
        index = random.randint(0, len(seqs) -1)
        while (index in topindices):
            index = random.randint(0, len(seqs) -1)
        seqs[index].updateCount(1)

def getUniqSeqRep(size, aaUsage, aa2codons):
    """Generate "size" number of unique nucleotide sequences with length (len(aaUsage))
    Using the amino-acid Usage from aaUsage, and codon usage from aa2codons
    """
    seqs = Seqs() #list of sequences
    while len(seqs) < size:
        aa = getAaSeq(aaUsage)
        nuc = aa2nuc(aa, aa2codons)
        seq = Seq(nuc, aa)
        seqs.add(seq)
    return seqs

def readStats( file ):
    cutoff2stats = {} #key = cutoff, val = [clones1, clones2, overlap, %1oc2, %2oc1, %1or2, %2or1]
    f = open(file, 'r')
    for line in f:
        if line[0] == "#":
            continue
        items = line.split("\t")
        if len(items) < 9:
            raise FileFormatError("Wrong file format, 9 fields are required, see %d. File %s\n" %(len(items), file))
        cutoff = float(items[0])
        #stats = [int(items[1]), int(items[2]), int(items[3]), int(items[4]), float(items[5]), float(items[6]), float(items[7]), float(items[8])]
        stats = [ [float(items[i])] for i in xrange(1,9) ]
        #stats = [float(items[1]), float(items[2]), float(items[3]), float(items[4]), float(items[5]), float(items[6]), float(items[7]), float(items[8])]
        cutoff2stats[cutoff] = stats
    f.close()
    return cutoff2stats

def getPc(count, total):
    if total == 0:
        return 0
    return 100.0*count/total

def getAaSeq(aaUsage):
    aa = ""
    seqlen = len(aaUsage)
    for i in xrange(seqlen):
        letters = aaUsage[i]
        j = random.randint(0, len(letters) -1)
        aa += letters[j]
    return aa

def aa2nuc(aaseq, aa2codons):
    nucseq = ''
    for aa in aaseq:
        codons = aa2codons[aa]
        i = random.randint( 0, len(codons) -1 )
        nucseq += codons[i]
    return nucseq

#================ ERROR CLASSES ============
class FileFormatError(Exception):
    pass

class TopFreqError(Exception):
    pass

class NonExistFileError(Exception):
    pass

#================ READ INPUT FILES ================
def readCodonUsage(file):
    """
    """
    aa2codons = {} #key = aa, val = {codon:freq}
    aa2codon2freq = {}
    f = open(file, 'r')
    for line in f:
        line.strip()
        if line == "" or line[0] == "#":
            continue
        items = line.split()
        if len(items) %4 > 0:
            raise FileFormatError("Wrong codonUsage format\n")
        for i in xrange(0, len(items), 4):
            codon = items[i]
            aa = items[i+1]
            freq = float( items[i+2] )
            if aa not in aa2codon2freq:
                aa2codon2freq[aa] = {codon: freq}
            else:
                aa2codon2freq[aa][codon] = freq
    f.close()

    #Normalize
    normlen = 100
    for aa in aa2codon2freq:
        aa2codons[aa] = []
        for codon, freq in aa2codon2freq[aa].iteritems():
            for i in xrange( int(freq*normlen) ):
                aa2codons[aa].append(codon)

    return aa2codons

def readAaUsage(file):
    '''
    Return: len2aaUsage = {} #key = length, val= position-array, each item represent a.a usage of that position.
    Example:
    9 1 C 10
    9 2 A 7
    9 2 R 1
    9 2 S 3
    9 2 V 1
    len2aaUsage = {9:["CCCCCCCCCC", "AAAAAAARSSSV"]}
    '''
    len2aaUsage = {} #key = length, val= position-array, each item represent a.a usage of that position.
    len2listLetter2freq = {}
    f = open(file, 'r')
    for line in f:
        line = line.strip()
        if line == "" or line[0] == "#":
            continue
        items = line.split()
        if len(items) != 4:
            raise FileFormatError("Wrong aaUsage file format. Required 4 fields, have %d. Line: %s\n" %(len(items), line) )
        l = int( items[0] )
        pos = int( items[1] ) - 1
        letter = items[2]
        count = int( items[3] )

        if l < 0 or pos < 0 or count < 0:
            raise ValueError("readAaUsage: length, position, count must >=0")
        if pos >= l:
            raise FileFormatError("readAaUsage: Column 2 value (Position) must <= col 1 value (length)\n")
        if l not in len2listLetter2freq:
            len2listLetter2freq[ l ] = [{} for i in xrange(l)]
        len2listLetter2freq[ l ][ pos ][ letter ] = count
    f.close()

    #Make sure that for each length, usage for all positions are provided from input file:
    for l in len2listLetter2freq:
        for i in xrange(l):
            if len( len2listLetter2freq[l][i].keys() ) == 0:
                raise FileFormatError("readAaUsage requires input file to have AA usage for all positions of each length. Could not find usage for length %d, position %d." %(l, i + 1))

    #Normalize counts:
    normLen = 10000 #equivalent the number of digits of the frequencies got taken into account
    for l in len2listLetter2freq:
        posList = len2listLetter2freq[l]
        len2aaUsage[ l ] = []
        for i, letter2count in enumerate(posList):
            total = sum([ letter2count[letter] for letter in letter2count ])
            if total > 0:
                for letter, count in letter2count.iteritems():
                    letter2count[letter] = normLen*count/total
            posLetters = "".join([ letter*count for letter, count in letter2count.iteritems()])
            len2aaUsage[l].append(posLetters)

    return len2aaUsage

def getUniAaUsage(lens):
    aas = "ARNDCEQGHILKFPSTWYV"
    len2aaUsage = {}
    for l in lens:
        if not isinstance(l, int):
            raise ValueError("getUniAaUsage, input lengths contain non-integer items\n")
        len2aaUsage[l] = []
        for i in xrange(l):#each position
            len2aaUsage[l].append(aas)
    return len2aaUsage

def readCdr3LenDist(file):
    '''Get length distribution of CDR3 amino-acid sequences. Format of input file: <length>\t<count>
    Return a dictionary len2freq, where key=length and val = freq of that length
    '''
    len2freq = {} #key = length, val=frequency
    f = open(file, 'r')
    total = 0
    for line in f:
        line = line.strip()
        if line == "" or line[0] == "#":
            continue
        #items = line.split("\t")
        items = line.split()
        if len(items) < 2:
            raise FileFormatError("Wrong cdr3LenDist file format. Require two fields: <length>\\t<count>\n.%s\n" %line)
        if items[0] == "TOTAL":
            continue
        try:
            l = int( (items[0].split("."))[0] )
            c = int( items[1] )
        except ValueError:
            raise ValueError("First and Second columns in cdr3LenDist file %s must be integer\n" %file)

        total += c
        if l not in len2freq:
            len2freq[l] = c
        else:
            len2freq[l] += c
    #Normalize counts to frequencies:
    if total > 0:
        for l, c in len2freq.iteritems():
            len2freq[l] = float(c)/total
    f.close()
    return len2freq

def getUniCdr3LenDist():
    len2freq = {}
    minlen = 9
    maxlen = 23
    r = maxlen + 1 - minlen
    for l in xrange(minlen, maxlen + 1):
        len2freq[l] = 1.0/r
    return len2freq

def getTopFreqs( topFreqs, totalClones, totalSeqs ):
    f = open(topFreqs, "r")
    freqs = []
    for line in f:
        line = line.strip()
        freq = float(line)
        if not isinstance(freq, float):
            raise ValueError("Wrong format topFreqs file. Requires floats.\n")
        if freq <= 0:
            raise ValueError("Frequencies in topFreqs file (%s) must be positive. Got %f\n" %(topFreqs, freq))

        freqs.append( freq )
    #Make sure that sum of freqs <= 1 - (totalClones - len(freqs))/totalSeqs
    if sum(freqs) > 1 - float((totalClones - len(freqs)))/totalSeqs:
        raise TopFreqError("Frequencies of top clones add up to be larger than permitted")
    topFreqs = freqs
    f.close()
    return topFreqs

################### OPTIONS #########################
def checkOptions( parser, options, args ):
    #Check output directory
    if not os.path.exists( options.outdir ):
        system( "mkdir %s\n" %(options.outdir) )

    #Number of samples:
    if not isinstance(options.numSamples, int):
        parser.error("Number of samples must be in integer\n")
    if options.numSamples < 2:
        parser.error("Number of samples must > 2. %d was given\n" %options.numSamples)

    #Make sure the number of totalClones < totalSeqs:
    if options.totalClones > options.totalSeqs:
        raise ValueError("Total number of clones must be smaller than total number of sequences. %d clones and %d sequences were given" %(options.totalClones, options.totalSeqs) )

    #sampling sizes:
    options.samSize = [long(s) for s in options.samSize.split(",")]

    #Top Freqs:
    if options.topFreqs:
        if not os.path.exists( options.topFreqs ):
            parser.error("TopFreqs file %s does not exists\n" %options.topFreqs)
        options.topFreqs = getTopFreqs( options.topFreqs, options.totalClones, options.totalSeqs )

    #Cutoffs:
    options.cutoffs = [ float(c) for c in options.cutoffs.split(",") ]

    #CDR3 length distribution
    if options.cdr3LenDist:
        if not os.path.exists(options.cdr3LenDist):
            parser.error("cdr3LenDist file %s does not exist\n" %options.cdr3LenDist)
        options.len2freq = readCdr3LenDist(options.cdr3LenDist)
    else:
        options.len2freq = getUniCdr3LenDist()

    if options.aaUsage:
        if not os.path.exists(options.aaUsage):
            parser.error("aaUsage file %s does not exist\n" %options.aaUsage)
        options.len2aaUsage = readAaUsage(options.aaUsage)
    else:
        options.len2aaUsage = getUniAaUsage( range(9, 24) )

    if not options.codonUsage:
        parser.error("Required argument codonUsage. None was given\n")
    if not os.path.exists(options.codonUsage):
        parser.error("codonUsage file %s does not exist\n" %options.codonUsage)
    options.aa2codons = readCodonUsage(options.codonUsage)

def addOptions( parser ):
    #parser.add_option("-i", "--indir", dest = 'indir', help="Required argument. Input directory that contains input fastq files (each represents a sample)\n")
    #parser.add_option("-n", "--popCdr3", dest='popCdr3', type="long", default=5*(10**11), help="Estimated number of unique CDR3 nucleotide sequences in the population. Default=5x10^11")
    parser.add_option("-o", "--outdir", dest = 'outdir', default=".", help="Output directory. Default: current directory\n")
    parser.add_option("-N", "--numSim", dest="numSim", type="int", default=100, help="Number of repertoire simulations. Default = 100")
    parser.add_option("--numSamplingPerSim", dest="numSamplingPerSim", type="int", default=100, help="Number of samplings per repertoire simulation. Default = %default")
    parser.add_option("-m", "--totalClones", dest="totalClones", type="long", default=3*(10**6), help="Estimated number of unique CDR3 nucleotide sequences in one individual. Default=3x10^6")
    parser.add_option("-M", "--totalSeqs", dest="totalSeqs", type="long", default=6*(10**7), help="Estimated number of CDR3 nucleotide sequences in one individual. Default=6x10^7 (number of CD8+ T cells per person)")
    parser.add_option("-S", "--samplingSize", dest="samSize", default="50000,100000,150000,200000", help="Comma separated string of sampling sizes. Default=%default")
    #parser.add_option("-a", "--aaUsage", dest="aaUsage", help="Amino acid usage of the CDR3 sequences. Format:<aa-cdr3-length>\t<position>\t<letter>\t<count>. If not specified, use uniform distribution")
    parser.add_option("-a", "--aaUsage", dest="aaUsage", help="Amino acid usage of the CDR3 sequences. Format:<aa-cdr3-length>\t<position>\t<letter>\t<count>. If not specified, use lognormal distribution")
    parser.add_option("-l", "--cdr3LenDist", dest="cdr3LenDist", help="CDR3 length distribution. Format: <aa-cdr3-length>\t<count>")
    parser.add_option("-c", "--codonUsage", dest="codonUsage", default="~/codonUsage.txt", help="Codon usage table")
    #parser.add_option("-f", "--topFreqs", dest="topFreqs", help="Frequencies of x top clones. Sum of these freqs must <= (M - (m-x))/M. If not specified, use uniform distribution")
    parser.add_option("-f", "--topFreqs", dest="topFreqs", help="Frequencies of x top clones. Sum of these freqs must <= (M - (m-x))/M. If not specified, use lognormal distribution")
    parser.add_option("-s", "--numSamples", dest="numSamples", type="int", default=2, help="Number of samples (individuals/repertoires) to compare. Default=2.")
    parser.add_option("-p", "--cutoffs", dest="cutoffs", default="0,0.001,0.01,0.05,0.1,0.5,1,5,10", help="Comma separated clonesize percentage cutoffs. Default='0,0.001,0.01,0.05,0.1,0.5,1,5'")
    parser.add_option("--maxSimPerRun", dest="maxSimPerRun", type='int', default=100, help="Maximum number of simulations per batch. Default=100")

################# MAIN ##############################
def main():
    parser = OptionParser()
    Stack.addJobTreeOptions( parser )

    addOptions( parser )

    options, args = parser.parse_args()
    checkOptions( parser, options, args )

    i = Stack( Setup(options, 0) ).startJobTree( options )
    if i:
        raise RuntimeError("The jobtree contains %d failed jobs.\n" %i)

if __name__ == "__main__":
    #from tcrRepSim import *
    from immunoseq.src.tcrRepSim import *
    main()


## Requirements

### Packages
import os
import json
import gzip
import numpy as np

from bandits_to_rank.environment import Environment_PBM, PositionsRanking
from bandits_to_rank.opponents import greedy
from bandits_to_rank.opponents.pbm_pie import PBM_PIE_Greedy_SVD, PBM_PIE_semi_oracle, PBM_PIE_Greedy_MLE
from bandits_to_rank.opponents.pbm_ucb import PBM_UCB_Greedy_SVD, PBM_UCB_semi_oracle, PBM_UCB_Greedy_MLE
from bandits_to_rank.opponents.pbm_ts import PBM_TS_Greedy_SVD, PBM_TS_semi_oracle, PBM_TS_Greedy_MLE
from bandits_to_rank.opponents.bc_mpts import BC_MPTS_Greedy_SVD, BC_MPTS_semi_oracle,BC_MPTS_Greedy_MLE
#from bandits_to_rank.opponents.pmed import PMED   # loaded only before usage to load tensorflow library only when required
from bandits_to_rank.opponents.top_rank import TOP_RANK
from bandits_to_rank.opponents.grab import GRAB
from bandits_to_rank.opponents.f_grab import sGRAB
from bandits_to_rank.opponents.combucb import CombUCB1, KL_CombUCB1
from bandits_to_rank.opponents.pb_mhb import *
from bandits_to_rank.referee import Referee

# set.seed(123)

# Path to bandits-to-rank module
import bandits_to_rank

packagedir = os.path.dirname(bandits_to_rank.__path__[0])

class NdArrayEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        if isinstance(obj, (np.int_, np.intc, np.intp, np.int8,
                            np.int16, np.int32, np.int64, np.uint8,
                            np.uint16, np.uint32, np.uint64)):
            return int(obj)
        # Let the base class default method raise the TypeError
        return json.JSONEncoder.default(self, obj)

def record_zip(filename, dico):
    print(type(dico))
    print('file', filename)
    json_str = json.dumps(dico, cls=NdArrayEncoder)
    json_bytes = json_str.encode('utf-8')
    with gzip.GzipFile(filename, 'w') as fout:
        fout.write(json_bytes)
    return 'done'

class Parameters():
    """ Parameters used for the experiment

    # Environement
        env
        env_name        str used for name of files
        logs_env_name   (only for merge)

    # Player
        player
        player_name

    # Rules
        rules_name
        referee

    # Sub-experiment
        first_game      (only for play)
        end_game        (only for play)
        input_path      (only for merge)
        output_path
        force           (only for play)
    """

    def __init__(self):
        self.env = Environment_PBM([1], [1], label="fake")
        self.positions_ranking = PositionsRanking.SHUFFLE_EXCEPT_FIRST  # default: shuffle kappas before each game
        self.nb_relevant_positions = None  # default: compute reward at each position
        self.rng = np.random.default_rng()

    #########" PBM_Setting
    def set_positions_ranking(self, positions_ranking):
        self.positions_ranking = positions_ranking

        # tag for file names and logs
        if positions_ranking == PositionsRanking.FIXED:
            raise ValueError('fixed ranking of positions should be set by the player')
        elif positions_ranking == PositionsRanking.DECREASING:
            # TODO: better naming for PBM '__decreasing_kappa'
            # TODO: better naming for CM '__std_order_on_views'
            tag = '__sorted_kappa' if type(self.env) == Environment_PBM else ''
        elif positions_ranking == PositionsRanking.SHUFFLE:
            # TODO: better naming for CM '__random_order_on_views'
            tag = '__shuffled_kappa' if type(self.env) == Environment_PBM else '_order_view_shuffle'
        elif positions_ranking == PositionsRanking.SHUFFLE_EXCEPT_FIRST:
            # TODO: better naming for PBM __shuffled_kappa_except_first
            tag = '' if type(self.env) == Environment_PBM else '__random_order_on_views_except_first'
        elif positions_ranking == PositionsRanking.INCREASING:
            tag = ('__increasing_kappa' if type(self.env) == Environment_PBM else '__reverse_order_on_views')
        elif positions_ranking == PositionsRanking.INCREASING_EXCEPT_FIRST:
            tag = ('__increasing_kappa_except_first' if type(self.env) == Environment_PBM
                   else '__reverse_order_on_views_except_first')
        else:
            raise ValueError(f'unhandled ranking on positions: {positions_ranking}')
        self.env_name += tag
        self.logs_env_name += tag
        self.env.label += tag

    def set_env_KDD_all(self):
        """!!! only to merge logs from several queries of KDD data !!!"""
        self.env_name = f'KDD_all'
        self.logs_env_name = f'KDD_[0-9]*_query'

    def set_env_KDD(self, query):
        # load KDD data
        # todo: to put in bandits_to_rank.data
        with open(packagedir + '/data/param_KDD.txt', 'r') as file:
            dict_theta_query = json.load(file)
        query_name, query_params = list(dict_theta_query.items())[query]

        # set environement
        self.env_name = f'KDD_{query}_query'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(query_params['thetas'], query_params['kappas']
                                   , label='%s (%d for us)' % (query_name, query))

    def set_env_Yandex_all(self):
        """!!! only to merge logs from several queries of Yandex data !!!"""
        self.env_name = f'Yandex_all'
        self.logs_env_name = f'Yandex_[0-9]*_query'

    def set_env_Yandex(self, query):
        # load Yandex data
        # todo: to put in bandits_to_rank.data
        with open(packagedir + '/data/param_Yandex.txt', 'r') as file:
            dict_theta_query = json.load(file)
        query_name, query_params = list(dict_theta_query.items())[query]

        # reduce to 10 products, 5 positions
        thetas = np.sort(query_params['thetas'])[:-11:-1]
        kappas = np.sort(query_params['kappas'])[:-6:-1]

        # set environement
        self.env_name = f'Yandex_{query}_query'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas
                                   , label='%s (%d for us)' % (query_name, query))

    def set_env_Yandex_equi_all(self, K):
        """!!! only to merge logs from several queries of Yandex data !!!"""
        self.env_name = f'Yandex_equi_{K}_K_all'
        self.logs_env_name = f'Yandex_equi_{K}_K__[0-9]*_query'

    def set_env_Yandex_equi(self, query, K):
        # load Yandex data
        with open(packagedir + '/data/param_Yandex.txt', 'r') as file:
            dict_theta_query = json.load(file)
        query_name, query_params = list(dict_theta_query.items())[query]

        # reduce to 10 products, 10 positions
        index_max = 1 + K
        thetas = np.sort(query_params['thetas'])[:-index_max:-1]
        kappas = np.sort(query_params['kappas'])[:-index_max:-1]

        # set environement
        self.env_name = f'Yandex_equi_{K}_K__{query}_query'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label=f'Yandex equi K={K} {query} ({query_name} for Yandex)')

    def set_env_test(self):
        """Purely simulated environment with standard click's probabilities"""
        kappas = [1, 0.6, 0.3]
        thetas = [0.1, 0.5, 0.1, 0.6, 0.1, 0.4, 0.1, 0.1, 0.1, 0.1]
        self.env_name = f'purely_simulated__test'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label="purely simulated, test")

    def set_env_std(self):
        """Purely simulated environment with standard click's probabilities"""
        kappas = [1, 0.75, 0.6, 0.3, 0.1]
        thetas = [0.3, 0.2, 0.15, 0.15, 0.15, 0.10, 0.05, 0.05, 0.01, 0.01]
        self.env_name = 'purely_simulated__std'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label="purely simulated, std")

    def set_env_small(self):
        """Purely simulated environment with click's probabilities close to 0"""
        kappas = [1, 0.75, 0.6, 0.3, 0.1]
        thetas = [0.15, 0.1, 0.1, 0.05, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01]
        self.env_name = 'purely_simulated__small'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label="purely simulated, small")

    def set_env_big(self):
        """Purely simulated environment with click's probabilities close to 1"""
        kappas = [1, 0.75, 0.6, 0.3, 0.1]
        thetas = [0.99, 0.95, 0.9, 0.85, 0.8, 0.75, 0.75, 0.75, 0.75, 0.75]
        self.env_name = 'purely_simulated__big'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label="purely simulated, big")

    def set_env_extra_small(self):
        """Purely simulated environment with click's probabilities close to 0"""
        kappas = [1, 0.75, 0.6, 0.3, 0.1]
        thetas = [0.10, 0.05, 0.01, 0.005, 0.001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001]
        self.env_name = 'purely_simulated__xsmall'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label="purely simulated, extra small")

    def set_env_xx_small(self):
        """Purely simulated environment with click's probabilities close to 0"""
        kappas = [1, 0.75, 0.6, 0.3, 0.1]
        thetas = [1e-3, 5e-4, 1e-4, 5e-5, 1e-5, 1e-6, 1e-6, 1e-6, 1e-6, 1e-6]
        self.env_name = 'purely_simulated__xxsmall'
        self.logs_env_name = self.env_name
        self.env = Environment_PBM(thetas, kappas, label="purely simulated, xx small")

    def set_env_simul(self, label):
        if label == "std":
            self.set_env_std()
        elif label == "big":
            self.set_env_big()
        elif label == "small":
            self.set_env_small()
        elif label == "xsmall":
            self.set_env_extra_small()
        elif label == "xxsmall":
            self.set_env_xx_small()
        else:
            raise ValueError("unknown label of environment")

    def set_rules(self, nb_trials, nb_records=1000):
        # Check inputs
        if nb_records > nb_trials:
            nb_records = -1

        self.rules_name = f'games_{nb_trials}_nb_trials_{nb_records}_record_length'
        self.referee = Referee(self.env, nb_trials, all_time_record=False, len_record_short=nb_records)

    def set_player_eGreedy(self, c, update=100, noSVD=False):
        nb_prop, nb_place = self.env.get_setting()
        if noSVD:
            self.player_name = f'Bandit_EGreedy_EM_{c}_c_{update}_update'
            self.player = greedy.greedy_EGreedy_EM(c, nb_prop, nb_place, update)
        else:
            self.player_name = f'Bandit_EGreedy_SVD_{c}_c_{update}_update'
            self.player = greedy.greedy_EGreedy(c, nb_prop, nb_place, update)

    def set_player_PBM_TS(self, type="oracle"):
        nb_prop, nb_place = self.env.get_setting()
        if type =="oracle":
            self.player_name = 'Bandit_PBM-TS_oracle'
            self.player = PBM_TS_semi_oracle(nb_prop, nb_place, discount_factor=self.env.kappas, count_update=1)
            self.positions_ranking = PositionsRanking.FIXED
        elif type =="greedyMLE":
            self.player_name = 'Bandit_PBM_TS_greedy_MLE'
            self.player = PBM_TS_Greedy_MLE(nb_prop, nb_place, count_update=1)
        elif type =="greedySVD":
            self.player_name = 'Bandit_PBM_TS_greedy_SVD'
            self.player = PBM_TS_Greedy_SVD(nb_prop, nb_place, count_update=1)
        else:
            self.player_name = 'Bandit_PBM-TS_greedy_SVD'
            self.player = PBM_TS_Greedy_SVD(nb_prop, nb_place, count_update=1)

    def set_player_PBM_PIE(self, epsilon, T, type ="oracle"):
        nb_prop, nb_place = self.env.get_setting()
        if type =="oracle":
            self.player_name = f'Bandit_PBM-PIE_oracle_{epsilon}_epsilon'
            self.player = PBM_PIE_semi_oracle(nb_prop, epsilon, T, nb_place, discount_factor=self.env.kappas, count_update=1)
            self.positions_ranking = PositionsRanking.FIXED
        elif type =="greedyMLE":
            self.player_name = 'Bandit_PBM_PIE_greedy_MLE'
            self.player = PBM_PIE_Greedy_MLE(nb_prop, epsilon, nb_place, count_update=1)
        elif type =="greedySVD":
            self.player_name = 'Bandit_PBM_PIE_greedy_SVD'
            self.player = PBM_PIE_Greedy_SVD(nb_prop, epsilon, nb_place, count_update=1)
        else:
            self.player_name = f'Bandit_PBM-PIE_greedy_SVD_{epsilon}_epsilon'
            self.player = PBM_PIE_Greedy_SVD(nb_prop, epsilon, T, nb_place, count_update=1)

    def set_player_PBM_UCB(self, epsilon, type ="oracle"):
        nb_prop, nb_place = self.env.get_setting()
        if type =="oracle":
            self.player_name = f'Bandit_PBM_UCB_oracle_{epsilon}_epsilon'
            self.player = PBM_UCB_semi_oracle(nb_prop, epsilon, nb_place, discount_factor=self.env.kappas, count_update=1)
            self.positions_ranking = PositionsRanking.FIXED
        elif type =="greedyMLE":
            self.player_name = 'Bandit_PBM_UCB_greedy_MLE'
            self.player = PBM_UCB_Greedy_MLE(nb_prop, epsilon, nb_place, count_update=1)
        elif type =="greedySVD":
            self.player_name = 'Bandit_PBM_UCB_greedy_SVD'
            self.player = PBM_UCB_Greedy_SVD(nb_prop, epsilon, nb_place, count_update=1)
        else:
            self.player_name = f'Bandit_PBM_UCB_greedy_SVD_{epsilon}_epsilon'
            self.player = PBM_UCB_Greedy_SVD(nb_prop, epsilon, nb_place, count_update=1)

    def set_player_BC_MPTS(self, type ="oracle"):
        nb_prop, nb_place = self.env.get_setting()
        if type =="oracle":
            self.player_name = 'Bandit_BC-MPTS_oracle'
            self.player = BC_MPTS_semi_oracle(nb_prop, nb_place, self.env.kappas)
            self.positions_ranking = PositionsRanking.FIXED
        elif type =="greedyMLE":
            self.player_name = 'Bandit_BC-MPTS_greedy_MLE'
            self.player = BC_MPTS_Greedy_MLE(nb_prop, nb_place, count_update=1)
        elif type =="greedySVD":
            self.player_name = 'Bandit_BC-MPTS_greedy_SVD'
            self.player = BC_MPTS_Greedy_SVD(nb_prop, nb_place, count_update=1)
        else:
            self.player_name = 'Bandit_BC-MPTS_greedy_SVD'
            self.player = BC_MPTS_Greedy_SVD(nb_prop, nb_place, count_update=1)

    def set_player_PMED(self, alpha, gap_MLE, gap_q, run=True):
        nb_prop, nb_place = self.env.get_setting()

        self.player_name = f'Bandit_PMED_{alpha}_alpha_{gap_MLE}_gap_MLE_{gap_q}_gap_q'

        if run:
            from bandits_to_rank.opponents.pmed import PMED
            self.player = PMED(nb_prop, nb_place, alpha, gap_MLE, gap_q)

    def set_player_CombUCB1(self, exploration_factor=2.):
        nb_prop, nb_place = self.env.get_setting()

        self.player_name = f'Bandit_CombUCB1_{exploration_factor}_exploration'
        self.player = CombUCB1(nb_arms=nb_prop, nb_positions=nb_place, exploration_factor=exploration_factor)

    def set_player_KL_COMB(self, horizon):
        nb_prop, nb_place = self.env.get_setting()

        self.player_name = f'Bandit_KL-COMB_{horizon}_horizon'
        self.player = KL_CombUCB1(nb_arms=nb_prop, nb_positions=nb_place, horizon=horizon)

    def set_player_PB_MHB(self, nb_steps, random_start=False):
        nb_prop, nb_place = self.env.get_setting()
        if random_start:
            self.player_name = f'Bandit_PB-MHB_random_start_{nb_steps}_step_{self.proposal_name}_proposal'
            self.player = PB_MHB(nb_prop, nb_place, proposal_method=self.proposal, step=nb_steps,
                                 part_followed=False)
        else:
            self.player_name = f'Bandit_PB-MHB_warm-up_start_{nb_steps}_step_{self.proposal_name}_proposal'
            self.player = PB_MHB(nb_prop, nb_place, proposal_method=self.proposal, step=nb_steps,
                                 part_followed=True)

    def set_player_TopRank(self, T, horizon_time_known=True, doubling_trick=False, oracle=False):
        nb_prop, nb_place = self.env.get_setting()
        if oracle:
            self.player_name = f'Bandit_TopRank_oracle_{T}_delta_{"TimeHorizonKnown" if horizon_time_known else ""}_{"doubling_trick" if doubling_trick else ""}'
            self.player = TOP_RANK(nb_arms=nb_prop,
                                   T=T, horizon_time_known=horizon_time_known,doubling_trick_active=doubling_trick,
                                   discount_factor=self.env.kappas)
            self.positions_ranking = PositionsRanking.FIXED
        else:
            self.player_name = f'Bandit_TopRank_{T}_delta_{"TimeHorizonKnown" if horizon_time_known else ""}_{"doubling_trick" if doubling_trick else ""}'
            self.player = TOP_RANK(nb_arms=nb_prop,
                                   T=T, horizon_time_known=horizon_time_known, doubling_trick_active=doubling_trick,
                                   discount_factor=np.arange(nb_place - 1, -1, -1))
        """
            self.player_name = f'Bandit_TopRank_greedy_{T}_delta_{"TimeHorizonKnown" if horizon_time_known else ""}_{"doubling_trick" if doubling_trick else ""}'
            self.player = TOP_RANK(nb_arms=nb_prop,
                                   T=T, horizon_time_known=horizon_time_known,doubling_trick_active=doubling_trick,
                                   nb_positions=nb_place, lag=1)
        """

    def set_player_SGRAB(self, T, gamma, forced_initiation):
        nb_prop, nb_place = self.env.get_setting()
        if gamma == 0:
            gamma_use = nb_prop * nb_place
        else:
            gamma_use = gamma
        self.player_name = f'Bandit_SGRAB_{T}_T_{gamma_use}_gamma{"_forced" if forced_initiation else ""}'
        self.player = sGRAB(nb_arms=nb_prop, nb_positions=nb_place, T=T)

    def set_player_GRAB(self, T, gamma, forced_initiation):
        nb_prop, nb_place = self.env.get_setting()
        if gamma == 0:
            gamma_use = nb_prop - nb_place
        else:
            gamma_use = gamma
        self.player_name = f'Bandit_GRAB_{T}_T_{gamma_use}_gamma{"_forced" if forced_initiation else ""}'
        self.player = GRAB(nb_arms=nb_prop, nb_positions=nb_place, T=T, gamma=gamma_use,
                                forced_initiation=forced_initiation)

    def set_proposal_TGRW(self, c, vari_sigma=True):
        self.proposal_name = f'TGRW_{c}_c{"_vari_sigma" if vari_sigma else ""}'
        self.proposal = propos_trunk_GRW(c, vari_sigma)

    def set_proposal_LGRW(self, c, vari_sigma=True):
        self.proposal_name = f'LGRW_{c}_c{"_vari_sigma" if vari_sigma else ""}'
        self.proposal = propos_logit_RW(c, vari_sigma)

    def set_proposal_RR(self, c, str_proposal_possible, vari_sigma=True):
        list_proposal_possible = list(str_proposal_possible.split("-"))
        self.proposal_name = f'RR_{c}_c_{len(list_proposal_possible)}_proposals'
        self.proposal = propos_Round_Robin(c, vari_sigma, list_proposal_possible)

    def set_proposal_MaxPos(self):
        self.proposal_name = f'MaxPos'
        self.proposal = propos_max_position()

    def set_proposal_PseudoView(self):
        self.proposal_name = f'PseudoView'
        self.proposal = propos_pseudo_view()

    def set_exp(self, first_game=-1, nb_games=-1, nb_checkpoints=10, input_path=None, output_path=None, force=True):
        self.first_game = first_game
        self.end_game = first_game + nb_games
        self.nb_checkpoints = nb_checkpoints
        self.input_path = input_path if input_path is not None else output_path
        self.output_path = output_path if output_path is not None else input_path
        self.force = force


"""
Code for top level interface.

This code is added to the main package level in __init__.py
"""
import numpy as np
import abcgan.constants as const
import abcgan.transforms as trans
from abcgan import persist
from abcgan.mask import mask_altitude
from abcgan.transforms import compute_valid, compute_valid_hfp
import torch
import h5py
from tqdm import tqdm
from warnings import warn

def generate(drivers, measurements=None,
             driver_names=const.driver_names, n_alt=const.max_alt,
             bv_model='bv_gan', hfp_model='hfp_gan', bv_type='radar',
             generate_hfps=False, return_z_scale=False, cuda_index=None,
             verbose=1):
    """
    Generate background variable profiles and HFP waves
    consistent with the historical distribution.

    Parameters
    -------------
    drivers: np.ndarray
        n_samples x n_drivers input list of driving parameters (not z-scaled).
    driver_names: list
        list of names of driving parameters
    measurements: np.ndarray, optional
        n_samples x n_alt_in x n_meas input list of altitude measurements,
        n_alt_in should be less than n_alt. These represent fixed
        measurements for the lowest altitudes to condition on.
        Usually left as default (None)
    n_alt: int, optional
        number of altitude measurements to draw, defaults to max_alt
    return_z_scale: bool, optional
        set to have the function return z scaled feature data
    bv_model: str, optional
        name of bv GAN to use
    bv_type: str. optional
        name of the type of background variables to use (lidar or radar)
    hfp_model: str, optional
        name of hfp GAN to use
    generate_hfps: bool, optional
        Set to generate hfps and return generated hfps
    cuda_index: int, optional
        GPU index to use when generating BVs and HFPs
    verbose: bool, optional
        set to show loading bar
    Returns
    -------------
    samples: (np.ndarray, np.ndarray, np.ndarray)
        1) n_samples x n_alt x n_bvs output measurements at each requested altitude.
        2) n_samples x n_hfps generated hfp waves
        3) n_sample probabilities that the generated wave is present
    """

    with torch.no_grad():
        disable1 = bool(verbose < 1)
        if n_alt > const.max_alt and bv_type == 'radar':
            raise ValueError(f"Requested {n_alt} altitudes but only {const.max_alt}"
                             f" can be simulated for {bv_type}.")
        elif n_alt > const.max_alt_lidar and bv_type == 'lidar':
            raise ValueError(f"Requested {n_alt} altitudes but only {const.max_alt_lidar}"
                             f" can be simulated for {bv_type}.")
        if bv_type == 'lidar':
            n_bv = const.n_lidar_bv
        else:
            n_bv = const.n_bv

        if measurements is None:
            # put placeholder measurements if none provided
            measurements = np.zeros((drivers.shape[0], 0, n_bv))

        n_batch = drivers.shape[0]
        n_alt_start = measurements.shape[1]

        # verify the correct shapes for the inputs
        if drivers.shape != (n_batch, len(driver_names)):
            raise ValueError(f"driver and driver_names must have the "
                             f"same length ({drivers.shape[-1]} != {len(driver_names)}")
        if measurements.shape != (n_batch, n_alt_start, n_bv):
            raise ValueError(f"Measurement shape must be n_batch x "
                             f"{n_alt_start} x {n_bv}.")

        # z scale inputs and place into tensors
        driver_feat = trans.scale_driver(drivers, driver_names=driver_names)
        bv_feat, valid_mask = trans.scale_bv(measurements, bv_type)
        driver_feat = torch.tensor(driver_feat, dtype=torch.float)
        bv_feat = torch.tensor(bv_feat, dtype=torch.float)
        bv_feat, alt_mask = mask_altitude(bv_feat)
        hfp_feats = torch.zeros(len(drivers), const.n_waves, const.n_hfp_feat)
        G_b = torch.zeros(len(drivers))

        if torch.cuda.is_available() and cuda_index is not None:
            device = torch.device('cuda:' + str(cuda_index))
            if len(drivers) >= 500:
                batch_size = 500
            else:
                batch_size = len(drivers)
        else:
            device = torch.device('cpu')
            if len(drivers) >= 100:
                batch_size = 100
            else:
                batch_size = len(drivers)

        batch_idxs = np.hstack((np.arange(0, len(drivers), step=batch_size), [len(drivers)]))

        # Load bv models
        bv_gen, _ = persist.recreate(name=bv_model)
        if bv_gen.transformer.dr_emb.shape[0] != driver_feat.shape[-1]:
            raise ValueError(f"Model must be trained with "
                             f"{driver_feat.shape[-1]} drivers.")
        bv_gen.to(device)
        bv_gen.eval()

        # iteratively build altitude profile
        for i in tqdm(range(len(batch_idxs) - 1), desc='Generate BV Profile', disable=disable1):
            for i_alt in range(n_alt_start, n_alt):
                dr_src = driver_feat[batch_idxs[i]:batch_idxs[i + 1]].to(device)
                bv_src = bv_feat[batch_idxs[i]:batch_idxs[i + 1]].to(device)

                bv_out = bv_gen(dr_src, bv_src)
                # fill in next altitude
                bv_feat[batch_idxs[i]:batch_idxs[i + 1], i_alt, :] = bv_out[:, i_alt, :].detach().cpu()
        G_bv_feats = bv_feat[:, :n_alt, :]
        G_bvs = trans.get_bv(bv_feat[:, :n_alt, :].numpy(), bv_type)
        if measurements is not None:
            G_bvs[:, :n_alt_start, :] = measurements
        bv_gen.cpu()

        if generate_hfps:
            # Load hfp GAN
            hfp_gen, _ = persist.recreate(name=hfp_model)
            hfp_gen.to(device)
            hfp_gen.eval()

            # Generate waves
            for i in tqdm(range(len(batch_idxs) - 1), desc='Generate HFP Batches', disable=disable1):
                dr_src = driver_feat[batch_idxs[i]:batch_idxs[i + 1]].to(device)
                gbv_src = G_bv_feats[batch_idxs[i]:batch_idxs[i + 1]].to(device)
                hfp_tgt = hfp_feats[batch_idxs[i]:batch_idxs[i + 1]].to(device)

                hfp_out, gb = hfp_gen(dr_src, gbv_src, hfp_tgt)
                hfp_feats[batch_idxs[i]:batch_idxs[i + 1], ...] = hfp_out.detach().cpu()
                G_b[batch_idxs[i]:batch_idxs[i + 1]] = gb.detach().cpu()
            G_hfps_feats = hfp_feats.numpy()
            G_hfps = trans.get_hfp(hfp_feats.numpy())
            G_b = G_b.numpy()

        if generate_hfps:
            if return_z_scale:
                return G_bv_feats.numpy(), G_hfps_feats, G_b
            else:
                return G_bvs, G_hfps, G_b
        else:
            if return_z_scale:
                return G_bv_feats.numpy()
            else:
                return G_bvs

def generate_multi(drivers, bvs=None, n_repeat=10,
                   bv_model='bv_gan', hfp_model='hfp_gan',
                   bv_type='radar', generate_hfps=False,
                   cuda_index=None, verbose=1):
    """
    Generate multiple background variable profiles and HFP waves
    consistent with the historical distribution for each driver sample

    Parameters
    -------------
    drivers: np.ndarray
        n_samples x n_drivers input list of driving parameters (not z-scaled).
    bvs: np.ndarray, optional
        n_samples x n_alt_in x n_meas input list of altitude measurements,
        n_alt_in should be less than n_alt. These represent fixed
        measurements for the lowest altitudes to condition on.
        Usually left as default (None)
    n_repeat: int, optional
        number of bv profiles/waves to generate for each driver sample
    bv_model: str, optional
        name of bv GAN to use
    bv_type: str. optional
        name of the type of background variables to use (lidar or radar)
    hfp_model: str, optional
        name of hfp GAN to use
    generate_hfps: bool, optional
        Set to generate hfps and return generated hfps
    cuda_index: int, optional
        GPU index to use when generating BVs and HFPs
    verbose: bool, optional
        set to show loading bar
    Returns
    -------------
    samples: (np.ndarray, np.ndarray, np.ndarray)
        1) (n_samples x n_repeat x n_alt x n_bvs) output measurements at each requested altitude.
        2) (n_samples x n_repeat x 1 x n_hfps) generated hfp waves
        3) (n_sample x n_repeat) probabilities that the generated wave is present
    """
    disable = bool(verbose < 1)
    G_bvs = np.zeros((len(drivers), n_repeat, const.max_alt, const.n_bv))
    G_hfps = np.zeros((len(drivers), n_repeat, const.n_waves, const.n_hfp))
    G_b = np.zeros((len(drivers), n_repeat))

    for i in tqdm(range(len(drivers)), desc='Generating Samples', disable=disable):
        sampled_driver = drivers[[i], ...].repeat(n_repeat, 0)
        if bvs is not None:
            sampled_bv = bvs[[i], ...].repeat(n_repeat, 0)
        else:
            sampled_bv = None
        if generate_hfps:
            gen_data = generate(sampled_driver,
                                measurements=sampled_bv,
                                bv_model=bv_model,
                                hfp_model=hfp_model,
                                bv_type=bv_type,
                                generate_hfps=True,
                                cuda_index=cuda_index,
                                verbose=False)
            G_bvs[i, ...] = gen_data[0]
            G_hfps[i, ...] = gen_data[1]
            G_b[i, ...] = gen_data[2]
        else:
            G_bvs[i, ...] = generate(sampled_driver,
                                     measurements=sampled_bv,
                                     bv_model=bv_model,
                                     hfp_model=hfp_model,
                                     bv_type=bv_type,
                                     generate_hfps=False,
                                     cuda_index=cuda_index,
                                     verbose=False)
    if generate_hfps:
        return G_bvs, G_hfps, G_b
    else:
        return G_bvs

def discriminate(drivers, bvs, hfps=None,
                 driver_names=const.driver_names,
                 bv_model='bv_gan', hfp_model='hfp_gan',
                 bv_type='radar'):
    """
    Score how well the measurements match with historical observations.

    Parameters
    -------------
    drivers: np.ndarray
        n_samples x n_drivers input list of driving parameters (not z-scaled).
    driver_names: list
        list of names of driving parameters
    bvs: np.ndarray
        n_samples x n_alt_in x n_meas input list of altitude measurements,
        n_alt_in should be less than max_alt.
    hfps: np.ndarray, optional
        n_samples x n_wave x n_hfps input list of wave measurements,
    bv_model: str, optional
        name of bv model to use
    hfp_model: str, optional
        name of model hfp to use
    bv_type: str. optional
        name of the type of background variables to use (lidar or radar)
    Returns
    -------------
    scores: (np.ndarray, np.ndarray)
        1) n_samples x n_alt bv normalcy scores in the range [0, 1.0].
        2) n_samples hfp wave normalcy scores in the range [0, 1.0].
    """
    warn(f'The discriminate function is deprecated. Does not produce credible results '
         f'outside of the training process. Please use the anomaly modules in place of'
         f'the discriminator', DeprecationWarning, stacklevel=2)
    with torch.no_grad():
        n_batch, n_alt = bvs.shape[:2]
        _, bv_crit = persist.recreate(name=bv_model)
        bv_crit.eval()

        driver_feat = trans.scale_driver(drivers, driver_names)
        bv_feat, _ = trans.scale_bv(bvs, bv_type)

        if bv_crit.transformer.dr_emb.shape[0] != driver_feat.shape[-1]:
            raise ValueError(f"Model must be trained with "
                             f"{driver_feat.shape[-1]} drivers.")

        driver_feat = torch.tensor(driver_feat, dtype=torch.float)
        bv_feat = torch.tensor(bv_feat, dtype=torch.float)
        bv_feat, alt_mask = mask_altitude(bv_feat)

        # Get bvs scores
        bv_scores = bv_crit(bv_feat, driver_feat, bv_feat, ~alt_mask)
        bv_scores = bv_scores.view(n_batch, -1)[:, :n_alt].cpu().numpy()

        if hfps is not None:
            if hfps.shape[-1] != const.n_hfp:
                raise ValueError(f"HFPs data must have "
                                 f"{const.n_hfp} features.")

            _, hfp_crit = persist.recreate(name=hfp_model)
            hfp_crit.eval()
            hfp_feat, _ = trans.scale_hfp(hfps)
            hfp_feat = torch.tensor(hfp_feat, dtype=torch.float)

            # Get hfp scores
            hfp_scores = hfp_crit(driver_feat, bv_feat, hfp_feat, hfp_feat, ~alt_mask)
            hfp_scores = hfp_scores.view(n_batch, -1).cpu().numpy()

        if hfps is None:
            return bv_scores
        else:
            return bv_scores, hfp_scores

def estimate_drivers(drivers, model='dr_gan'):
    """
    Predict drivers 2 hours into the future driver GAN model. Used for real-time
    background predictions using drivers from 2 hours ago.

    Parameters
    -------------
    drivers: np.ndarray
        n_samples x n_drivers input list of driving parameters (not z-scaled).
    model: str, optional
        name of model to use

    Returns
    -------------
    predicted_drivers: np.ndarray
        estimation of driver features two hours from the drivers inputted
    """
    driver_feats = torch.tensor(trans.scale_driver(drivers), dtype=torch.float)
    dr_gen, _ = persist.recreate(name=model)
    dr_gen.eval()

    with torch.no_grad():
        predicted_feats = dr_gen(driver_feats)

    predicted_drivers = trans.get_driver(predicted_feats.cpu().numpy())
    return predicted_drivers

def hellinger_scores_bv(real, fake, mask=None, bins=None, filter_length=None,
                        return_hist_info=False, z_scale=True, bv_type='radar'):
    """
    Returns the hellinger distance score that measures how similarity between
    real and generated background variable profiles.
    ----------------
    real: np.ndarray
        tensor of real values for a particular alt and bv feat
    fake: np.ndarray
        tensor of generated values for a particular alt and bv feat
    bins: int
        number of bins to use in histogram calculations
        (If None # of bins will be calculated based on number of samples)
    filter_length: int
        averaging filter length to smooth out noise in histograms
        (If None filter length will be calculated based on number of samples)
    return_hist_info: bool
        set to have function return the histograms and bin edges used which
        were used to calculate the hellinger distance metric.
    bv_type:
        type of data (radar or lidar)
    Returns
    -------------
    dist:
        the hellinger distance (n_alts x n_feats)
    """
    if mask is None:
        mask = np.ones((real.shape[0], real.shape[1]), dtype=bool)

    if bins is None:
        bins = max(15, int((real.shape[0])**const.bin_exp))
    if filter_length is None:
        filter_length = max(2, int(len(real)**const.filter_exp))

    dists = np.zeros((real.shape[1], real.shape[2]))
    r_hists = np.zeros((bins, real.shape[1], real.shape[2]))
    f_hists = np.zeros((bins, real.shape[1], real.shape[2]))
    edges = np.zeros((bins + 1, real.shape[1], real.shape[2]))

    for i in range(real.shape[1]):
        for j in range(real.shape[2]):
            if z_scale:
                r = trans.scale_bv(real, bv_type)[0][mask[:, i], i, j]
                f = trans.scale_bv(fake, bv_type)[0][mask[:, i], i, j]
                args = {'bins': bins, 'range': (-3, 3), 'density': True}
            else:
                r = real[mask[:, i], i, j]
                f = fake[mask[:, i], i, j]
                args = {'bins': bins, 'range': (const.bv_thresholds[j, 0], const.bv_thresholds[j, 1]),
                        'density': True}
            r_hist, edg = np.histogram(r, **args)
            f_hist, edg = np.histogram(f, **args)
            if filter_length:
                r_hist = np.convolve(r_hist, np.ones(filter_length), mode='same') / filter_length
                f_hist = np.convolve(f_hist, np.ones(filter_length), mode='same') / filter_length
            r_area = r_hist * np.diff(edg)
            f_area = f_hist * np.diff(edg)
            dists[i, j] = (1 / np.sqrt(2)) * np.sqrt(np.sum((np.sqrt(r_area) - np.sqrt(f_area)) ** 2))
            r_hists[:, i, j] = r_hist
            f_hists[:, i, j] = f_hist
            edges[:, i, j] = edg

    if return_hist_info:
        return dists, (r_hists, f_hists, edges)
    else:
        return dists

def hellinger_scores_hfp(real, fake, r_mask=None, f_mask=None,
                         bins=None, filter_length=None, z_scale=True,
                         return_hist_info=False):
    """
    Returns the hellinger distance score that measures the similarity between
    real and generated background variable profiles.
    ----------------
    real:
        tensor of real values for a particular alt and bv feat
    fake:
        tensor of generated values for a particular alt and bv feat
    bins:
        tensor of real values for a particular alt and bv feat
    filter_length:
        averaging filter length to smooth out histograms
    z_scale: bool
        used z-scaled values (recommended)
    return_hist_info: bool
        set to have function return the real hist,
        fake hist, and bin edges used in calculation
    Returns
    -------------
    dist:
        the hellinger distance (n_alts or n_waves x n_feats)
    """
    if r_mask is None:
        r_mask = np.ones((real.shape[0], real.shape[1]), dtype=bool)
    if f_mask is None:
        f_mask = np.ones((fake.shape[0], fake.shape[1]), dtype=bool)

    if bins is None:
        bins = max(15, int(r_mask.sum() ** const.bin_exp))
    if filter_length is None:
        filter_length = max(2, int(r_mask.sum() ** const.filter_exp))

    dists = np.zeros((real.shape[1], real.shape[2]))
    r_hists = np.zeros((bins, real.shape[1], real.shape[2]))
    f_hists = np.zeros((bins, real.shape[1], real.shape[2]))
    edges = np.zeros((bins + 1, real.shape[1], real.shape[2]))

    for i in range(real.shape[1]):
        for j in range(real.shape[2]):
            if z_scale:
                r = trans.scale_hfp(real)[0][r_mask[:, i], i, j]
                f = trans.scale_hfp(fake)[0][f_mask[:, i], i, j]
                args = {'bins': bins, 'range': const.hfp_z_ranges[j], 'density': True}
            else:
                r = real[r_mask[:, i], i, j]
                f = fake[f_mask[:, i], i, j]
                args = {'bins': bins,
                        # 'range': (min(np.min(r), np.min(f)), max(np.max(r), np.max(f))),
                        'range': (const.hfp_thresholds[j, 0], const.hfp_thresholds[j, 1]),
                        'density': True}
            r_hist, edg = np.histogram(r, **args)
            f_hist, edg = np.histogram(f, **args)
            if filter_length:
                r_hist = np.convolve(r_hist, np.ones(filter_length), mode='same') / filter_length
                f_hist = np.convolve(f_hist, np.ones(filter_length), mode='same') / filter_length
            r_area = r_hist * np.diff(edg)
            f_area = f_hist * np.diff(edg)
            dists[i, j] = (1 / np.sqrt(2)) * np.sqrt(np.sum((np.sqrt(r_area) -
                                                             np.sqrt(f_area)) ** 2))
            r_hists[:, i, j] = r_hist
            f_hists[:, i, j] = f_hist
            edges[:, i, j] = edg

    if return_hist_info:
        return dists, (r_hists, f_hists, edges)
    else:
        return dists

def stack_drivers(driver_dict, driver_names=const.driver_names):
    """
    Stacks drivers in appropriate format.

    This function is provided for convenience.

    Parameters
    ----------------
    driver_dict: dict
        Dictionary mapping names of drivers to the numpy arrays
        with values for those drivers. Each array has a single
        dimension of the same length n_samples.
        Can also use an `h5py.Group`.
    driver_names: list
        names of the drivers to load

    Valid names for drivers can be found at `abcgan.driver_names`

    Raises
    ------------------
    ValueError:
        If the driver values have the wrong type or shape.
    KeyError:
        If one of the required drivers is missing.
    """
    if isinstance(driver_dict, h5py.Group):
        driver_dict = {k: v[()] for k, v in driver_dict.items()}
    shp = None
    for v in driver_dict.values():
        if not isinstance(v, np.ndarray):
            raise ValueError(f"Values in driver_dict must have"
                             f" type np.ndarray not {type(v)}.")
        if shp is None:
            shp = v.shape
            if len(shp) != 1:
                raise ValueError("Driver dict values must have only one"
                                 " dimension for the number of samples.")
        if shp != v.shape:
            raise ValueError("All values in driver_dict must have"
                             " the same length.")
    return np.stack([driver_dict[k] for k in driver_names],
                    axis=-1)

def stack_bvs(bv_dict, bv_type='radar'):
    """
    Stacks drivers in appropriate format.

    This function is provided for convenience.

    Parameters
    ----------------
    bv_dict: dict
        Dictionary mapping names of background variables
        to numpy arrays with values for those bvs. Each
        array should have shape n_sapmles x n_altitudes.
        Can also use `h5py.Group`.
    bv_type: str
        string specifying weather to stack radar or
        lidar data

    Valid names for drivers can be found at `abcgan.bv_names`

    Raises
    ------------------
    ValueError:
        If the input shape of the bv dict values is not corrects
    KeyError:
        If one of the required bvs is missing.
    """
    if isinstance(bv_dict, h5py.Group):
        bv_dict = {k: v[()] for k, v in bv_dict.items()}
    shp = None
    for v in bv_dict.values():
        if not isinstance(v, np.ndarray):
            raise ValueError(f"Values in bv_dict must have"
                             f" type np.ndarray not {type(v)}.")
        if shp is None:
            shp = v.shape
            if len(shp) != 2:
                raise ValueError("BV dict values must have 2 "
                                 "dimensions [n_samples x n_altitudes].")
        if shp != v.shape:
            raise ValueError("All values in bv_dict must have the"
                             " same shape.")
    if bv_type == 'lidar':
        bvs = np.stack([bv_dict[k] for k in const.lidar_bv_names],
                       axis=-1)[:, :const.max_alt_lidar, :]
    else:
        bvs = np.stack([bv_dict[k] for k in const.bv_names],
                       axis=-1)[:, 26:26 + const.max_alt, :]
    return bvs

def load_h5_data(fname, bv_type='radar', load_hfp=False, n_samples=None):
    """
    loads and returns external drivers, background variables, HFP waves and data
    mask all aligned in time with outlier/invalid data filtered out

    Parameters
    -------------
    fname: str
        name of h5 file to load the data from
    bv_type: str. optional
        name of the type of background variables to use (lidar or radar)
    load_hfp: bool. optional
        set to load HFP waves along with the background variables
    n_samples: int. optional
        number of samples to load (None to load all samples)
    Returns
    -------------
    drivers: np.ndarray
        (n_samples x n_dr) external drivers.
    bvs: np.ndarray
        (n_samples x n_alt x n_bv) background variables.
    alt_mask: np.ndarray
        (n_samples x n_alt) background variables alt mask.
    hfps: np.ndarray
        (n_samples x 1 x n_hpf) HPF waves.
    wave_mask: np.ndarray
        (n_samples x 1) HFP wave mask.
    unix_time: np.ndarray
        (n_samples, ) time stamp of each sample
    """
    with h5py.File(fname, 'r') as f:
        # --------------------------------------
        # Read driver and time stamps from file
        # --------------------------------------
        dr_dict = f['Drivers']
        unix_time = f['UnixTime'][()]
        drivers = np.stack([dr_dict[driver_name][:]
                            for driver_name in const.driver_names
                            if driver_name in dr_dict.keys()],
                           axis=-1)
        # --------------------------------------
        # Read background variable data from file
        # --------------------------------------
        bv_dict = f['BackgroundValues']
        if bv_type == 'lidar':
            bvs = np.stack([bv_dict[bv_name][:]
                            for bv_name in const.lidar_bv_names],
                           axis=-1)[:, :const.max_alt_lidar, :]
            bv_thresholds = const.lidar_thresholds
        else:
            bvs = np.stack([bv_dict[bv_name][:]
                            for bv_name in const.bv_names],
                           axis=-1)[:, 26:26 + const.max_alt]
            bv_thresholds = const.bv_thresholds

        if load_hfp:
            # --------------------------------------
            # Read HFP data from file
            # --------------------------------------
            hfp_dict = f['HFPValues']
            hfps = np.stack([hfp_dict[hfp_name]
                             for hfp_name in const.hfp_names
                             if hfp_name in hfp_dict.keys()],
                            axis=-1)
            hfps = hfps[:, None, :]

    # Get valid bvs and altitude mask
    valid_bv_mask = compute_valid(bvs, bv_thresholds)

    if load_hfp:
        valid_hfp_mask = compute_valid_hfp(hfps)
        valid_mask = valid_bv_mask & valid_hfp_mask & ~(np.isnan(drivers).any(-1))
    else:
        valid_mask = valid_bv_mask & ~(np.isnan(drivers).any(-1))

    # Filter out any invalid samples
    if n_samples is not None:
        drivers = drivers[valid_mask][:n_samples]
        bvs = bvs[valid_mask][:n_samples]
        unix_time = unix_time[valid_mask][:n_samples]
    else:
        drivers = drivers[valid_mask]
        bvs = bvs[valid_mask]
        unix_time = unix_time[valid_mask]

    # Get altitude mask for bvs
    _, alt_mask = mask_altitude(torch.tensor(bvs, dtype=torch.float))
    alt_mask = alt_mask.detach().numpy()

    if load_hfp:
        # Filter out invalid samples and get wave mask
        if n_samples is not None:
            hfps = hfps[valid_mask][:n_samples]
        else:
            hfps = hfps[valid_mask]
        wave_mask = ~(np.isnan(hfps).any(-1))

        return drivers, bvs, alt_mask, hfps, wave_mask, unix_time

    else:
        return drivers, bvs, alt_mask, unix_time