<commit_before># USBKeyboard.py
#
# Contains class definitions to implement a USB keyboard.

import struct
import time
from umap2.core.usb import DescriptorType
from umap2.core.usb_class import USBClass
from umap2.core.usb_device import USBDevice
from umap2.core.usb_configuration import USBConfiguration
from umap2.core.usb_interface import USBInterface
from umap2.core.usb_endpoint import USBEndpoint
from umap2.fuzz.helpers import mutable

class Requests(object):
    GET_REPORT = 0x01  # Mandatory
    GET_IDLE = 0x02
    GET_PROTOCOL = 0x03  # Ignored - only for boot device
    SET_REPORT = 0x09
    SET_IDLE = 0x0A
    SET_PROTOCOL = 0x0B  # Ignored - only for boot device

class USBKeyboardClass(USBClass):
    name = 'KeyboardClass'

    def setup_local_handlers(self):
        self.local_handlers = {
            Requests.GET_REPORT: self.handle_get_report,
            Requests.GET_IDLE: self.handle_get_idle,
            Requests.SET_REPORT: self.handle_set_report,
            Requests.SET_IDLE: self.handle_set_idle,
        }

    @mutable('hid_get_report_response')
    def handle_get_report(self, req):
        response = b'\xff' * req.length
        return response

    @mutable('hid_get_idle_response')
    def handle_get_idle(self, req):
        return b''

    @mutable('hid_set_report_response')
    def handle_set_report(self, req):
        return b''

    @mutable('hid_set_idle_response')
    def handle_set_idle(self, req):
        return b''

class USBKeyboardInterface(USBInterface):
    name = 'KeyboardInterface'

    def __init__(self, app, phy):
        super(USBKeyboardInterface, self).__init__(
            app=app,
            phy=phy,
            interface_number=0,
            interface_alternate=0,
            interface_class=USBClass.HID,
            interface_subclass=0,
            interface_protocol=0,
            interface_string_index=0,
            endpoints=[
                USBEndpoint(
                    app=app,
                    phy=phy,
                    number=2,
                    direction=USBEndpoint.direction_in,
                    transfer_type=USBEndpoint.transfer_type_interrupt,
                    sync_type=USBEndpoint.sync_type_none,
                    usage_type=USBEndpoint.usage_type_data,
                    max_packet_size=0x40,
                    interval=0x40,
                    handler=self.handle_buffer_available
                )
            ],
            descriptors={
                DescriptorType.hid: self.get_hid_descriptor,
                DescriptorType.report: self.get_report_descriptor
            },
            usb_class=USBKeyboardClass(app, phy)
        )

        empty_preamble = [0x00] * 10
        text = [0x0f, 0x00, 0x16, 0x00, 0x28, 0x00]

        self.keys = [chr(x) for x in empty_preamble + text]
        self.call_count = 0
        self.first_call = None

    @mutable('hid_descriptor')
    def get_hid_descriptor(self, *args, **kwargs):
        report_descriptor = self.get_report_descriptor()
        bDescriptorType = b'\x21'  # HID
        bcdHID = b'\x10\x01'
        bCountryCode = b'\x00'
        bNumDescriptors = b'\x01'
        bDescriptorType2 = b'\x22'  # REPORT
        desclen = len(report_descriptor)
        wDescriptorLength = struct.pack('<H', desclen)
        hid_descriptor = (
            bDescriptorType +
            bcdHID +
            bCountryCode +
            bNumDescriptors +
            bDescriptorType2 +
            wDescriptorLength
        )
        bLength = struct.pack('<B', len(hid_descriptor) + 1)
        hid_descriptor = bLength + hid_descriptor
        return hid_descriptor

    @mutable('hid_report_descriptor')
    def get_report_descriptor(self, *args, **kwargs):
        self.first_call = None
        usage_page_generic_desktop_controls = b'\x05\x01'
        # usage_page_generic_desktop_controls = b'\xb1\x01'
        usage_keyboard = b'\x09\x06'
        collection_application = b'\xA1\x01'
        usage_page_keyboard = b'\x05\x07'
        usage_minimum1 = b'\x19\xE0'
        usage_maximum1 = b'\x29\xE7'
        logical_minimum1 = b'\x15\x00'
        logical_maximum1 = b'\x25\x01'
        report_size1 = b'\x75\x01'
        report_count1 = b'\x95\x08'
        input_data_variable_absolute_bitfield = b'\x81\x02'
        report_count2 = b'\x95\x01'
        report_size2 = b'\x75\x08'
        input_constant_array_absolute_bitfield = b'\x81\x01'
        usage_minimum2 = b'\x19\x00'
        usage_maximum2 = b'\x29\x65'
        logical_minimum2 = b'\x15\x00'
        logical_maximum2 = b'\x25\x65'
        report_size3 = b'\x75\x08'
        report_count3 = b'\x95\x01'
        input_data_array_absolute_bitfield = b'\x81\x00'
        end_collection = b'\xc0'

        report_descriptor = (
            usage_page_generic_desktop_controls +
            usage_keyboard +
            collection_application +
            usage_page_keyboard +
            usage_minimum1 +
            usage_maximum1 +
            logical_minimum1 +
            logical_maximum1 +
            report_size1 +
            report_count1 +
            input_data_variable_absolute_bitfield +
            report_count2 +
            report_size2 +
            input_constant_array_absolute_bitfield +
            usage_minimum2 +
            usage_maximum2 +
            logical_minimum2 +
            logical_maximum2 +
            report_size3 +
            report_count3 +
            input_data_array_absolute_bitfield +
            end_collection
        )
        return report_descriptor

    def handle_buffer_available(self):
        #
        # this is really ugly, but sometimes the host expects
        # (during initialization) to get the report on ep0 and
        # ignores the actual ep (2 in this case), we'll just
        # wait for a little while... (see section 7.2.1 in HID spec)
        #
        if self.first_call is None:
            self.first_call = time.time()
        if time.time() - self.first_call > 2:
            self.usb_function_supported()
            if self.keys:
                letter = self.keys.pop(0)
                self.type_letter(letter)

    def type_letter(self, letter, modifiers=0):
        data = struct.pack('<BBB', 0, 0, ord(letter))
        self.verbose('sending keypress 0x%02x' % ord(letter))
        self.send_on_endpoint(2, data)

class USBKeyboardDevice(USBDevice):
    name = 'KeyboardDevice'

    def __init__(self, app, phy, vid=0x610b, pid=0x4653, rev=0x1234, **kwargs):
        super(USBKeyboardDevice, self).__init__(
            app=app,
            phy=phy,
            device_class=USBClass.Unspecified,
            device_subclass=0,
            protocol_rel_num=0,
            max_packet_size_ep0=64,
            vendor_id=vid,
            product_id=pid,
            device_rev=rev,
            manufacturer_string='Dell',
            product_string='Dell USB Entry Keyboard',
            serial_number_string='00001',
            configurations=[
                USBConfiguration(
                    app=app,
                    phy=phy,
                    index=1,
                    string='Emulated Keyboard',
                    interfaces=[
                        USBKeyboardInterface(app, phy)
                    ]
                )
            ],
        )

usb_device = USBKeyboardDevice
<commit_msg>[KeyboardDevice] keep sending reports when done with text<commit_after># USBKeyboard.py
#
# Contains class definitions to implement a USB keyboard.

import struct
import time
from umap2.core.usb import DescriptorType
from umap2.core.usb_class import USBClass
from umap2.core.usb_device import USBDevice
from umap2.core.usb_configuration import USBConfiguration
from umap2.core.usb_interface import USBInterface
from umap2.core.usb_endpoint import USBEndpoint
from umap2.fuzz.helpers import mutable

class Requests(object):
    GET_REPORT = 0x01  # Mandatory
    GET_IDLE = 0x02
    GET_PROTOCOL = 0x03  # Ignored - only for boot device
    SET_REPORT = 0x09
    SET_IDLE = 0x0A
    SET_PROTOCOL = 0x0B  # Ignored - only for boot device

class USBKeyboardClass(USBClass):
    name = 'KeyboardClass'

    def setup_local_handlers(self):
        self.local_handlers = {
            Requests.GET_REPORT: self.handle_get_report,
            Requests.GET_IDLE: self.handle_get_idle,
            Requests.SET_REPORT: self.handle_set_report,
            Requests.SET_IDLE: self.handle_set_idle,
        }

    @mutable('hid_get_report_response')
    def handle_get_report(self, req):
        response = b'\xff' * req.length
        return response

    @mutable('hid_get_idle_response')
    def handle_get_idle(self, req):
        return b''

    @mutable('hid_set_report_response')
    def handle_set_report(self, req):
        return b''

    @mutable('hid_set_idle_response')
    def handle_set_idle(self, req):
        return b''

class USBKeyboardInterface(USBInterface):
    name = 'KeyboardInterface'

    def __init__(self, app, phy):
        super(USBKeyboardInterface, self).__init__(
            app=app,
            phy=phy,
            interface_number=0,
            interface_alternate=0,
            interface_class=USBClass.HID,
            interface_subclass=0,
            interface_protocol=0,
            interface_string_index=0,
            endpoints=[
                USBEndpoint(
                    app=app,
                    phy=phy,
                    number=2,
                    direction=USBEndpoint.direction_in,
                    transfer_type=USBEndpoint.transfer_type_interrupt,
                    sync_type=USBEndpoint.sync_type_none,
                    usage_type=USBEndpoint.usage_type_data,
                    max_packet_size=0x40,
                    interval=0x40,
                    handler=self.handle_buffer_available
                )
            ],
            descriptors={
                DescriptorType.hid: self.get_hid_descriptor,
                DescriptorType.report: self.get_report_descriptor
            },
            usb_class=USBKeyboardClass(app, phy)
        )

        empty_preamble = [0x00] * 10
        text = [0x0f, 0x00, 0x16, 0x00, 0x28, 0x00]

        self.keys = [chr(x) for x in empty_preamble + text]
        self.call_count = 0
        self.first_call = None

    @mutable('hid_descriptor')
    def get_hid_descriptor(self, *args, **kwargs):
        report_descriptor = self.get_report_descriptor()
        bDescriptorType = b'\x21'  # HID
        bcdHID = b'\x10\x01'
        bCountryCode = b'\x00'
        bNumDescriptors = b'\x01'
        bDescriptorType2 = b'\x22'  # REPORT
        desclen = len(report_descriptor)
        wDescriptorLength = struct.pack('<H', desclen)
        hid_descriptor = (
            bDescriptorType +
            bcdHID +
            bCountryCode +
            bNumDescriptors +
            bDescriptorType2 +
            wDescriptorLength
        )
        bLength = struct.pack('<B', len(hid_descriptor) + 1)
        hid_descriptor = bLength + hid_descriptor
        return hid_descriptor

    @mutable('hid_report_descriptor')
    def get_report_descriptor(self, *args, **kwargs):
        self.first_call = None
        usage_page_generic_desktop_controls = b'\x05\x01'
        # usage_page_generic_desktop_controls = b'\xb1\x01'
        usage_keyboard = b'\x09\x06'
        collection_application = b'\xA1\x01'
        usage_page_keyboard = b'\x05\x07'
        usage_minimum1 = b'\x19\xE0'
        usage_maximum1 = b'\x29\xE7'
        logical_minimum1 = b'\x15\x00'
        logical_maximum1 = b'\x25\x01'
        report_size1 = b'\x75\x01'
        report_count1 = b'\x95\x08'
        input_data_variable_absolute_bitfield = b'\x81\x02'
        report_count2 = b'\x95\x01'
        report_size2 = b'\x75\x08'
        input_constant_array_absolute_bitfield = b'\x81\x01'
        usage_minimum2 = b'\x19\x00'
        usage_maximum2 = b'\x29\x65'
        logical_minimum2 = b'\x15\x00'
        logical_maximum2 = b'\x25\x65'
        report_size3 = b'\x75\x08'
        report_count3 = b'\x95\x01'
        input_data_array_absolute_bitfield = b'\x81\x00'
        end_collection = b'\xc0'

        report_descriptor = (
            usage_page_generic_desktop_controls +
            usage_keyboard +
            collection_application +
            usage_page_keyboard +
            usage_minimum1 +
            usage_maximum1 +
            logical_minimum1 +
            logical_maximum1 +
            report_size1 +
            report_count1 +
            input_data_variable_absolute_bitfield +
            report_count2 +
            report_size2 +
            input_constant_array_absolute_bitfield +
            usage_minimum2 +
            usage_maximum2 +
            logical_minimum2 +
            logical_maximum2 +
            report_size3 +
            report_count3 +
            input_data_array_absolute_bitfield +
            end_collection
        )
        return report_descriptor

    def handle_buffer_available(self):
        #
        # this is really ugly, but sometimes the host expects
        # (during initialization) to get the report on ep0 and
        # ignores the actual ep (2 in this case), we'll just
        # wait for a little while... (see section 7.2.1 in HID spec)
        #
        if self.first_call is None:
            self.first_call = time.time()
        if time.time() - self.first_call > 2:
            self.usb_function_supported()
            if self.keys:
                letter = self.keys.pop(0)
            else:
                letter = '\x00'
            self.type_letter(letter)

    def type_letter(self, letter, modifiers=0):
        data = struct.pack('<BBB', 0, 0, ord(letter))
        self.send_on_endpoint(2, data)

class USBKeyboardDevice(USBDevice):
    name = 'KeyboardDevice'

    def __init__(self, app, phy, vid=0x610b, pid=0x4653, rev=0x1234, **kwargs):
        super(USBKeyboardDevice, self).__init__(
            app=app,
            phy=phy,
            device_class=USBClass.Unspecified,
            device_subclass=0,
            protocol_rel_num=0,
            max_packet_size_ep0=64,
            vendor_id=vid,
            product_id=pid,
            device_rev=rev,
            manufacturer_string='Dell',
            product_string='Dell USB Entry Keyboard',
            serial_number_string='00001',
            configurations=[
                USBConfiguration(
                    app=app,
                    phy=phy,
                    index=1,
                    string='Emulated Keyboard',
                    interfaces=[
                        USBKeyboardInterface(app, phy)
                    ]
                )
            ],
        )

usb_device = USBKeyboardDevice
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#=======================================================================
# huffman.py
# ---------
# Test implementation of huffman encoder and decoder. Just to see
# that I actually understand it. Also good practice to learn to write
# bit patterns from Python.
#
#
# (c) 2012 Secworks
# JoachimS
#=======================================================================

#-------------------------------------------------------------------
# Imports.
#-------------------------------------------------------------------
import argparse
import random
import sys

#-------------------------------------------------------------------
# Constants.
#-------------------------------------------------------------------
VERBOSE = False
DEBUG = True

VERSION = '0.1 Beta'

#-------------------------------------------------------------------
# class Node
#
# A basic node class (struct) used to create the prefix tree.
#-------------------------------------------------------------------
class Node:
    def __init__(self, char, weight):
        self.char = char
        self.weight = weight
        self.lchild = None
        self.rchild = None
        self.children = 0

    def print_fields(self):
        if self.char != None:
            print("Char: %c, weight = %d" % (self.char, self.weight))
        else:
            print("Inner node with %d children and total weight %d" %\
                  (self.children, self.weight))

        if self.lchild != None:
            print("Left child contains subtree")
        else:
            print("Left child contains no subtree")

        if self.rchild != None:
            print("Right child contains subtree")
        else:
            print("Right child contains no subtree")

#-------------------------------------------------------------------
# get_node_codes()
#
# Get a list of tuples with the prefix codes for a char in a
# given node as well as for its subnodes.
#-------------------------------------------------------------------
def get_node_codes(prefix, ptree):
    if ptree.char == None:
        left_list = []
        right_list = []
        if ptree.lchild != None:
          left_list = get_node_codes(prefix + '1', ptree.lchild)
        if ptree.rchild != None:
          right_list = get_node_codes(prefix + '0', ptree.rchild)
        return left_list + right_list
    else:
        return [(ptree.char, prefix, ptree.weight)]

#-------------------------------------------------------------------
# extract_prefix_codes()
#
# Given a prefix tree, extracts the prefix codes for all
# leaves. The prefix codes are returned as a dictionary.
#-------------------------------------------------------------------
def extract_prefix_codes(ptree):
    left_list = get_node_codes('1', ptree.lchild)
    right_list = get_node_codes('0', ptree.rchild)
    prefix_dict = {}

    for (char, prefix, weight) in left_list:
        prefix_dict[char] = (prefix, weight)

    for (char, prefix, weight) in right_list:
        prefix_dict[char] = (prefix, weight)

    return prefix_dict

#-------------------------------------------------------------------
# gen_prefix_tree()
#
# Given a list of tuples with frequency (weight) for a character
# the function returns the corresponding tuple based prefix tree.
#-------------------------------------------------------------------
def gen_prefix_tree(nlist):
    while (len(nlist) > 1):
        if VERBOSE:
            print("Length of prefix list: %d" % len(nlist))

        nlist = sort_node_list(nlist)
        node1 = nlist.pop()
        node2 = nlist.pop()

        if VERBOSE:
            print("node1 weigth: %d" % node1.weight)
            print("node2 weigth: %d" % node2.weight)

        new_node = Node(None, (node1.weight + node2.weight))
        new_node.lchild = node1
        new_node.rchild = node2
        new_node.children = 2 + node1.children + node2.children
        nlist.append(new_node)

    return nlist.pop()

#-------------------------------------------------------------------
# sort_node_list()
#
# Given a list of nodes returns a list with the nodes
# sorted by weight in decreasing order.
#-------------------------------------------------------------------
def sort_node_list(nlist):
    nodes = 0
    new_list = []
    if VERBOSE:
        print("Length of given list: %d" % len(nlist))
    for i in range(len(nlist)):
        node = nlist.pop()
        nodes += 1
        if len(new_list) == 0:
            new_list.append(node)
        else:
            i = 0
            while ((i < len(new_list) and (node[1] < new_list[i][1]))):
                i += 1
            new_list = new_list[:i] + [node] + new_list[i:]

    return new_list

#-------------------------------------------------------------------
# print_prefix_codes()
#
# Given a db with prefix codes for a set of keys (chars) will
# print the contents of the db as well as some interesting
# statistics.
#-------------------------------------------------------------------
def print_prefix_codes(prefix_codes):
    min_len = 100000000
    max_len = 0
    num_raw_bits = 0
    num_prefix_bits = 0

    if VERBOSE:
        print("The generated prefix codes:")
    for key in prefix_codes:
        char = key
        (weight, prefix) = prefix_codes[key]
        num_prefix_bits += len(prefix) * weight
        num_raw_bits +=  weight * 8
        if VERBOSE:
            print("Char: %03d, prefix: %015s, prefix length: %02d, weight %010d" %\
                  (char, prefix, len(prefix), weight))

        if len(prefix) < min_len:
            min_len = len(prefix)
            min_char = char
            min_prefix = prefix
            min_weight = weight

        if len(prefix) > max_len:
            max_len = len(prefix)
            max_char = char
            max_prefix = prefix
            max_weight = weight

    reduction = int((1 - (num_prefix_bits / num_raw_bits)) * 100)

    print("")
    print("Minimum prefix length %02d for char %03d and prefix %015s with weight %010d" %\
          (min_len, min_char, min_prefix, min_weight))
    print("Maximum prefix length %02d for char %03d and prefix %015s with weight %010d" %\
          (max_len, max_char, max_prefix, max_weight))

    print("")
    print("Total number of bits in raw data:             %016d" % num_raw_bits)
    print("Total number of bits to prefix code all data: %016d" % num_prefix_bits)
    print("Reduction:                                    %02d percent" % reduction)
    print("")

#-------------------------------------------------------------------
# huffman_encode()
#
# Huffman encode a file with a given filename.
# Note: We assume that the file contains bytes.
#-------------------------------------------------------------------
def huffman_encode(filename):
    '''Encode the contents of the file using Huffman coding.'''

    # 1. First pass. Scan the file and build byte
    # frequency statistics.
    byte_freq = [0] * 256
    with open(filename, 'rb') as my_file:
        byte = my_file.read(1)
        bytectr = 1
        while byte:
            byteval = ord(byte)
            byte_freq[byteval] = byte_freq[byteval] + 1
            byte = my_file.read(1)
            bytectr = bytectr + 1

    if DEBUG:
        print "Encode step one."
        print "Number of bytes read: %d" % bytectr
        print byte_freq
        print ""

    # 2. Build weighted binary tree for all bytes in the file.
    alphabet = []
    for i in range(256):
        if (byte_freq[i] > 0):
            alphabet.append([i, byte_freq[i], 0, 0, 0])

    if DEBUG:
        print "Number of symbols in the alphabet: %d" % len(alphabet)
        print "The collected alphabet with weighs:"
        print alphabet

    elements_left = len(alphabet)
    while elements_left:
        pass
        # Find the two elements with the least weight

    # 3. Create symbol table based on the tree

    # 4. Emit symbol table as header for file.

    # 5. Second pass. Read file and emit symbols.

    # 6. Done!
    return 0

#-------------------------------------------------------------------
# huffman_decode()
#
# Huffman decode a file with the given filename.
#-------------------------------------------------------------------
def huffman_decode(filename):
    pass

#-------------------------------------------------------------------
# gen_node_list()
#
# Generates a list of max_types numbder of nodes. Each tuple
# contains a randomly selected frequency of a character from the
# set of characters [chr(0) .. chr((max_types - 1))]
#
# The nodes are tuples with the contents:
# (char, weight, left subtree, right subtree)
#
# For leaf nodes the subtrees are None. For all other nodes
# the char is None."
#-------------------------------------------------------------------
def gen_node_list(max_types, max_nums):
    return [(i, random.randint(0,max_nums), 0, None, None) for i in range(max_types)]

#-------------------------------------------------------------------
# encdec_huffman()
#
# Perform huffman encoding or decoding for the given infile
# as given by the arguments.
#-------------------------------------------------------------------
def encdec_huffman(args):
    pass

#-------------------------------------------------------------------
# test_huffman()
#
# Test mode. Generates data, builds prefix tree, performs
# encoding and then decoding and presents the results.
#-------------------------------------------------------------------
def test_huffman():
    print("Generating a tuple based prefix tree")
    print("====================================")
    print

    max_types = 256
    max_nums  = int(1E8)

    print("Generating %d nodes with up to %d instances." %\
          (max_types, max_nums))
    my_list = gen_node_list(max_types, max_nums)

    if VERBOSE:
        print("List of nodes before sort:")

    my_list = sort_node_list(my_list)

    if VERBOSE:
        print("List of nodes after sort:")
        print(my_list)
        print("")

    print("Generating prefix tree and extracting db for prefix codes.")
    my_tree = gen_prefix_tree(my_list)
    if VERBOSE:
        print("The generated prefix tree:")
        print(my_tree)

    my_codes = extract_prefix_codes(my_tree)
    if VERBOSE:
        print("The generated code db:")
        print(my_codes)

    print_prefix_codes(my_codes)

#-------------------------------------------------------------------
# main()
#
# Create an argument parser as needed to get input and output
# filenames as well as commands and options. Perform huffman
# encoding or decoding based on the given arguments.
#-------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('-i', '--infile',
                        help='The input file to be encoded or decoded.')

    parser.add_argument('-o', '--outfile',
                        help='The outputfile the processed file data will be saved to.')

    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Enable verbose processing.')

    parser.add_argument('-t', '--test', action='store_true',
                        help='Perform test processing.')

    parser.add_argument('-d', '--decode', action='store_true',
                        help='Perform Huffman decoding.')

    parser.add_argument('-e', '--encode', action='store_true',
                        help='Perform Huffman encoding.')

    parser.add_argument('--version', action='version', version=VERSION)

    args = parser.parse_args()

    if args.infile==None and not args.test:
        print "Error: No input file given and not in test mode."
        exit(1)

    if args.test:
        test_huffman()
    else:
        huffman_encdec(args)

#-------------------------------------------------------------------
# __name__
# Python thingy which allows the file to be run standalone as
# well as parsed from within a Python interpreter.
#-------------------------------------------------------------------
if __name__=="__main__":
    # Run the main function.
    sys.exit(main())

#=======================================================================
# EOF huffman.py
#=======================================================================
<commit_msg>Moved the new prefix tree generator.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#=======================================================================
# huffman.py
# ---------
# Test implementation of huffman encoder and decoder. Just to see
# that I actually understand it. Also good practice to learn to write
# bit patterns from Python.
#
#
# (c) 2012 Secworks
# JoachimS
#=======================================================================

#-------------------------------------------------------------------
# Imports.
#-------------------------------------------------------------------
import argparse
import random
import sys

#-------------------------------------------------------------------
# Constants.
#-------------------------------------------------------------------
VERBOSE = False
DEBUG = True

VERSION = '0.1 Beta'

#-------------------------------------------------------------------
# class Node
#
# A basic node class (struct) used to create the prefix tree.
#-------------------------------------------------------------------
class Node:
    def __init__(self, char, weight):
        self.char = char
        self.weight = weight
        self.lchild = None
        self.rchild = None
        self.children = 0

    def print_fields(self):
        if self.char != None:
            print("Char: %c, weight = %d" % (self.char, self.weight))
        else:
            print("Inner node with %d children and total weight %d" %\
                  (self.children, self.weight))

        if self.lchild != None:
            print("Left child contains subtree")
        else:
            print("Left child contains no subtree")

        if self.rchild != None:
            print("Right child contains subtree")
        else:
            print("Right child contains no subtree")

#-------------------------------------------------------------------
# get_node_codes()
#
# Get a list of tuples with the prefix codes for a char in a
# given node as well as for its subnodes.
#-------------------------------------------------------------------
def get_node_codes(prefix, ptree):
    if ptree.char == None:
        left_list = []
        right_list = []
        if ptree.lchild != None:
          left_list = get_node_codes(prefix + '1', ptree.lchild)
        if ptree.rchild != None:
          right_list = get_node_codes(prefix + '0', ptree.rchild)
        return left_list + right_list
    else:
        return [(ptree.char, prefix, ptree.weight)]

#-------------------------------------------------------------------
# extract_prefix_codes()
#
# Given a prefix tree, extracts the prefix codes for all
# leaves. The prefix codes are returned as a dictionary.
#-------------------------------------------------------------------
def extract_prefix_codes(ptree):
    left_list = get_node_codes('1', ptree.lchild)
    right_list = get_node_codes('0', ptree.rchild)
    prefix_dict = {}

    for (char, prefix, weight) in left_list:
        prefix_dict[char] = (prefix, weight)

    for (char, prefix, weight) in right_list:
        prefix_dict[char] = (prefix, weight)

    return prefix_dict

#-------------------------------------------------------------------
# gen_prefix_tree()
#
# Given a list of tuples with frequency (weight) for a character
# the function returns the corresponding tuple based prefix tree.
#-------------------------------------------------------------------
def gen_prefix_tree(nlist):
    while (len(nlist) > 1):
        if VERBOSE:
            print("Length of prefix list: %d" % len(nlist))

        nlist = sort_node_list(nlist)
        node1 = nlist.pop()
        node2 = nlist.pop()

        if VERBOSE:
            print("node1 weigth: %d" % node1[1])
            print("node2 weigth: %d" % node2[1])

        new_weight = node1[1] + node2[1]
        new_children = 2 + node1[2] + node2[2]
        new_node = (None, new_weight, new_children, node1, node2)
        nlist.append(new_node)

    return nlist.pop()

#-------------------------------------------------------------------
# sort_node_list()
#
# Given a list of nodes returns a list with the nodes
# sorted by weight in decreasing order.
#-------------------------------------------------------------------
def sort_node_list(nlist):
    nodes = 0
    new_list = []
    if VERBOSE:
        print("Length of given list: %d" % len(nlist))
    for i in range(len(nlist)):
        node = nlist.pop()
        nodes += 1
        if len(new_list) == 0:
            new_list.append(node)
        else:
            i = 0
            while ((i < len(new_list) and (node[1] < new_list[i][1]))):
                i += 1
            new_list = new_list[:i] + [node] + new_list[i:]

    return new_list

#-------------------------------------------------------------------
# print_prefix_codes()
#
# Given a db with prefix codes for a set of keys (chars) will
# print the contents of the db as well as some interesting
# statistics.
#-------------------------------------------------------------------
def print_prefix_codes(prefix_codes):
    min_len = 100000000
    max_len = 0
    num_raw_bits = 0
    num_prefix_bits = 0

    if VERBOSE:
        print("The generated prefix codes:")
    for key in prefix_codes:
        char = key
        (weight, prefix) = prefix_codes[key]
        num_prefix_bits += len(prefix) * weight
        num_raw_bits +=  weight * 8
        if VERBOSE:
            print("Char: %03d, prefix: %015s, prefix length: %02d, weight %010d" %\
                  (char, prefix, len(prefix), weight))

        if len(prefix) < min_len:
            min_len = len(prefix)
            min_char = char
            min_prefix = prefix
            min_weight = weight

        if len(prefix) > max_len:
            max_len = len(prefix)
            max_char = char
            max_prefix = prefix
            max_weight = weight

    reduction = int((1 - (num_prefix_bits / num_raw_bits)) * 100)

    print("")
    print("Minimum prefix length %02d for char %03d and prefix %015s with weight %010d" %\
          (min_len, min_char, min_prefix, min_weight))
    print("Maximum prefix length %02d for char %03d and prefix %015s with weight %010d" %\
          (max_len, max_char, max_prefix, max_weight))

    print("")
    print("Total number of bits in raw data:             %016d" % num_raw_bits)
    print("Total number of bits to prefix code all data: %016d" % num_prefix_bits)
    print("Reduction:                                    %02d percent" % reduction)
    print("")

#-------------------------------------------------------------------
# huffman_encode()
#
# Huffman encode a file with a given filename.
# Note: We assume that the file contains bytes.
#-------------------------------------------------------------------
def huffman_encode(filename):
    '''Encode the contents of the file using Huffman coding.'''

    # 1. First pass. Scan the file and build byte
    # frequency statistics.
    byte_freq = [0] * 256
    with open(filename, 'rb') as my_file:
        byte = my_file.read(1)
        bytectr = 1
        while byte:
            byteval = ord(byte)
            byte_freq[byteval] = byte_freq[byteval] + 1
            byte = my_file.read(1)
            bytectr = bytectr + 1

    if DEBUG:
        print "Encode step one."
        print "Number of bytes read: %d" % bytectr
        print byte_freq
        print ""

    # 2. Build weighted binary tree for all bytes in the file.
    alphabet = []
    for i in range(256):
        if (byte_freq[i] > 0):
            alphabet.append([i, byte_freq[i], 0, 0, 0])

    if DEBUG:
        print "Number of symbols in the alphabet: %d" % len(alphabet)
        print "The collected alphabet with weighs:"
        print alphabet

    elements_left = len(alphabet)
    while elements_left:
        pass
        # Find the two elements with the least weight

    # 3. Create symbol table based on the tree

    # 4. Emit symbol table as header for file.

    # 5. Second pass. Read file and emit symbols.

    # 6. Done!
    return 0

#-------------------------------------------------------------------
# huffman_decode()
#
# Huffman decode a file with the given filename.
#-------------------------------------------------------------------
def huffman_decode(filename):
    pass

#-------------------------------------------------------------------
# gen_node_list()
#
# Generates a list of max_types numbder of nodes. Each tuple
# contains a randomly selected frequency of a character from the
# set of characters [chr(0) .. chr((max_types - 1))]
#
# The nodes are tuples with the contents:
# (char, weight, left subtree, right subtree)
#
# For leaf nodes the subtrees are None. For all other nodes
# the char is None."
#-------------------------------------------------------------------
def gen_node_list(max_types, max_nums):
    return [(i, random.randint(0,max_nums), 0, None, None) for i in range(max_types)]

#-------------------------------------------------------------------
# encdec_huffman()
#
# Perform huffman encoding or decoding for the given infile
# as given by the arguments.
#-------------------------------------------------------------------
def encdec_huffman(args):
    pass

#-------------------------------------------------------------------
# test_huffman()
#
# Test mode. Generates data, builds prefix tree, performs
# encoding and then decoding and presents the results.
#-------------------------------------------------------------------
def test_huffman():
    print("Generating a tuple based prefix tree")
    print("====================================")
    print

    max_types = 256
    max_nums  = int(1E8)

    print("Generating %d nodes with up to %d instances." %\
          (max_types, max_nums))
    my_list = gen_node_list(max_types, max_nums)

    if VERBOSE:
        print("List of nodes before sort:")

    my_list = sort_node_list(my_list)

    if VERBOSE:
        print("List of nodes after sort:")
        print(my_list)
        print("")

    print("Generating prefix tree and extracting db for prefix codes.")
    my_tree = gen_prefix_tree(my_list)
    if VERBOSE:
        print("The generated prefix tree:")
        print(my_tree)

    my_codes = extract_prefix_codes(my_tree)
    if VERBOSE:
        print("The generated code db:")
        print(my_codes)

    print_prefix_codes(my_codes)

#-------------------------------------------------------------------
# main()
#
# Create an argument parser as needed to get input and output
# filenames as well as commands and options. Perform huffman
# encoding or decoding based on the given arguments.
#-------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('-i', '--infile',
                        help='The input file to be encoded or decoded.')

    parser.add_argument('-o', '--outfile',
                        help='The outputfile the processed file data will be saved to.')

    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Enable verbose processing.')

    parser.add_argument('-t', '--test', action='store_true',
                        help='Perform test processing.')

    parser.add_argument('-d', '--decode', action='store_true',
                        help='Perform Huffman decoding.')

    parser.add_argument('-e', '--encode', action='store_true',
                        help='Perform Huffman encoding.')

    parser.add_argument('--version', action='version', version=VERSION)

    args = parser.parse_args()

    if args.infile==None and not args.test:
        print "Error: No input file given and not in test mode."
        exit(1)

    if args.test:
        test_huffman()
    else:
        huffman_encdec(args)

#-------------------------------------------------------------------
# __name__
# Python thingy which allows the file to be run standalone as
# well as parsed from within a Python interpreter.
#-------------------------------------------------------------------
if __name__=="__main__":
    # Run the main function.
    sys.exit(main())

#=======================================================================
# EOF huffman.py
#=======================================================================
<|endoftext|>


<commit_before># Copyright 2016 Queen's University Biological Station

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from django.db import models

class Sensor(models.Model):
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	name = models.CharField(max_length=100)
	data_id = models.CharField(max_length=20) # For downloaded data, sometimes the full name would be unwieldy

	decimals = models.PositiveSmallIntegerField()

class Station(models.Model):
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	name = models.CharField(max_length=100)
	goes_id = models.CharField(max_length=8)

	sensors = models.ManyToManyField(Sensor)

class Reading(models.Model):
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	read_time = models.DateTimeField()
	value = models.IntegerField()

	sensor = models.ForeignKey("Sensor", on_delete=models.SET_NULL, null=True)
	station = models.ForeignKey("Station", on_delete=models.SET_NULL, null=True)

<commit_msg>Add a message model (which relies on PostgreSQL array fields).<commit_after># Copyright 2016 Queen's University Biological Station

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from django.db import models
from django.contrib.postgres.fields import ArrayField

class Sensor(models.Model):
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	name = models.CharField(max_length=100)
	data_id = models.CharField(max_length=20) # For downloaded data, sometimes the full name would be unwieldy

	decimals = models.PositiveSmallIntegerField()

class Station(models.Model):
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	name = models.CharField(max_length=100)
	goes_id = models.CharField(max_length=8)

	# Foreign keys
	sensors = models.ManyToManyField(Sensor)

class Reading(models.Model):
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	read_time = models.DateTimeField()
	value = models.IntegerField()

	# Foreign keys
	sensor = models.ForeignKey("Sensor", on_delete=models.SET_NULL, null=True)
	station = models.ForeignKey("Station", on_delete=models.SET_NULL, null=True)

class Message(models.Model):
	# See http://eddn.usgs.gov/dcpformat.html for format details.

	NORMAL = "N"

	LOW = "L"
	HIGH = "H"

	FAIR = "F"
	POOR = "P"

	EAST = "E"
	WEST = "W"

	MODULATION_INDEX_CHOICES = (
		(NORMAL, "Normal (60 degrees +/- 5)"),
		(LOW, "Low (50 degrees)"),
		(HIGH, "High (70 degrees)"),
	)

	DATA_QUALITY_CHOICES = (
		(NORMAL, "Normal (error rate < 10^-6)"),
		(FAIR, "Fair (10^-6 < error rate < 10^-4)"),
		(POOR, "Poor (error rate > 10^-4)"),
	)

	GOES_SPACECRAFT_CHOICES = (
		(EAST, "East"),
		(WEST, "West"),
	)

	# From http://eddn.usgs.gov/dataSourceCodes.html
	DATA_SOURCE_CHOICES = (
		("LE", "Cincinnati East; USACE LRD Cincinnati"),
		("d1", "NIFC West Boise ID - Unit 1; NIFC Boise"),
		("d2", "NIFC West Boise ID - Unit 2; NIFC Boise"),
		("OW", "Omaha West; USACE NWO"),
		("RE", "Rock Island East; USACE MVR"),
		("RW", "Rock Island West; USACE MVR"),
		("SF", "West Palm Beach East; SFWMD"),
		("UB", "Ucom Backup @ WCDA; NOAA Wallops CDA"),
		("UP", "Ucom Primary @ WCDA; NOAA Wallops CDA"),
		("XE", "Sioux Falls, East; USGS EROS"),
		("XW", "Sioux Falls, West; USGS EROS"),
		("XL", "Sioux Falls, LRIT; USGS EROS"),
		("RL", "Reston, LRIT; Reston, Virginia"),
	)

	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)

	goes_id = models.CharField(max_length=8)
	goes_channel = models.PositiveSmallIntegerField()
	goes_spacecraft = models.CharField(max_length=1, choices=GOES_SPACECRAFT_CHOICES, default=EAST)

	arrival_time = models.DateTimeField()
	failure_code = models.CharField(max_length=1)

	signal_strength = models.PositiveSmallIntegerField()
	frequency_offset = models.CharField(max_length=2)
	modulation_index = models.CharField(max_length=1, choices=MODULATION_INDEX_CHOICES, default=NORMAL)

	data_quality = models.CharField(max_length=1, choices=DATA_QUALITY_CHOICES, default=NORMAL)
	data_source = models.CharField(max_length=2, choices=DATA_SOURCE_CHOICES)

	recorded_message_length = models.PositiveSmallIntegerField()

	values = ArrayField(models.IntegerField())
	message_text = models.TextField()

	# Foreign keys
	station = models.ForeignKey("Station", on_delete=models.SET_NULL, null=True)
<|endoftext|>


<commit_before>from django.db import models
from django.db.models.signals import post_save
from django.contrib.auth.models import User
from django.forms import ModelForm
from django.conf import settings

from ummeli.base import email_copy
from ummeli.base.utils import render_to_pdf
from django.core.mail import send_mail,  EmailMessage

class Certificate (models.Model):
    name = models.CharField(max_length=45)
    institution = models.CharField(max_length=200, null=True, blank=True)
    year = models.IntegerField(default=0, null=True, blank=True)

    def __unicode__(self):  # pragma: no cover
        return self.name + " @ " + self.institution

class Language (models.Model):
    language = models.CharField(max_length=45)
    readWrite = models.BooleanField(default=False)

    def __unicode__(self):  # pragma: no cover
        return self.language

class WorkExperience (models.Model):
    title = models.CharField(max_length=45)
    company = models.CharField(max_length=45, null=True, blank=True)
    startYear = models.IntegerField(default=0, null=True, blank=True)
    endYear = models.IntegerField(default=0, null=True, blank=True)

    def __unicode__(self):  # pragma: no cover
        return self.title + " @ " + self.company

class Reference (models.Model):
    fullname = models.CharField(max_length=45)
    relationship = models.CharField(max_length=45, null=True, blank=True)
    contactNo = models.CharField(max_length=45, null=True, blank=True)

    def __unicode__(self):  # pragma: no cover
        return self.fullname

class CurriculumVitae(models.Model):
    firstName = models.CharField(max_length=45, null=True, blank=True)
    surname = models.CharField(max_length=45, null=True, blank=True)
    gender = models.CharField(max_length=45, null=True, blank=True)
    email = models.CharField(max_length=45, null=True, blank=True)
    telephoneNumber = models.CharField(max_length=45, null=True, blank=True)
    location = models.CharField(max_length=45, null=True, blank=True)
    streetName = models.CharField(max_length=45, null=True, blank=True)
    school = models.CharField(max_length=45, null=True, blank=True)
    highestGrade = models.CharField(max_length=45, null=True, blank=True)
    highestGradeYear = models.IntegerField(default=0, null=True, blank=True)
    dateOfBirth = models.CharField(max_length=45, null=True, blank=True)
    houseNumber = models.CharField(max_length=45, null=True, blank=True)
    certificates = models.ManyToManyField(Certificate, blank=True)
    languages = models.ManyToManyField(Language, blank=True)
    workExperiences = models.ManyToManyField(WorkExperience, blank=True)
    references = models.ManyToManyField(Reference, blank=True)
    user = models.OneToOneField('auth.User')
    nr_of_faxes_sent = models.IntegerField(default=0,  editable=False)

    def fullname(self):
        return '%s %s' % (self.firstName,  self.surname)

    def can_send_fax(self):
        return self.nr_of_faxes_sent < settings.MAX_LAUNCH_FAXES_COUNT

    def fax_cv(self,  fax_nr, article_text = None):
        if(self.can_send_fax()):
            self.nr_of_faxes_sent += 1
            self.save()
            return self.email_cv('%s@faxfx.net' % fax_nr.replace(' ', ''),  article_text)
        return None

    def email_cv(self, email_address, article_text = None):
        email_text = ''
        if article_text:
            email_text = email_copy.APPLY_COPY % {'sender': self.fullname(),
                                                                            'job_ad':article_text}
        else:
            email_text = email_copy.SEND_COPY % {'sender': self.fullname(),
                                                                           'job_ad':article_text}

        email = EmailMessage('CV for %s' % self.fullname(), email_text,
                                            'no-reply@ummeli.org',
                                            [email_address])
        pdf = render_to_pdf('pdf_template.html', {'model': self})
        email.attach('curriculum_vitae_for_%s_%s' % (self.firstName, self.surname),
                            pdf,  'application/pdf')
        return email.send(fail_silently=False)

    def __unicode__(self):  # pragma: no cover
        return u"CurriculumVitae %s - %s" % (self.pk, self.firstName)

class CurriculumVitaeForm(ModelForm):
    class Meta:
        model = CurriculumVitae
        exclude = ('user')

def create_cv(sender, instance, created, **kwargs):
    if created:
        cv = CurriculumVitae.objects.create(firstName=instance.first_name,
                surname=instance.last_name, email=instance.email,
                user=instance)

post_save.connect(create_cv, sender = User,
                  dispatch_uid = "users-profilecreation-signal")
<commit_msg>added faxes_remaining method<commit_after>from django.db import models
from django.db.models.signals import post_save
from django.contrib.auth.models import User
from django.forms import ModelForm
from django.conf import settings

from ummeli.base import email_copy
from ummeli.base.utils import render_to_pdf
from django.core.mail import send_mail,  EmailMessage

class Certificate (models.Model):
    name = models.CharField(max_length=45)
    institution = models.CharField(max_length=200, null=True, blank=True)
    year = models.IntegerField(default=0, null=True, blank=True)

    def __unicode__(self):  # pragma: no cover
        return self.name + " @ " + self.institution

class Language (models.Model):
    language = models.CharField(max_length=45)
    readWrite = models.BooleanField(default=False)

    def __unicode__(self):  # pragma: no cover
        return self.language

class WorkExperience (models.Model):
    title = models.CharField(max_length=45)
    company = models.CharField(max_length=45, null=True, blank=True)
    startYear = models.IntegerField(default=0, null=True, blank=True)
    endYear = models.IntegerField(default=0, null=True, blank=True)

    def __unicode__(self):  # pragma: no cover
        return self.title + " @ " + self.company

class Reference (models.Model):
    fullname = models.CharField(max_length=45)
    relationship = models.CharField(max_length=45, null=True, blank=True)
    contactNo = models.CharField(max_length=45, null=True, blank=True)

    def __unicode__(self):  # pragma: no cover
        return self.fullname

class CurriculumVitae(models.Model):
    firstName = models.CharField(max_length=45, null=True, blank=True)
    surname = models.CharField(max_length=45, null=True, blank=True)
    gender = models.CharField(max_length=45, null=True, blank=True)
    email = models.CharField(max_length=45, null=True, blank=True)
    telephoneNumber = models.CharField(max_length=45, null=True, blank=True)
    location = models.CharField(max_length=45, null=True, blank=True)
    streetName = models.CharField(max_length=45, null=True, blank=True)
    school = models.CharField(max_length=45, null=True, blank=True)
    highestGrade = models.CharField(max_length=45, null=True, blank=True)
    highestGradeYear = models.IntegerField(default=0, null=True, blank=True)
    dateOfBirth = models.CharField(max_length=45, null=True, blank=True)
    houseNumber = models.CharField(max_length=45, null=True, blank=True)
    certificates = models.ManyToManyField(Certificate, blank=True)
    languages = models.ManyToManyField(Language, blank=True)
    workExperiences = models.ManyToManyField(WorkExperience, blank=True)
    references = models.ManyToManyField(Reference, blank=True)
    user = models.OneToOneField('auth.User')
    nr_of_faxes_sent = models.IntegerField(default=0,  editable=False)

    def fullname(self):
        return '%s %s' % (self.firstName,  self.surname)

    def can_send_fax(self):
        return self.nr_of_faxes_sent < settings.MAX_LAUNCH_FAXES_COUNT

    def faxes_remaining(self):
        return settings.MAX_LAUNCH_FAXES_COUNT - self.nr_of_faxes_sent

    def fax_cv(self,  fax_nr, article_text = None):
        if(self.can_send_fax()):
            self.nr_of_faxes_sent += 1
            self.save()
            return self.email_cv('%s@faxfx.net' % fax_nr.replace(' ', ''),  article_text)
        return None

    def email_cv(self, email_address, article_text = None):
        email_text = ''
        if article_text:
            email_text = email_copy.APPLY_COPY % {'sender': self.fullname(),
                                                                            'job_ad':article_text}
        else:
            email_text = email_copy.SEND_COPY % {'sender': self.fullname(),
                                                                           'job_ad':article_text}

        email = EmailMessage('CV for %s' % self.fullname(), email_text,
                                            'no-reply@ummeli.org',
                                            [email_address])
        pdf = render_to_pdf('pdf_template.html', {'model': self})
        email.attach('curriculum_vitae_for_%s_%s' % (self.firstName, self.surname),
                            pdf,  'application/pdf')
        return email.send(fail_silently=False)

    def __unicode__(self):  # pragma: no cover
        return u"CurriculumVitae %s - %s" % (self.pk, self.firstName)

class CurriculumVitaeForm(ModelForm):
    class Meta:
        model = CurriculumVitae
        exclude = ('user')

def create_cv(sender, instance, created, **kwargs):
    if created:
        cv = CurriculumVitae.objects.create(firstName=instance.first_name,
                surname=instance.last_name, email=instance.email,
                user=instance)

post_save.connect(create_cv, sender = User,
                  dispatch_uid = "users-profilecreation-signal")
<|endoftext|>


<commit_before>#!/usr/bin/python
import os
import re
import vcf
import glob
import numpy as np

#GENE FORMAT
##chr    start    stop    name
#3    178866311    178952497    PIK3CA

from optparse import OptionParser
# -------------------------------------------------
parser = OptionParser()
parser.add_option("--vcfdir",   dest="vcfdir",     help="Path to directory containing VCF files",  default=False)
parser.add_option("--outdir",   dest="outdir",     help="Path to directory to write output to",    default="./DriverProfile/")
parser.add_option("--genelist", dest="genelist",   help="File containing Genes to test/plot)",     default=False)

parser.add_option("--bgzip",    dest="bgzip",      help="Path to bgzip binary",                    default="bgzip")
parser.add_option("--tabix",    dest="tabix",      help="Path to tabix binary",                    default="tabix")

parser.add_option("--t",        dest="nrcpus",     help="Number of CPUs to use per sample",        default=2)

parser.add_option("--dp",       dest="mindepth",   help="Minimum read depth to consider reliable", default=10)
parser.add_option("--af",       dest="minvaf",     help="Minimum variant allele fraction",         default=0.25)
parser.add_option("--pf",       dest="popfreq",    help="Maximum popultaion frequency",            default=0.05)

parser.add_option("--debug",    dest="debug",      help="Flag for debug logging",                  default=False)
parser.add_option("--format",   dest="format",     help="VCF output format [GATK/FREEB/..]",       default="GATK")
(options, args) = parser.parse_args()
# -------------------------------------------------

# PREP FILES AND ENVIRONMENT
# for i in *.vcf; do sed -r -i 's/;dbNSFP_clinvar_clnsig=(.+);/;/g' $i; done
# os.system("module load tabix")

# -------------------------------------------------
vocabulary = {"None":-1, "clean":0, "sequence_feature":0, "synonymous_variant":0, "intron_variant":0, "3_prime_UTR_variant":0.5, "5_prime_UTR_variant":0.5, "non_coding_exon_variant":0.5, "TF_binding_site_variant":1.0, "splice_region_variant":1.1, "missense_variant":1.5, "splice_donor_variant":2, "splice_acceptor_variant":2, "inframe_deletion":2.1, "inframe_insertion":2.1, "disruptive_inframe_deletion":2.5, "disruptive_inframe_insertion":2.5, "5_prime_UTR_premature_start_codon_gain_variant":3, "stop_gained":4, "nonsense_mediated_decay":4, "frameshift_variant":5}
toselect = [k for k,v in vocabulary.items() if v >= 1.5]

mapping = {"missense_variant":"Missense_Mutation", "disruptive_inframe_deletion":"Frame_Shift_Del", "disruptive_inframe_insertion":"Frame_Shift_Ins", "5_prime_UTR_premature_start_codon_gain_variant":"Nonsense_Mutation", "stop_gained":"Nonsense_Mutation", "nonsense_mediated_decay":"Nonsense_Mutation", "frameshift_variant":"Frame_Shift_???"}
lolipop = ["Hugo_Symbol","Sample_ID","Protein_Change","Mutation_Type","Chromosome","Start_Position","End_Position","Reference_Allele","Variant_Allele"]

FREQ_FIELDS = ["dbNSFP_ExAC_AF", "dbNSFP_ExAC_Adj_AF", "GoNLv5_Freq", "GoNLv5_AF"]
# -------------------------------------------------
debug = options.debug
DEPTH_KEY=""
VAF_KEY=""
# -------------------------------------------------
def check_arguments():
    global DEPTH_KEY
    global VAF_KEY

    if not os.path.exists(options.vcfdir):
        print("Invalid VCF folder %s"%(options.vcfdir))
        return False

    if not os.path.exists(options.outdir):
        print("Creating output folder %s"%(options.outdir))
        try:
            os.mkdir(options.outdir)
        except OSError:
            print("Invalid / unable to create, output folder %s"%(options.outdir))
            return(False)

    if options.format == "GATK":
        DEPTH_KEY="AD"
        VAF_KEY="AD"

    if options.format == "FREEB":
        DEPTH_KEY="DP"
        VAF_KEY="DPR"

    print("Running with the following settings:")
    print("------------------------------------")
    print(options)
    print("DEPTH FIELD:"+DEPTH_KEY)
    print("ALLELE FIELD:"+VAF_KEY)
    print("------------------------------------")
    return(True)

# -------------------------------------------------

# Extract population frequency from VCF record
# Annoation assumed to be in SNPeff formatting
def find_popfreq(vcf_record):
    popfreq=[0.0]
    for field in FREQ_FIELDS:
        if field in vcf_record.INFO:
            #if debug: print(vcf_record.INFO[field])
            for x in vcf_record.INFO[field]:
                if x is None:
                    popfreq.append(0.0)
                else:
                    popfreq.append(float(x))
    return(popfreq)

# Determine the most damaging effect of the variant
def find_effects(vcf_record):
    maxeffect="None"
    if debug: print(vcf_record.INFO)

    if "ANN" not in vcf_record.INFO:
        return maxeffect

    # TRAVERSE ALL ANNOTATIONS
    for pred in vcf_record.INFO["ANN"]:
        # SPLIT THE SEPERATE FIELDS WITHIN THE ANNOTATION
        items = pred.split("|")
        allele = items[0]
        effects = items[1].split("&")
        for effect in effects:
            if effect not in vocabulary:
                # A NEW MUTATION EFFECT WAS FOUND
                if debug:
                    print("NEW Mutation effect identified:")
                    print(pred)
                    print(effect)

            else:
                # STORE THE MOST DELETERIOUS EFFECT
                if vocabulary[effect] > vocabulary[maxeffect]:
                    maxeffect = effect
    if debug: print(maxeffect)
    return(maxeffect)

# ETRACT THE MOST DELETERIOUS MUTATIONS IN A GENE
def select_maximum_effect(effects):
    effectvalues = [vocabulary[eff] for eff in effects]
    if debug: print(effectvalues)
    indices = np.argmax(effectvalues)
    return(indices)

# CHECK AND GENERATE GZ AND TBI
def zip_and_index(vcffile):
    if not os.path.exists(vcffile+".gz"):
        os.system(options.bgzip+" -c "+vcffile+" > "+vcffile+".gz")
    if not os.path.exists(vcffile+".gz"+".tbi"):
        os.system(options.tabix+" "+vcffile+".gz")

# GENE FORMAT
# Gene name + location + variants or not
# VARIANT FORMAT
# Variant + DEPTH + POP FREQ + MLEAF + EFFECT

def check_ad(sample_vcf):
    try:
        ad_item = sample_vcf[DEPTH_KEY]
    except AttributeError as e:
        return(False)
    if sample_vcf[DEPTH_KEY] is None:
        return(False)
    return(True)

#sample_vcf == vcf_record.genotype(sample)
def check_depth(sample_vcf):
    #single depth field
    if isinstance(sample_vcf[DEPTH_KEY], int):
        # SKIP LOW DEPTH POSITIONS
        if sample_vcf[DEPTH_KEY] < int(options.mindepth):
            return(False)
    #multi depth field
    else:
        # SKIP LOW DEPTH POSITIONS
        if sum(sample_vcf[DEPTH_KEY]) < int(options.mindepth):
            return(False)
    return(True)

def check_vaf(sample_vcf):
    #single depth field
    if isinstance(sample_vcf[DEPTH_KEY], int):
        # CHECK VAF
        if (sum(sample_vcf[VAF_KEY][1:])*1.0/sample_vcf[DEPTH_KEY]) < float(options.minvaf):
            return(False)
    #multi depth field
    else:
        # CHECK VAF
        if (sum(sample_vcf[VAF_KEY][1:])*1.0/sum(sample_vcf[DEPTH_KEY])) < float(options.minvaf):
            return(False)
    return(True)

# -------------------------------------------------

def main():
    global DEPTH_KEY
    global VAF_KEY

    file_list = glob.glob(os.path.join(options.vcfdir, "*.vcf"))
    for vcf_file in file_list:
        zip_and_index(vcf_file)

    genelist=open(options.genelist, 'r').read().split('\n')

    df = {}
    rdf= {}

    # FOR ALL VCF FILES
    for vcf_file in file_list:
        if (debug):
            print("------")
            print(vcf_file)
        vcfread = vcf.Reader(open(vcf_file+".gz",'r'), compressed="gz")
        if debug: print(vcfread.samples)

        # FOR EACH SAMPLE
        for i,sample in enumerate(vcfread.samples):
            samplename = False
            if (debug): print(options.format)

            if options.format == "GATK":
                samplename = sample
            elif options.format == "FREEB":
                if (debug): print("++ "+vcfread.samples[1])
                sample = vcfread.samples[i+1]
                samplename = vcf_file.split(".")[1].split("_")[1]
            df[samplename] = {}
            rdf[samplename] = {}

        # FOR EACH GENE OF INTREST
        for gene in genelist:
            if len(gene)<=0:
                continue
            thisgene = dict(zip(["Chr","Start","Stop","SYMBOL"], gene.strip().split('\t')))

            #if debug: print(")
            vcf_records=False
            try:
                vcf_records = vcfread.fetch(thisgene["Chr"], int(thisgene["Start"])-20, int(thisgene["Stop"])+20)
            except ValueError as e:
                if debug: print("-- {}\tNO RECORDS FOUND".format(thisgene))
                for samplename in df:
                    df[samplename][thisgene["SYMBOL"]] = "None"
                continue

            for samplename in df:
                effects = []
                records = []

                # FILTER NON-QC RECORDS
                for vcf_record in vcf_records:
                    #if debug: print("** {}".format(vcf_record))
                    #if debug: print("DEPTH {}".format(vcf_record.genotype(sample)[DEPTH_KEY]))
                    #if debug: print("VAF {}".format(sum(vcf_record.genotype(sample)[VAF_KEY][1:])*1.0/vcf_record.genotype(sample)[DEPTH_KEY]))
                    #if debug: print("SKIPPING due to MLEAF")

                    #CHECK IF SAMPLE GENOTYPE AVAILABLE
                    sgenot = None
                    try:
                        sgenot = vcf_record.genotype(sample)
                    except AttributeError as e:
                        if debug: print("-- {}\t{}\tNO GT FOUND".format(thisgene, sample))
                        continue

                    PASS = False
                    log = "++ {}\t{}\t{}".format(thisgene,sample,vcf_record)
                    # CHEK IF AD FIELD PRESENT
                    if check_ad(sgenot):
                        log += "\tAD:PASS"
                        log += "\tDEPTH:{}".format(vcf_record.genotype(sample)[DEPTH_KEY])
                        # CHECK TOTAL COVERAGE OF IDENTIFIED ALLELLES
                        if check_depth(sgenot):
                            log += ":PASS"
                            log += "\tVAF:{}".format(sum(vcf_record.genotype(sample)[VAF_KEY][1:])*1.0/vcf_record.genotype(sample)[DEPTH_KEY])
                            # CHECK VARIANT ALLELE FREQUENCY
                            if check_vaf(sgenot):
                                log +=":PASS"
                                log +="\tPOP:{}".format([vcf_record.INFO[rf] for rf in FREQ_FIELDS])
                                # CHECK POPULATION FREQUENCY
                                if max(find_popfreq(vcf_record)) <= float(options.popfreq):
                                    log += ":PASS"
                                    log += "\tMLEAF:{}".format(vcf_record.INFO["MLEAF"])
                                    # CHECK OCCURENCE IN TOTAL POOL
                                    if vcf_record.INFO["MLEAF"] <= 0.4:
                                        log +=":PASS"
                                        PASS = True

                    if debug: print(log)
                    if PASS:
                        effects.append(find_effects(vcf_record))
                        records.append(vcf_record)

                # ON GENE+SAMPLE LEVEL determine maximum mutation effect
                if len(effects) <= 0:
                    df[samplename][thisgene["SYMBOL"]] = "None"
                else:
                    loc = select_maximum_effect(effects)
                    eff = effects[loc]
                    if eff in toselect:
                        df[samplename][thisgene["SYMBOL"]] = eff
                        if eff in mapping:
                            rdf[samplename][thisgene["SYMBOL"]] = {}
                            rdf[samplename][thisgene["SYMBOL"]]["REC"] = records[loc]
                            rdf[samplename][thisgene["SYMBOL"]]["EFF"] = eff
                    else:
                        df[samplename][thisgene["SYMBOL"]] = "None"
                if debug: print("** {}\t{}\t{}\t{}".format(thisgene, sample, df[samplename][thisgene["SYMBOL"]], ",".join(effects)))

        # WIRTE OUTPUT
        outfile = open(options.outdir+"/"+"MutationOverview.txt",'w')
        outfile.write("Sample\t"+'\t'.join(df[samplename].keys())+"\n")
        #print "##############################"
        for sp in df:
            outfile.write(sp+'\t'+'\t'.join(df[sp].values())+"\n")
        #    print sp+'\t'+'\t'.join(df[sp].values())
        #print "##############################"
        outfile.close()

        outfile = open(options.outdir+"/"+"MutationChart.txt",'w')
        outfile.write('\t'.join(lolipop)+"\n")
        #print "##############################"
        for samplename in rdf:
            for gene in rdf[samplename]:
                thisrec = rdf[samplename][gene]["REC"]

                proteffect=None
                for pred in thisrec.INFO["ANN"]:
                    if rdf[samplename][gene]["EFF"] in pred.split("|")[1].split("&"):
                        proteffect=pred.split("|")[10]
                        break
                if (debug):
                    print("%%\t",thisrec.INFO["ANN"])
                    print("%%\t",gene, samplename, proteffect, mapping[rdf[samplename][gene]["EFF"]], str(thisrec.CHROM), str(thisrec.POS), str(thisrec.POS+len(thisrec.ALT[0])), thisrec.REF, str(thisrec.ALT[0]))
                outfile.write("\t".join([gene, samplename, proteffect, mapping[rdf[samplename][gene]["EFF"]], str(thisrec.CHROM), str(thisrec.POS), str(thisrec.POS+len(thisrec.ALT[0])), thisrec.REF, str(thisrec.ALT[0])] )+"\n")
        #print "##############################"
        outfile.close()

# -------------------------------------------------

print("Starting Analysis")

if __name__ == '__main__':
    if check_arguments():
        main()
    else:
        print("Error in provided arguments")

print("DONE")

# -------------------------------------------------
<commit_msg>Adding sum to account for lists in AD<commit_after>#!/usr/bin/python
import os
import re
import vcf
import glob
import numpy as np

#GENE FORMAT
##chr    start    stop    name
#3    178866311    178952497    PIK3CA

from optparse import OptionParser
# -------------------------------------------------
parser = OptionParser()
parser.add_option("--vcfdir",   dest="vcfdir",     help="Path to directory containing VCF files",  default=False)
parser.add_option("--outdir",   dest="outdir",     help="Path to directory to write output to",    default="./DriverProfile/")
parser.add_option("--genelist", dest="genelist",   help="File containing Genes to test/plot)",     default=False)

parser.add_option("--bgzip",    dest="bgzip",      help="Path to bgzip binary",                    default="bgzip")
parser.add_option("--tabix",    dest="tabix",      help="Path to tabix binary",                    default="tabix")

parser.add_option("--t",        dest="nrcpus",     help="Number of CPUs to use per sample",        default=2)

parser.add_option("--dp",       dest="mindepth",   help="Minimum read depth to consider reliable", default=10)
parser.add_option("--af",       dest="minvaf",     help="Minimum variant allele fraction",         default=0.25)
parser.add_option("--pf",       dest="popfreq",    help="Maximum popultaion frequency",            default=0.05)

parser.add_option("--debug",    dest="debug",      help="Flag for debug logging",                  default=False)
parser.add_option("--format",   dest="format",     help="VCF output format [GATK/FREEB/..]",       default="GATK")
(options, args) = parser.parse_args()
# -------------------------------------------------

# PREP FILES AND ENVIRONMENT
# for i in *.vcf; do sed -r -i 's/;dbNSFP_clinvar_clnsig=(.+);/;/g' $i; done
# os.system("module load tabix")

# -------------------------------------------------
vocabulary = {"None":-1, "clean":0, "sequence_feature":0, "synonymous_variant":0, "intron_variant":0, "3_prime_UTR_variant":0.5, "5_prime_UTR_variant":0.5, "non_coding_exon_variant":0.5, "TF_binding_site_variant":1.0, "splice_region_variant":1.1, "missense_variant":1.5, "splice_donor_variant":2, "splice_acceptor_variant":2, "inframe_deletion":2.1, "inframe_insertion":2.1, "disruptive_inframe_deletion":2.5, "disruptive_inframe_insertion":2.5, "5_prime_UTR_premature_start_codon_gain_variant":3, "stop_gained":4, "nonsense_mediated_decay":4, "frameshift_variant":5}
toselect = [k for k,v in vocabulary.items() if v >= 1.5]

mapping = {"missense_variant":"Missense_Mutation", "disruptive_inframe_deletion":"Frame_Shift_Del", "disruptive_inframe_insertion":"Frame_Shift_Ins", "5_prime_UTR_premature_start_codon_gain_variant":"Nonsense_Mutation", "stop_gained":"Nonsense_Mutation", "nonsense_mediated_decay":"Nonsense_Mutation", "frameshift_variant":"Frame_Shift_???"}
lolipop = ["Hugo_Symbol","Sample_ID","Protein_Change","Mutation_Type","Chromosome","Start_Position","End_Position","Reference_Allele","Variant_Allele"]

FREQ_FIELDS = ["dbNSFP_ExAC_AF", "dbNSFP_ExAC_Adj_AF", "GoNLv5_Freq", "GoNLv5_AF"]
# -------------------------------------------------
debug = options.debug
DEPTH_KEY=""
VAF_KEY=""
# -------------------------------------------------
def check_arguments():
    global DEPTH_KEY
    global VAF_KEY

    if not os.path.exists(options.vcfdir):
        print("Invalid VCF folder %s"%(options.vcfdir))
        return False

    if not os.path.exists(options.outdir):
        print("Creating output folder %s"%(options.outdir))
        try:
            os.mkdir(options.outdir)
        except OSError:
            print("Invalid / unable to create, output folder %s"%(options.outdir))
            return(False)

    if options.format == "GATK":
        DEPTH_KEY="AD"
        VAF_KEY="AD"

    if options.format == "FREEB":
        DEPTH_KEY="DP"
        VAF_KEY="DPR"

    print("Running with the following settings:")
    print("------------------------------------")
    print(options)
    print("DEPTH FIELD:"+DEPTH_KEY)
    print("ALLELE FIELD:"+VAF_KEY)
    print("------------------------------------")
    return(True)

# -------------------------------------------------

# Extract population frequency from VCF record
# Annoation assumed to be in SNPeff formatting
def find_popfreq(vcf_record):
    popfreq=[0.0]
    for field in FREQ_FIELDS:
        if field in vcf_record.INFO:
            #if debug: print(vcf_record.INFO[field])
            for x in vcf_record.INFO[field]:
                if x is None:
                    popfreq.append(0.0)
                else:
                    popfreq.append(float(x))
    return(popfreq)

# Determine the most damaging effect of the variant
def find_effects(vcf_record):
    maxeffect="None"
    if debug: print(vcf_record.INFO)

    if "ANN" not in vcf_record.INFO:
        return maxeffect

    # TRAVERSE ALL ANNOTATIONS
    for pred in vcf_record.INFO["ANN"]:
        # SPLIT THE SEPERATE FIELDS WITHIN THE ANNOTATION
        items = pred.split("|")
        allele = items[0]
        effects = items[1].split("&")
        for effect in effects:
            if effect not in vocabulary:
                # A NEW MUTATION EFFECT WAS FOUND
                if debug:
                    print("NEW Mutation effect identified:")
                    print(pred)
                    print(effect)

            else:
                # STORE THE MOST DELETERIOUS EFFECT
                if vocabulary[effect] > vocabulary[maxeffect]:
                    maxeffect = effect
    if debug: print(maxeffect)
    return(maxeffect)

# ETRACT THE MOST DELETERIOUS MUTATIONS IN A GENE
def select_maximum_effect(effects):
    effectvalues = [vocabulary[eff] for eff in effects]
    if debug: print(effectvalues)
    indices = np.argmax(effectvalues)
    return(indices)

# CHECK AND GENERATE GZ AND TBI
def zip_and_index(vcffile):
    if not os.path.exists(vcffile+".gz"):
        os.system(options.bgzip+" -c "+vcffile+" > "+vcffile+".gz")
    if not os.path.exists(vcffile+".gz"+".tbi"):
        os.system(options.tabix+" "+vcffile+".gz")

# GENE FORMAT
# Gene name + location + variants or not
# VARIANT FORMAT
# Variant + DEPTH + POP FREQ + MLEAF + EFFECT

def check_ad(sample_vcf):
    try:
        ad_item = sample_vcf[DEPTH_KEY]
    except AttributeError as e:
        return(False)
    if sample_vcf[DEPTH_KEY] is None:
        return(False)
    return(True)

#sample_vcf == vcf_record.genotype(sample)
def check_depth(sample_vcf):
    #single depth field
    if isinstance(sample_vcf[DEPTH_KEY], int):
        # SKIP LOW DEPTH POSITIONS
        if sample_vcf[DEPTH_KEY] < int(options.mindepth):
            return(False)
    #multi depth field
    else:
        # SKIP LOW DEPTH POSITIONS
        if sum(sample_vcf[DEPTH_KEY]) < int(options.mindepth):
            return(False)
    return(True)

def check_vaf(sample_vcf):
    #single depth field
    if isinstance(sample_vcf[DEPTH_KEY], int):
        # CHECK VAF
        if (sum(sample_vcf[VAF_KEY][1:])*1.0/sample_vcf[DEPTH_KEY]) < float(options.minvaf):
            return(False)
    #multi depth field
    else:
        # CHECK VAF
        if (sum(sample_vcf[VAF_KEY][1:])*1.0/sum(sample_vcf[DEPTH_KEY])) < float(options.minvaf):
            return(False)
    return(True)

# -------------------------------------------------

def main():
    global DEPTH_KEY
    global VAF_KEY

    file_list = glob.glob(os.path.join(options.vcfdir, "*.vcf"))
    for vcf_file in file_list:
        zip_and_index(vcf_file)

    genelist=open(options.genelist, 'r').read().split('\n')

    df = {}
    rdf= {}

    # FOR ALL VCF FILES
    for vcf_file in file_list:
        if (debug):
            print("------")
            print(vcf_file)
        vcfread = vcf.Reader(open(vcf_file+".gz",'r'), compressed="gz")

        if (debug): print(vcfread.samples)
        if (debug): print(options.format)

        # FOR EACH SAMPLE
        for i,sample in enumerate(vcfread.samples):
            samplename = False

            if options.format == "GATK":
                samplename = sample
            elif options.format == "FREEB":
                if (debug): print("++ "+vcfread.samples[1])
                sample = vcfread.samples[i+1]
                samplename = vcf_file.split(".")[1].split("_")[1]
            df[samplename] = {}
            rdf[samplename] = {}

        # FOR EACH GENE OF INTREST
        for gene in genelist:
            if len(gene)<=0:
                continue
            thisgene = dict(zip(["Chr","Start","Stop","SYMBOL"], gene.strip().split('\t')))

            #if debug: print(")
            vcf_records=False
            try:
                vcf_records = vcfread.fetch(thisgene["Chr"], int(thisgene["Start"])-20, int(thisgene["Stop"])+20)
            except ValueError as e:
                if debug: print("-- {}\tNO RECORDS FOUND".format(thisgene))
                for samplename in df:
                    df[samplename][thisgene["SYMBOL"]] = "None"
                continue

            for samplename in df:
                effects = []
                records = []

                # FILTER NON-QC RECORDS
                for vcf_record in vcf_records:
                    #if debug: print("** {}".format(vcf_record))
                    #if debug: print("DEPTH {}".format(vcf_record.genotype(sample)[DEPTH_KEY]))
                    #if debug: print("VAF {}".format(sum(vcf_record.genotype(sample)[VAF_KEY][1:])*1.0/vcf_record.genotype(sample)[DEPTH_KEY]))
                    #if debug: print("SKIPPING due to MLEAF")

                    #CHECK IF SAMPLE GENOTYPE AVAILABLE
                    sgenot = None
                    try:
                        sgenot = vcf_record.genotype(sample)
                    except AttributeError as e:
                        if debug: print("-- {}\t{}\tNO GT FOUND".format(thisgene, sample))
                        continue

                    PASS = False
                    log = "++ {}\t{}\t{}".format(thisgene,sample,vcf_record)
                    # CHEK IF AD FIELD PRESENT
                    if check_ad(sgenot):
                        log += "\tAD:PASS"
                        log += "\tDEPTH:{}".format(vcf_record.genotype(sample)[DEPTH_KEY])
                        # CHECK TOTAL COVERAGE OF IDENTIFIED ALLELLES
                        if check_depth(sgenot):
                            log += ":PASS"
                            log += "\tVAF:{}".format(sum(vcf_record.genotype(sample)[VAF_KEY][1:])*1.0/sum(vcf_record.genotype(sample)[DEPTH_KEY]))
                            # CHECK VARIANT ALLELE FREQUENCY
                            if check_vaf(sgenot):
                                log +=":PASS"
                                log +="\tPOP:{}".format([vcf_record.INFO[rf] for rf in FREQ_FIELDS])
                                # CHECK POPULATION FREQUENCY
                                if max(find_popfreq(vcf_record)) <= float(options.popfreq):
                                    log += ":PASS"
                                    log += "\tMLEAF:{}".format(vcf_record.INFO["MLEAF"])
                                    # CHECK OCCURENCE IN TOTAL POOL
                                    if vcf_record.INFO["MLEAF"] <= 0.4:
                                        log +=":PASS"
                                        PASS = True

                    if debug: print(log)
                    if PASS:
                        effects.append(find_effects(vcf_record))
                        records.append(vcf_record)

                # ON GENE+SAMPLE LEVEL determine maximum mutation effect
                if len(effects) <= 0:
                    df[samplename][thisgene["SYMBOL"]] = "None"
                else:
                    loc = select_maximum_effect(effects)
                    eff = effects[loc]
                    if eff in toselect:
                        df[samplename][thisgene["SYMBOL"]] = eff
                        if eff in mapping:
                            rdf[samplename][thisgene["SYMBOL"]] = {}
                            rdf[samplename][thisgene["SYMBOL"]]["REC"] = records[loc]
                            rdf[samplename][thisgene["SYMBOL"]]["EFF"] = eff
                    else:
                        df[samplename][thisgene["SYMBOL"]] = "None"
                if debug: print("** {}\t{}\t{}\t{}".format(thisgene, sample, df[samplename][thisgene["SYMBOL"]], ",".join(effects)))

        # WIRTE OUTPUT
        outfile = open(options.outdir+"/"+"MutationOverview.txt",'w')
        outfile.write("Sample\t"+'\t'.join(df[samplename].keys())+"\n")
        #print "##############################"
        for sp in df:
            outfile.write(sp+'\t'+'\t'.join(df[sp].values())+"\n")
        #    print sp+'\t'+'\t'.join(df[sp].values())
        #print "##############################"
        outfile.close()

        outfile = open(options.outdir+"/"+"MutationChart.txt",'w')
        outfile.write('\t'.join(lolipop)+"\n")
        #print "##############################"
        for samplename in rdf:
            for gene in rdf[samplename]:
                thisrec = rdf[samplename][gene]["REC"]

                proteffect=None
                for pred in thisrec.INFO["ANN"]:
                    if rdf[samplename][gene]["EFF"] in pred.split("|")[1].split("&"):
                        proteffect=pred.split("|")[10]
                        break
                if (debug):
                    print("%%\t",thisrec.INFO["ANN"])
                    print("%%\t",gene, samplename, proteffect, mapping[rdf[samplename][gene]["EFF"]], str(thisrec.CHROM), str(thisrec.POS), str(thisrec.POS+len(thisrec.ALT[0])), thisrec.REF, str(thisrec.ALT[0]))
                outfile.write("\t".join([gene, samplename, proteffect, mapping[rdf[samplename][gene]["EFF"]], str(thisrec.CHROM), str(thisrec.POS), str(thisrec.POS+len(thisrec.ALT[0])), thisrec.REF, str(thisrec.ALT[0])] )+"\n")
        #print "##############################"
        outfile.close()

# -------------------------------------------------

print("Starting Analysis")

if __name__ == '__main__':
    if check_arguments():
        main()
    else:
        print("Error in provided arguments")

print("DONE")

# -------------------------------------------------
<|endoftext|>