"""
Contains numerical kernel for Seismic FDFD class
"""

import numpy
import scipy
import scipy.sparse
import SimPEG
import shutil, os, errno
from IPython.parallel import require, interactive, Reference

DEFAULT_FREESURF_BOUNDS = [False, False, False, False]
DEFAULT_PML_SIZE = 10
DEFAULT_IREG = 4
DEFAULT_DTYPE = 'double'

try:
    from pymatsolver import MumpsSolver
    DEFAULT_SOLVER = MumpsSolver
except:
    DEFAULT_SOLVER = SimPEG.SolverWrapD(scipy.sparse.linalg.splu)

class SeisFDFDKernel(object):

    # source array ref

    # receiver array ref

    mesh = None
    freq = None
    Solver = lambda: None

    def __init__(self, systemConfig, **kwargs):

        if systemConfig.get('cache', False):
            try:
                from tempfile import mkdtemp
                from joblib import Memory
            except ImportError:
                pass
            else:
                if 'cacheDir' in systemConfig:
                    cacheDir = systemConfig['cacheDir']
                    try:
                        os.makedirs(cacheDir)
                    except OSError as e:
                        if e.errno == errno.EEXIST and os.path.isdir(cacheDir):
                            pass
                        else:
                            raise
                else:
                    cacheDir = mkdtemp()

                self._mem = Memory(cachedir=cacheDir, verbose=0)

                # Cache outputs of these methods
                self.forward = self._mem.cache(self.forward)
                self.backprop = self._mem.cache(self.backprop)

        hx = [(systemConfig['dx'], systemConfig['nx']-1)]
        hz = [(systemConfig['dz'], systemConfig['nz']-1)]
        self.mesh = SimPEG.Mesh.TensorMesh([hx, hz], '00')

        self.mesh.ireg = systemConfig.get('ireg', DEFAULT_IREG)
        self.mesh.freeSurf = systemConfig.get('freeSurf', DEFAULT_FREESURF_BOUNDS)

        initMap = {
        #   Argument        Rename to Property
            'c':            'cR',
            'Q':            None,
            'rho':          None,
            'nPML':         None,
            'freeSurf':     None,
            'freq':         None,
            'ky':           None,
            'kyweight':     None,
            'Solver':       None,
            'dx':           None,
            'dz':           None,
            'dtype':        None,
        }

        for key in initMap.keys():
            if key in systemConfig:
                if initMap[key] is None:
                    setattr(self, key, systemConfig[key])
                else:
                    setattr(self, initMap[key], systemConfig[key])

    def __del__(self):
        if hasattr(self, '_mem'):
            self._mem.clear()
            cacheDir = self._mem.cachedir
            del self._mem
            shutil.rmtree(cacheDir)

    # Model properties

    @property
    def c(self):
        return self.cR + self.cI
    @c.setter
    def c(self, value):
        self._cR = value.real
        self._cI = value.imag
        self._invalidateMatrix()

    @property
    def rho(self):
        if getattr(self, '_rho', None) is None:
            self._rho = 310 * self.c**0.25
        return self._rho
    @rho.setter
    def rho(self, value):
        self._rho = value
        self._invalidateMatrix()

    @property
    def Q(self):
        if getattr(self, '_Q', None) is None:
            self._Q = numpy.inf
        return self._Q
    @Q.setter
    def Q(self, value):
        self._Q = value
        self._invalidateMatrix()

    @property
    def cR(self):
        return self._cR
    @cR.setter
    def cR(self, value):
        self._cR = value
        self._invalidateMatrix()

    @property
    def cI(self):
        if self.Q is numpy.inf:
            return 0
        else:
            return 1j * self.cR / (2*self.Q)
    @cI.setter
    def cI(self, value):
        if (value == 0).all():
            self._Q = numpy.inf
        else:
            self._Q = 1j * self.cR / (2*value)
        self._invalidateMatrix()

    # Modelling properties

    @property
    def nPML(self):
        if getattr(self, '_nPML', None) is None:
            self._nPML = DEFAULT_PML_SIZE
        return self._nPML
    @nPML.setter
    def nPML(self, value):
        self._nPML = value
        self._invalidateMatrix()

    @property
    def ky(self):
        if getattr(self, '_ky', None) is None:
            self._ky = 0.
        return self._ky
    @ky.setter
    def ky(self, value):
        self._ky = value
        self._invalidateMatrix()

    @property
    def kyweight(self):
        if getattr(self, '_kyweight', None) is None:
            self._kyweight = 1.
        return self._kyweight
    @kyweight.setter
    def kyweight(self, value):
        self._kyweight = value
        self._invalidateMatrix()

    # Clever matrix setup properties

    @property
    def Solver(self):
        if getattr(self, '_Solver', None) is None:
            self._Solver = DEFAULT_SOLVER
        return self._Solver
    @Solver.setter
    def Solver(self, value):
        self._Solver = value

    @property
    def A(self):
        if getattr(self, '_A', None) is None:
            self._A = self._initHelmholtzNinePoint()
        return self._A

    @property
    def Ainv(self):
        if getattr(self, '_Ainv', None) is None:
            self._mfact()
        return self._Ainv

    def _invalidateMatrix(self):
        if getattr(self, '_A', None) is not None:
            del(self._A)
        if getattr(self, '_Ainv', None) is not None:
            del(self._Ainv)
        if getattr(self, '_mem', None) is not None:
            self._mem.clear()

    @property
    def dtypeReal(self):
        if self.dtype == 'float':
            return numpy.float32
        elif self.dtype == 'double':
            return numpy.float64
        else:
            raise NotImplementedError('Unknown dtype: %s'%self.dtype)

    @property
    def dtypeComplex(self):
        if self.dtype == 'float':
            return numpy.complex64
        elif self.dtype == 'double':
            return numpy.complex128
        else:
            raise NotImplementedError('Unknown dtype: %s'%self.dtype)

    @property
    def dtype(self):
        return getattr(self, '_dtype', DEFAULT_DTYPE)
    @dtype.setter
    def dtype(self, value):
        # Currently this doesn't work because all the solvers assume doubles
        # if value in ['float', 'double']:
        if value in ['double']:
            self._dtype = value
        else:
            raise NotImplementedError('Unknown dtype: %s'%value)

    # ------------------------------------------------------------------------
    # Matrix setup

    def _mfact(self):
        self._Ainv = self.Solver(self.A)

    def _initHelmholtzNinePoint(self):
        """
        An attempt to reproduce the finite-difference stencil and the
        general behaviour of OMEGA by Pratt et al. The stencil is a 9-point
        second-order version based on work by a number of people in the mid-90s
        including Ivan Stekl. The boundary conditions are based on the PML
        implementation by Steve Roecker in fdfdpml.f.
        """

        # Set up SimPEG mesh
        dims = (self.mesh.nNy, self.mesh.nNx)
        # mAve = self.mesh.aveN2CC

        # c = (mAve.T * self.c.ravel()).reshape(dims)
        # rho = (mAve.T * self.rho.ravel()).reshape(dims)

        c = self.c
        rho = self.rho

        # fast --> slow is x --> y --> z as Fortran

        # Set up physical properties in matrices with padding
        omega   = 2 * numpy.pi * self.freq
        cPad    = numpy.pad(c, pad_width=1, mode='edge')
        rhoPad  = numpy.pad(rho, pad_width=1, mode='edge')

        aky = 2*numpy.pi*self.ky

        # Model parameter M
        K = ((omega**2 / cPad**2) - aky**2) / rhoPad

        # Horizontal, vertical and diagonal geometry terms
        dx  = self.mesh.hx[0]
        dz  = self.mesh.hy[0]
        dxx = dx**2
        dzz = dz**2
        dxz = dx*dz
        dd  = numpy.sqrt(dxz)

        # PML decay terms
        # NB: Arrays are padded later, but 'c' in these lines
        #     comes from the original (un-padded) version

        nPML    = self.nPML

        pmldx   = dx*(nPML - 1)
        pmldz   = dz*(nPML - 1)
        pmlr    = 1e-3
        pmlfx   = 3.0 * numpy.log(1/pmlr)/(2*pmldx**3)
        pmlfz   = 3.0 * numpy.log(1/pmlr)/(2*pmldz**3)

        dpmlx   = numpy.zeros(dims, dtype=self.dtypeComplex)
        dpmlz   = numpy.zeros(dims, dtype=self.dtypeComplex)
        isnx    = numpy.zeros(dims, dtype=self.dtypeReal)
        isnz    = numpy.zeros(dims, dtype=self.dtypeReal)

        # Only enable PML if the free surface isn't set

        freeSurf = self.mesh.freeSurf

        if freeSurf[0]:
            isnz[-nPML:,:] = -1 # Top

        if freeSurf[1]:
            isnx[:,-nPML:] = -1 # Right Side

        if freeSurf[2]:
            isnz[:nPML,:] = 1 # Bottom

        if freeSurf[3]:
            isnx[:,:nPML] = 1 # Left side

        dpmlx[:,:nPML] = (numpy.arange(nPML, 0, -1)*dx).reshape((1,nPML))
        dpmlx[:,-nPML:] = (numpy.arange(1, nPML+1, 1)*dx).reshape((1,nPML))
        dnx     = pmlfx*c*dpmlx**2
        ddnx    = 2*pmlfx*c*dpmlx
        denx    = dnx + 1j*omega
        r1x     = 1j*omega / denx
        r1xsq   = r1x**2
        r2x     = isnx*r1xsq*ddnx/denx

        dpmlz[:nPML,:] = (numpy.arange(nPML, 0, -1)*dz).reshape((nPML,1))
        dpmlz[-nPML:,:] = (numpy.arange(1, nPML+1, 1)*dz).reshape((nPML,1))
        dnz     = pmlfz*c*dpmlz**2
        ddnz    = 2*pmlfz*c*dpmlz
        denz    = dnz + 1j*omega
        r1z     = 1j*omega / denz
        r1zsq   = r1z**2
        r2z     = isnz*r1zsq*ddnz/denz

        # Visual key for finite-difference terms
        # (per Pratt and Worthington, 1990)
        #
        #   This         Original
        # AF FF CF  vs.  AD DD CD
        # AA BE CC  vs.  AA BE CC
        # AD DD CD  vs.  AF FF CF

        # Set of keys to index the dictionaries
        keys = ['AD', 'DD', 'CD', 'AA', 'BE', 'CC', 'AF', 'FF', 'CF']

        # Diagonal offsets for the sparse matrix formation
        offsets = {
            'AD':   (-1) * dims[1] + (-1),
            'DD':   (-1) * dims[1] + ( 0),
            'CD':   (-1) * dims[1] + (+1),
            'AA':   ( 0) * dims[1] + (-1),
            'BE':   ( 0) * dims[1] + ( 0),
            'CC':   ( 0) * dims[1] + (+1),
            'AF':   (+1) * dims[1] + (-1),
            'FF':   (+1) * dims[1] + ( 0),
            'CF':   (+1) * dims[1] + (+1),
        }

        # Buoyancies
        bMM = 1. / rhoPad[0:-2,0:-2] # bottom left
        bME = 1. / rhoPad[0:-2,1:-1] # bottom centre
        bMP = 1. / rhoPad[0:-2,2:  ] # bottom right
        bEM = 1. / rhoPad[1:-1,0:-2] # middle left
        bEE = 1. / rhoPad[1:-1,1:-1] # middle centre
        bEP = 1. / rhoPad[1:-1,2:  ] # middle right
        bPM = 1. / rhoPad[2:  ,0:-2] # top    left
        bPE = 1. / rhoPad[2:  ,1:-1] # top    centre
        bPP = 1. / rhoPad[2:  ,2:  ] # top    right

        # Initialize averaged buoyancies on most of the grid
        bMM = (bEE + bMM) / 2 # a2
        bME = (bEE + bME) / 2 # d1
        bMP = (bEE + bMP) / 2 # d2
        bEM = (bEE + bEM) / 2 # a1
        # ... middle
        bEP = (bEE + bEP) / 2 # c1
        bPM = (bEE + bPM) / 2 # f2
        bPE = (bEE + bPE) / 2 # f1
        bPP = (bEE + bPP) / 2 # c2

        # Reset the buoyancies on the outside edges
        # bMM[ 0, :] = bEE[ 0, :]
        # bMM[ :, 0] = bEE[ :, 0]
        # bME[ 0, :] = bEE[ 0, :]
        # bMP[ 0, :] = bEE[ 0, :]
        # bMP[ :,-1] = bEE[ :,-1]
        # bEM[ :, 0] = bEE[ :, 0]
        # bEP[ :,-1] = bEE[ :,-1]
        # bPM[-1, :] = bEE[-1, :]
        # bPM[ :, 0] = bEE[ :, 0]
        # bPE[-1, :] = bEE[-1, :]
        # bPP[-1, :] = bEE[-1, :]
        # bPP[ :,-1] = bEE[ :,-1]

        # K = omega^2/(c^2 . rho)
        kMM = K[0:-2,0:-2] # bottom left
        kME = K[0:-2,1:-1] # bottom centre
        kMP = K[0:-2,2:  ] # bottom centre
        kEM = K[1:-1,0:-2] # middle left
        kEE = K[1:-1,1:-1] # middle centre
        kEP = K[1:-1,2:  ] # middle right
        kPM = K[2:  ,0:-2] # top    left
        kPE = K[2:  ,1:-1] # top    centre
        kPP = K[2:  ,2:  ] # top    right

        # 9-point fd star
        acoef   = 0.5461
        bcoef   = 0.4539
        ccoef   = 0.6248
        dcoef   = 0.09381
        ecoef   = 0.000001297

        # 5-point fd star
        # acoef = 1.0
        # bcoef = 0.0
        # ecoef = 0.0

        # NB: bPM and bMP here are switched relative to S. Roecker's version
        #     in OMEGA. This is because the labelling herein is always ?ZX.

        diagonals = {
            'AD':   ecoef*kMM
                    + bcoef*bMM*((r1zsq+r1xsq)/(4*dxz) - (r2z+r2x)/(4*dd)),
            'DD':   dcoef*kME
                    + acoef*bME*(r1zsq/dz - r2z/2)/dz
                    + bcoef*(r1zsq-r1xsq)*(bMP+bMM)/(4*dxz),
            'CD':   ecoef*kMP
                    + bcoef*bMP*((r1zsq+r1xsq)/(4*dxz) - (r2z+r2x)/(4*dd)),
            'AA':   dcoef*kEM
                    + acoef*bEM*(r1xsq/dx - r2x/2)/dx
                    + bcoef*(r1xsq-r1zsq)*(bPM+bMM)/(4*dxz),
            'BE':   ccoef*kEE
                    + acoef*(r2x*(bEM-bEP)/(2*dx) + r2z*(bME-bPE)/(2*dz) - r1xsq*(bEM+bEP)/dxx - r1zsq*(bME+bPE)/dzz)
                    + bcoef*(((r2x+r2z)*(bMM-bPP) + (r2z-r2x)*(bMP-bPM))/(4*dd) - (r1xsq+r1zsq)*(bMM+bPP+bPM+bMP)/(4*dxz)),
            'CC':   dcoef*kEP
                    + acoef*bEP*(r1xsq/dx + r2x/2)/dx
                    + bcoef*(r1xsq-r1zsq)*(bMP+bPP)/(4*dxz),
            'AF':   ecoef*kPM
                    + bcoef*bPM*((r1zsq+r1xsq)/(4*dxz) - (r2z+r2x)/(4*dd)),
            'FF':   dcoef*kPE
                    + acoef*bPE*(r1zsq/dz - r2z/2)/dz
                    + bcoef*(r1zsq-r1xsq)*(bPM+bPP)/(4*dxz),
            'CF':   ecoef*kPP
                    + bcoef*bPP*((r1zsq+r1xsq)/(4*dxz) - (r2z+r2x)/(4*dd)),
        }

        diagonals['AD'] = diagonals['AD'].ravel()[dims[1]+1:          ]
        diagonals['DD'] = diagonals['DD'].ravel()[dims[1]  :          ]
        diagonals['CD'] = diagonals['CD'].ravel()[dims[1]-1:          ]
        diagonals['AA'] = diagonals['AA'].ravel()[        1:          ]
        diagonals['BE'] = diagonals['BE'].ravel()[         :          ]
        diagonals['CC'] = diagonals['CC'].ravel()[         :-1        ]
        diagonals['AF'] = diagonals['AF'].ravel()[         :-dims[1]+1]
        diagonals['FF'] = diagonals['FF'].ravel()[         :-dims[1]  ]
        diagonals['CF'] = diagonals['CF'].ravel()[         :-dims[1]-1]

        # self._setupBoundary(diagonals, freeSurf)
        if any(freeSurf):
            raise NotImplementedError('Free surface not implemented!')

        # for key in diagonals.keys():
        #     print('%s:\t%d\t%d'%(key, diagonals[key].size, offsets[key]))

        diagonals = [diagonals[key] for key in keys]
        offsets = [offsets[key] for key in keys]

        A = scipy.sparse.diags(diagonals, offsets, shape=(self.mesh.nN, self.mesh.nN), format='csr', dtype=self.dtypeComplex)#, shape=(self.mesh.nN, self.mesh.nN))#, self.mesh.nN, self.mesh.nN, format='csr')

        return A

    # def _setupBoundary(self, diagonals, freeSurf):
    #     """
    #     Function to set up boundary regions for the Seismic FDFD problem
    #     using the 9-point finite-difference stencil from OMEGA/FULLWV.
    #     """

    #     keys = diagonals.keys()
    #     pickDiag = lambda x: -1. if freeSurf[x] else 1.

    #     # Left
    #     for key in keys:
    #         if key is 'BE':
    #             diagonals[key][:,0] = pickDiag(3)
    #         else:
    #             diagonals[key][:,0] = 0.

    #     # Right
    #     for key in keys:
    #         if key is 'BE':
    #             diagonals[key][:,-1] = pickDiag(1)
    #         else:
    #             diagonals[key][:,-1] = 0.

    #     # Bottom
    #     for key in keys:
    #         if key is 'BE':
    #             diagonals[key][0,:] = pickDiag(2)
    #         else:
    #             diagonals[key][0,:] = 0.

    #     # Top
    #     for key in keys:
    #         if key is 'BE':
    #             diagonals[key][-1,:] = pickDiag(0)
    #         else:
    #             diagonals[key][-1,:] = 0.

    # ------------------------------------------------------------------------
    # Externally-callable functions

    def clear(self):
        self._invalidateMatrix()

    # What about @caching decorators?
    def forward(self, src, dOnly=True):

        q = self.kyweight * src.getq(self.mesh)
        u = self.Ainv * q

        d = numpy.array([numpy.dot(P,u) for P in src.getP(self.mesh, self.ky)]).ravel()

        if dOnly:
            return d
        else:
            return u, d

    def backprop(self, src, dresid=1.):

        qr = self.kyweight * src.getqback(self.mesh, dresid, self.ky)
        u = self.Ainv * qr

        return u

    # def gradient(self, isrc, sterm, dresid):

    #     uF, d = self.forward(isrc, False, sterm)
    #     uB = self.backprop(isrc, dresid)

    #     return uF * uB


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun  6 09:37:12 2019

@author: bressler
"""

import SBCcode as sbc
from os import listdir
from os.path import isfile,join
import numpy as np
import matplotlib.pyplot as plt
import scipy
from gaincalc import get_gain
import pulse_integrator as pi
from runlistscatalogue import *
import gc

CONVERSION_TO_CHARGE = (125.0/128)*(1/50.0)*(1/1000.0)*(1/(1.602e-19))

def trig_difference(runs):
    pmtdiffs = []
    pmtnobubdiffs = []
    dubbubdiffs = []
    for run in runs:
        print(run)
        runrawpath = "/bluearc/storage/SBC-17-data/%s/"%run
        runreconpath = "/pnfs/coupp/persistent/grid_output/SBC-17/output/%s/"%run
        acousticfilename = runreconpath+"AcousticAnalysis_%s.bin"%run
        getbubfile = "/coupp/data/home/coupp/HumanGetBub_output_SBC-17/HumanGetBub_%s.bin"%run
        a = sbc.DataHandling.ReadBinary.ReadBlock(acousticfilename)
        c = sbc.DataHandling.ReadBinary.ReadBlock(getbubfile)
        eventn = c["ev"]
        bubt0 = a["bubble_t0"]
        #events = [evnt for evnt in listdir(runrawpath) if not isfile(join(runrawpath,evnt))]
        #for x in events:
        for x in range(101):
            gc.collect()
            try:
                e = sbc.DataHandling.GetSBCEvent.GetEvent(runrawpath,x,'fastDAQ','PMTtraces')
                cgate = e["fastDAQ"]["CAMgate"]
                dcam = np.diff(cgate)
                fdt = e["fastDAQ"]["time"]
                camOffTimes = [fdt[i] for i in range(len(dcam)) if dcam[i] > 0.5]
                pmttracetime = e["PMTtraces"]["t0_sec"][:,0]+e["PMTtraces"]["t0_frac"][:,0]
                d=sbc.AnalysisModules.PMTfastDAQalignment.PMTandFastDAQalignment(e)
                pmtalign = d["PMT_trigt0_sec"]+d["PMT_trigt0_frac"]
                tracetimes = pmttracetime - pmtalign

                at0 = bubt0[int(x),0]
                for t in (tracetimes-at0):
                    if t<0 and t>-500e-6:
                        lastCamOff = 0
                        for k in range(len(camOffTimes)):
                            if t+at0 > camOffTimes[k]:
                                lastCamOff = camOffTimes[k]
                            elif t+at0 < camOffTimes[k]:
                                break
                        if t+at0-lastCamOff > 25e-6:
                            if list(eventn).count(int(x)) == 2:
                                pmtdiffs.append(t)
                            elif list(eventn).count(int(x)) == 1:
                                pmtnobubdiffs.append(t)
                            elif list(eventn).count(int(x)) == 3:
                                print(3)
                            elif list(eventn).count(int(x)) == 4:
                                dubbubdiffs.append(t)
            except:
                print("Last event: %d"%(x-1))
                break

    return [pmtnobubdiffs,pmtdiffs,dubbubdiffs]

def zdependence(runs, m):
    #m = 1e7

    #m=get_gain("/bluearc/storage/SBC-17-data/",runs[0])

    Ncoinc = [0,0]
    ntotcoinc = [0,0]
    totevents = [0,0]
    totbub = [0,0]
    diffs = [[],[]]
    goodz=[[],[]]
    pmtdiffs = [[],[]]
    coincspec = [[],[]]

    allxyzfname = "/pnfs/coupp/persistent/grid_output/SBC-17/output/SimpleXYZ_all.bin"
    xyzf = sbc.DataHandling.ReadBinary.ReadBlock(allxyzfname)
    for run in runs:
        print("zdependence processing run "+run)
        indices = [i for i,x in enumerate(xyzf["runid"]) if str(x[0])+"_"+str(x[1]) == run]
        runposreco = {"ev":[xyzf["ev"][indices]],"x":[xyzf["bubX"][indices]],
                      "y":[xyzf["bubY"][indices]],"z":[xyzf["bubZ"][indices]]}
        runrawpath = "/bluearc/storage/SBC-17-data/%s/"%run
        runreconpath = "/pnfs/coupp/persistent/grid_output/SBC-17/output/%s/"%run
        acousticfilename = runreconpath+"AcousticAnalysis_%s.bin"%run
        a = sbc.DataHandling.ReadBinary.ReadBlock(acousticfilename)
        #c = sbc.DataHandling.ReadBinary.ReadBlock(getbubfile)
        bubt0 = a["bubble_t0"]
        events = [evnt for evnt in listdir(runrawpath) if not isfile(join(runrawpath,evnt))]
        for j in [0,1]:
            with open("/nashome/b/bressler/sbcoutput/%s_PMTmatching_ch%s.txt"%(run,str(j)),"w+") as f, open("/nashome/b/bressler/sbcoutput/%s_muonCoincidences.txt"%run,'w+') as fmu:
                f.write("run event PMT_t0_index PMT_t0_-at0_us PMT_t0 at0 phe z\n")
                fmu.write("run event phe\n")
                for x in events:
                    yes = False
                    if int(x)<len(runposreco["z"][0])-1:
                        yes=True
                    totevents[j] += 1
                    if yes and not np.isnan(runposreco["z"][0][int(x)]):
                        totbub[j] += 1
                        e = sbc.DataHandling.GetSBCEvent.GetEvent(runrawpath,x)
                        veto = e["fastDAQ"]["VetoCoinc"]
                        cgate = e["fastDAQ"]["CAMgate"]
                        dcam = np.diff(cgate)
                        fdt = e["fastDAQ"]["time"]
                        camOffTimes = [fdt[i] for i in range(len(dcam)) if dcam[i] > 0.5]
                        muon = False
                        pmttracetime = e["PMTtraces"]["t0_sec"][:,0]+e["PMTtraces"]["t0_frac"][:,0]
                        d=sbc.AnalysisModules.PMTfastDAQalignment.PMTandFastDAQalignment(e)
                        pmtalign = d["PMT_trigt0_sec"]+d["PMT_trigt0_frac"]
                        tracetimes = pmttracetime - pmtalign
                        at0 = bubt0[int(x),j]
                        i=0 # to match the indexing of the pre-made code I had 1???
                        candidate = 0
                        candidate_time = 0
                        candidate_PMTtime = 0
                        candidate_index = 0
                        for t in (tracetimes-at0):
                            # loop through every PMT trace for the event

                            if t<-150e-6 and t>-600e-6:
                                # if the trace time is within 500 microsec before acoustic t0
                                if max(veto)>0.1:
                                    if fdt[list(veto).index(max(veto))]-at0<0 and fdt[list(veto).index(max(veto))]-at0>-500e-6:
                                        print("Veto Coincidence: event "+run+"-"+str(x))
                                        muon = True

                                #lastCamOff = 0
                                #for k in range(len(camOffTimes)):
                                #    if t+at0 > camOffTimes[k]:
                                #        lastCamOff = camOffTimes[k]
                                #    elif t+at0 < camOffTimes[k]:
                                #        break
                                #if t+at0-lastCamOff > 25e-6:
                                    # if the trace time is more than 25 microseconds away from a camera gate rise
                                    #but not doing it this way anymore because we'll check for the LED being on later during the merge
                                ntotcoinc[j]+=1
                                pmtdiffs.append(t)

                                #take abs to get positive area:
                                trace = np.fabs(e["PMTtraces"]["traces"][i][0])
                                #if ch0 saturated, stitch in low res channel:
                                if max(trace) == 128:
                                    trace = pi.stitchTraces(trace,np.fabs(e["PMTtraces"]["traces"][i][1]))
                                dt = e["PMTtraces"]["dt"][i][0]

                                #subtract baseline:
                                #Actually this gets done in pulse_integrator anyway
                                #baseline = np.mean(trace[0:50])
                                #trace -= baseline

                                #integrate and convert to phe:
                                [phe,n,totInt,pktimes] = pi.SBC_pulse_integrator_bressler(trace,dt)
                                if phe != None:
                                    phe /= m
                                    #keep track of largest candidate:
                                    if phe > candidate:
                                        candidate = phe
                                        candidate_time = t
                                        candidate_PMTtime = t+at0
                                        candidate_index = i
                                #else:
                                #    candidate = -1.0
                            i+=1
                        #i.e. if there is a candidate PMT trace with area greater than zero
                        if candidate > 0:
                            Ncoinc[j] += 1
                            ind = candidate_index
                            pmtt = candidate_PMTtime
                            diffs[j].append(candidate_time)
                            goodz[j].append(runposreco["z"][0][int(x)])
                            coincspec[j].append(candidate)
                            f.write("%s %s %d %f %f %f %f %f\n"%(run,x,ind,
                                                              candidate_time*1e6,
                                                              pmtt,at0,candidate,
                                                              runposreco["z"][0][int(x)]))
                        if muon:
                            fmu.write("%s %s %f\n"%(run,x,candidate))
                    gc.collect()
            print("run "+run+" file %s written"%str(j))
                        #pmtdiffs.append(candidate_times[candidates.index(max(candidates))])
    print("total number of events: "+str(totevents))
    print("total number of bubbles: "+str(totbub))
    print("total coincident triggers: "+str(ntotcoinc))
    print("total coincident bubbles with scintillation greater than 0phe: "+str(Ncoinc))
    print("fraction of bubbles with a coincident scintillation signal greater than 0phe: "+str(sum(Ncoinc)*100/sum(totbub))+"%")

    return [goodz,diffs,coincspec,Ncoinc,ntotcoinc,totevents,totbub]

def main():
    bgruns = bgOct10and11

    biberuns = BiBeSept23and24

    """
    ch1 files made for:
        "20171003_4","20171003_5",
        "20171004_0","20171004_1","20171004_2","20171004_3","20171004_4",
        "20171005_0","20171005_1","20171005_2","20171005_3","20171005_4",
        "20171006_0","20171006_1"

    ch0 files made for:
        "20171003_4","20171003_5",
        "20171004_0","20171004_1","20171004_2","20171004_3","20171004_4",
        "20171005_0","20171005_1","20171005_2","20171005_3","20171005_4",
        "20171006_0","20171006_1"

    """

    BiAlruns = []

    """
    ch1 files made for:
        "20171006_2","20171006_3","20171006_4","20171006_5","20171007_0",
        "20171007_1","20171007_2","20171007_4","20171007_5",
        "20171007_6","20171008_0","20171008_1","20171008_4","20171008_6",
        "20171008_7","20171009_0","20171009_1","20171009_2"

    ch0 files made for:
        "20171006_2","20171006_3","20171006_4","20171006_5","20171007_0",
        "20171007_1","20171007_2","20171007_4","20171007_5",
        "20171007_6","20171008_0","20171008_1","20171008_4","20171008_6",
        "20171008_7","20171009_0","20171009_1","20171009_2"
    """

    cfruns = ["20170711_15","20170711_16"]
    """
    ch0 files made for:
        "20170707_6","20170707_7","20170707_8","20170707_9","20170707_10","20170708_0",
            "20170708_1","20170708_3","20170708_5","20170708_6","20170708_7",
            "20170708_8","20170708_9","20170709_0","20170709_1","20170709_2",
            "20170709_3","20170709_4","20170709_6","20170709_7","20170709_8",
            "20170710_0","20170710_1","20170710_2","20170710_3","20170710_4",
            "20170710_5","20170710_6","20170710_7","20170710_8","20170710_9",
            "20170711_0","20170711_14","20170711_16"

    ch1 files made for:

            "20170707_6","20170707_7","20170707_8","20170707_9","20170707_10","20170708_0",
            "20170708_1","20170708_3","20170708_4","20170708_5","20170708_6",
            "20170708_7","20170708_8","20170708_9","20170709_0","20170709_1","20170709_2",
            "20170709_3","20170709_4","20170709_6",20170709_7","20170709_8",
            "20170710_0","20170710_1","20170710_2","20170710_3","20170710_4",
            "20170710_5","20170710_6","20170710_7","20170710_8", "20170710_9","20170711_0",
            "20170711_14","20170711_15","20170711_16"

    bad AcousticAnalysis_ .bin files:
        "20170708_2","20170708_4",
    """

    pmtnobubdiffs,pmtdiffs,dubbubdiffs = trig_difference(biberuns)
    goodz,diffs,coincspec,Ncoinc,ntotcoinc,totevents,totbub = zdependence(biberuns)

    bgpmtnobubdiffs,bgpmtdiffs,bgdubbubdiffs = trig_difference(bgruns)
    bggoodz,bgdiffs,bgcoincspec,bgNcoinc,bgntotcoinc,bgtotevents,bgtotbub = zdependence(bgruns)

    """
    plt.figure()
    _,bins,_=plt.hist(pmtdiffs,150,histtype='step',label="one bubble",lw=4)
    plt.hist(diffs,200,histtype='step')
    plt.hist(pmtnobubdiffs,bins,histtype='step', label = "no bubble",lw=4)
    plt.hist(dubbubdiffs,bins,histtype='step', label="two bubbles",lw=4)
    plt.xlabel("PMT trigger times minus acoustic t_0",fontsize=25)
    #plt.xlim([-100e-6,300e-6])
    plt.yscale('log')
    plt.legend(fontsize=18)
    plt.show
    """
    def ft(x,m,b):
        return m*x +b

    params,params_cov = scipy.optimize.curve_fit(ft,goodz,diffs)
    p1 = params[0]
    p0 = params[1]
    print("the slope is"+str(p1))
    print("The speed of sound is "+str((1/np.fabs(p1))/100)+" m/s")

    plt.figure()
    plt.scatter(goodz,diffs)
    plt.plot(np.arange(-3,0.1,0.1),p0+p1*np.arange(-3,0.1,0.1),lw=3)
    plt.ylim([-500e-6, 0])
    plt.xlabel("z position (cm)")
    plt.ylabel("time difference (seconds)")
    plt.show

    plt.figure()
    vals,bins,_=plt.hist(coincspec,np.ceil(max(coincspec)),histtype='step',color='r')
    bgvals,_,_ = plt.hist(bgcoincspec,bins,histtype='step')
    plt.xlabel("spectrum of PMT pulse areas within 500 microseconds before acoustic t_0 (photoelectrons)")
    plt.show

    plt.figure()
    plt.bar(bins[:(len(vals))],[v/totbub for v in vals],1,color='r',linewidth=0,label="californium")
    plt.bar(bins[:len(vals)],[v/bgtotbub for v in bgvals],0.7,color='b',linewidth = 0,label="background")
    plt.xlabel("spectrum of PMT pulse areas within 500 microseconds before acoustic t_0 (photoelectrons)",fontsize=25)
    plt.ylabel("probability of a scintillation pulse of this area per bubble",fontsize=25)
    plt.legend(fontsize=18)
    plt.show

if __name__ == "__main__":
    main()


from __future__ import absolute_import
from __future__ import print_function

import copy
from collections import defaultdict
import numpy as np
import tensorflow as tf
from tqdm import tqdm
from six.moves import xrange
import sys
sys.path.append('../../.')
from cleverhans.utils import other_classes
from cleverhans.utils_tf import model_argmax
from cleverhans.evaluation import batch_eval
from cleverhans.attacks_tf import (jacobian_graph, jacobian,
                                   apply_perturbations, saliency_map)
import keras.backend as K
import os
import pickle

def adaptive_fgsm(x, predictions, eps, clip_min=None, clip_max=None,
                  log_dir=None, y=None, model_logits = None,
                  alpha = None, dataset=None
                  ):
    """
    Computes symbolic TF tensor for the adversarial samples. This must
    be evaluated with a session.run call.
    :param x: the input placeholder
    :param predictions: the model's output tensor
    :param eps: the epsilon (input variation parameter)
    :param clip_min: optional parameter that can be used to set a minimum
                    value for components of the example returned
    :param clip_max: optional parameter that can be used to set a maximum
                    value for components of the example returned
    :param y: the output placeholder. Use None (the default) to avoid the
            label leaking effect.
    :return: a tensor for the adversarial example
    """

    # Compute loss]
    logits, = predictions.op.inputs

    fingerprint_dir = log_dir
    fixed_dxs = pickle.load(open(os.path.join(fingerprint_dir, "fp_inputs_dx.pkl"), "rb"))
    fixed_dys = pickle.load(open(os.path.join(fingerprint_dir, "fp_outputs.pkl"), "rb"))

    if y is None:
        # In this case, use model predictions as ground truth
        y = tf.to_float(
            tf.equal(predictions,
                     tf.reduce_max(predictions, 1, keep_dims=True)))

    output = logits
    pred_class = tf.argmax(y,axis=1)
    loss_fp = 0
    [a,b,c] = np.shape(fixed_dys)
    num_dx = b
    target_dys = tf.convert_to_tensor(fixed_dys)
    target_dys = (tf.gather(target_dys,pred_class))
    norms = tf.sqrt(tf.reduce_sum(tf.square(output), axis=1, keep_dims=True))
    norm_logits = output/norms

    for i in range(num_dx):
        logits_p = model_logits(x + fixed_dxs[i])
        p_norm = tf.sqrt(tf.reduce_sum(tf.square(logits_p), axis=1, keep_dims=True))
        logits_p_norm = logits_p/p_norm
        loss_fp = loss_fp + tf.losses.mean_squared_error((logits_p_norm - norm_logits),target_dys[:,i,:])
        #self appropriate fingerprint

    y = y / tf.reduce_sum(y, 1, keep_dims=True)
    loss_ce = tf.reduce_mean(
        tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=y)
    )
    ## Tune this alpha!!

    loss = loss_ce - alpha*loss_fp

    # Define gradient of loss wrt input
    grad, = tf.gradients(loss, x)

    # Take sign of gradient
    signed_grad = tf.sign(grad)

    # Multiply by constant epsilon
    scaled_signed_grad = eps * signed_grad

    # Add perturbation to original example to obtain adversarial example
    adv_x = tf.stop_gradient(x + scaled_signed_grad)

    # If clipping is needed, reset all values outside of [clip_min, clip_max]
    if (clip_min is not None) and (clip_max is not None):
        adv_x = tf.clip_by_value(adv_x, clip_min, clip_max)

    return adv_x

def adaptive_fast_gradient_sign_method(sess, model, X, Y, eps, clip_min=None,
                              clip_max=None, batch_size=256, log_dir = None,
                                       model_logits = None, binary_steps = 2,
                                        dataset="cifar"):
    """
    TODO
    :param sess:
    :param model: predictions or after-softmax
    :param X:
    :param Y:
    :param eps:
    :param clip_min:
    :param clip_max:
    :param batch_size:
    :return:
    """
    # Define TF placeholders for the input and output
    x = tf.placeholder(tf.float32, shape=(None,) + X.shape[1:])
    y = tf.placeholder(tf.float32, shape=(None,) + Y.shape[1:])
    alpha = tf.placeholder(tf.float32, shape=(None,) + (1,))
    num_samples = np.shape(X)[0]
    ALPHA = 0.1*np.ones((num_samples,1))
    ub = 10.0*np.ones(num_samples)
    lb = 0.0*np.ones(num_samples)
    Best_X_adv = None
    for i in range(binary_steps):
        adv_x = adaptive_fgsm(
            x, model(x), eps=eps,
            clip_min=clip_min,
            clip_max=clip_max, y=y,
            log_dir= log_dir,
            model_logits = model_logits,
            alpha = alpha
        )

        X_adv = batch_eval(
            sess, [x, y, alpha], [adv_x],
            [X, Y, ALPHA], feed={},
            args={'batch_size': batch_size}
        )
        X_adv = np.array(X_adv[0])
        if(i==0):
            Best_X_adv = X_adv

        ALPHA, Best_X_adv = binary_refinement(sess,Best_X_adv,
                      X_adv, Y, ALPHA, ub, lb, model, dataset)

    return Best_X_adv

def binary_refinement(sess,Best_X_adv,
                      X_adv, Y, ALPHA, ub, lb, model, dataset='cifar'):
    num_samples = np.shape(X_adv)[0]
    print(dataset)
    if(dataset=="mnist"):
        X_place = tf.placeholder(tf.float32, shape=[1, 1, 28, 28])
    else:
        X_place = tf.placeholder(tf.float32, shape=[1, 3, 32, 32])

    pred = model(X_place)
    for i in range(num_samples):
        logits_op = sess.run(pred,feed_dict={X_place:X_adv[i:i+1,:,:,:]})
        if(not np.argmax(logits_op) == np.argmax(Y[i,:])):
            # Success, increase alpha
            Best_X_adv[i,:,:,:] = X_adv[i,:,:,]
            lb[i] = ALPHA[i,0]
        else:
            ub[i] = ALPHA[i,0]
        ALPHA[i] = 0.5*(lb[i] + ub[i])
    return ALPHA, Best_X_adv

def adaptive_basic_iterative_method(sess, model, X, Y, eps, eps_iter, nb_iter=50,
                           clip_min=None, clip_max=None, batch_size=256,
                           log_dir = None, model_logits = None,
                                     binary_steps =2, attack_type = "bim-b",
                                     dataset="cifar"):
    """
    TODO
    :param sess:
    :param model: predictions or after-softmax
    :param X:
    :param Y:
    :param eps:
    :param eps_iter:
    :param nb_iter:
    :param clip_min:
    :param clip_max:
    :param batch_size:
    :return:
    """
    print("nb_iter",nb_iter)
    # Define TF placeholders for the input and output
    x = tf.placeholder(tf.float32, shape=(None,)+X.shape[1:])
    y = tf.placeholder(tf.float32, shape=(None,)+Y.shape[1:])
    alpha = tf.placeholder(tf.float32, shape=(None,) + (1,))
    num_samples = np.shape(X)[0]
    ALPHA = 0.1*np.ones((num_samples,1))
    ub = 10.0*np.ones(num_samples)
    lb = 0.0*np.ones(num_samples)
    Best_X_adv = None

    results = np.zeros((nb_iter, X.shape[0],) + X.shape[1:])
    # Initialize adversarial samples as the original samples, set upper and
    # lower bounds
    X_adv = X
    X_min = X_adv - eps
    X_max = X_adv + eps
    print('Running BIM iterations...')
    # "its" is a dictionary that keeps track of the iteration at which each
    # sample becomes misclassified. The default value will be (nb_iter-1), the
    # very last iteration.
    def f(val):
        return lambda: val
    its = defaultdict(f(nb_iter-1))
    # Out keeps track of which samples have already been misclassified
    out = set()
    for j in range(binary_steps):

        for i in tqdm(range(nb_iter)):
            adv_x = adaptive_fgsm(
                x, model(x), eps=eps_iter,
                clip_min=clip_min, clip_max=clip_max, y=y,
                log_dir= log_dir,
                model_logits = model_logits,
                alpha = alpha
            )
            X_adv, = batch_eval(
                sess, [x, y, alpha], [adv_x],
                [X_adv, Y, ALPHA], feed={K.learning_phase(): 0},
                args={'batch_size': batch_size}
            )
            X_adv = np.maximum(np.minimum(X_adv, X_max), X_min)
            results[i] = X_adv
            # check misclassifieds
            predictions = model.predict_classes(X_adv, batch_size=512, verbose=0)
            misclassifieds = np.where(predictions != Y.argmax(axis=1))[0]
            for elt in misclassifieds:
                if elt not in out:
                    its[elt] = i
                    out.add(elt)
            print(i)

        X_adv = results[-1]
        if(j==0):
            Best_X_adv = X_adv
        ALPHA, Best_X_adv = binary_refinement(sess,Best_X_adv,
                      X_adv, Y, ALPHA, ub, lb, model, dataset)
    return Best_X_adv


#!/usr/bin/env python
# -*- coding: utf-8 -*-

### Example of use
###, L. Darme 02/07/2019

import matplotlib.pyplot as plt
import numpy as np
# Importing additional user-defined function

import UsefulFunctions as uf
import Amplitudes as am
import Production as br
import Detection as de
import LimitsList as lim
#############################################################################
###############        Several limits example    ##############################
#############################################################################

# This example loads a list of limits, recast them and save them in the Output folder, then plot some of them

if __name__ == "__main__":

    geffem={"gu11":2/3.,"gu22":2/3.,"gd11":-1/3.,"gd22":-1/3.,"gd33":-1/3.,"gl11":-1.,"gl22":-1.}

    ExperimentsList=np.array(["lsnd_decay","charm_decay","seaquest_phase2_decay","faser_decay","ship_decay", \
                              "babar_monogam","belle2_monogam", \
                              "miniboone_scattering","sbnd_scattering","ship_scattering","nova_scattering", \
                              "sn1987_low_cooling","sn1987_high_cooling", \
                              "na62_invisibledecayPi0","bes_invisibledecayJPsi","babar_invisibledecayUpsilon",\
                              "atlas_monojet_down","atlas_monojet_up","lep_monogam"])

    Lim,LabelLimit = lim.GetLimits(ExperimentsList,10,geffem,"V",True)

    fig=plt.figure(1)
    s = fig.add_subplot(1, 1, 1)
    yup = 1e4;xup=5
    ydown = 10;xdown=0.005
    s.set_xlim((xdown,xup))
    s.set_ylim((ydown,yup))
    xbasic=np.linspace(xdown,xup,75)
    s.set_xscale("log", nonposx='clip')
    s.set_yscale("log", nonposy='clip')

    s.loglog(Lim['lsnd_decay'][0],Lim['lsnd_decay'][1],linestyle='-',linewidth=1.5,color='xkcd:green',zorder=15)
    s.fill_between(Lim['lsnd_decay'][0],ydown,Lim['lsnd_decay'][1],color='xkcd:green',alpha=0.5, zorder=15)
    s.loglog(Lim['ship_decay'][0],Lim['ship_decay'][1],linestyle='--',linewidth=1.5,color='xkcd:orange',zorder=15)
    s.loglog(Lim['charm_decay'][0],Lim['charm_decay'][1],linestyle='-',linewidth=1.5,color='xkcd:darkgreen',zorder=15)
    s.fill_between(Lim['charm_decay'][0],ydown,Lim['charm_decay'][1],color='xkcd:darkgreen',alpha=0.5, zorder=15)
    s.loglog(Lim['seaquest_phase2_decay'][0],Lim['seaquest_phase2_decay'][1],linestyle='--',linewidth=1.5,color='xkcd:red',zorder=15)
    s.loglog(Lim['faser_decay'][0],Lim['faser_decay'][1],linestyle='-.',linewidth=1.5,color='xkcd:indigo',zorder=15)

    # -------- Mono photon at LEP
    s.fill_between(Lim['lep_monogam'][0],200,Lim['lep_monogam'][1],color='xkcd:blue',alpha=0.25, zorder=15)
    s.axhline(200,linestyle='--',linewidth=1.,color='xkcd:blue',zorder=15)
    s.loglog(Lim['lep_monogam'][0],Lim['lep_monogam'][1],linestyle='--',linewidth=1.5,color='xkcd:blue',zorder=15)

    # -------- Missing energy searches
    s.loglog(Lim['babar_monogam'][0],Lim['babar_monogam'][1],linestyle='-',linewidth=2,color='xkcd:grey',zorder=15)
    s.fill_between(Lim['babar_monogam'][0],ydown,Lim['babar_monogam'][1],color='xkcd:grey',alpha=0.5, zorder=15)
    s.loglog(Lim['belle2_monogam'][0],Lim['belle2_monogam'][1],linestyle='--',linewidth=1.5,color='xkcd:black',zorder=15)

    # ------ Self consistency
    s.loglog(xbasic,2*xbasic,linestyle='-',linewidth=1.5,color='xkcd:grey',zorder=15)
    s.fill_between(xbasic,ydown,2*xbasic,color='xkcd:grey',alpha=0.5, zorder=15)

    # -------- SN bounds
    gp=np.logical_and(Lim['sn1987_high_cooling'][1]>1.1*Lim['sn1987_low_cooling'][1],Lim['sn1987_high_cooling'][1]>1)
    s.loglog(Lim['sn1987_high_cooling'][0][gp],Lim['sn1987_high_cooling'][1][gp],linestyle='-',linewidth=1,color='xkcd:purple',zorder=15)
    s.loglog(Lim['sn1987_low_cooling'][0][gp],Lim['sn1987_low_cooling'][1][gp],linestyle='--',linewidth=0.5,color='xkcd:purple',zorder=15)
    s.fill_between(Lim['sn1987_low_cooling'][0][gp],Lim['sn1987_low_cooling'][1][gp],Lim['sn1987_high_cooling'][1][gp],color='xkcd:purple',alpha=0.15, zorder=15)

    # ----- Invisible decay of Upsilon

    s.loglog(Lim['babar_invisibledecayUpsilon'][0],Lim['babar_invisibledecayUpsilon'][1],linestyle='-',linewidth=1.1,color='xkcd:dark grey',zorder=15)
    s.fill_between(Lim['babar_invisibledecayUpsilon'][0],ydown,Lim['babar_invisibledecayUpsilon'][1],color='xkcd:grey',alpha=0.25, zorder=15)

    s.text(0.08,220,r'FASER',fontsize=10,color='xkcd:indigo', zorder=40,rotation=40)
    s.text(1.1,52,r'BaBar' ,fontsize=10, zorder=50)
    s.text(1.1,110,r'Belle II ($50$ ab${}^{-1}$)' ,fontsize=10, zorder=50)
    s.text(1.1,220,r'LEP (DELPHI)' ,color='xkcd:darkblue',fontsize=10, zorder=50)
    s.text(0.01,2500,r'SN1987 (Cooling)' ,color='xkcd:purple',fontsize=10, zorder=50,rotation=0)
    s.text(0.015,70,r'LSND',fontsize=10,color='xkcd:darkgreen', zorder=50,rotation=20)
    s.text(0.15,610,r'SeaQuest',fontsize=10,color='xkcd:red', zorder=40,rotation=35)
    s.text(0.28,510,r'CHARM',fontsize=10,color='xkcd:darkgreen', zorder=40,rotation=35)
    s.text(1.,4.e3,'SHIP',fontsize=10,color='xkcd:rust', zorder=50,rotation=30)
    s.text(1.1,160,r'BaBar ($\Upsilon \to$ inv)' ,fontsize=10, zorder=50)
    s.text(0.02,0.94,r'$g_e:g_d:g_u = -1:-\frac{1}{3}:\frac{2}{3}$', style='italic',fontsize=14,transform = s.transAxes, zorder=20)
    #----Adjusting the labels and color bar
    s.set_xlabel(r'$M_{\chi_2}$  [GeV]',fontsize=18)
    s.set_ylabel(r'$\Lambda$  / $\sqrt{g}$ [GeV]',fontsize=18)

    #---- Saving and showing on screen the Figure
    plt.tight_layout()
    plt.savefig('Output/Example2.pdf')
    plt.show()


import os
import pickle
import warnings

import numpy as np
from shutil import copyfile

from astropy import units as u
from astropy.table import Table, Column

from .kinematic_distance import KinematicDistance

class BayesianDistance(object):
    def __init__(self, filename=None):
        """
        initializes the BayesianDistance class

        Parameters
        ----------
        path_to_bdc : file path to the Bayesian distance program
        path_to_input_table : file path of a table containing information of the
            Gaussian decompositions
        verbose : The default is 'True'. Prints status messages to the
            terminal.
        """
        self.path_to_bdc = None
        self.version = '2.4'
        self.path_to_input_table = None
        self.path_to_output_table = None
        self.input_table = None
        self.verbose = True
        self.add_kinematic_distance = True
        self.add_galactocentric_distance = True
        self.check_for_kda_solutions = True
        self.prior_velocity_dispersion = False
        self.colname_lon, self.colname_lat, self.colname_vel,\
            self.colname_e_vel, self.colname_kda, self.colname_name,\
            self.colname_vel_disp = (None for i in range(7))
        self.colnr_lon, self.colnr_lat, self.colnr_vel,\
            self.colnr_e_vel, self.colnr_kda, self.colnr_name,\
            self.colnr_vel_disp = (None for i in range(7))
        self.prob_sa, self.prob_kd, self.prob_gl, self.prob_ps, self.prob_pm =\
            (None for _ in range(5))
        self.table_format = 'ascii'
        self.save_temporary_files = False
        self.max_e_vel = 5.0
        self.default_e_vel = 5.0
        self.kda_info_tables = []
        self.exclude_kda_info_tables = []
        self.kda_weight = 1

        self.random_seed = 177
        self.sample = 1000
        self.beam = None
        self.size_linewidth_index = 0.5
        self.size_linewidth_e_index = 0.1
        self.size_linewidth_sigma_0 = 0.7
        self.size_linewidth_e_sigma_0 = 0.1

        self.use_ncpus = None
        self.plot_probability = False

        self._p = {
            '1.0': {
                'bdc_fortran': 'Bayesian_distance_v1.0.f',
                'summary_suffix': '.prt',
                'fct_extract': self.extract_results_v1p0,
                'R_0': 8.34},
            '2.4': {
                'bdc_fortran': 'Bayesian_distance_2019_fromlist_v2.4.f',
                'summary_suffix': 'summary.prt',
                'fct_extract': self.extract_results_v2p4,
                'R_0': 8.15}
        }

    def say(self, message, end=None):
        """Diagnostic messages."""
        if self.verbose:
            print(message, end=end)

    def check_settings(self):
        self.initialize_bdc()
        self.initialize_table()
        self.set_probability_controls()
        if self.check_for_kda_solutions:
            self.initialize_kda_tables()
        if self.prior_velocity_dispersion:
            self.initialize_prior_velocity_dispersion()

        text = 'Python wrapper for Bayesian distance calculator v{}'.format(
            self.version)
        border = len(text) * '='
        heading = '\n{a}\n{b}\n{a}\n'.format(a=border, b=text)
        self.say(heading)

    def initialize_bdc(self):
        if self.version is None:
            raise Exception("Need to specify 'version'")

        path_script = os.path.dirname(
            os.path.dirname(os.path.realpath(__file__)))

        self.path_to_bdc = os.path.join(
            path_script, 'BDC', 'v' + self.version)
        path_to_file = os.path.join(
            self.path_to_bdc, self._p[self.version]['bdc_fortran'])

        with open(path_to_file, "r") as fin:
            self.bdc_script = fin.readlines()

    def initialize_table(self):
        if self.path_to_output_table is not None:
            self.path_to_table = self.path_to_output_table

        if self.path_to_table is None:
            errorMessage = str("specify 'path_to_output_table'")
            raise Exception(errorMessage)

        self.dirname_table = os.path.dirname(self.path_to_table)
        if len(self.dirname_table) == 0:
            self.dirname_table = os.getcwd()
        self.table_file = os.path.basename(self.path_to_table)
        self.table_filename, self.table_file_extension =\
            os.path.splitext(self.table_file)
        if not os.path.exists(self.dirname_table):
            os.makedirs(self.dirname_table)

    def initialize_kda_tables(self):
        dirname = os.path.dirname(
            os.path.dirname(os.path.realpath(__file__)))
        if not self.kda_info_tables:
            files = os.listdir(os.path.join(dirname, 'KDA_info'))
            self.kda_info_tables = [
                name[:-4] for name in files if name.endswith('.ini')]

        if self.exclude_kda_info_tables:
            self.kda_info_tables = [
                table for table in self.kda_info_tables
                if table not in self.exclude_kda_info_tables]

        self._kda_tables = []
        keys = ['GLON', 'GLAT', 'VLSR', 'd_VLSR', 'p_far',
                'cos_pa', 'sin_pa', 'aa', 'bb']
        for tablename in self.kda_info_tables:
            table = Table.read(os.path.join(
                dirname, 'KDA_info', tablename + '.dat'), format='ascii')
            table = table[keys]

            self._kda_tables.append(table)

    def initialize_prior_velocity_dispersion(self):
        try:
            self.beam = self.beam.to(u.rad).value
        except AttributeError:
            err_msg = "'beam' needs to be specified as valid astropy unit"
            raise Exception(err_msg)

        self.kd = KinematicDistance()
        self.kd.initialize()

        np.random.seed = self.random_seed
        self._indices = self.size_linewidth_index + np.random.randn(
            self.sample) * self.size_linewidth_e_index
        self._sigma_0 = self.size_linewidth_sigma_0 + np.random.randn(
            self.sample) * self.size_linewidth_e_sigma_0

    def set_probability_controls(self):
        s = '      '

        default_vals = {
            '1.0': {'SA': 0.5, 'KD': 1.0, 'GL': 1.0, 'PS': 0.25, 'PM': None},
            '2.4': {'SA': 0.85, 'KD': 0.85, 'GL': 0.85, 'PS': 0.15, 'PM': 0.85}
            }

        if self.prob_sa is None:
            self.prob_sa = default_vals[self.version]['SA']
        if self.prob_kd is None:
            self.prob_kd = default_vals[self.version]['KD']
        if self.prob_gl is None:
            self.prob_gl = default_vals[self.version]['GL']
        if self.prob_ps is None:
            self.prob_ps = default_vals[self.version]['PS']
        if self.prob_pm is None:
            self.prob_pm = default_vals[self.version]['PM']

        cwd = os.getcwd()
        os.chdir(self.path_to_bdc)

        with open(os.path.join(
                self.path_to_bdc, 'probability_controls.inp'), 'r') as fin:
            file_content = fin.readlines()
        with open(os.path.join(
                self.path_to_bdc, 'probability_controls.inp'), 'w') as fout:
            for line in file_content:
                if not line.startswith('!'):
                    line = '{s}{a}{s}{b}{s}{c}{s}{d}'.format(
                        s=s, a=self.prob_sa, b=self.prob_kd, c=self.prob_gl,
                        d=self.prob_ps)
                    if self.prob_pm is not None:
                        line += '{s}{a}'.format(s=s, a=self.prob_pm)
                fout.write(line)
        os.chdir(cwd)

        string = str("prob_sa: {a}\nprob_kd: {b}\n"
                     "prob_gl: {c}\nprob_ps: {d}\n".format(
                         a=self.prob_sa, b=self.prob_kd, c=self.prob_gl,
                         d=self.prob_ps))
        if self.version == '2.4':
            string += 'prob_pm: {}\n'.format(self.prob_pm)
        self.say("setting probability controls to the following values:")
        self.say(string)

    def determine_column_indices(self):
        self.colnr_lon = self.input_table.colnames.index(self.colname_lon)
        self.colnr_lat = self.input_table.colnames.index(self.colname_lat)
        self.colnr_vel = self.input_table.colnames.index(self.colname_vel)
        if self.colname_e_vel is not None:
            if not isinstance(self.colname_e_vel, list):
                self.colname_e_vel = [self.colname_e_vel]
            self.colnr_e_vel = [self.input_table.colnames.index(colname)
                                for colname in self.colname_e_vel]
        if self.colname_kda is not None:
            self.colnr_kda = self.input_table.colnames.index(self.colname_kda)
        if self.colname_vel_disp is not None:
            self.colnr_vel_disp = self.input_table.colnames.index(
                self.colname_vel_disp)
        if self.colname_name is not None:
            self.colnr_name = self.input_table.colnames.index(self.colname_name)

    def make_fortran_out(self, source):
        """Create a fortran executable for the source.

        Replaces the default input file in the fortran script of the Bayesian
        distance calculator with the input file of the source, then creates a
        Fortran executable file.
        """
        with open("{}.f".format(self.path_to_source), "w") as fout:
            for line in self.bdc_script:
                fout.write(line.replace('sources_info.inp',
                                        '{}_sources_info.inp'.format(source)))
        os.system('gfortran {}.f -o {}.out'.format(
                self.path_to_source, self.path_to_source))

    def extract_string(self, s, first, last, incl=False):
        """Search for a substring inside a string.

        Parameters
        ----------
        s : string that is searched for the substring
        first : first characters of the substring
        last : last characters of the substring
        incl : defines if the `first` and `last` characters are still part of
            the substring that will be returned. The default is `False`
            (`first` and `last` are not part of the returned substring)

        Returns
        -------
        substring of s

        """
        try:
            if incl is True:
                start = s.index(first)
                end = s.index(last) + len(last)
            else:
                start = s.index(first) + len(first)
                end = s.index(last, start)
            return s[start:end]
        except ValueError:
            return ""

    def extract_probability_info(self, line, lon, lat, p_far):
        """
        Extract the distance results from the corresponding string in
        the output file of the Bayesian distance calculator tool.
        """
        deleteString = self.extract_string(
                line, 'Probability component', ':', incl=True)
        replaceString = self.extract_string(
                line, 'Probability component', ':')
        line = line.replace(deleteString, replaceString)
        line = line.replace('\n', '')
        comp, dist, err, prob, arm = line.split()
        c_u, c_v, c_w = self.get_cartesian_coords(lon, lat, float(dist))
        # if np.isnan(dist) is True:
        #     dist, err, prob = (0.0 for i in range(3))
        return [comp, dist, err, prob, arm, c_u, c_v, c_w, p_far]

    def extract_results_v1p0(self, input_file_content, result_file_content,
                             kin_dist=None, kda_ref=None):
        """
        Loop through the lines of the output file of the Bayesian distance
        calculator tool and search for the distance results.

        Parameters
        ----------
        result_file_content : List containing read-in lines of the output file
            ({source_name}.prt) of the Bayesian distance calculator tool
        """
        results = []
        flag = False
        for line in result_file_content:
            if flag:
                params = line.split()
                lon, lat, p_far =\
                    float(params[1]), float(params[2]), float(params[4])
                flag = False
            if 'Extra_info' in line:
                flag = True
            searchString = 'Probability component'
            if searchString in line:
                result = self.extract_probability_info(line, lon, lat, p_far)

                if kda_ref is not None:
                    result += [kda_ref]

                if kin_dist is not None:
                    result += kin_dist

                results.append(result)
        return results

    def extract_kinematic_distances(self, result_file_content):
        """"""
        kinDist = [np.NAN, np.NAN]

        flag = 'one'
        for line in result_file_content:
            searchString = 'Kinematic distance(s):'
            if searchString in line:
                if flag == 'one':
                    kinDist[0] = self.extract_kinematic_info(line)
                    flag = 'two'
                elif flag == 'two':
                    kinDist[1] = self.extract_kinematic_info(line)
        return kinDist

    def extract_kinematic_info(self, line):
        """
        Extract the distance results from the corresponding string in
        the output file of the Bayesian distance calculator tool.
        """
        line = line.replace('Kinematic distance(s):', '')
        line = line.replace('\n', '')
        return float(line)

    def extract_results_v2p4(self, input_file_content, result_file_content,
                             kin_dist=None, kda_ref=None):
        for line in input_file_content:
            if line.startswith('!'):
                continue
            params = line.split()
            p_far = params[5]

        for line in result_file_content:
            if line.startswith('!'):
                continue
            params = line.split()

            n_params = len(params)

            lon, lat, vlsr, e_vlsr = params[:4]

            results = []

            for i in range(1, int(n_params / 4)):
                comp = int(n_params / 4) - 1
                dist, e_dist, prob, arm = params[i*4:(i + 1)*4]
                c_u, c_v, c_w = self.get_cartesian_coords(
                    float(lon), float(lat), float(dist))

                result = [comp, dist, e_dist, prob, arm, c_u, c_v, c_w, p_far]

                if kda_ref is not None:
                    result += [kda_ref]

                if kin_dist is not None:
                    result += kin_dist

                results.append(result)
        return results

    def delete_all_temporary_files(self, source):
        for filename in [f for f in os.listdir(self.path_to_bdc) if f.startswith(source)]:
            os.remove(os.path.join(self.path_to_bdc, filename))

    def get_results(self, source, kda_ref=None, name=None):
        """
        Extract the distance results from the output file ({source_name}.prt)
        of the Bayesian distance calculator tool.
        """
        suffix = self._p[self.version]['summary_suffix']
        for filename in [f for f in os.listdir(self.path_to_bdc)
                         if f.startswith(source) and f.endswith(suffix)]:
            with open(os.path.join(self.path_to_bdc, filename), 'r') as fin:
                result_file_content = fin.readlines()

        for filename in [f for f in os.listdir(self.path_to_bdc)
                         if f.startswith(source) and f.endswith("info.inp")]:
            with open(os.path.join(self.path_to_bdc, filename), 'r') as fin:
                input_file_content = fin.readlines()

        if self.add_kinematic_distance:
            if self.version == '1.0':
                kd_content = result_file_content.copy()
            elif self.version == '2.4':
                with open(os.path.join(self.path_to_bdc, source + '.prt'), 'r') as fin:
                    kd_content = fin.readlines()
            kinDist = self.extract_kinematic_distances(kd_content)
        else:
            kinDist = None

        results = self._p[self.version]['fct_extract'](
            input_file_content, result_file_content, kin_dist=kinDist, kda_ref=kda_ref)

        if self.plot_probability:
            if self.save_temporary_files:
                for filename in [f for f in os.listdir(self.path_to_bdc) if f.startswith(source)]:
                    src = os.path.join(self.path_to_bdc, filename)
                    if name is not None:
                        filename = filename.replace(source, name)
                    dst = os.path.join(
                        os.path.dirname(self.path_to_output_table), filename)
                    copyfile(src, dst)

            self.plot_probability_density(
                source, results, input_file_content, name=name)

        self.delete_all_temporary_files(source)

        return results

    def run_bdc_script(self, source, input_string):
        self.path_to_source = os.path.join(self.path_to_bdc, source)
        filepath = '{}_sources_info.inp'.format(self.path_to_source)
        with open(filepath, 'w') as fin:
            fin.write(input_string)
        self.make_fortran_out(source)
        cwd = os.getcwd()
        os.chdir(self.path_to_bdc)
        os.system('./{}.out'.format(source))
        os.chdir(cwd)

    def bdc_calculation_ok(self, source):
        """Check if BDC yielded any distance results."""
        suffix = self._p[self.version]['summary_suffix']
        for filename in [f for f in os.listdir(self.path_to_bdc)
                         if f.startswith(source) and f.endswith(suffix)]:
            with open(os.path.join(self.path_to_bdc, filename), 'r') as fin:
                result_file_content = fin.readlines()
        for line in result_file_content:
            if not line.startswith('!'):
                return True

        return False

    def determine_e_vel(self, row):
        """Determine uncertainty for vlsr value."""
        e_vel = None

        if self.colnr_e_vel is not None:
            e_vel = max([row[colnr] for colnr in self.colnr_e_vel])
            if abs(float(e_vel)) > self.max_e_vel:  # abs(float(vel)):
                e_vel = None

        if self.version == '1.0':
            plusminus = ''
        elif self.version == '2.4':
            # TODO: implement minimum error for velocity
            if e_vel is not None:
                plusminus = '{}\t'.format(e_vel)
            else:
                plusminus = '{}\t'.format(self.default_e_vel)

        return plusminus

    def determine_p_far_and_kda_ref(self, row, lon, lat, vel):
        """Determine KDA prior and corresponding literature reference."""
        p_far = 0.5
        kda_ref = None

        if self.colnr_kda is not None:
            if row[self.colnr_kda] == 'F':
                p_far = 0.5 + 0.5 * self.kda_weight
            elif row[self.colnr_kda] == 'N':
                p_far = 0.5 - 0.5 * self.kda_weight
            elif isinstance(row[self.colnr_kda], float):
                warnings.warn(
                    "KDA solutions need to be given as strings ('N', 'F')")
                p_far = row[self.colnr_kda] * self.kda_weight
        elif self.check_for_kda_solutions:
            p_far, kda_ref = self.check_KDA(lon, lat, vel)

        return round(p_far, 2), kda_ref

    def get_expected_vel_disp(self, distance):
        size = self.beam * distance * 1e3
        sigma_exp = self._sigma_0 * (size)**(self._indices)
        return np.mean(sigma_exp), np.std(sigma_exp)

    def normalized_gauss(self, mean, sigma, x):
        return np.exp(-0.5 * ((x - mean) / sigma)**2)

    def check_limit(self, mean, sigma, vel_disp, prob, limit=0.01):
        limit_prob = None
        if prob < limit:
            if (vel_disp - mean) > 0:
                limit_prob = 'higher'
            else:
                limit_prob = 'lower'
        return limit_prob

    def determine_pfar_from_vel_disp(self, dist_n, dist_f, vel_disp):
        sigma_exp_mean, sigma_exp_std = self.get_expected_vel_disp(dist_n)
        prob_near = self.normalized_gauss(
            sigma_exp_mean, sigma_exp_std, vel_disp)
        limit_n = self.check_limit(
            sigma_exp_mean, sigma_exp_std, vel_disp, prob_near)

        sigma_exp_mean, sigma_exp_std = self.get_expected_vel_disp(dist_f)
        prob_far = self.normalized_gauss(
            sigma_exp_mean, sigma_exp_std, vel_disp)
        limit_f = self.check_limit(
            sigma_exp_mean, sigma_exp_std, vel_disp, prob_far)

        if (limit_f == 'lower') and (limit_n != 'higher'):
            pfar = 0
        else:
            pfar = (-prob_near + prob_far) * 0.5 + 0.5

        return pfar

    def determine_p_far_from_velocity_dispersion(self, row, lon, lat, vel):
        dist_n, dist_f = self.kd.calc_kinematic_distance(lon, lat, vel)
        vel_disp = row[self.colnr_vel_disp]
        p_far = self.determine_pfar_from_vel_disp(dist_n, dist_f, vel_disp)
        return round(p_far, 2)

    def determine(self, row, idx):
        """Determine distance of lbv data point via the BDC."""
        row = list(row)

        source = "SRC{}".format(str(idx).zfill(9))
        lon, lat, vel =\
            row[self.colnr_lon], row[self.colnr_lat], row[self.colnr_vel]

        name = None
        if self.colnr_name is not None:
            name = row[self.colnr_name]

        plusminus = self.determine_e_vel(row)
        p_far, kda_ref = self.determine_p_far_and_kda_ref(row, lon, lat, vel)
        condition = ((p_far == 0.5) and
                     self.prior_velocity_dispersion and
                     (self.colnr_vel_disp is not None))
        if condition:
            p_far = self.determine_p_far_from_velocity_dispersion(
                row, lon, lat, vel)

        input_string = "{a}\t{b}\t{c}\t{d}\t{e}{f}\t-\n".format(
            a=source, b=lon, c=lat, d=vel, e=plusminus, f=p_far)

        self.run_bdc_script(source, input_string)

        #  rerun BDC calculation with p_far = 0.5 if chosen p_far value did not yield distance results
        if (self.version == '2.4') and (p_far != 0.5):
            if not self.bdc_calculation_ok(source):
                self.delete_all_temporary_files(source)

                p_far = 0.5
                input_string = "{a}\t{b}\t{c}\t{d}\t{e}{f}\t-\n".format(
                    a=source, b=lon, c=lat, d=vel, e=plusminus, f=p_far)
                self.run_bdc_script(source, input_string)

        rows = []
        results = self.get_results(source, kda_ref=kda_ref, name=name)
        for result in results:
            rows.append(row + result)

        return rows

    def get_values_from_init_file(self, init_file):
        """Read in values from init file."""
        import ast
        import configparser
        config = configparser.ConfigParser()
        config.read(init_file)

        for key, value in config['DEFAULT'].items():
            try:
                setattr(self, '_' + key, ast.literal_eval(value))
            except ValueError:
                raise Exception('Could not parse parameter {} from config file'.format(key))

    def gaussian_weight(self, x, std=False):
        """Calculate the Gaussian weight.

        Gaussian function: amp * np.exp(-4. * np.log(2) * (x-mean)**2 / fwhm**2)
        mean = 0

        Renormalization factor for amplitude, so that Gaussian function is 1 at the fwhm/2; scale = 1 / (np.exp(-np.log(2)) = np.exp(np.log(2)
        fwhm_factor = 2 * np.sqrt(2 * np.log(2)) = 2.354820045
        """
        if std:
            return np.exp((1 - x**2) / 2)  # = np.exp(0.5) * np.exp(-4. * np.log(2) * (x / 2.354820045)**2)
        else:
            return np.exp(np.log(2) * (1 - 4. * x**2))  # np.exp(np.log(2)) * np.exp(-4. * np.log(2) * x**2)

    def point_in_ellipse(self, table, lon, lat):
        """Adapted from: https://stackoverflow.com/questions/7946187/
        See also: https://math.stackexchange.com/questions/426150/"""
        cos_pa = table['cos_pa'].data
        sin_pa = table['sin_pa'].data
        glon = table['GLON'].data
        glat = table['GLAT'].data
        aa = table['aa'].data
        bb = table['bb'].data

        a = (cos_pa * (lon - glon) + sin_pa * (lat - glat))**2
        b = (sin_pa * (lon - glon) - cos_pa * (lat - glat))**2
        epsilon = (a / aa) + (b / bb)

        std = False
        if self._size == 'std':
            std = True
        else:
            epsilon = epsilon / 4

        weight = self.gaussian_weight(np.sqrt(epsilon), std=std)
        weight[weight > 1] = 1

        return weight >= self._threshold_spatial, weight

    def get_weight_velocity(self, table, vel):
        """Calculate the weight for the velocity association.

        Parameters
        ----------
        table : astropy.table.table.Table
            Table containing sources with solved kinematic distance ambiguities.
        vel : float
            vlsr position of the coordinate.

        Returns
        -------
        mask : numpy.ndarray
            Mask that is true for each weight that exceeds _threshold_spectral.
        weight : numpy.ndarray
            Array of the weight values.

        """
        vlsr = table['VLSR'].data.data
        dvlsr = table['d_VLSR'].data.data

        x = np.abs(vlsr - vel) / dvlsr

        # fwhm_factor = 2.354820045
        # if fwhm:
        #     x = np.abs((vlsr - vel) / (dvlsr))
        # else:
        #     x = np.abs((vlsr - vel) / (dvlsr)) / (2 * fwhm_factor)

        std = False
        if self._linewidth == 'std':
            std = True

        weight = self.gaussian_weight(x, std=std)
        weight[weight > 1] = 1

        return weight >= self._threshold_spectral, weight

    def get_kda(self, weights_kda, refs):
        if len(weights_kda) == 0:
            return 0, '--'

        if len(weights_kda) == 1:
            return weights_kda[0], refs[0]

        weights_kda_abs = [abs(x) for x in weights_kda]
        max_weight = max(weights_kda_abs)

        indices = np.argwhere(weights_kda_abs == max_weight).flatten().tolist()

        if len(indices) == 1:
            i = indices[0]
            return weights_kda[i], refs[i]

        if sum(weights_kda[indices]) == 0:
            return 0, '--'
        else:
            list_weights_kda = weights_kda[indices].tolist()
            weight = max(list_weights_kda, key=list_weights_kda.count)
            i = np.argwhere(weights_kda == weight).flatten()[0]
            return weight, refs[i]

    def check_KDA(self, lon, lat, vel):
        weights_kda, refs = np.array([]), []
        dirname = os.path.dirname(
            os.path.dirname(os.path.realpath(__file__)))

        for table, tablename in zip(self._kda_tables, self.kda_info_tables):
            path_to_ini = os.path.join(dirname, 'KDA_info', tablename + '.ini')
            self.get_values_from_init_file(path_to_ini)
            mask_pp, weight_pp = self.point_in_ellipse(table, lon, lat)
            mask_vlsr, weight_vlsr = self.get_weight_velocity(table, vel)

            mask_total = np.logical_and(mask_pp, mask_vlsr)
            weight_total = weight_pp * weight_vlsr
            weight_total = weight_total[mask_total]
            p_far_values = table['p_far'].data
            p_far_values = p_far_values[mask_total]

            n_values = np.count_nonzero(mask_total)
            if n_values == 0:
                continue
            elif n_values == 1:
                weights_kda = np.append(
                    weights_kda, self._weight_cat * p_far_values * weight_total)
                refs.append(self._reference)
            else:
                weights_kda = np.append(weights_kda, self._weight_cat * (np.average(
                    p_far_values * weight_total, weights=weight_total)))
                refs.append(self._reference)

        weight_kda, ref = self.get_kda(weights_kda, refs)
        p_far = 0.5 + weight_kda

        return round(float(p_far), 2), ref

    def get_cartesian_coords(self, lon, lat, dist):
        from astropy.coordinates import SkyCoord
        from astropy import units as u

        c = SkyCoord(l=lon*u.degree,
                     b=lat*u.degree,
                     distance=dist*u.kpc,
                     frame='galactic')
        c.representation_type = 'cartesian'
        c_u = round(c.u.value, 4)
        c_v = round(c.v.value, 4)
        c_w = round(c.w.value, 4)

        return c_u, c_v, c_w

    def calculate_distances(self):
        self.check_settings()
        self.say('calculating Bayesian distance...')

        if self.input_table is None:
            self.input_table = Table.read(
                self.path_to_input_table, format=self.table_format)
            #  TESTING:
            # self.input_table = self.input_table[62000:62001]
        self.determine_column_indices()

        condition = (self.prior_velocity_dispersion and
                     (self.colnr_vel_disp is None))
        if condition:
            self.colnr_vel_disp = False
            warnings.warn(str("Did not specify 'colnr_vel_disp' or 'colname_vel_disp'. Setting 'prior_velocity_dispersion=False'."))

        from . import BD_multiprocessing
        BD_multiprocessing.init([self, self.input_table])
        results_list = BD_multiprocessing.func(use_ncpus=self.use_ncpus)
        print('SUCCESS\n')

        for i, item in enumerate(results_list):
            if not isinstance(item, list):
                self.say("Error for distance with index {}: {}".format(i, item))
                del results_list[i]
                continue

        results_list = np.array([item for sublist in results_list
                                 for item in sublist])

        if self.save_temporary_files:
            filepath = os.path.join(
                os.path.dirname(self.path_to_table),
                '_bdc_results_list.pickle')
            with open(filepath, 'wb') as p_file:
                pickle.dump(results_list, p_file)

        self.create_astropy_table(results_list)

    def galactocentric_distance(self, glon, dist_los, glat=None):
        """Calculate galactocentric distance.

        Parameters
        ----------
        glon : float [radians]
            Galactic longitude angle of the line of sight. Has to be supplied in [radians].
        dist_los : float [kpc]
            Distance along the line of sight. Has to be supplied in [kpc].
        glat : float [radians]
            Galactic latitude angle of the line of sight. Has to be supplied in [radians].

        Returns
        -------
        Galactocentric distance in [kpc].

        """
        if glat is not None:
            dist_los = dist_los * np.cos(glat)
        R_0 = self._p[self.version]['R_0']
        return np.sqrt(R_0**2 + dist_los**2 - 2*R_0*dist_los*np.cos(glon))

    def create_astropy_table(self, results):
        self.say('creating Astropy table...')

        added_colnames = ['comp', 'dist', 'e_dist', 'prob', 'arm',
                          'c_u', 'c_v', 'c_w', 'p_far']

        dtypeinput_table = []
        for name, dtype in self.input_table.dtype.descr:
            dtypeinput_table.append(dtype)
        added_dtype = ['i4', 'f4', 'f4', 'f4', 'object',
                       'f4', 'f4', 'f4', 'f4']

        if self.check_for_kda_solutions and (self.colname_kda is None):
            added_colnames += ['KDA_ref']
            added_dtype += ['object']
        if self.add_kinematic_distance:
            added_colnames += ['kDist_1', 'kDist_2']
            added_dtype += ['f4', 'f4']

        names = self.input_table.colnames + added_colnames
        dtype = dtypeinput_table + added_dtype

        self.table_results = Table(data=results, names=names, dtype=dtype)

        if self.add_galactocentric_distance:
            rgal = self.galactocentric_distance(
                np.radians(self.table_results[self.colname_lon].data),
                self.table_results['dist'].data,
                glat=np.radians(self.table_results[self.colname_lat].data))
            self.table_results.add_column(Column(data=rgal, name='rgal'))

        for key in ['c_u', 'c_v', 'c_w', 'rgal']:
            if key in self.table_results.colnames:
                self.table_results[key].format = "{0:.3f}"
        for key in ['dist', 'e_dist', 'prob', 'p_far', 'kDist_1', 'kDist_2']:
            if key in self.table_results.colnames:
                self.table_results[key].format = "{0:.2f}"

        self.say(">> saved table '{}' in {}\n".format(
                 self.table_file, self.dirname_table))

        self.table_results.write(self.path_to_table, format=self.table_format,
                                 overwrite=True)

    def choose_distance(self, probabilities, distances, dist_errors):
        """Choose distance from alternative solutions.

        Flags for the chosen distance:
        - 0: only 1 distance solution existed
        - 1: only distance solution for which associated Gaussian fit had amplitude above three standard deviations of flat distance probability density
        - 2: distance solution had the highest probability
        - 3: distance solutions were tied in their probabilites; chosen distance had the lowest distance error
        - 4: distance solutions were tied in their probabilites and distance errors; chosen distance is the near distance

        fwhm_factor = 2 * np.sqrt(2 * np.log(2)) = 2.354820045

        Calculate the integrated area of the Gaussian function:
        area_gauss = amp * fwhm / ((1. / np.sqrt(2*np.pi)) * 2*np.sqrt(2*np.log(2)))

        combining all constants yields a factor of 0.93943727869965132
        """
        if len(probabilities) == 1:
            return [], 0

        #  check if one of the components had a probability of 1; this implies
        #  that the remaining components have a probability of zero. This can #  happen as v2.4 of the BDC by default always returns two components
        if 1 in probabilities:
            remove = np.where(probabilities != 1)[0]
            return remove, 0

        #  to get from integrated intensity (= probabilities) and std (= dist_errors) to amplitude
        amps = probabilities * 0.93943727869965132 / (2.354820045 * dist_errors)
        remove = np.where(amps < 3 * 0.04)[0]
        if len(remove) == 1:
            return remove, 1

        remove = np.where(probabilities == min(probabilities))[0]
        if len(remove) == 1:
            return remove, 2

        remove = np.where(dist_errors == max(dist_errors))[0]
        if len(remove) == 1:
            return remove, 3

        remove = np.argmax(distances)
        return remove, 4

    def get_table_distance_max_probability(self, save=True):
        from tqdm import tqdm
        self.say('creating Astropy table containing only distance results '
                 'with the highest probability...')

        remove_rows, choice_flags = np.array([]), np.array([])

        if self.version == '1.0':
            for idx, component in tqdm(enumerate(self.table_results['comp'])):
                if idx == 0:
                    comps_indices = np.array([idx])
                else:
                    if (component == 1):
                        if comps_indices.size > 1:
                            sort_indices_highest_probability = np.argsort(
                                self.table_results['prob'][comps_indices])[::-1]
                            remove = sort_indices_highest_probability[1:]
                            remove_rows = np.append(remove_rows, comps_indices[remove])
                        comps_indices = np.array([idx])
                    else:
                        comps_indices = np.append(comps_indices, idx)

            #  take care of the last distance results in the list
            sort_indices_highest_probability = np.argsort(
                self.table_results['prob'][comps_indices])[::-1]
            remove = sort_indices_highest_probability[1:]
            remove_rows = np.append(remove_rows, comps_indices[remove])
        elif self.version == '2.4':
            comps_indices = np.array([], dtype='int')

            for idx, component in tqdm(enumerate(self.table_results['comp'])):
                comps_indices = np.append(comps_indices, idx)

                if len(comps_indices) == component:
                    remove, flag = self.choose_distance(
                        self.table_results['prob'][comps_indices],
                        self.table_results['dist'][comps_indices],
                        self.table_results['e_dist'][comps_indices]
                        )
                    remove_rows = np.append(remove_rows, comps_indices[remove])
                    choice_flags = np.append(choice_flags, flag)
                    comps_indices = np.array([], dtype='int')

        remove_rows = remove_rows.astype(int)
        self.table_results.remove_rows(remove_rows)

        if self.version == '2.4':
            self.table_results.add_column(
                Column(data=choice_flags, name='flag', dtype='int'))

        if save:
            self.table_file = '{}{}{}'.format(self.table_filename, '_p_max',
                                              self.table_file_extension)
            self.path_to_table = os.path.join(
                self.dirname_table, self.table_file)

            self.say(">> saved table '{}' in {}".format(
                     self.table_file, self.dirname_table))

            self.table_results.write(self.path_to_table,
                                     format=self.table_format,
                                     overwrite=True)

    def find_index_max_probability(self, indices, arm=False):
        idx = [i for i in indices]
        prob = [self.table['prob'][i] for i in indices]
        max_idx = prob.index(max(prob))

        if arm:
            arms = [self.table['arm'][i] for i in indices]
            return idx[max_idx], arms[max_idx]
        else:
            return idx[max_idx]

    def order_distances(self, results):
        indices = list(range(len(results)))

        distances = np.array([float(result[1]) for result in results])
        dist_errors = np.array([float(result[2]) for result in results])
        probabilities = np.array([float(result[3]) for result in results])

        remove, _ = self.choose_distance(probabilities, distances, dist_errors)

        first_choice = [i for i in indices if i not in remove]
        choices = first_choice + remove.tolist()
        results = [results[i] for i in choices]
        return results

    def plot_probability_density(self, source, results, input_file_content,
                                 name=None):
        import matplotlib.pyplot as plt

        def get_maximum_distance(distance, probability, max_dist=None):
            try:
                prob_threshold = 0.05
                distance_cutoff = distance[probability > prob_threshold][-1]
                return max(max_dist, int(distance_cutoff + 1))
            except IndexError:
                return max_dist

        distKD, probKD = np.loadtxt(
            os.path.join(self.path_to_bdc, '{}_kinematic_distance_pdf.dat'.format(source)),
            usecols=(0, 1), skiprows=2, unpack=True)

        if self.version == '1.0':
            distGL, probGL = np.loadtxt(
                os.path.join(self.path_to_bdc, '{}_latitude_pdf.dat'.format(source)),
                usecols=(0, 1), skiprows=2, unpack=True)

            distSA, probSA = np.loadtxt(
                os.path.join(self.path_to_bdc, '{}_spiral_arm_pdf.dat'.format(source)),
                usecols=(0, 1), skiprows=2, unpack=True)
        elif self.version == '2.4':
            distSA, probSA = np.loadtxt(
                os.path.join(self.path_to_bdc, '{}_arm_latitude_pdf.dat'.format(source)),
                usecols=(0, 1), skiprows=2, unpack=True)

            distGL, probGL = distSA, probSA

        arm_ranges_lower, arm_ranges_upper = np.loadtxt(
            os.path.join(self.path_to_bdc, '{}_arm_ranges.dat'.format(source)),
            usecols=(0, 1), skiprows=2, unpack=True)

        spiral_arms = np.genfromtxt(
            os.path.join(self.path_to_bdc, '{}_arm_ranges.dat'.format(source)),
            skip_header=2, usecols=2, dtype='str')

        skip_spiral_arm_ranges = False
        if spiral_arms.size == 2:
            arm_ranges_lower = [arm_ranges_lower[1]]
            arm_ranges_upper = [arm_ranges_upper[1]]
            spiral_arms = [spiral_arms[1]]
        elif spiral_arms.size > 2:
            arm_ranges_lower = arm_ranges_lower[1:]
            arm_ranges_upper = arm_ranges_upper[1:]
            spiral_arms = spiral_arms[1:]
        else:
            skip_spiral_arm_ranges = True

        distPS, probPS = np.loadtxt(
            os.path.join(self.path_to_bdc, '{}_parallaxes_pdf.dat'.format(source)),
            usecols=(0, 1), skiprows=2, unpack=True)

        distFD, probFD = np.loadtxt(
            os.path.join(self.path_to_bdc, '{}_final_distance_pdf.dat'.format(source)),
            usecols=(0, 1), skiprows=2, unpack=True)

        max_dist = 0
        for dist, prob in zip(
                [distKD, distGL, distSA, distPS], [probKD, probGL, probSA, probPS]):
            max_dist = get_maximum_distance(dist, prob, max_dist=max_dist)

        fig = plt.figure(figsize=(10, 7.5))
        ax = fig.add_subplot(1, 1, 1)

        ax.set_xlabel('Distance [kpc]', size=20)
        ax.set_ylabel('Probability density [kpc$^{-1}$]', size=20)

        ax.tick_params(axis='both', labelsize=16, pad=8)
        ax.tick_params(axis='both', which='major', direction='out',
                       width=1.25, length=10, pad=8)
        ax.tick_params(axis='both', which='minor', direction='out',
                       width=1.25, length=5)

        ax.plot(distKD, probKD, label='KD', lw=2.5, ls='solid', c='dodgerblue', alpha=0.75)
        ax.plot(distGL, probGL, label='GL', lw=2.5, ls='dotted', c='orange', alpha=1.0)
        ax.plot(distSA, probSA, label='SA', lw=2.5, ls='--', c='indianred', alpha=1.0)
        ax.plot(distPS, probPS, label='PS', lw=2.5, ls='-.', c='forestgreen', alpha=1.0)
        ax.plot(distFD, probFD, label='combined', lw=2, ls='solid', c='black', alpha=1.0)

        if not skip_spiral_arm_ranges:
            for lower, upper, text in zip(
                    arm_ranges_lower, arm_ranges_upper, spiral_arms):
                if lower > max_dist:
                    continue
                horizontalalignment = 'center'
                if text == 'AqR':
                    horizontalalignment = 'left'
                ax.axvspan(lower, upper, alpha=0.15, color='indianred')
                ax.text((upper + lower)/2, ax.get_ylim()[1] * 0.99, text, size=14, color='indianred', horizontalalignment=horizontalalignment, verticalalignment='top')

        box = ax.get_position()
        ax.set_position([box.x0, box.y0,
                         box.width, box.height * 0.9])
        # Put a legend below current axis
        leg1 = ax.legend(
            loc='upper center', bbox_to_anchor=(0.5, 1.08),
            fancybox=False, shadow=False, ncol=5,
            fontsize=14, numpoints=1, frameon=0)

        markers, texts = [], []

        if self.version == '2.4':
            results = self.order_distances(results)
        for i, result in enumerate(results):
            dist = float(result[1])
            if dist <= 0:
                continue
            e_dist = float(result[2])
            prob = float(result[3])
            index = 'D$_{{\\mathregular{{{}}}}}$'.format(i + 1)
            text = '{a}={b:.1f}$\\pm${c:.1f} kpc ({d:.0%})'.format(
                a=index, b=dist, c=e_dist, d=prob)
            marker = ax.scatter(dist, 0 - i*0.01)
            markers.append(marker)
            texts.append(text)
            ax.errorbar(dist, 0 - i*0.01, xerr=e_dist)

        leg2 = ax.legend(
            markers, texts,
            loc='upper center', bbox_to_anchor=(0.5, 1.135),
            fancybox=False, shadow=False, ncol=len(results),
            fontsize=14, numpoints=1, frameon=0)

        for line in input_file_content:
            if line.startswith('!'):
                continue
            params = line.split()
            glon, glat, vlsr, e_vlsr, p_far = params[1:6]
            break

        text = str('$\\ell$={} deg, $b$={} deg, '
                   'V$_{{\\mathregular{{LSR}}}}$={} $\\pm$ {} km/s, '
                   'P$_{{\\mathregular{{far}}}}$={}'. format(
                       glon, glat, vlsr, e_vlsr, p_far))
        plt.title(text, fontsize=14, pad=50)

        ax.set_xlim([0, max_dist])

        ax.add_artist(leg1)

        if name is not None:
            source = name

        path_to_file = os.path.join(self.dirname_table, source + '.pdf')
        plt.savefig(path_to_file, bbox_inches='tight')
        plt.close()