#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
# CODE NAME HERE

# CODE DESCRIPTION HERE

Created on 2020-05-07

@author: cook
"""
from astropy.io import fits
from astropy.table import Table
from astropy import units as uu
from astropy.time import Time
import numpy as np
import sys
import os
import argparse
from scipy.interpolate import InterpolatedUnivariateSpline
import matplotlib.pyplot as plt
import itertools

# =============================================================================
# Define variables
# =============================================================================
__VERSION__ = '0.2.000'
__DATE__ = '2020-05-07'
# constants
SPEED_OF_LIGHT = 299792.458  # [km/s]
IMAGE_PIXEL_SIZE = 2.28  # IMAGE_PIXEL_SIZE

# =============================================================================
# Argument functions
# =============================================================================
class Arguments:
    def __init__(self, rawargs):
        self.version = __VERSION__
        self.date = __DATE__
        # static
        self.CASE = 1
        self.OUTDIR = '.'
        self.PLOT = True
        self.DEBUG = False
        self.kernelparams = None
        # can be dynamic
        self.KERNEL = None
        self.KERNEL_WIDTH = None
        self.KERNEL_EWIDTH = None
        self.KERNEL_BETA = None
        self.IN_FILE = None
        self.BLAZE_FILE = None
        self.WAVE_FILE = None
        self.MASK_FILE = None
        self.MASK_WIDTH = None
        self.MASK_MIN_WEIGHT = None
        self.CCF_STEP = None
        self.CCF_WIDTH = None
        self.CCF_RV_NULL = None
        self.IN_RV = None
        self.CCF_N_ORD_MAX = None
        self.BLAZE_NORM_PERCENTILE = None
        self.BLAZE_THRESHOLD = None
        self.NOISE_SIGDET = None
        self.NOISE_SIZE = None
        self.NOISE_THRES = None
        # get raw args
        self.rawargs = rawargs[1:]
        self.cmdargs = None
        # deal with version argument
        if '--version' in self.rawargs:
            print('Version={0} ({1})'.format(__VERSION__, __DATE__))
            sys.exit(1)
        # do we get arguments from args?
        if len(self.rawargs) > 0:
            self.parseargs()

    def cmdkeys(self):
        keys = dict()
        keys['KERNEL'] = '--kernel={0} '
        keys['KERNEL_WIDTH'] = '--kernelwid={0} '
        keys['KERNEL_EWIDTH'] = '--kernelewid={0} '
        keys['KERNEL_BETA'] = '--kernelbeta={0} '
        keys['IN_FILE'] = '--infile={0} '
        keys['BLAZE_FILE'] = '--blaze={0} '
        keys['WAVE_FILE'] = '--wave={0} '
        keys['MASK_WIDTH'] = '--maskwidth={0} '
        keys['MASK_MIN_WEIGHT'] = '--maskminweight={0} '
        keys['CCF_STEP'] = '--ccfstep={0} '
        keys['CCF_WIDTH'] = '--ccfwidth={0} '
        keys['CCF_RV_NULL'] = '--ccf_rv_null={0} '
        keys['IN_RV'] = '--targetrv={0} '
        keys['CCF_N_ORD_MAX'] = '--ccf_ord_max={0} '
        keys['BLAZE_NORM_PERCENTILE'] = '--blazenormper={0} '
        keys['BLAZE_THRESHOLD'] = '--blaze_thres={0} '
        keys['NOISE_SIGDET'] = '--noise_sigdet={0} '
        keys['NOISE_SIZE'] = '--noise_size={0} '
        keys['NOISE_THRES'] = '--noise_thres={0} '
        return keys

    def parseargs(self):
        # get parser
        desc = ('Stand alone CCF code from the DRS. Version {0}. '
                'Note any arguments left out will default to the default '
                'defined in the --case argument '
                '(if left out this defaults to --case=1)')
        parser = argparse.ArgumentParser(description=desc.format(self.version))

        parser.add_argument('--case', action='store', default=self.CASE,
                            dest='CASE', choices=[1, 2], required=True,
                            type=int,
                            help='The case to use for all undefined keywords. '
                                 'case=1   default setup for datatype=OBJ. '
                                 'case=2   default setup for datatype=FP')
        parser.add_argument('--plot', action='store_true', default=self.PLOT,
                            dest='PLOT', help='If set plots graphics')
        parser.add_argument('--debug', action='store_true', default=False,
                            dest='DEBUG', help='Debug mode')
        parser.add_argument('--outdir', action='store', default=self.OUTDIR,
                            dest='OUTDIR',
                            help='Output path for files (absolute path) if'
                                 'blank uses the input file directory.')
        parser.add_argument('--kernel', action='store', default=self.KERNEL,
                            dest='KERNEL',
                            choices=['None', 'boxcar', 'gaussian',
                                     'supergaussian'],
                            help='Kernal name. '
                                 'If boxcar must define --kernelwid '
                                 'If gaussian must define --kernelewid '
                                 'If supergaussian must define '
                                 '--kernelewid --kernelbeta.')
        parser.add_argument('--kernelwid', action='store',
                            default=self.KERNEL_WIDTH, dest='KERNEL_WIDTH',
                            help='The kernel width if --kernel=boxcar')
        parser.add_argument('--kernelewid', action='store',
                            default=self.KERNEL_EWIDTH, dest='KERNEL_EWIDTH',
                            help='The kernel e-width if --kernel=gaussian or'
                                 '--kernel=supergaussian')
        parser.add_argument('--kernelbeta', action='store',
                            default=self.KERNEL_BETA, dest='KERNEL_BETA',
                            help='The kernel beta value if '
                                 '--kernel=supergaussian')
        parser.add_argument('--infile', action='store',
                            default=self.IN_FILE, dest='IN_FILE',
                            help='The input e2dsff file (absolute path)')
        parser.add_argument('--blaze', action='store',
                            default=self.BLAZE_FILE, dest='BLAZE_FILE',
                            help='The input blaze file (absolute path)')
        parser.add_argument('--wave', action='store',
                            default=self.WAVE_FILE, dest='WAVE_FILE',
                            help='The input wave solution (absolute path)'
                                 'if one wishes to use the header set this to'
                                 '"None" or "header" i.e. --wave="None" or'
                                 ' --wave="header"')
        parser.add_argument('--mask', action='store', default=self.MASK_FILE,
                            dest='MASK_FILE',
                            help='The ccf mask file to use (absolute path)')
        parser.add_argument('--maskwidth', action='store', type=float,
                            default=self.MASK_WIDTH, dest='MASK_WIDTH',
                            help='The size of the mask lines to use')
        parser.add_argument('--maskminweight', action='store', type=float,
                            default=self.MASK_MIN_WEIGHT,
                            dest='MASK_MIN_WEIGHT',
                            help='The minimum line weighting to use. If set to'
                                 '1 forces all weights to be equal')
        parser.add_argument('--ccfstep', action='store', default=self.CCF_STEP,
                            dest='CCF_STEP', type=float,
                            help='CCF step size [km/s]')
        parser.add_argument('--ccfwidth', action='store',
                            default=self.CCF_WIDTH, dest='CCF_WIDTH',
                            help='CCF step width [km/s]')
        parser.add_argument('--ccf_rv_null', action='store',
                            default=self.CCF_RV_NULL, dest='CCF_RV_NULL',
                            help='The largest absolute RV value to accept from '
                                 'header. Above or below this values are '
                                 'rejected')
        parser.add_argument('--targetrv', action='store',
                            default=self.IN_RV, dest='IN_RV',
                            help='The input target RV [km/s]')
        parser.add_argument('--ccf_ord_max', action='store',
                            default=self.CCF_N_ORD_MAX, dest='CCF_N_ORD_MAX',
                            help='The reddest order to use (bluest = 0)')
        parser.add_argument('--blazenormper', action='store',
                            default=self.BLAZE_NORM_PERCENTILE,
                            dest='BLAZE_NORM_PERCENTILE',
                            help='The blaze percentile to normalise by [0-100]')
        parser.add_argument('--blaze_thres', action='store',
                            default=self.BLAZE_THRESHOLD,
                            dest='BLAZE_THRESHOLD',
                            help='The blaze threshold to cut at '
                                 '(below this blaze level flux is ignored)')
        parser.add_argument('--noise_sigdet', action='store',
                            default=self.NOISE_SIGDET, dest='NOISE_SIGDET',
                            help='The noise level used to calculate dv rms '
                                 'and ccf snr')
        parser.add_argument('--noise_size', action='store',
                            default=self.NOISE_SIZE, dest='NOISE_SIZE',
                            help='The size around saturated pixels to flag as '
                                 'unusable for dv rms')
        parser.add_argument('--noise_thres', action='store',
                            default=self.NOISE_THRES, dest='NOISE_THRES',
                            help='The maximum flux for a good unsaturated pixel'
                                 'for dv rms')

        # parse arguments
        self.cmdargs = parser.parse_args()
        # update case (as other parameters are defined by default from this)
        if self.cmdargs.CASE is not None:
            self.CASE = self.cmdargs.CASE

    def args_from_cmd(self):
        """
        Push commands from arguments into constants in self
        i.e. access through Argument.VARIABLE

        :return:
        """
        if self.cmdargs is None:
            return
        # update self with args
        for key in self.__dict__:
            if hasattr(self.cmdargs, key):
                argval = getattr(self.cmdargs, key)
                if argval is not None:
                    setattr(self, key, argval)

    def kernel_args(self):
        # get the kernel name
        name = self.KERNEL.lower()
        # deal with kernel cases
        if name in [None, 'None', '']:
            self.kernelparams = None
        elif 'box' in name:
            # get the width
            width = self.KERNEL_WIDTH
            if width is None:
                raise ValueError('ERROR: kernel width must be set')
            try:
                width = float(width)
            except:
                raise ValueError('ERROR: kernel width must be a valid float')
            # set kernel params
            self.kernelparams = ['boxcar', width]
        elif ('gauss' in name) and ('super' not in name):
            # get the e-width
            ewidth = self.KERNEL_EWIDTH
            if ewidth is None:
                raise ValueError('ERROR: kernel e-width must be set')
            try:
                ewidth = float(ewidth)
            except:
                raise ValueError('ERROR: kernel e-width must be a valid float')
            # set kernel params
            self.kernelparams = ['gaussian', ewidth]
        elif ('gauss' in name):
            # get the e-width
            ewidth = self.KERNEL_EWIDTH
            if ewidth is None:
                raise ValueError('ERROR: kernel e-width must be set')
            try:
                ewidth = float(ewidth)
            except:
                raise ValueError('ERROR: kernel e-width must be a valid float')
            # get the beta value
            beta = self.KERNEL_BETA
            if beta is None:
                raise ValueError('ERROR: kernel beta must be set')
            try:
                beta = float(beta)
            except:
                raise ValueError('ERROR: kernel beta must be a valid float')
            # set kernel params
            self.kernelparams = ['supergaussian', ewidth, beta]
        else:
            self.kernelparams = None

def get_combinations(names, variables, newdirs, outdir):
    combinations = []
    # generate time string
    timestr = Time.now().iso.replace(' ', '_')
    # get all combinations
    combs = list(itertools.product(*variables))
    # loop around combinations and create combination
    for c_it, comb in enumerate(combs):
        cargs = [c_it, comb, names, newdirs, len(combs), outdir, timestr]
        combinations.append(Combination(*cargs))
    # return all combination instances
    return combinations

class Combination:
    def __init__(self, number, variables, names, newdirs, total, outpath,
                 timestr):
        # get some stats (for ease of use)
        self.number = number
        self.total = total
        # get output directory
        self.outdir = str(timestr)
        # setup parameter dictionary
        self.parameters = dict()
        # append parameters and outdir
        for n_it, name in enumerate(names):
            value = variables[n_it]
            self.parameters[name] = value
            # deal with paths
            if os.sep in value:
                strvalue = os.path.basename(value)
                strvalue.split('.')[0]
            # add to the outdir
            if newdirs[n_it]:
                self.outdir += '_{0}={1}'.format(name, strvalue)
        # set output dir
        self.outpath = os.path.join(outpath, self.outdir)
        # need to make the directory
        if not os.path.exists(self.outpath):
            print('Creating directory {0}'.format(self.outpath))
            os.mkdir(self.outpath)

# =============================================================================
# Reading functions
# =============================================================================
def read_mask(mask_file, mask_cols):
    table = Table.read(mask_file, format='ascii')
    # get column names
    oldcols = list(table.colnames)
    # rename columns
    for c_it, col in enumerate(mask_cols):
        table[oldcols[c_it]].name = col
    # return table
    return table

def read_wave(image, iheader, wavefile):
    if wavefile not in [None, 'None', '', 'False', False, 'header']:
        # load the wave file and return it
        wavemap, waveheader = fits.getdata(wavefile, header=True)
        # return the wave map and the wave header
        return wavemap, waveheader, wavefile
    else:
        # get image dimensions
        nbo, nbx = image.shape
        deg = iheader['WAVEDEGN']
        # populate the wave coefficients
        wave_coeffs = np.zeros((nbo, deg + 1))
        for it in range(nbo):
            for jt in range(deg + 1):
                kt = it * (deg + 1 + jt)
                wave_coeffs[it][jt] = iheader['WAVE{0:04d}'.format(kt)]
        # set up storage
        wavemap = np.zeros((nbo, nbx))
        xpixels = np.arange(nbx)
        # loop aroun each order and make the wave map
        for order_num in range(nbo):
            # get this order coefficients
            ocoeffs = wave_coeffs[order_num][::-1]
            # calculate polynomial values and push into wavemap
            wavemap[order_num] = np.polyval(ocoeffs, xpixels)
        # get wave keys from header
        waveheader = fits.Header()
        waveheader['MJDMID'] = iheader['WAVETIME']
        # return the wave map and the wave header
        return wavemap, waveheader, iheader['WAVEFILE']

def get_mask(table, mask_width, mask_min, mask_units='nm'):
    ll_mask_e = np.array(table['ll_mask_e']).astype(float)
    ll_mask_s = np.array(table['ll_mask_s']).astype(float)
    ll_mask_d = ll_mask_e - ll_mask_s
    ll_mask_ctr = ll_mask_s + ll_mask_d * 0.5
    # if mask_width > 0 ll_mask_d is multiplied by mask_width/c
    if mask_width > 0:
        ll_mask_d = mask_width * ll_mask_s / SPEED_OF_LIGHT
    # make w_mask an array
    w_mask = np.array(table['w_mask']).astype(float)
    # use w_min to select on w_mask or keep all if w_mask_min >= 1
    if mask_min < 1.0:
        mask = w_mask > mask_min
        ll_mask_d = ll_mask_d[mask]
        ll_mask_ctr = ll_mask_ctr[mask]
        w_mask = w_mask[mask]
    # else set all w_mask to one (and use all lines in file)
    else:
        w_mask = np.ones(len(ll_mask_d))
    # ----------------------------------------------------------------------
    # deal with the units of ll_mask_d and ll_mask_ctr
    # must be returned in nanometers
    # ----------------------------------------------------------------------
    # get unit object from mask units string
    unit = getattr(uu, mask_units)
    # add units
    ll_mask_d = ll_mask_d * unit
    ll_mask_ctr = ll_mask_ctr * unit
    # convert to nanometers
    ll_mask_d = ll_mask_d.to(uu.nm).value
    ll_mask_ctr = ll_mask_ctr.to(uu.nm).value
    # ----------------------------------------------------------------------
    # return the size of each pixel, the central point of each pixel
    #    and the weight mask
    return ll_mask_d, ll_mask_ctr, w_mask

# =============================================================================
# Math functions
# =============================================================================
def relativistic_waveshift(dv, units='km/s'):
    """
    Relativistic offset in wavelength

    default is dv in km/s

    :param dv: float or numpy array, the dv values
    :param units: string or astropy units, the units of dv
    :return:
    """
    # get c in correct units
    # noinspection PyUnresolvedReferences
    if units == 'km/s' or units == uu.km / uu.s:
        c = SPEED_OF_LIGHT
    # noinspection PyUnresolvedReferences
    elif units == 'm/s' or units == uu.m / uu.s:
        c = SPEED_OF_LIGHT * 1000
    else:
        raise ValueError("Wrong units for dv ({0})".format(units))
    # work out correction
    corrv = np.sqrt((1 + dv / c) / (1 - dv / c))
    # return correction
    return corrv

def iuv_spline(x, y, **kwargs):
    # copy x and y
    x, y = np.array(x), np.array(y)
    # find all NaN values
    nanmask = ~np.isfinite(y)

    if np.sum(~nanmask) < 2:
        y = np.zeros_like(x)
    elif np.sum(nanmask) == 0:
        pass
    else:
        # replace all NaN's with linear interpolation
        badspline = InterpolatedUnivariateSpline(x[~nanmask], y[~nanmask],
                                                 k=1, ext=1)
        y[nanmask] = badspline(x[nanmask])
    # return spline
    return InterpolatedUnivariateSpline(x, y, **kwargs)

def gauss_function(x, a, x0, sigma, dc):
    """
    A standard 1D gaussian function (for fitting against)]=

    :param x: numpy array (1D), the x data points
    :param a: float, the amplitude
    :param x0: float, the mean of the gaussian
    :param sigma: float, the standard deviation (FWHM) of the gaussian
    :param dc: float, the constant level below the gaussian

    :return gauss: numpy array (1D), size = len(x), the output gaussian
    """
    return a * np.exp(-0.5 * ((x - x0) / sigma) ** 2) + dc

def fwhm(sigma=1.0):
    """
    Get the Full-width-half-maximum value from the sigma value (~2.3548)

    :param sigma: float, the sigma, default value is 1.0 (normalised gaussian)
    :return: 2 * sqrt(2 * log(2)) * sigma = 2.3548200450309493 * sigma
    """
    return 2 * np.sqrt(2 * np.log(2)) * sigma

# =============================================================================
# Plot functions
# =============================================================================
def plotloop(looplist):
    # check that looplist is a valid list
    if not isinstance(looplist, list):
        # noinspection PyBroadException
        try:
            looplist = list(looplist)
        except Exception as _:
            print('PLOT ERROR: looplist must be a list')
    # define message to give to user
    message = ('Plot loop navigation: Go to \n\t [P]revious plot '
               '\n\t [N]ext plot \n\t [E]nd plotting '
               '\n\t Number from [0 to {0}]: \t')
    message = message.format(len(looplist) - 1)
    # start the iterator at zero
    it = 0
    first = True
    # loop around until we hit the length of the loop list
    while it < len(looplist):
        # deal with end of looplist
        if it == len(looplist):
            # break out of while
            break
        # if this is the first iteration do not print message
        if first:
            # yield the first iteration value
            yield looplist[it]
            # increase the iterator value
            it += 1
            first = False
        # else we need to ask to go to previous, next or end
        else:
            # get user input
            userinput = input(message)
            # try to cast into a integer
            # noinspection PyBroadException
            try:
                userinput = int(userinput)
            except Exception as _:
                userinput = str(userinput)
            # if 'p' in user input we assume they want to go to previous
            if 'P' in str(userinput).upper():
                yield looplist[it - 1]
                it -= 1
            # if 'n' in user input we assume they want to go to next
            elif 'N' in str(userinput).upper():
                yield looplist[it + 1]
                it += 1
            elif isinstance(userinput, int):
                it = userinput
                # deal with it too low
                if it < 0:
                    it = 0
                # deal with it too large
                elif it >= len(looplist):
                    it = len(looplist) - 1
                # yield the value of it
                yield looplist[it]
            # else we assume the loop is over and we want to exit
            else:
                # break out of while
                break

def plot_individual_ccf(props, nbo):
    # get the plot loop generator
    generator = plotloop(range(nbo))
    # loop around orders
    for order_num in generator:
        plt.close()
        fig, frame = plt.subplots(ncols=1, nrows=1)
        frame.plot(props['RV_CCF'], props['CCF'][order_num], color='b',
                   marker='+', ls='None', label='data')
        frame.plot(props['RV_CCF'], props['CCF_FIT'][order_num], color='r', )
        rvorder = props['CCF_FIT_COEFFS'][order_num][1]
        frame.set(title='Order {0}  RV = {1} km/s'.format(order_num, rvorder),
                  xlabel='RV [km/s]', ylabel='CCF')
        plt.show()
        plt.close()

def plot_mean_ccf(props):
    plt.close()
    fig, frame = plt.subplots(ncols=1, nrows=1)
    frame.plot(props['RV_CCF'], props['MEAN_CCF'], color='b', marker='+',
               ls='None')
    frame.plot(props['RV_CCF'], props['MEAN_CCF_FIT'], color='r')
    frame.set(title='Mean CCF   RV = {0} km/s'.format(props['MEAN_RV']),
              xlabel='RV [km/s]', ylabel='CCF')
    plt.show()
    plt.close()

# =============================================================================
# Writing functions
# =============================================================================
def write_file(props, infile, maskname, header, wheader, wave_file,
               outpath=None):
    # ----------------------------------------------------------------------
    # construct out file name
    inbasename = os.path.basename(infile).split('.')[0]
    maskbasename = os.path.basename(maskname).split('.')[0]
    inpath = os.path.dirname(infile)
    outfile = 'CCFTABLE_{0}_{1}.fits'.format(inbasename, maskbasename)
    # deal with no outpath
    if outpath is None:
        outpath = os.path.join(inpath, outfile)
    # ----------------------------------------------------------------------
    # produce CCF table
    table1 = Table()
    table1['RV'] = props['RV_CCF']
    for order_num in range(len(props['CCF'])):
        table1['ORDER{0:02d}'.format(order_num)] = props['CCF'][order_num]
    table1['COMBINED'] = props['MEAN_CCF']
    # ----------------------------------------------------------------------
    # produce stats table
    table2 = Table()
    table2['ORDERS'] = np.arange(len(props['CCF'])).astype(int)
    table2['NLINES'] = props['CCF_LINES']
    # get the coefficients
    coeffs = props['CCF_FIT_COEFFS']
    table2['CONTRAST'] = np.abs(100 * coeffs[:, 0])
    table2['RV'] = coeffs[:, 1]
    table2['FWHM'] = coeffs[:, 2]
    table2['DC'] = coeffs[:, 3]
    table2['SNR'] = props['CCF_SNR']
    table2['NORM'] = props['CCF_NORM']

    # ----------------------------------------------------------------------
    # add to the header
    # ----------------------------------------------------------------------
    # add results from the CCF
    header['CCFMNRV'] = (props['MEAN_RV'],
                         'Mean RV calc. from the mean CCF [km/s]')
    header['CCFMCONT'] = (props['MEAN_CONTRAST'],
                          'Mean contrast (depth of fit) from mean CCF')
    header['CCFMFWHM'] = (props['MEAN_FWHM'],
                          'Mean FWHM from mean CCF')
    header['CCFMRVNS'] = (props['MEAN_RV_NOISE'],
                          'Mean RV Noise from mean CCF')
    header['CCFTLINE'] = (props['TOT_LINE'],
                          'Total no. of mask lines used in CCF')
    # ----------------------------------------------------------------------
    # add constants used to process
    header['CCFMASK'] = (props['CCF_MASK'], 'CCF mask file used')
    header['CCFSTEP'] = (props['CCF_STEP'], 'CCF step used [km/s]')
    header['CCFWIDTH'] = (props['CCF_WIDTH'], 'CCF width used [km/s]')
    header['CCFTRGRV'] = (props['TARGET_RV'],
                          'CCF central RV used in CCF [km/s]')
    header['CCFSIGDT'] = (props['CCF_SIGDET'],
                          'Read noise used in photon noise calc. in CCF')
    header['CCFBOXSZ'] = (props['CCF_BOXSIZE'],
                          'Size of bad px used in photon noise calc. in CCF')
    header['CCFMAXFX'] = (props['CCF_MAXFLUX'],
                          'Flux thres for bad px in photon noise calc. in CCF')
    header['CCFORDMX'] = (props['CCF_NMAX'],
                          'Last order used in mean for mean CCF')
    header['CCFMSKMN'] = (props['MASK_MIN'],
                          'Minimum weight of lines used in the CCF mask')
    header['CCFMSKWD'] = (props['MASK_WIDTH'],
                          'Width of lines used in the CCF mask')
    header['CCFMUNIT'] = (props['MASK_UNITS'], 'Units used in CCF Mask')
    # ----------------------------------------------------------------------
    header['RV_WAVFN'] = (os.path.basename(wave_file), 'RV wave file used')
    header['RV_WAVTM'] = (wheader['MJDMID'],
                          'RV wave file time [mjd]')
    header['RV_WAVTD'] = (header['MJDMID'] - wheader['MJDMID'],
                          'RV timediff [days] btwn file and wave solution')
    header['RV_WAVFP'] = ('None', 'RV measured from wave sol FP CCF [km/s]')
    header['RV_SIMFP'] = ('None', 'RV measured from simultaneous FP CCF [km/s]')
    header['RV_DRIFT'] = ('None',
                          'RV drift between wave sol and sim. FP CCF [km/s]')
    header['RV_OBJ'] = (props['MEAN_RV'],
                        'RV calc in the object CCF (non corr.) [km/s]')
    header['RV_CORR'] = ('None', 'RV corrected for FP CCF drift [km/s]')
    # ----------------------------------------------------------------------
    # log where we are writing the file to
    print('Writing file to {0}'.format(outpath))
    # construct hdus
    hdu = fits.PrimaryHDU()
    t1 = fits.BinTableHDU(table1, header=header)
    t2 = fits.BinTableHDU(table2, header=header)
    # construct hdu list
    hdulist = fits.HDUList([hdu, t1, t2])
    # write hdulist
    hdulist.writeto(outpath, overwrite=True)


from __future__ import print_function

import numpy as np

from astropy import units as u
from glue.external import six
from astropy.io.fits import PrimaryHDU, ImageHDU, Header

from .utils.wcs_utils import get_spatial_scale, sanitize_wcs
from .geometry import extract_slice
from .geometry import path as paths
from .utils.wcs_slicing import slice_wcs

def extract_pv_slice(cube, path, wcs=None, spacing=1.0, order=3,
                     respect_nan=True):
    """
    Given a position-position-velocity cube with dimensions (nv, ny, nx), and
    a path, extract a position-velocity slice.

    Alternative implementations:
        gipsy::sliceview
        karma::kpvslice
        casaviewer::slice

    Parameters
    ----------
    cube : :class:`~numpy.ndarray` or :class:`~spectral_cube.SpectralCube` or str or HDU
        The cube to extract a slice from. If this is a plain
        :class:`~numpy.ndarray` instance, the WCS information can optionally
        be specified with the ``wcs`` parameter. If a string, it should be
        the name of a file containing a spectral cube.
    path : `Path` or list of 2-tuples
        The path along which to define the position-velocity slice. The path
        can contain coordinates defined in pixel or world coordinates.
    wcs : :class:`~astropy.wcs.WCS`, optional
        The WCS information to use for the cube. This should only be
        specified if the ``cube`` parameter is a plain
        :class:`~numpy.ndarray` instance.
    spacing : float
        The position resolution in the final position-velocity slice. This
        can be given in pixel coordinates or as a
        :class:`~astropy.units.Quantity` instance with angle units.
    order : int, optional
        Spline interpolation order when using paths with zero width. Does not
        have any effect for paths with a non-zero width.
    respect_nan : bool, optional
        If set to `False`, NaN values are changed to zero before computing
        the slices. If set to `True`, in the case of line paths a second
        computation is performed to ignore the NaN value while interpolating,
        and set the output values of NaNs to NaN.

    Returns
    -------
    slice : `PrimaryHDU`
        The position-velocity slice, as a FITS HDU object
    """

    if isinstance(cube, (six.string_types, ImageHDU, PrimaryHDU)):
        try:
            from spectral_cube import SpectralCube
            cube = SpectralCube.read(cube)
        except ImportError:
            raise ImportError("spectral_cube package required for working "
                              "with fits data. Install spectral_cube or "
                              "use NumPy arrays")

    if _is_spectral_cube(cube):
        wcs = cube.wcs
        # The fits HEADER will preserve the UNIT, but pvextractor does not care
        # what the flux units are
        cube = cube.filled_data[...].value

    if wcs is not None:
        wcs = sanitize_wcs(wcs)

    if not isinstance(cube, np.ndarray) or wcs is not None:
        scale = get_spatial_scale(wcs)
        if isinstance(spacing, u.Quantity):
            pixel_spacing = (spacing / scale).decompose()
            world_spacing = spacing
        else:
            pixel_spacing = spacing
            world_spacing = spacing * scale
    else:
        if isinstance(spacing, u.Quantity):
            raise TypeError("No WCS has been specified, so spacing should be given in pixels")
        else:
            pixel_spacing = spacing
            world_spacing = None

    # Allow path to be passed in as list of 2-tuples
    if not isinstance(path, paths.Path):
        path = paths.Path(path)

    pv_slice = extract_slice(cube, path, wcs=wcs, spacing=pixel_spacing,
                             order=order, respect_nan=respect_nan)

    # Generate output header
    if wcs is None:
        header = Header()
    else:
        header = slice_wcs(wcs, spatial_scale=world_spacing).to_header()

    # TODO: write path to BinTableHDU

    return PrimaryHDU(data=pv_slice, header=header)

def _is_spectral_cube(obj):
    try:
        from spectral_cube import SpectralCube
        return isinstance(obj, SpectralCube)
    except ImportError:
        return False


import copy
import cmath
import h5py
import math
import numpy
import scipy.linalg
import sys
import time
from pauxy.walkers.multi_ghf import MultiGHFWalker
from pauxy.walkers.single_det import SingleDetWalker
from pauxy.walkers.multi_det import MultiDetWalker
from pauxy.walkers.multi_coherent import MultiCoherentWalker
from pauxy.walkers.thermal import ThermalWalker
from pauxy.walkers.stack import FieldConfig
from pauxy.utils.io import get_input_value
from pauxy.utils.misc import update_stack

class Walkers(object):
    """Container for groups of walkers which make up a wavefunction.

    Parameters
    ----------
    system : object
        System object.
    trial : object
        Trial wavefunction object.
    nwalkers : int
        Number of walkers to initialise.
    nprop_tot : int
        Total number of propagators to store for back propagation + itcf.
    nbp : int
        Number of back propagation steps.
    """

    def __init__(self, system, trial, qmc, walker_opts={}, verbose=False,
                 comm=None, nprop_tot=None, nbp=None):
        self.nwalkers = qmc.nwalkers
        self.ntot_walkers = qmc.ntot_walkers
        if verbose:
            print("# nwalkers = {}".format(self.nwalkers))
            print("# ntot_walkers = {}".format(self.ntot_walkers))
        self.write_freq = walker_opts.get('write_freq', 0)
        self.write_file = walker_opts.get('write_file', 'restart.h5')
        self.use_log_shift = walker_opts.get('use_log_shift', False)
        self.shift_counter = 1
        self.read_file = walker_opts.get('read_file', None)
        if comm is None:
            rank = 0
        else:
            rank = comm.rank
        if verbose:
            print("# Setting up wavefunction object.")
        if trial.name == 'MultiSlater':
            self.walker_type = 'MSD'
            # TODO: FDM FIXTHIS
            if trial.ndets == 1:
                if verbose:
                    print("# Usinge single det walker with msd wavefunction.")
                self.walker_type = 'SD'
                trial.psi = trial.psi[0]
                self.walkers = [SingleDetWalker(system, trial, walker_opts=walker_opts,
                                                index=w, nprop_tot=nprop_tot,
                                                nbp=nbp)
                                for w in range(qmc.nwalkers)]
            else:
                self.walkers = [
                        MultiDetWalker(system, trial, walker_opts=walker_opts,
                                       verbose=(verbose and w == 0))
                        for w in range(qmc.nwalkers)
                        ]
            self.buff_size = self.walkers[0].buff_size
            if nbp is not None:
                self.buff_size += self.walkers[0].field_configs.buff_size
            self.walker_buffer = numpy.zeros(self.buff_size,
                                             dtype=numpy.complex128)
        elif trial.name == 'thermal':
            self.walker_type = 'thermal'
            self.walkers = [ThermalWalker(system, trial,
                                          walker_opts=walker_opts,
                                          verbose=(verbose and w==0))
                            for w in range(qmc.nwalkers)]
            self.buff_size = self.walkers[0].buff_size + self.walkers[0].stack.buff_size
            self.walker_buffer = numpy.zeros(self.buff_size,
                                             dtype=numpy.complex128)
            stack_size = self.walkers[0].stack_size
            if system.name == "Hubbard":
                if stack_size % qmc.nstblz != 0 or qmc.nstblz < stack_size:
                    if verbose:
                        print("# Stabilisation frequency is not commensurate "
                              "with stack size.")
                        print("# Determining a better value.")
                    if qmc.nstblz < stack_size:
                        qmc.nstblz = stack_size
                        if verbose:
                            print("# Updated stabilization frequency: "
                                  " {}".format(qmc.nstblz))
                    else:
                        qmc.nstblz = update_stack(qmc.nstblz, stack_size,
                                                  name="nstblz", verbose=verbose)
        elif trial.name == "coherent_state" and trial.symmetrize:
            self.walker_type = 'MSD'
            self.walkers = [MultiCoherentWalker(system, trial, walker_opts=walker_opts,
                                        index=w, nprop_tot=nprop_tot,
                                        nbp=nbp)
                        for w in range(qmc.nwalkers)]
            self.buff_size = self.walkers[0].buff_size
            if nbp is not None:
                if verbose:
                    print("# Performing back propagation.")
                    print("# Number of steps in imaginary time: {:}.".format(nbp))
                self.buff_size += self.walkers[0].field_configs.buff_size
            self.walker_buffer = numpy.zeros(self.buff_size,
                                             dtype=numpy.complex128)
        else:
            self.walker_type = 'SD'
            self.walkers = [SingleDetWalker(system, trial, walker_opts=walker_opts,
                                            index=w, nprop_tot=nprop_tot,
                                            nbp=nbp)
                            for w in range(qmc.nwalkers)]
            self.buff_size = self.walkers[0].buff_size
            if nbp is not None:
                if verbose:
                    print("# Performing back propagation.")
                    print("# Number of steps in imaginary time: {:}.".format(nbp))
                self.buff_size += self.walkers[0].field_configs.buff_size
            self.walker_buffer = numpy.zeros(self.buff_size,
                                             dtype=numpy.complex128)
        if system.name == "Generic" or system.name == "UEG":
            dtype = complex
        else:
            dtype = int
        self.pcont_method = get_input_value(walker_opts, 'population_control',
                                            default='comb')
        self.min_weight = walker_opts.get('min_weight', 0.1)
        self.max_weight = walker_opts.get('max_weight', 4.0)
        if verbose:
            print("# Using {} population control "
                  "algorithm.".format(self.pcont_method))
            mem = float(self.walker_buffer.nbytes) / (1024.0**3)
            print("# Buffer size for communication: {:13.8e} GB".format(mem))
            if mem > 2.0:
                # TODO: FDM FIX THIS
                print(" # Warning: Walker buffer size > 2GB. May run into MPI"
                      "issues.")
        if not self.walker_type == "thermal":
            walker_size = 3 + self.walkers[0].phi.size
        if self.write_freq > 0:
            self.write_restart = True
            self.dsets = []
            with h5py.File(self.write_file,'w',driver='mpio',comm=comm) as fh5:
                for i in range(self.ntot_walkers):
                    fh5.create_dataset('walker_%d'%i, (walker_size,),
                                       dtype=numpy.complex128)

        else:
            self.write_restart = False
        if self.read_file is not None:
            if verbose:
                print("# Reading walkers from %s file series."%self.read_file)
            self.read_walkers(comm)
        self.target_weight = qmc.ntot_walkers
        self.nw = qmc.nwalkers
        self.set_total_weight(qmc.ntot_walkers)

    def orthogonalise(self, trial, free_projection):
        """Orthogonalise all walkers.

        Parameters
        ----------
        trial : object
            Trial wavefunction object.
        free_projection : bool
            True if doing free projection.
        """
        for w in self.walkers:
            detR = w.reortho(trial)
            if free_projection:
                (magn, dtheta) = cmath.polar(detR)
                w.weight *= magn
                w.phase *= cmath.exp(1j*dtheta)

    def add_field_config(self, nprop_tot, nbp, system, dtype):
        """Add FieldConfig object to walker object.

        Parameters
        ----------
        nprop_tot : int
            Total number of propagators to store for back propagation + itcf.
        nbp : int
            Number of back propagation steps.
        nfields : int
            Number of fields to store for each back propagation step.
        dtype : type
            Field configuration type.
        """
        for w in self.walkers:
            w.field_configs = FieldConfig(system.nfields, nprop_tot, nbp, dtype)

    def copy_historic_wfn(self):
        """Copy current wavefunction to psi_n for next back propagation step."""
        for (i,w) in enumerate(self.walkers):
            numpy.copyto(self.walkers[i].phi_old, self.walkers[i].phi)

    def copy_bp_wfn(self, phi_bp):
        """Copy back propagated wavefunction.

        Parameters
        ----------
        phi_bp : object
            list of walker objects containing back propagated walkers.
        """
        for (i, (w,wbp)) in enumerate(zip(self.walkers, phi_bp)):
            numpy.copyto(self.walkers[i].phi_bp, wbp.phi)

    def copy_init_wfn(self):
        """Copy current wavefunction to initial wavefunction.

        The definition of the initial wavefunction depends on whether we are
        calculating an ITCF or not.
        """
        for (i,w) in enumerate(self.walkers):
            numpy.copyto(self.walkers[i].phi_right, self.walkers[i].phi)

    def pop_control(self, comm):
        if self.ntot_walkers == 1:
            return
        if self.use_log_shift:
           self.update_log_ovlp(comm)
        weights = numpy.array([abs(w.weight) for w in self.walkers])
        global_weights = numpy.empty(len(weights)*comm.size)
        comm.Allgather(weights, global_weights)
        total_weight = sum(global_weights)
        # Rescale weights to combat exponential decay/growth.
        scale = total_weight / self.target_weight
        if total_weight < 1e-8:
            if comm.rank == 0:
                print("# Warning: Total weight is {:13.8e}: "
                      .format(total_weight))
                print("# Something is seriously wrong.")
            sys.exit()
        self.set_total_weight(total_weight)
        # Todo: Just standardise information we want to send between routines.
        for w in self.walkers:
            w.unscaled_weight = w.weight
            w.weight = w.weight / scale
        if self.pcont_method == "comb":
            global_weights = global_weights / scale
            self.comb(comm, global_weights)
        elif self.pcont_method == "pair_branch":
            self.pair_branch(comm)
        else:
            if comm.rank == 0:
                print("Unknown population control method.")

    def comb(self, comm, weights):
        """Apply the comb method of population control / branching.

        See Booth & Gubernatis PRE 80, 046704 (2009).

        Parameters
        ----------
        comm : MPI communicator
        """
        # Need make a copy to since the elements in psi are only references to
        # walker objects in memory. We don't want future changes in a given
        # element of psi having unintended consequences.
        # todo : add phase to walker for free projection
        if comm.rank == 0:
            parent_ix = numpy.zeros(len(weights), dtype='i')
        else:
            parent_ix = numpy.empty(len(weights), dtype='i')
        if comm.rank == 0:
            total_weight = sum(weights)
            cprobs = numpy.cumsum(weights)
            r = numpy.random.random()
            comb = [(i+r) * (total_weight/self.target_weight) for i in
                    range(self.target_weight)]
            iw = 0
            ic = 0
            while ic < len(comb):
                if comb[ic] < cprobs[iw]:
                    parent_ix[iw] += 1
                    ic += 1
                else:
                    iw += 1
            data = {'ix': parent_ix}
        else:
            data = None

        data = comm.bcast(data, root=0)
        parent_ix = data['ix']
        # Keep total weight saved for capping purposes.
        # where returns a tuple (array,), selecting first element.
        kill = numpy.where(parent_ix == 0)[0]
        clone = numpy.where(parent_ix > 1)[0]
        reqs = []
        walker_buffers = []
        # First initiate non-blocking sends of walkers.
        comm.barrier()
        for i, (c, k) in enumerate(zip(clone, kill)):
            # Sending from current processor?
            if c // self.nw == comm.rank:
                # Location of walker to clone in local list.
                clone_pos = c % self.nw
                # copying walker data to intermediate buffer to avoid issues
                # with accessing walker data during send. Might not be
                # necessary.
                dest_proc = k // self.nw
                # with h5py.File('before_{}.h5'.format(comm.rank), 'a') as fh5:
                    # fh5['walker_{}_{}_{}'.format(c,k,dest_proc)] = self.walkers[clone_pos].get_buffer()
                buff = self.walkers[clone_pos].get_buffer()
                reqs.append(comm.Isend(buff, dest=dest_proc, tag=i))
        # Now receive walkers on processors where walkers are to be killed.
        for i, (c, k) in enumerate(zip(clone, kill)):
            # Receiving to current processor?
            if k // self.nw == comm.rank:
                # Processor we are receiving from.
                source_proc = c // self.nw
                # Location of walker to kill in local list of walkers.
                kill_pos = k % self.nw
                comm.Recv(self.walker_buffer, source=source_proc, tag=i)
                # with h5py.File('walkers_recv.h5', 'w') as fh5:
                    # fh5['walk_{}'.format(k)] = self.walker_buffer.copy()
                self.walkers[kill_pos].set_buffer(self.walker_buffer)
                # with h5py.File('after_{}.h5'.format(comm.rank), 'a') as fh5:
                    # fh5['walker_{}_{}_{}'.format(c,k,comm.rank)] = self.walkers[kill_pos].get_buffer()
        # Complete non-blocking send.
        for rs in reqs:
            rs.wait()
        # Necessary?
        # if len(kill) > 0 or len(clone) > 0:
            # sys.exit()
        comm.Barrier()
        # Reset walker weight.
        # TODO: check this.
        for w in self.walkers:
            w.weight = 1.0

    def pair_branch(self, comm):
        walker_info = [[abs(w.weight),1,comm.rank,comm.rank] for w in self.walkers]
        glob_inf = comm.gather(walker_info, root=0)
        # Want same random number seed used on all processors
        if comm.rank == 0:
            # Rescale weights.
            glob_inf = numpy.array([item for sub in glob_inf for item in sub])
            total_weight = sum(w[0] for w in glob_inf)
            sort = numpy.argsort(glob_inf[:,0], kind='mergesort')
            isort = numpy.argsort(sort, kind='mergesort')
            glob_inf = glob_inf[sort]
            s = 0
            e = len(glob_inf) - 1
            tags = []
            isend = 0
            while s < e:
                if glob_inf[s][0] < self.min_weight or glob_inf[e][0] > self.max_weight:
                    # sum of paired walker weights
                    wab = glob_inf[s][0] + glob_inf[e][0]
                    r = numpy.random.rand()
                    if r < glob_inf[e][0] / wab:
                        # clone large weight walker
                        glob_inf[e][0] = 0.5 * wab
                        glob_inf[e][1] = 2
                        # Processor we will send duplicated walker to
                        glob_inf[e][3] = glob_inf[s][2]
                        send = glob_inf[s][2]
                        # Kill small weight walker
                        glob_inf[s][0] = 0.0
                        glob_inf[s][1] = 0
                        glob_inf[s][3] = glob_inf[e][2]
                    else:
                        # clone small weight walker
                        glob_inf[s][0] = 0.5 * wab
                        glob_inf[s][1] = 2
                        # Processor we will send duplicated walker to
                        glob_inf[s][3] = glob_inf[e][2]
                        send = glob_inf[e][2]
                        # Kill small weight walker
                        glob_inf[e][0] = 0.0
                        glob_inf[e][1] = 0
                        glob_inf[e][3] = glob_inf[s][2]
                    tags.append([send])
                    s += 1
                    e -= 1
                else:
                    break
            nw = self.nwalkers
            glob_inf = glob_inf[isort].reshape((comm.size,nw,4))
        else:
            data = None
            total_weight = 0
        data = comm.scatter(glob_inf, root=0)
        # Keep total weight saved for capping purposes.
        walker_buffers = []
        reqs = []
        for iw, walker in enumerate(data):
            if walker[1] > 1:
                tag = comm.rank*len(walker_info) + walker[3]
                self.walkers[iw].weight = walker[0]
                buff = self.walkers[iw].get_buffer()
                reqs.append(comm.Isend(buff,
                                       dest=int(round(walker[3])),
                                       tag=tag))
        for iw, walker in enumerate(data):
            if walker[1] == 0:
                tag = walker[3]*len(walker_info) + comm.rank
                comm.Recv(self.walker_buffer,
                          source=int(round(walker[3])),
                          tag=tag)
                self.walkers[iw].set_buffer(self.walker_buffer)
        for r in reqs:
            r.wait()

    def recompute_greens_function(self, trial, time_slice=None):
        for w in self.walkers:
            w.greens_function(trial, time_slice)

    def set_total_weight(self, total_weight):
        for w in self.walkers:
            w.total_weight = total_weight
            w.old_total_weight = w.total_weight

    def reset(self, trial):
        for w in self.walkers:
            w.stack.reset()
            w.stack.set_all(trial.dmat)
            w.greens_function(trial)
            w.weight = 1.0
            w.phase = 1.0 + 0.0j

    def get_write_buffer(self, i):
        w = self.walkers[i]
        buff = numpy.concatenate([[w.weight], [w.phase], [w.ot], w.phi.ravel()])
        return buff

    def set_walker_from_buffer(self, i, buff):
        w = self.walkers[i]
        w.weight = buff[0]
        w.phase = buff[1]
        w.ot = buff[2]
        w.phi = buff[3:].reshape(self.walkers[i].phi.shape)

    def write_walkers(self, comm):
        start = time.time()
        with h5py.File(self.write_file,'r+',driver='mpio',comm=comm) as fh5:
            for (i,w) in enumerate(self.walkers):
                ix = i + self.nwalkers*comm.rank
                buff = self.get_write_buffer(i)
                fh5['walker_%d'%ix][:] = self.get_write_buffer(i)
        if comm.rank == 0:
            print(" # Writing walkers to file.")
            print(" # Time to write restart: {:13.8e} s"
                  .format(time.time()-start))

    def update_log_ovlp(self, comm):
        send = numpy.zeros(3, dtype=numpy.complex128)
        # Overlap log factor
        send[0] = sum(abs(w.ot) for w in self.walkers)
        # Det R log factor
        send[1] = sum(abs(w.detR) for w in self.walkers)
        send[2] = sum(abs(w.log_detR) for w in self.walkers)
        global_av = numpy.zeros(3, dtype=numpy.complex128)
        comm.Allreduce(send, global_av)
        log_shift = numpy.log(global_av[0]/self.ntot_walkers)
        detR_shift = numpy.log(global_av[1]/self.ntot_walkers)
        log_detR_shift = global_av[2]/self.ntot_walkers
        # w.log_shift = -0.5
        n = self.shift_counter
        nm1 = self.shift_counter - 1
        for w in self.walkers:
            w.log_shift = (w.log_shift*nm1 + log_shift)/n
            w.log_detR_shift = (w.log_detR_shift*nm1 + log_detR_shift)/n
            w.detR_shift = (w.detR_shift*nm1 + detR_shift)/n
        self.shift_counter += 1

    def read_walkers(self, comm):
        with h5py.File(self.read_file, 'r') as fh5:
            for (i,w) in enumerate(self.walkers):
                try:
                    ix = i + self.nwalkers*comm.rank
                    self.set_walker_from_buffer(i, fh5['walker_%d'%ix][:])
                except KeyError:
                    print(" # Could not read walker data from:"
                          " %s"%(self.read_file))


#!/usr/bin/env python

import copy, math, numpy, os, pdb, sys
import chipseq_analysis, cmdlineProgs, random, selexDb, seq_utility, seqUtils, utility
import numpy

# SUMMARY: I have created this class to encapsulate motifs, for use in my
# chipseq qc work. Should possibly have done this earlier, as there are
# various motif-related functions etc. scattered throughout my python libraries.

# Introduced on June 18th 2013, for use in snv analysis. Should be useful
# elsewhere too:
def motifPrefixes2meme(outputFile, motifPrefixes, baseDir, suffix):
    for motifPrefix in motifPrefixes:
        currMotif = pwm(open(baseDir + motifPrefix + suffix))
        currMotif.setName(motifPrefix)
        currMotif.writeToMEME(outputFile)
        print >> outputFile, ""

# Introduced on June 18th 2013, for use in snv analysis. Should be useful
# elsewhere too:
def getTFsForMotifPrefix(motifPrefix, motifsCursor):
    """Retrieves all TFs corresponding to the specified motif prefix,
    by identity and sequence similarity, according to the motifs database."""

    # Get the TFs linked by gene ID:
    cmdStr = "select genes.Name from motifs.genes inner join motifs.tfMotifs on motifs.genes.geneID = motifs.tfMotifs.geneID where motifs.tfMotifs.motifPrefix = \"" + str(motifPrefix) + "\";"
    motifsCursor.execute(cmdStr)
    rows = motifsCursor.fetchall()

    tfsByGeneID = map(lambda row: row[0], rows)

    # Get the TFs linked by DBD sequence similarity:
    cmdStr = "select distinct g2.Name from motifs.genes inner join motifs.tfMotifs on motifs.genes.geneID = motifs.tfMotifs.geneID inner join geneProteins on genes.GeneID = geneProteins.GeneID inner join proteins on geneProteins.proteinID = proteins.proteinID inner join dbdBlocks on proteins.proteinID = dbdBlocks.proteinID inner join aaSeq on dbdBlocks.aaSeqID = aaSeq.seqID inner join aaSeqSim on aaSeq.seqID = aaSeqSim.seq1id inner join aaSeq as2 on aaSeqSim.seq2id = as2.seqID inner join dbdBlocks db2 on as2.seqID = db2.aaSeqID inner join proteins p2 on db2.proteinID = p2.proteinID inner join geneProteins gp2 on p2.proteinID = gp2.proteinID inner join genes g2 on gp2.geneID = g2.geneID where motifs.tfMotifs.motifPrefix = \"" + str(motifPrefix) + "\" and aaSeqSim.similarity >= 1;"
    motifsCursor.execute(cmdStr)
    rows = motifsCursor.fetchall()

    tfsBySeqSim = map(lambda row: row[0], rows)

    if (len(rows) > 0):
        tfName = rows[0][0]
    else:
        tfName = None # There are no chip-seq TFs for that motif

    # Fixed bug on July 29th 2013: Uniquify the resulting list before returning:
    tfsForMotifDict = {}
    for tf in tfsByGeneID + tfsBySeqSim:
        tfsForMotifDict[tf] = 1
    return tfsForMotifDict.keys()

def cons2countMatrix(consString, count):
    """Takes a consensus sequence and a count value as input, and generates a
    numpy matrix representing the equivelant pwm. Returns that matrix object
    (rather than a pwm object)."""

    # Set up the alphabet mapping dictionary, in order to map from letters to
    # indeces in the matrix:
    mappingDict = {'A':0, 'C':1, 'G':2, 'T':3, 'a':0, 'c':1, 'g':2, 't':3}
    aLen = 4

    # Generate a new zeros matrix as long as the alphabet and as wide as the
    # consensus string, to build the count matrix on:
    countMat = numpy.zeros((len(consString), aLen))

    # Set the count value in each column of the matrix, based on the consensus
    # sequence letter at that position, and the specified input count value:
    for seqIdx in range(len(consString)):
        lettIdx = mappingDict[consString[seqIdx]]
        countMat[seqIdx][lettIdx] = count
    return countMat

class pwm(object):
    """A single pwm motif model."""

    def __init__(self, motifData, dataType="memeFile"):
        self.name = None
        self.matrix = None
        self.memeFilePath = None
        self.eValue = None
        self.consensusCount = None # Optional, for count data such as selex
        if (dataType == "memeFile"):
            assert (isinstance(motifData, file))
            # Input data is a MEME file => Initialise accordingly:
            self.initFromMEME(motifData)
            if (self.matrix == []):
                # Matrix was still empty after trying to initialise from
                # MEME file -> Report this as an exception:
                raise ValueError("MEME file had no more motifs in it:" +
                                 motifData.name)
        if (dataType == "xxMotifFile"):
            assert (isinstance(motifData, file))
            # Input data is a xxMotif file => Initialise accordingly:
            self.initFrom_xxMotif(motifData)
            if (self.matrix == []):
                # Matrix was still empty after trying to initialise from
                # xxMotif file -> Report this as an exception:
                raise ValueError("xxMotif file had no more motifs in it:" +
                                 motifData.name)
        elif (dataType == "countsFile"):
            assert (isinstance(motifData, file))
            self.initFromCounts(motifData)
        elif (dataType == "freqMatrix"):
            self.initFromFreqMatrix(motifData)
        elif (dataType == "seqAln"):
            # Input data is a list of strings representing aligned DNA sequences
            # from which this pwm should be constructed:
            self.initFromAlign(motifData)
        elif (dataType == "iniMotifFile"):
            self.initFromInimotifFile(motifData)
        self.logoFilename = None

    def getScoreProf(self, dnaSeq, bgFreqs = [0.25,0.25,0.25,0.25,0.25]):
        """Introduced April 15th 2011.

        Calculates the motif LLR score contribution at each position in the
        specified DNA sequence. Returns the resulting LLR score contributions
        as an array. Sequence positions with the letter "N" result in a zero
        value contribution.

        NOTE: This only works/makes sense when a zero-order background
        model is used. Thus, the bgFreqs is assumed to be an array describing
        such a background model (for A, C, G, T and N)."""

        assert len(dnaSeq) == self.getWidth()

        # Generate a data structure mapping from letter to pwm (and background
        # model) index:
        lett2colIdx = {'A':0, 'C':1, 'G':2, 'T':3, 'N':4, \
                           'a':0, 'c':1, 'g':2, 't':3, 'n':4}

        # The array showing contributions of each letter to the total LLR:
        llrScoreContribs = []

        # Consider each position in the motif...
        for columnIdx in range(self.getWidth()):
            # Current column; bgFreqs[-1] gives the "N" frequency:
            motifColumn = self.getMatrix()[columnIdx] + [bgFreqs[-1]]

            # Get motif likelihood, bg likelihood, and then calculate llr:
            letter = dnaSeq[columnIdx]
            letterColIdx = lett2colIdx[letter]
            motifProb = motifColumn[letterColIdx]
            bgProb = bgFreqs[letterColIdx]
            llrScoreContribs.append(math.log(motifProb/bgProb, 10))

        return llrScoreContribs

    def getName(self):
        return self.name

    def setName(self, name):
        self.name = name

    def getIC_arr(self):
        """Returns an array storing the information content of each column of
        this pwm, in order of motif position."""
        matrix = self.matrix
        ICs = []
        for col in matrix:
            ICs.append(seq_utility.calc_IC(col))
        return ICs

    def getMatrix(self):
        return self.matrix

    def getWidth(self):
        return len(self.matrix)

    def getLogoFilename(self):
        return self.logoFilename

    def getMemeFilePath(self):
        return self.memeFilePath

    def getEValue(self):
        return self.eValue

    def trimLowIC(self, icThresh=0.5, copy=False):
        """Trims off low information-content flanking columns from the motif."""
        trimmedMotif = self.getMatrix()

        # Trim the leading low IC columns...
        while ((len(trimmedMotif) > 0) and
               (seq_utility.calc_IC(trimmedMotif[0]) < icThresh)):
            trimmedMotif = trimmedMotif[1:]

        # Trim the tailing low IC columns...
        while ((len(trimmedMotif) > 0) and
               (seq_utility.calc_IC(trimmedMotif[-1]) < icThresh)):
            trimmedMotif = trimmedMotif[:-1]
        if (not copy):
            # Editing the original motif => Set matrix:
            self.matrix = trimmedMotif
            return None
        else:
            # Making a copy matrix:
            copyMotif = pwm(trimmedMotif, dataType="freqMatrix")
            return copyMotif

    def makeSeqLogo(self, outFilePrefix, format="eps"):
        """Generates an eps sequence logo for this motif, writing it out to the
        specified filename."""

        self.logoFilename = outFilePrefix + "." + format

        tmpMemeFilename = \
            utility.makeTempFilename(outFilePrefix + "_tmp_MEME_file_ceqlogo",
                                     fileSuffix = ".meme")
        tmpMemeFile = open(tmpMemeFilename, 'w')
        self.writeToMEME(tmpMemeFile)
        print >> tmpMemeFile, ""
        tmpMemeFile.flush()
        tmpMemeFile.close()

        # FIXME: Currently, I assume the matrix file exists, and that it is
        # in MEME format. Will need to adapt this in the future when
        # dynamically-generated motifs are run instead:
        seqUtils.run_ceqlogo(tmpMemeFilename, outFilePrefix, format=format)
        cmdlineProgs.deleteFiles([tmpMemeFilename])

    def writeToMAST(self, outFile, pseudo=0.01):
        """Writes the motif out the specified filehandle in MAST format."""

        # The motif matrix data must have been set before this method can
        # be called:
        assert (self.matrix != None)

        lengthStr = str(len(self.matrix))

        hdrString = """MEME version 4.5

ALPHABET= ACGT

strands: + -

Background letter frequencies (from dataset with add-one prior applied):
A 0.25 C 0.25 G 0.25 T 0.25

MOTIF """
        extra = str(self.name) + "\nBL   MOTIF " + str(self.name) + " width= " + lengthStr + " seqs=0\n\nlog-odds matrix: alength= 4 w= " + lengthStr + "\n"
        hdrString = hdrString + extra

        outFile.write(hdrString)

        # Currently just assume uniform background model:
        for col in self.matrix:
            # Add pseudo-counts to the elements of the matrix...
            # FIXME: This could be done better. What pseudo-count
            # value is best? Shouldn't this be in a psfm2llr function?
            newCol = map(lambda prob: (prob+pseudo)/(1+(pseudo*4)), col)

            # From mast source code, it seems the log base is 10, although I'm
            # not 100% sure!:
            try:
                llrCol = map(lambda prob: math.log(prob/0.25, 10), newCol)
            except ValueError, e:
                print >> sys.stderr, "Invalid probability column:", newCol
                raise e
            outFile.write(reduce(lambda p1, p2: str(p1) + " " + str(p2),
                                 llrCol, "") + "\n")
        print >> outFile, ""

    def writeToTRANSFAC(self, outFile, nSeqs=1000, pseudo=0.01):
        """Writes the motif out the specified filehandle in TRANSFAC format."""

        # The motif matrix data must have been set before this method can
        # be called:
        assert (self.matrix != None)

        lengthStr = str(len(self.matrix))

        hdrString = "AC " + self.name + """
XX
TY Motif
ID """ + self.name + """
BF undef
P0\tA\tC\tG\tT\n"""

        outFile.write(hdrString)

        # Currently just assume uniform background model:
        colNum = 0
        for col in self.matrix:
            # Add pseudo-counts to the elements of the matrix...
            newCol = map(lambda prob: (prob+pseudo)/(1+(pseudo*4)), col)
            newColCounts = map(lambda freq: int(freq*1000), newCol)

            colNumStr = ("%2d" % colNum).replace(" ", "0")
            outFile.write(colNumStr +
                          reduce(lambda p1, p2: str(p1) + "\t" + str(p2),
                                 newColCounts, "") + "\n")
            colNum += 1

        print >> outFile, "XX\n//"

    def addPseudoCounts(self, pseudo=0.01):
        """Modifies the psfm for this motif, by adding a specified pseudo count
        at each position."""

        newMatrix = []
        for col in self.matrix:
            newCol = map(lambda prob: (prob+pseudo)/(1+(pseudo*4)), col)
            newMatrix.append(newCol)
        self.matrix = newMatrix

    def writeCountMatrix(self, outFile, nSeqs=100, pseudo=0.01):
        """Writes the motif out as a count matrix, using the specified
        number of sequences."""
        assert (self.matrix != None)

        hdrString = ">" + self.getName()
        outFile.write(hdrString + "\n")
        for col in self.matrix:
            probCol = map(lambda prob: (prob+pseudo)/(1+(pseudo*4)), col)
            countCol = map(lambda prob: int(prob*nSeqs), probCol)
            outFile.write(reduce(lambda p1, p2: str(p1) + " " + str(p2),
                                 countCol, "") + "\n")

    def writeFreqMatrix(self, outFile, nSeqs=100, pseudo=0.01):
        """Writes the motif out as a frequency matrix."""
        assert (self.matrix != None)

        hdrString = ">" + self.getName()
        outFile.write(hdrString + "\n")
        for col in self.matrix:
            probCol = map(lambda prob: (prob+pseudo)/(1+(pseudo*4)), col)
            outFile.write(reduce(lambda p1, p2: str(p1) + " " + str(p2),
                                 probCol, "") + "\n")

    def writeToMEME(self, outFile, pseudo=0.01):
        """Writes the motif out the specified filehandle in MEME format."""

        # The motif matrix data must have been set before this method can
        # be called:
        assert (self.matrix != None)

        hdrString = """MEME version 4.5

ALPHABET= ACGT

strands:  + -

Background letter frequencies (from dataset with add-one prior applied):
A 0.25 C 0.25 G 0.25 T 0.25

MOTIF """ + str(self.name) + """
BL   MOTIF """ + str(self.name) + """ width= """ + str(len(self.matrix)) + """ seqs=0
letter-probability matrix: alength= 4 w= """ + str(len(self.matrix)) + " nsites= 10 E= 0\n"

        outFile.write(hdrString)

        # Currently just assume uniform background model:
        for col in self.matrix:
            newCol = map(lambda prob: (prob+pseudo)/(1+(pseudo*4)), col)
            outFile.write(reduce(lambda p1, p2: str(p1) + " " + str(p2),
                                 newCol, "") + "\n")
        outFile.flush()

    def initFromFreqMatrix(self, matrix, name="Unknown"):
        """Initialise the matrix by simply assigning to the specified frequency
        matrix."""

        self.matrix = matrix

    def initFromInimotifFile(self, iniMotifFile, name="Unknown"):
        """Initialise the matrix from an iniMotif file."""

        # Read in the count of the consensus sequence...
        currLine = iniMotifFile.readline()
        elems = currLine.split()
        while ((len(elems) < 2) or (elems[1] != "consensus")):
            currLine = iniMotifFile.readline()
            elems = currLine.split()
        self.consensusCount = int(elems[-1])

        # Skip to the start of the frequency matrix...
        while (currLine[:16] != "Frequency Matrix"):
            currLine = iniMotifFile.readline()

        # Read the matrix information:
        currLine = iniMotifFile.readline()
        freqsTranspose = []
        while (currLine != "\n"):
            elems = currLine.split()
            freqs = map(lambda tok: float(tok), elems[1:])
            freqsTranspose.append(freqs)
            currLine = iniMotifFile.readline()

        freqMatrixTranspose = numpy.matrix(freqsTranspose)
        freqMatrix = freqMatrixTranspose.transpose()
        freqMatrixAsList = freqMatrix.tolist()
        self.matrix = freqMatrixAsList

    def initFromCounts(self, countsInfile, name="Unknown", pseudo=0.01):
        """Initialise the matrix from an input text file of counts, where each
        row has four values specifying the A, C, G, T counts for a given column
        of the psfm."""

        self.name = name

        matrix = []
        for line in countsInfile.readlines():
            elems = line.strip().split()
            countsColumn = map(lambda tok: int(tok), elems)
            countsSum = reduce(lambda count1, count2: count1 + count2,
                               countsColumn)
            freqsColumn = \
                map(lambda count: float(count + pseudo)/float(countsSum + pseudo*4),
                              countsColumn)
            matrix.append(freqsColumn)
        self.matrix = matrix

    def initFromAlign(self, seqAlnList, name="Unknown", pseudo=0.01):
        """Initialise the matrix from a list of aligned DNA sequences."""

        lett_to_idx = \
            {'A':0, 'C':1, 'G':2, 'T':3, \
                 'a':0, 'c':1, 'g':2, 't':3}

        # Initialise count matrix:
        seq_len = len(seqAlnList[0])
        count_matrix = []
        seq_idx = 0
        while (seq_idx < seq_len):
            count_matrix.append([0,0,0,0])
            seq_idx = seq_idx + 1

        for sequence in seqAlnList:
            # Add counts to count_matrix for current sequence:
            seq_idx = 0;
            while (seq_idx < seq_len):
                curr_lett = sequence[seq_idx]
                # If 'N' is found, add a count of 0.25 to each count
                # matrix letter:
                if ((curr_lett == 'N') or (curr_lett == 'n')):
                    count_matrix[seq_idx][0] = count_matrix[seq_idx][0] + 0.25
                    count_matrix[seq_idx][1] = count_matrix[seq_idx][1] + 0.25
                    count_matrix[seq_idx][2] = count_matrix[seq_idx][2] + 0.25
                    count_matrix[seq_idx][3] = count_matrix[seq_idx][3] + 0.25
                else:
                    lett_idx = lett_to_idx[curr_lett]
                    count_matrix[seq_idx][lett_idx] = \
                        count_matrix[seq_idx][lett_idx] + 1
                seq_idx = seq_idx + 1
                self.name = name

        freqMatrix = []
        for countsColumn in count_matrix:
            countsSum = reduce(lambda count1, count2: count1 + count2,
                               countsColumn)
            freqsColumn = \
                map(lambda count: float(count + pseudo)/float(countsSum + pseudo*4),
                              countsColumn)
            freqMatrix.append(freqsColumn)
        self.matrix = freqMatrix

    def initFrom_xxMotif(self, xxMotif_infile):
        # Store name of file as attribute of motif:

        self.xxMotifFilePath = xxMotif_infile.name

        # Parse a single motif from the assumed open file...

        # Throw away lines until the next motif line is reached...
        currLine = xxMotif_infile.readline()
        if (currLine == ""):
            raise ValueError(xxMotif_infile.name + " contains no more motifs!")

        while ((currLine != "") and (currLine[:5] != "Motif")):
            currLine = xxMotif_infile.readline()

        if (currLine == ""):
            raise ValueError(xxMotif_infile.name + " contains no more motifs!")

        # Parse the motif line, to obtain the motif name:
        elems = currLine.strip().split()
        motifName = elems[1][:-1]
        motifEvalue = float(elems[-1])
        self.name = motifName
        self.eValue = motifEvalue

        # Parse the frequency matrix:
        freqMatrixList = []
        for lettIdx in [1,2,3,4]:
            currLine = xxMotif_infile.readline()
            elems = currLine.strip().split()
            freqsCol = map(lambda tok: float(tok), elems[1:])
            freqMatrixList.append(freqsCol)

        freqMatrix = numpy.matrix(freqMatrixList)
        freqMatrixTranspose = freqMatrix.transpose()
        freqMatrixTransposeList = freqMatrixTranspose.tolist()
        self.matrix = freqMatrixTransposeList

    def initFromMEME(self, meme_infile):
        # Initialise this motif from the specified motifData input file...

        # Store the name of the file as an attribute of this motif:
        self.memeFilePath = meme_infile.name

        # Parse a single motif from an open MEME file...

        # While the currline doesn't start with "MOTIF", discard the line and
        # get the next...
        currLine = meme_infile.readline()
        if (currLine == ""):
            raise ValueError(meme_infile.name + " contains no more motifs!")

        # Ignore lines until "MOTIF" is reached:
        while ((currLine != "") and (currLine[:5] != "MOTIF")):
            currLine = meme_infile.readline()

        if (currLine == ""):
            raise ValueError(meme_infile.name + " contains no more motifs!")

        # Parse the motif line, to obtain the motif name:
        elems = currLine.strip().split()
        motifName = elems[1]

        # Discard lines until the  "letter-probability" line is reached...
        while ((currLine != "") and (currLine[:18] != "letter-probability")):
            currLine = meme_infile.readline()

        if (currLine == ""):
            raise ValueError(meme_infile.name + " contains no matrix info" + \
                                 " for matrix " + motifName)

        # Parse the motif stats; extract and save the E-value of the motif:
        elems = currLine.strip().split()
        eValue = float(elems[9])
        nsites = int(elems[7])

        # Instantiate the motif with the parsed values:
        self.name = motifName
        self.eValue = eValue
        self.nsites = nsites

        # Parse the actual matrix data...
        currLine = meme_infile.readline()
        matrix = []
        # Modified April 17th 2013; allowing new line character to indicate
        # end of current motif:
        while ((currLine != "") and (currLine[:3] != "---") and (currLine != "\n")):
            elems = currLine.strip().split()
            matrixCol = map(lambda tok: float(tok), elems)
            matrix.append(matrixCol)
            currLine = meme_infile.readline()
        self.matrix = matrix

    def getRevCompMatrix(self):
        """Returns a raw frequency matrix representing the reverse complement of this matrix."""
        matrixCopy = copy.deepcopy(self.matrix)
        matrixCopy.reverse()
        for col in matrixCopy:
            col.reverse()
        return matrixCopy

    def getConsensus(self):
        """Returns the consensus sequence of this motif."""

        letters = ["A","C","G","T"]

        maxPositions = numpy.array(self.matrix).argmax(1)
        consensus = reduce(lambda lett1, lett2: lett1+lett2,
                           map(lambda maxPos:letters[maxPos], maxPositions))
        return consensus

    def loadIntoDb(self, analysisID, cycle, db):
        """Loads this motif into the specified SELEX database, with the
        specified analysisID and SELEX cycle number that it was generated
        from."""

        # Obtain the ensemble ID of this factor, using the analysisID and a
        # mysql query...

        # Find the Barcode and Batch studied with the specified analysis:
        cursor=db.cursor()
        cmdStr = "select BarcodeAnalysed, BarcodeBatchAnalysed from SELEX_Analyses where (AnalysisID = " + str(analysisID) + ");"
        cursor.execute(cmdStr)
        rows = cursor.fetchall()

        # Return the first element of the first result from the query run:
        (barcode, batch) = rows[0]

        # Obtain the ensembleID, by using a SelexSample object:
        selexSample = selexDb.SelexSample(barcode, batch, cycle)
        ensembleID = selexSample.getEnsembleID()

        # Obtain the width of the motif:
        motWidth=self.getWidth()

        # Obtain the consensus sequence:
        consensus=self.getConsensus()
        consensusCount=self.consensusCount

        # Insert a new entry into the Motifs table using a new mysql command,
        # inserting analysisID, ensembleID, width, cycle, consensus sequence
        # and consensus count:
        if (ensembleID != None):
            cmdStr = "INSERT INTO Motifs (AnalysisID, EnsembleID, Width, SelexCycle, ConsensusSequence, ConsensusCount) VALUES (" + str(analysisID) + ",\"" + str(ensembleID) + "\"," + str(motWidth) + "," + str(cycle) + ",\"" + consensus  + "\"," + str(consensusCount) + ");"
        else:
            cmdStr = "INSERT INTO Motifs (AnalysisID, Width, SelexCycle, ConsensusSequence, ConsensusCount) VALUES (" + str(analysisID) + "," + str(motWidth) + "," + str(cycle) + ",\"" + consensus  + "\"," + str(consensusCount) + ");"

        # Run that command:
        cursor=db.cursor()
        cursor.execute(cmdStr)

        # Obtain the new autoincremented motifID value:
        newMotifID = db.insert_id()

        # Insert data for the columns...
        for columnIdx in range(len(self.matrix)):
            currColumn = self.matrix[columnIdx]

            # Calculate information content of that column:
            ic = seq_utility.calc_IC(currColumn)

            # Insert a new entry into MotifColumns, with the above motifID,
            # columnIdx, information content, and A, C, G and T frequencies:
            cmdStr = "INSERT INTO MotifColumns (MotifID, ColumnNumber, InformationContent, A, C, G, T) VALUES (" + str(newMotifID) + "," + str(columnIdx) + "," + str(ic) + "," + str(currColumn[0]) + "," + str(currColumn[1]) + "," + str(currColumn[2]) + "," + str(currColumn[3]) + ");"

            # Run that command:
            cursor=db.cursor()
            cursor.execute(cmdStr)

    def permuteMatrix(self):
        """Permutes the columns of this matrix."""
        random.shuffle(self.matrix)

    def isCrap(self):
        """Returns True if this motif is dodgy, False otherwise."""

        # A bit tricky; Get a trimmed version of the motif, and then compute
        # whether the motif is crap based on the trimmed version and the
        # original version...

        trimmedMot = self.trimLowIC(icThresh=0.5, copy=True)
        if (self.eValue > 0.01):
            return True
        if (trimmedMot.getWidth() <= 10):
            return False
        if (self.eValue < 10**-30):
            return False
        motICs = trimmedMot.getIC_arr()
        letters = trimmedMot.getConsensus()
        lettersAtHighIC = {}
        columnIdx = 0
        while (columnIdx < trimmedMot.getWidth()):
            currLetter = letters[columnIdx]
            ic = motICs[columnIdx]
            if (ic >= 1.5):
                if (not lettersAtHighIC.has_key(currLetter)):
                    lettersAtHighIC[currLetter] = 1
            columnIdx += 1
        if (len(lettersAtHighIC.keys()) >= 3):
            return False
        return True

class JASPAR_pwm(pwm):
    """A pwm derived from a high-throughput SELEX experiment."""
    def __init__(self, db, motifID):
        super(JASPAR_pwm, self).__init__(None, dataType="")

        # Obtain the JASPAR ID and collection values for this motifID:
        cmdStr = "SELECT COLLECTION, BASE_ID, NAME, VERSION from MATRIX where ID = " + str(motifID) + ";"
        cursor=db.cursor()
        cursor.execute(cmdStr)
        rows = cursor.fetchall()

        assert(len(rows) > 0) # Otherwise motifID is an invalid input.

        self.collection = rows[0][0]
        self.jasparID = rows[0][1]
        tfName = rows[0][2]
        version = rows[0][3]

        # Set the name of the matrix to the jaspar code and tf name:
        self.setName(self.jasparID + "_" + str(version) + "_" + tfName)

        # Obtain the motif data:
        cmdStr = "SELECT col, row, val from MATRIX_DATA where ID = \"" + str(motifID) + "\";"
        cursor=db.cursor()
        cursor.execute(cmdStr)
        rows = cursor.fetchall()

        positions = map(lambda tup: tup[0], rows)
        maxColNum = max(positions)

        # Parse the output into the matrix attribute...

        # Obtain the count matrix:
        countMatrix = numpy.zeros((maxColNum, 4))
        lett2col = {'A':0, 'C':1, 'G':2, 'T':3}
        for tup in rows:
            row = tup[0] - 1
            col = lett2col[tup[1]]
            countMatrix[row][col] = tup[2]

        # Divide all entries by row sums to obtain freq matrix:
        freqMatrix = numpy.zeros((maxColNum, 4))
        for rowIdx in range(maxColNum):
            rowSum = float(sum(countMatrix[rowIdx]))
            for colIdx in range(4):
                freqMatrix[rowIdx][colIdx] = \
                    (countMatrix[rowIdx][colIdx])/rowSum

        self.matrix = freqMatrix.tolist()

class SELEX_pwm(pwm):
    """A pwm derived from a high-throughput SELEX experiment."""
    def __init__(self, db, motifID):
        super(SELEX_pwm, self).__init__(None, dataType="")

        # Obtain the cycle and consensusCount values for this motifID:
        cmdStr = "SELECT SelexCycle, ConsensusCount from Motifs where MotifID = \"" + str(motifID) + "\";"
        cursor=db.cursor()
        cursor.execute(cmdStr)
        rows = cursor.fetchall()

        self.cycle = rows[0][0]
        self.consensusCount = rows[0][1]

        # Obtain the motif data:
        cmdStr = "SELECT ColumnNumber, A, C, G, T from MotifColumns where MotifID = \"" + str(motifID) + "\";"
        cursor=db.cursor()
        cursor.execute(cmdStr)
        rows = cursor.fetchall()
        matrixData = list(rows)
        matrixData.sort(cmpTups1)

        # Obtain a name for the motif:
        cmdStr = "select tf.HGNC_Name, \"w\", m.Width, \"cyc\", m.SelexCycle, sa.BarcodeBatchAnalysed, sa.BarcodeAnalysed from Motifs m inner join SELEX_Analyses sa on m.AnalysisID = sa.AnalysisID inner join TranscriptionFactors tf on m.EnsembleID = tf.EnsembleID where m.MotifID = " + str(motifID) + ";"
        cursor=db.cursor()
        cursor.execute(cmdStr)

        rows = cursor.fetchall()
        name = reduce(lambda tok1, tok2: str(tok1) + "_" + str(tok2), rows[0])

        self.setName(name)

        # Parse the output into the matrix attribute:
        matrix = map(lambda col: col[1:], matrixData)
        self.matrix = matrix

    def getConsCount(self):
        return self.consensusCount

def cmpTups1(tup1, tup2):
    if (tup1[0] < tup2[0]):
        return -1
    elif (tup1[0] == tup2[0]):
        return 0
    else:
        return 1

class motifLibrary:
    """A library of motifs. Introduced on 9th November 2010 in order to allow
    MEME-derived motifs to be added to an existing set of motifs more
    seamlessly."""

    def __init__(self, inputLibrary):
        self.motifs = {}
        if (isinstance(inputLibrary, list)):
            # Library specifies name of file containing motif locations =>
            # initialise as such:
            self.initFromMEME(inputLibrary)

    def getMotif(self, motifName):
        return self.motifs[motifName]

    def getMotifs(self):
        return self.motifs

    def initFromMEME(self, motifLibraryList):
        """motifLibraryList is a list of strings specifying the absolute
        paths of motifs in the input library."""

        # Set up the motif library from the input file...

        # For each motif filename in the input library...
        for currMotifFilename in motifLibraryList:
            if (not os.path.exists(currMotifFilename)):
                # Don't quit if a motif file isn't found; just skip it:
                print >> sys.stderr, "WARNING: Specified motif file \"" + \
                    currMotifFilename + "\" does not exist."
            else:
                # Attempt to generate a new motif object from that file, using
                # the motif constructor:
                try:
                    currMotif = pwm(open(currMotifFilename))
                    currMotifName = currMotif.getName()
                    if (self.motifs.has_key(currMotifName)):
                        print >> sys.stderr, "WARNING: motif \"" + \
                            currMotif.getName() + \
                            "\" was encountered more than once." + \
                            " Excluding second instance from the input library."
                    else:
                        # Non-redundant motif was successfully parsed => Add
                        # it to the library of motifs:
                        self.motifs[currMotifName] = currMotif
                except ValueError, e:
                    print >> sys.stderr, "WARNING: ValueError occured \"" + \
                        "whilst parsing motif file \"" + currMotifFilename + \
                        "\". Excluding that motif from the input library."
                    print >> sys.stderr, e

    def addMotif(self, motif):
        """Add a single specified motif to this library."""
        assert(not self.motifs.has_key(motif.getName()))
        self.motifs[motif.getName()] = motif

class scannedReg:
    """This class encapsulates the general idea of a genomic region that has
    been scanned by a set of motifs."""

    def __init__(self, region):
        # Store the attributes (in particular the genomic coordinates) of the
        # genomic region:
        self.reg = region # A BED_line object.

        # Set up an intially-empty dictionary of motif hits. It will be
        # populated with motif name as keys and motifOccurence as values:
        self.hits = {}

    def addHit(self, motifHit):
        """Stores the specified motifOccurence object under it's name.
        Throws an Exception if a motif by that name has already been
        registered for this sequence."""

        if (self.hits.has_key(motifHit.getMotif().getName())):
            raise Exception("Region already had motif instance of that name.")
        else:
            self.hits[motifHit.getMotif().getName()] = motifHit

    def getHit(self, motifName):
        """Returns the occurence of the specified motif for this region.
        Returns None if no hits are registered for the region."""
        if (self.hits.has_key(motifName)):
            return self.hits[motifName]
        else:
            return None

    def getDisp(self, mot1_name, mot2_name):
        """Returns a (displacement, sameStrand) tuple showing the
        displacement from motif1 to motif2 and whether they occur on the
        same strand of this sequence. Returns None if the two motifs overlap.
        Throws a KeyError exception if either motif is not present in the
        sequence."""
        mot1_hit = self.hits[mot1_name]
        mot2_hit = self.hits[mot2_name]
        disp1to2 = mot1_hit.get_disp(mot2_hit)

    def getBED(self):
        """Returns the BED_line object showing the coordinates of this
        region."""
        return self.reg

    def getMotCoords(self, motifName, flankWidth=0):
        """Returns a strand-specific BED_line object representing the
        *genomic* coordinates of the occurrence of the specific motif in this
        scanned region. Returns None if no motif hit exists in this region for
        the specified motif."""

        motifHit = self.getHit(motifName)
        if (motifHit == None):
            return None

        genomicHitLoc = chipseq_analysis.bestHit_scanner.getAbsHitCoords(self.getBED(), motifHit, flankWidth=flankWidth)
        return genomicHitLoc