#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AUTHOR

    Pedro Cerqueira
    github: @pedrorvc

DESCRIPTION

    This script serves to create xml files contaning the information necessary
    for the execution of BRIG (Blast Ring Image Generator), reducing the time
    performing the tedious task of setting up all the information on the GUI
    and provides a quick way to produce an image.

    The arguments for this script provide some (but not all)
    of the available options in BRIG, which were the ones I used to change the most.

    USAGE:

    brigaid.py -q reference_sequence.fna -rfd path/to/reference/dir -od path/to/output/dir -of path/to/output/dir/output_file
                                -oi path/to/output/BRIG/output_image -t Image_title -a annotation_file.gbk --genes genes_of_interest.txt
                                --contig-order contig_order.tsv

"""

import argparse
import csv
import os
import xml.etree.ElementTree as ET
from collections import OrderedDict
from xml.dom import minidom

from Bio import SeqIO
from matplotlib import cm

def listdir_fullpath(path):
    """ Gets the full path of the files from a directory

        Args:
            path (str): full path to a directory

        Returns:
            list containing the full path of every file contained in the input directory

    """

    return [os.path.join(path, f) for f in os.listdir(path)]

def ring_attributes(colour, name, position):
    """ Creates ring attributes.

        Args:
            colour (str): color of the ring.
            name (str): name of the ring.
            position (str): position of the ring.

        Returns:
            ring_attrs (dict): attributes of any regular ring of the BRIG xml.
    """

    ring_attrs = {"colour" : colour,
                 "name": name,
                 "position" : position,
                 "upperInt" : "90",
                 "lowerInt" : "70",
                 "legend" : "yes",
                 "size" : "30",
                 "labels" : "no",
                 "blastType" : "blastn"}

    return ring_attrs

def annotation_ring_attributes(position):
    """ Creates annotation ring attributes.

        Args:
            position (str): position of the ring.

        Returns:
            annotation_ring_attrs (dict): attributes of the annotation ring of the BRIG xml.
    """

    annotation_ring_attrs = {"colour" : '172,14,225',
                 "name": 'null',
                 "position" : position,
                 "upperInt" : "70",
                 "lowerInt" : "50",
                 "legend" : "yes",
                 "size" : "30",
                 "labels" : "no",
                 "blastType" : "blastn"}

    return annotation_ring_attrs

def create_feature_attrs(label, colour, decoration, start, stop):
    """ Create attributes for the Feature SubElements of the annotation ring.

        Args:
            label (str): name of the gene/CDS to annotate
            colour (str): colour of the decoration for the annotation
            decoration (str): shape of the gene/CDS to annotate, for example, 'clockwise-arrow'
            start (str): start of the gene/CDS to annotate
            stop (str): stop of the gene/CDS to annotate

        Results:
            feature_element_attrs (dict): attributes of the feature element.
            feature_range_element_attrs (dict): attributes of the feature range element
    """

    feature_element_attrs = {'label' : label,
                             'colour' : colour,
                             'decoration' : decoration}

    feature_range_element_attrs = {'start' : start,
                                   'stop' : stop}

    return feature_element_attrs, feature_range_element_attrs

def create_annotation_ring_tsv(annotation_ring, annotation_file):
    """ Uses a tsv file to annotate the reference genome.

        Args:
            annotation_ring: ElementTree SubElement object containing the 'ring' tag and its attributes.
            annotation_file (str): Full path to the file containing annotations for the reference genome.

    """

    with open(annotation_file) as tsvfile:
      reader = csv.DictReader(tsvfile, dialect='excel-tab')

      # Obtain the annotations from the file contents
      for row in reader:
          start = row['#START']
          stop = row['STOP']
          label = row['Label']
          colour = row['Colour']
          decoration = row['Decoration']

          # Create xml attributes
          feature_element_attrs, feature_range_element_attrs = create_feature_attrs(label, colour, decoration, start, stop)

          # Create xml elements
          feature_element = ET.SubElement(annotation_ring, 'feature', attrib=feature_element_attrs)
          feature_range_element = ET.SubElement(feature_element, 'featureRange', attrib=feature_range_element_attrs)

def annotation_ring_feature_elements_gbk_concat(annotation_ring, record, genome_size=False):
    """ Creates the annotation ring feature elements, using a concatenated Genbank annotation file.

        Args:
            annotation_ring: ElementTree SubElement object containing the 'ring' tag and its attributes.
            record (SeqRecord): Object of BioPython containing the information of the input Genbank.
            genome_size (bool): Size of genome. Integer when a Genbank divided by contigs is provided.
                                                Boolean (False) when a concatenated Genbank is provided.

    """

    #if type(genome_size) == int:

    # Obtain the features of the Genbank file records
    for fea in record.features:

        # Get the start and end position of the genome
        # Also get the strand
        if fea.type == 'CDS':
            start = str(fea.location.start.position)
            end = str(fea.location.end.position)
            strand = fea.location.strand

            # Get the label of the gene or product
            if 'gene' in fea.qualifiers:
                label = str(fea.qualifiers['gene'][0])
            elif 'product' in fea.qualifiers:
                product = fea.qualifiers['product'][0]
                label = str(product)
            else:
                continue

            # Define the decoration of the annotation based on the strand
            if strand == -1:
                decoration = 'counterclockwise-arrow'

            elif strand == 1:
                decoration = 'clockwise-arrow'

            # Create xml attributes
            feature_element_attrs, feature_range_element_attrs = create_feature_attrs(label, "black", decoration, start, end)

            # Create xml elements
            feature_element = ET.SubElement(annotation_ring, 'feature', attrib=feature_element_attrs)
            feature_range_element = ET.SubElement(feature_element, 'featureRange', attrib=feature_range_element_attrs)

        # If a genome size is provided, get the size of the records
        if type(genome_size) == int:

            if fea.type == 'source':
                size = fea.location.end.position

    try:
        size
        genome_size += size

        return genome_size

    except NameError:
        pass

def annotation_ring_feature_elements_genes_of_interest_gbk_concat(annotation_ring, record, genes, genome_size=False):
    """ Creates the annotation ring feature elements, using a concatenated Genbank annotation file
        and specific gene annotations.

        Args:
            annotation_ring: ElementTree SubElement object containing the 'ring' tag and its attributes.
            record (SeqRecord): Object of BioPython containing the information of the input Genbank.
            genome_size (bool): Size of genome. Integer when a Genbank divided by contigs is provided.
                                                Boolean (False) when a concatenated Genbank is provided.

    """

    for f in record.features:

        if f.type == 'CDS':

            # Find the 'gene' tag and determine if the gene belongs to the specified genes to be annotated
            if 'gene' in f.qualifiers and f.qualifiers['gene'][0] in genes:
                label = f.qualifiers['gene'][0]
            elif 'product' in f.qualifiers and f.qualifiers['product'][0] in genes:
                product = f.qualifiers['product'][0]
                label = product
            else:
                continue

            # Determine the start, stop and strand of the gene
            start = str(f.location.start.position + genome_size)
            end = str(f.location.end.position + genome_size)
            strand = f.location.strand

            # Define the decoration of the annotation based on the strand
            if strand == -1:
                decoration = 'counterclockwise-arrow'

            elif strand == 1:
                decoration = 'clockwise-arrow'

            # Create xml attributes
            feature_element_attrs, feature_range_element_attrs = create_feature_attrs(label, "black", decoration, start, end)

            # Create xml elements
            feature_element = ET.SubElement(annotation_ring, 'feature', attrib=feature_element_attrs)
            feature_range_element = ET.SubElement(feature_element, 'featureRange', attrib=feature_range_element_attrs)

        # If a genome size is provided, get the size of the records
        if type(genome_size) == int:

            if f.type == "source":
                size = f.location.end.position

    try:
        size
        genome_size += size

        return genome_size

    except NameError:
        pass

def create_annotation_ring_gbk_concat(annotation_ring, annotation_file, genes_of_interest, records):
    """ Create annotation ring using a concatenated Genbank annotation file.

        Args:
            annotation_ring: ElementTree SubElement object containing the 'ring' tag and its attributes.
            annotation_file (str): Full path to the file containing annotations for the reference genome.
            genes_of_interest (str): Full path to the file containing the genes to search for in the Genbank file.
            records (SeqRecord): Object of BioPython containing the information of the input Genbank.

    """

    if genes_of_interest != []:

        # Get the genes to serach in the Genbank file
        with open(genes_of_interest, "r") as f:
            genes = f.readlines()
            genes = [gene.rstrip() for gene in genes]

        # Create feature elements of the annotation ring
        for seq_record in records:
            annotation_ring_feature_elements_genes_of_interest_gbk_concat(annotation_ring, seq_record, genes)

    else:

        for seq_record in records:
            annotation_ring_feature_elements_gbk_concat(annotation_ring, seq_record)

def create_annotation_ring_gbk_contigs(annotation_ring, annotation_file, records, genes_of_interest, contig_order):
    """ Create annotation ring using a Genbank annotation file divided by contigs.

        Args:
            annotation_ring: ElementTree SubElement object containing the 'ring' tag and its attributes.
            annotation_file (str): Full path to the file containing annotations for the reference genome.
            genes_of_interest (str): Full path to the file containing the genes to search for in the Genbank file.
            records (SeqRecord): Object of BioPython containing the information of the input Genbank.
            contig_order (str): Full path to the file containing the order of the contigs.

    """

    if contig_order != []:

        with open(contig_order) as tsvfile:
            reader = csv.DictReader(tsvfile, dialect='excel-tab')

            # Create an OrderedDict with the contents of the file
            # The keys are the order are a number representing the order of the contig
            # The values are the names of the contigs
            content_dict = OrderedDict()
            for r in reader:
                content_dict[r["order"]] = r["contig"]

        # Create an OrderedDict with the content of each contig
        # The keys are the names of the contigs
        # The values are SeqRecord objects from BipPython
        seq_records_dict = OrderedDict()
        for record in records:
            seq_records_dict[record.id] = record

        if genes_of_interest != []:

            with open(genes_of_interest, "r") as f:
                genes = f.readlines()
                genes = [gene.rstrip() for gene in genes]

            genome_size = 0
            for i in range(1, len(records)+1):
                ord_record = seq_records_dict[content_dict[str(i)]]

                gsize = annotation_ring_feature_elements_genes_of_interest_gbk_concat(annotation_ring, ord_record, genes, genome_size)

                genome_size = gsize
        else:

            genome_size = 0
            for i in range(1, len(records)+1):
                ord_record = seq_records_dict[content_dict[str(i)]]

                gsize = annotation_ring_feature_elements_gbk_concat(annotation_ring, ord_record, genome_size)

                genome_size = gsize

    else:

        if genes_of_interest != []:

            with open(genes_of_interest, "r") as f:
                genes = f.readlines()
                genes = [gene.rstrip() for gene in genes]

            for seq_record in records:
                annotation_ring_feature_elements_genes_of_interest_gbk_concat(annotation_ring, seq_record, genes)

        else:
            for seq_record in records:
                annotation_ring_feature_elements_gbk_concat(annotation_ring, seq_record)

def write_xml(root_elem, output_file):
    """ Writes a xml file.

        Args:
            root_elem is a ElementTree Element object containing all the information
            required for the output file.
            output_file (str): full path to the output file

    """

    xml_file = ET.tostring(root_elem, encoding='utf8').decode('utf8')

    pretty_xml_file = minidom.parseString(xml_file).toprettyxml(indent='    ')

    output_file = output_file + ".xml"

    with open(output_file, "w") as f:
        f.write(pretty_xml_file)

####### Create xml elemnts

# Create root element

def create_root_element(blast_options, legend_position, query_file,
                               output_folder, image_output_file, title, image_format):
    """
    Creates the root element of the xml file and its attributes.

    Args:
        blast_options (str): additional options for blast, for example, -evalue or num_threads
        legend_position (str): position of the legend on the image
        query_file (str): full path to the query file
        output_folder (str): full path to the output folder
        image_output_file (str): full path to the image output file
        title (str): title of the output image
        image_format (str): format of the image output file

    Returns:
        root: ElementTree Element object containing the BRIG tag and its attributes

    """

    root_attrs = {"blastOptions" : blast_options,
                  "legendPosition" : legend_position,
                  "queryFile" : query_file,
                  "outputFolder" : output_folder,
                  "blastPlus" : "yes",
                  "outputFile" : os.path.join(output_folder, image_output_file),
                  "title" : title,
                  "imageFormat" : image_format,
                  "queryFastaFile" : query_file,
                  "cgXML" : os.path.join(output_folder + "/scratch", os.path.basename(query_file) + ".xml")}

    root = ET.Element('BRIG', attrib=root_attrs)

    return root

#### Create root children

# Create cgview_settings element

def create_cgview_settings_element(root, height, width):
    """ Creates the cgview_settings element of the xml file and its attributes.

        Args:
            root: ElementTree Element object containing the BRIG tag and its attributes.
            height (str): height of the output image in pixels
            width (str): width of the output image in pixels

        Returns:
            cgview_settings: ElementTree SubElement object containing the cgview settings tag and its attributes
    """

    cgview_settings_attrs = {"arrowheadLength" : "medium",
                             "backboneColor" : "black",
                             "backboneRadius" : "600",
                             "backboneThickness" : "medium",
                             "backgroundColor" : "white",
                             "borderColor" : "black",
                             "featureSlotSpacing" : "medium",
                             "featureThickness" : "30",
                             "giveFeaturePositions" : "false",
                             "globalLabel" : "true",
                             "height" : height,
                             "isLinear" : "false",
                             "labelFont" : "SansSerif,plain,25",
                             "labelLineLength" : "medium",
                             "labelLineThickness" : "medium",
                             "labelPlacementQuality" : "best",
                             "labelsToKeep" : "1000",
                             "longTickColor" : "black",
                             "minimumFeatureLength" : "medium",
                             "moveInnerLabelsToOuter" :"true",
                             "origin" : "12",
                             "rulerFont" : "SansSerif,plain,35",
                             "rulerFontColor" : "black",
                             "rulerPadding" : "40",
                             "rulerUnits" : "bases",
                             "shortTickColor" : "black",
                             "shortTickThickness" : "medium",
                             "showBorder" : "false",
                             "showShading" : "true",
                             "showWarning" : "false",
                             "tickDensity" : "0.2333",
                             "tickThickness" : "medium",
                             "titleFont" : "SansSerif,plain,45",
                             "titleFontColor" : "black",
                             "useColoredLabelBackgrounds" : "false",
                             "useInnerLabels" : "true",
                             "warningFont" : "Default,plain,35",
                             "warningFontColor" : "black",
                             "width" : width,
                             "zeroTickColor" : "black",
                             "tickLength" : "medium"}

    cgview_settings = ET.SubElement(root, 'cgview_settings', attrib=cgview_settings_attrs)

    return cgview_settings

# Create brig_settings element

def create_brig_settings_element(root, java_memory):
    """ Creates the brig_settings element of the xml file and its attributes.

        Args:
            root: ElementTree Element object containing the BRIG tag and its attributes.
            java_memory (str): amount of memory (in bytes) java is allowed to use for BRIG

        Returns:
            brig_settings: ElementTree SubElement object containing the brig settings tag and its attributes

    """

    brig_settings_attrs = {"Ring1" : "172,14,225",
                           "Ring2" : "222,149,220",
                           "Ring3" : "161,221,231",
                           "Ring4" : "49,34,221",
                           "Ring5" : "116,152,226",
                           "Ring6" : "224,206,38",
                           "Ring7" : "40,191,140",
                           "Ring8" : "158,223,139",
                           "Ring9" : "226,38,122",
                           "Ring10" :"211,41,77",
                           "defaultUpper" : "70",
                           "defaultLower" : "50",
                           "defaultMinimum" : "50",
                           "genbankFiles" : "gbk,gb,genbank",
                           "fastaFiles" : "fna,faa,fas,fasta,fa",
                           "emblFiles" : "embl",
                           "blastLocation" : "",
                           "divider" : "3",
                           "multiplier" : "3",
                           "memory" : java_memory,
                           "defaultSpacer" : "0"}

    brig_settings = ET.SubElement(root,
                                  "brig_settings",
                                  attrib=brig_settings_attrs)

    return brig_settings

## Create special element

def create_special_element(root):
    """Creates the 'special' element of the xml file and its attributes

        Args:
            root: ElementTree Element object containing the BRIG tag and its attributes.

        Returns:
            gc_content_special: ElementTree SubElement object containing the 'special' tag and its attributes
            gc_skew_special: ElementTree SubElement object containing the 'special' tag and its attributes

    """

    gc_content_special = ET.SubElement(root, 'special', attrib={'value' : 'GC Content'})
    gc_skew_special = ET.SubElement(root, 'special', attrib={'value' : 'GC Skew'})

    return gc_content_special, gc_skew_special

# Create reference dir element

def create_reference_directory_element(root, reference_directory):
    """ Creates the 'reference directory' element of the xml file and its attributes.

        Args:
            root: ElementTree Element object containing the 'BRIG' tag and its attributes.
            reference_directory (str): full path to the reference directory that contains
                                        the fasta files used to build the rings.

        Returns:
            ref_file: ElementTree SubElement object containing the 'refFile' tag and its attributes

    """

    ref_dir = ET.SubElement(root,
                            "refDir",
                            attrib={"location" : reference_directory})

    # Obtain the full path for all the files in the directory
    ref_dir_list = listdir_fullpath(reference_directory)

    for f in ref_dir_list:
        ref_file = ET.SubElement(ref_dir,
                                 "refFile",
                                 attrib={"location" : f})

    return ref_file

# Create the ring where the annotations are defined

def create_annotation_ring(root, reference_directory, annotation_file, genes_of_interest, contig_order):
    """ Creates the ring that will contain the annotations for the reference genome.

        Args:
            root: ElementTree Element object containing the 'BRIG' tag and its attributes.
            reference_directory (str): full path to the reference directory that contains
                                        the fasta files used to build the rings.
            annotation_file (str): Full path to the file containing annotations for the reference genome.
            genes_of_interest (str): Full path to the file containing a list of specific genes.
            contig_order (str): Full path to the tab-delimited file containing the order of the contigs.

    """

    # Determine the position of the annotation ring, which will be the position after the last reference genome
    ring_position = len(os.listdir(reference_directory)) + 2

    # Create the annotation ring element
    annotation_ring = ET.SubElement(root, 'ring', attrib=annotation_ring_attributes(str(ring_position)))

    # Check for tab-delimited annotation file input
    if list(SeqIO.parse(annotation_file, "genbank")) == []:
        create_annotation_ring_tsv(annotation_ring, annotation_file)

    else:

        # Get the records of the Genbank file
        records = [r for r in SeqIO.parse(annotation_file, "genbank")]

        ### Check if a contig order file has been provided

        if len(records) > 1:   # If more than 1 record exists, then the Genbank file is divided by contigs
            create_annotation_ring_gbk_contigs(annotation_ring, annotation_file, records, genes_of_interest, contig_order)
        else:
            create_annotation_ring_gbk_concat(annotation_ring, annotation_file, genes_of_interest, records)

## Create remaining rings

def create_ring_element(root, reference_directory, colormap):
    """ Creates the ring elements of the xml file, containing the position and color of the rings.

        Args:
            root: ElementTree Element object containing the 'BRIG' tag and its attributes.

            reference_directory (str): full path to the reference directory that contains
                                        the fasta files used to build the rings.
            colormap (str): name of the colormap (available in matplotlib) to use for the color of the rings

        Returns:
            ring_number_element: ElementTree SubElement object containing the 'ring' tag and its attributes
            ring_sequence_element: ElementTree SubElement object containing the 'sequence' tag and its attributes

    """

    ref_dir_list = listdir_fullpath(reference_directory)

    # Gets the colormap from matplotlib with as many colors as the number of files
    cmap = cm.get_cmap(colormap, len(ref_dir_list))

    list_colormap = cmap.colors.tolist()

    # Remove the fourth element (transparency) because it is not necessary
    colors_to_use = []
    for l in list_colormap:
        convert = [round(x * 255) for x in l]
        convert.pop()
        colors_to_use.append(convert)

    #reversed_colors_to_use = colors_to_use[::-1]

    # Check if the user provided an order for the rings
    has_digit = [os.path.basename(x).split("_")[0].isdigit() for x in ref_dir_list]

    if True in has_digit:
        # Obtain the ring positions
        ring_positions = [os.path.basename(x).split("_")[0] for x in ref_dir_list]

        # Reverse sort the positions of the rings, because they will be created
        # in a descending order of their positions
        ring_positions.sort(reverse=True)
        ref_dir_list.sort(reverse=True)

        for ring in range(len(ref_dir_list)):

            # The ring positions start at 2 due to the special rings (GC Content and GC Skew)
            ring_position = int(ring_positions[ring]) + 1

            # Select a color for the ring
            ring_color = ",".join([str(e) for e in colors_to_use[ring]])

            # Define the name of the ring
            ring_name = os.path.basename(ref_dir_list[ring]).split("_")[1]

            # Create the xml elements
            ring_number_element = ET.SubElement(root,
                                    'ring',
                                    ring_attributes(ring_color, ring_name, str(ring_position)))

            ring_sequence_element = ET.SubElement(ring_number_element,
                                      "sequence",
                                      attrib={"location" : ref_dir_list[ring]})

    else:
        # Sort files by lowercase
        ref_dir_list.sort(key=lambda y: y.lower())

        # The number of rings starts at 2 due to the GC Content and GC Skew
        ring_number = len(ref_dir_list) + 1
        for ring in range(len(ref_dir_list)):

            # Select a color for the ring
            ring_color = ",".join([str(e) for e in colors_to_use[ring]])

            # Define the name of the ring
            ring_name = os.path.basename(ref_dir_list[ring]).split("_")[0]

            # Create the xml elements
            ring_number_element = ET.SubElement(root,
                                                'ring',
                                                ring_attributes(ring_color, ring_name, str(ring_number)))

            ring_sequence_element = ET.SubElement(ring_number_element,
                                                  "sequence",
                                                  attrib={"location" : ref_dir_list[ring]})

            ring_number -= 1

    return ring_number_element, ring_sequence_element

## Create special rings

def create_special_ring_element(root):
    """ Create the 'special' ring element and its attributes.

        Args:
            root: ElementTree Element object containing the 'BRIG' tag and its attributes.

        Returns:
            gc_content_location: ElementTree SubElement object containing the 'sequence' tag and its attributes
            gc_skew_location: ElementTree SubElement object containing the 'sequence' tag and its attributes
    """

    # Create ring attributes
    gc_content_ring_attrs = ring_attributes('225,0,0', "GC Content", "0")
    gc_skew_ring_attrs = ring_attributes('225,0,0', "GC Skew", "1")

    # Add ring element to root
    gc_skew_ring = ET.SubElement(root, 'ring', attrib=gc_skew_ring_attrs)
    gc_content_ring = ET.SubElement(root, 'ring', attrib=gc_content_ring_attrs)

    # Add sequence element to ring
    gc_content_location = ET.SubElement(gc_content_ring, 'sequence', attrib={'location' : 'GC Content'})
    gc_skew_location = ET.SubElement(gc_skew_ring, 'sequence', attrib={'location' : 'GC Skew'})

    return gc_content_location, gc_skew_location

def main(query_file, reference_directory, output_folder, output_xml, image_output_file, title, annotation_file,
         genes_of_interest, contig_order, blast_options, legend_position, image_format, height, width, java_memory, colormap):

    root = create_root_element(blast_options, legend_position, query_file,
                               output_folder, image_output_file, title, image_format)

    cgview_settings = create_cgview_settings_element(root, height, width)

    brig_settings = create_brig_settings_element(root, java_memory)

    special = create_special_element(root)

    refdir = create_reference_directory_element(root, reference_directory)

    if annotation_file:
        create_annotation_ring(root, reference_directory, annotation_file, genes_of_interest, contig_order)

    rings = create_ring_element(root, reference_directory, colormap)

    special_ring = create_special_ring_element(root)

    write_xml(root, output_xml)

    print("\n File written to {}".format(output_xml))

def parse_arguments():

    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('-q', '--query', type=str, required=True, dest='query_file',
                        help='Path to the query/reference FASTA file.')

    parser.add_argument('-rfd', '--ref_dir', type=str, required=True, dest='reference_directory',
                        help='Path to the directory where the FASTA files to compare against the reference are located.')

    parser.add_argument('-od', '--out_dir', type=str, required=True, dest='output_folder',
                        help='Path to the output directory for the results of BRIG.')

    parser.add_argument('-of', '--out_xml', type=str, required=True, dest='output_file',
                        help='Path to the output of this script.')

    parser.add_argument('-oi', '--out_img', type=str, required=True, dest='image_output_file',
                        help='Path to the output file of the resulting image of BRIG.')

    parser.add_argument('-t', '--title', type=str, required=True, dest='title',
                        help='Title of the resulting image from BRIG.')

    parser.add_argument('-a', '--annotation', type=str, required=False, dest='annotation_file', default=False,
                        help='File containing annotations for the reference genome. '
                        'The annoation file can be a tab-delimited file (.tsv) or a Genbank format file (.gbk, .gb)')

    parser.add_argument('--genes', type=str, required=False, dest='genes_of_interest', default=[],
                        help='File containing a list of specific genes (one gene per line) to search when a Genbank annotation file is provided. ')

    parser.add_argument('--contig_order', type=str, required=False, dest='contig_order', default=[],
                        help='Tab-delimited file containing the order of the contigs when a Genbank (divided by contigs) annotation file is provided. '
                             'Example:  order    contig '
                                          '1     Contig8')

    parser.add_argument('-b', '--blast_options', type=str, required=False, dest="blast_options", default="-evalue 0.001 -num_threads 6",
                        help='Options for running BLAST.')

    parser.add_argument('-l', '--legend_pos', type=str, required=False, dest="legend_position", default="middle-right",
                        help='Positon of the legend on the resulting image.'
                        'The options available are upper, center or lower, '
                        'paired with left, center or right')

    parser.add_argument('-if', '--image_format', type=str, required=False, dest="image_format", default="jpg",
                        help='Format of the resulting image file.'
                        'The available options are: jpg, png, svg or svgz.')

    parser.add_argument('-ht', '--height', type=str, required=False, dest="height", default="3000",
                        help='Height (in pixels) of the resulting image.')

    parser.add_argument('-wd', '--width', type=str, required=False, dest="width", default="3000",
                        help='Width (in pixels) of the resulting image.')

    parser.add_argument('-jm', '--java_memory', type=str, required=False, dest="java_memory", default="1500",
                        help='Amount of memory (in bytes) that Java is allowed to use for BRIG.')

    parser.add_argument('-cm', '--colormap', type=str, required=False, dest="colormap", default="viridis",
                        help='Colormap from matplotlib to use for the color of the rings. '
                        'The available options are: viridis, plasma, inferno, magma and cividis.'
                        'More options for colormaps at: '
                        'https://matplotlib.org/users/colormaps.html')

    args = parser.parse_args()

    return [args.query_file, args.reference_directory, args.output_folder, args.output_file,
            args.image_output_file, args.title, args.annotation_file, args.genes_of_interest, args.contig_order,
            args.blast_options, args.legend_position, args.image_format, args.height, args.width, args.java_memory, args.colormap]

if __name__ == '__main__':

    args = parse_arguments()
    main(args[0], args[1], args[2], args[3], args[4], args[5], args[6],
         args[7], args[8], args[9], args[10], args[11], args[12], args[13],
         args[14], args[15])


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

from matplotlib.externals import six

from nose.tools import assert_equal

from matplotlib.testing.decorators import knownfailureif
from pylab import *

def test_simple():
    assert_equal(1 + 1, 2)

@knownfailureif(True)
def test_simple_knownfail():
    # Test the known fail mechanism.
    assert_equal(1 + 1, 3)

def test_override_builtins():
    ok_to_override = set([
        '__name__',
        '__doc__',
        '__package__',
        '__loader__',
        '__spec__',
        'any',
        'all',
        'sum'
    ])

    # We could use six.moves.builtins here, but that seems
    # to do a little more than just this.
    if six.PY3:
        builtins = sys.modules['builtins']
    else:
        builtins = sys.modules['__builtin__']

    overridden = False
    for key in globals().keys():
        if key in dir(builtins):
            if (globals()[key] != getattr(builtins, key) and
                    key not in ok_to_override):
                print("'%s' was overridden in globals()." % key)
                overridden = True

    assert not overridden

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)


# -*- coding: utf-8 -*-
from pandas.compat import range

import numpy as np
from numpy.random import RandomState

from pandas.core.api import Series, Categorical, CategoricalIndex
import pandas as pd

from pandas import compat
import pandas.core.algorithms as algos
import pandas.util.testing as tm
import pandas.hashtable as hashtable

class TestMatch(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_ints(self):
        values = np.array([0, 2, 1])
        to_match = np.array([0, 1, 2, 2, 0, 1, 3, 0])

        result = algos.match(to_match, values)
        expected = np.array([0, 2, 1, 1, 0, 2, -1, 0])
        self.assert_numpy_array_equal(result, expected)

        result = Series(algos.match(to_match, values, np.nan))
        expected = Series(np.array([0, 2, 1, 1, 0, 2, np.nan, 0]))
        tm.assert_series_equal(result, expected)

        s = pd.Series(np.arange(5),dtype=np.float32)
        result = algos.match(s, [2,4])
        expected = np.array([-1, -1, 0, -1, 1])
        self.assert_numpy_array_equal(result, expected)

        result = Series(algos.match(s, [2,4], np.nan))
        expected = Series(np.array([np.nan, np.nan, 0, np.nan, 1]))
        tm.assert_series_equal(result,expected)

    def test_strings(self):
        values = ['foo', 'bar', 'baz']
        to_match = ['bar', 'foo', 'qux', 'foo', 'bar', 'baz', 'qux']

        result = algos.match(to_match, values)
        expected = np.array([1, 0, -1, 0, 1, 2, -1])
        self.assert_numpy_array_equal(result, expected)

        result = Series(algos.match(to_match, values, np.nan))
        expected = Series(np.array([1, 0, np.nan, 0, 1, 2, np.nan]))
        tm.assert_series_equal(result,expected)

class TestFactorize(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_warn(self):

        s = Series([1, 2, 3])
        with tm.assert_produces_warning(FutureWarning):
            algos.factorize(s, order='A')

    def test_basic(self):

        labels, uniques = algos.factorize(['a', 'b', 'b', 'a',
                                           'a', 'c', 'c', 'c'])
        # self.assert_numpy_array_equal(labels, np.array([ 0, 1, 1, 0, 0, 2, 2, 2],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array(['a','b','c'], dtype=object))

        labels, uniques = algos.factorize(['a', 'b', 'b', 'a',
                                           'a', 'c', 'c', 'c'], sort=True)
        self.assert_numpy_array_equal(labels, np.array([ 0, 1, 1, 0, 0, 2, 2, 2],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array(['a','b','c'], dtype=object))

        labels, uniques = algos.factorize(list(reversed(range(5))))
        self.assert_numpy_array_equal(labels, np.array([0, 1, 2, 3, 4], dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array([ 4, 3, 2, 1, 0],dtype=np.int64))

        labels, uniques = algos.factorize(list(reversed(range(5))), sort=True)
        self.assert_numpy_array_equal(labels, np.array([ 4, 3, 2, 1, 0],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array([0, 1, 2, 3, 4], dtype=np.int64))

        labels, uniques = algos.factorize(list(reversed(np.arange(5.))))
        self.assert_numpy_array_equal(labels, np.array([0., 1., 2., 3., 4.], dtype=np.float64))
        self.assert_numpy_array_equal(uniques, np.array([ 4, 3, 2, 1, 0],dtype=np.int64))

        labels, uniques = algos.factorize(list(reversed(np.arange(5.))), sort=True)
        self.assert_numpy_array_equal(labels, np.array([ 4, 3, 2, 1, 0],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array([0., 1., 2., 3., 4.], dtype=np.float64))

    def test_mixed(self):

        # doc example reshaping.rst
        x = Series(['A', 'A', np.nan, 'B', 3.14, np.inf])
        labels, uniques = algos.factorize(x)

        self.assert_numpy_array_equal(labels, np.array([ 0,  0, -1,  1,  2,  3],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array(['A', 'B', 3.14, np.inf], dtype=object))

        labels, uniques = algos.factorize(x, sort=True)
        self.assert_numpy_array_equal(labels, np.array([ 2,  2, -1,  3,  0,  1],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array([3.14, np.inf, 'A', 'B'], dtype=object))

    def test_datelike(self):

        # M8
        v1 = pd.Timestamp('20130101 09:00:00.00004')
        v2 = pd.Timestamp('20130101')
        x = Series([v1,v1,v1,v2,v2,v1])
        labels, uniques = algos.factorize(x)
        self.assert_numpy_array_equal(labels, np.array([ 0,0,0,1,1,0],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array([v1.value,v2.value],dtype='M8[ns]'))

        labels, uniques = algos.factorize(x, sort=True)
        self.assert_numpy_array_equal(labels, np.array([ 1,1,1,0,0,1],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, np.array([v2.value,v1.value],dtype='M8[ns]'))

        # period
        v1 = pd.Period('201302',freq='M')
        v2 = pd.Period('201303',freq='M')
        x = Series([v1,v1,v1,v2,v2,v1])

        # periods are not 'sorted' as they are converted back into an index
        labels, uniques = algos.factorize(x)
        self.assert_numpy_array_equal(labels, np.array([ 0,0,0,1,1,0],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, pd.PeriodIndex([v1, v2]))

        labels, uniques = algos.factorize(x,sort=True)
        self.assert_numpy_array_equal(labels, np.array([ 0,0,0,1,1,0],dtype=np.int64))
        self.assert_numpy_array_equal(uniques, pd.PeriodIndex([v1, v2]))

    def test_factorize_nan(self):
        # nan should map to na_sentinel, not reverse_indexer[na_sentinel]
        # rizer.factorize should not raise an exception if na_sentinel indexes
        # outside of reverse_indexer
        key = np.array([1, 2, 1, np.nan], dtype='O')
        rizer = hashtable.Factorizer(len(key))
        for na_sentinel in (-1, 20):
            ids = rizer.factorize(key, sort=True, na_sentinel=na_sentinel)
            expected = np.array([0, 1, 0, na_sentinel], dtype='int32')
            self.assertEqual(len(set(key)), len(set(expected)))
            self.assertTrue(np.array_equal(pd.isnull(key), expected == na_sentinel))

        # nan still maps to na_sentinel when sort=False
        key = np.array([0, np.nan, 1], dtype='O')
        na_sentinel = -1
        ids = rizer.factorize(key, sort=False, na_sentinel=na_sentinel)
        expected = np.array([ 2, -1,  0], dtype='int32')
        self.assertEqual(len(set(key)), len(set(expected)))
        self.assertTrue(np.array_equal(pd.isnull(key), expected == na_sentinel))

    def test_vector_resize(self):
        # Test for memory errors after internal vector
        # reallocations (pull request #7157)

        def _test_vector_resize(htable, uniques, dtype, nvals):
            vals = np.array(np.random.randn(1000), dtype=dtype)
            # get_labels appends to the vector
            htable.get_labels(vals[:nvals], uniques, 0, -1)
            # to_array resizes the vector
            uniques.to_array()
            htable.get_labels(vals, uniques, 0, -1)

        test_cases = [
            (hashtable.PyObjectHashTable, hashtable.ObjectVector, 'object'),
            (hashtable.Float64HashTable,  hashtable.Float64Vector, 'float64'),
            (hashtable.Int64HashTable,    hashtable.Int64Vector, 'int64')]

        for (tbl, vect, dtype) in test_cases:
            # resizing to empty is a special case
            _test_vector_resize(tbl(), vect(), dtype, 0)
            _test_vector_resize(tbl(), vect(), dtype, 10)

class TestIndexer(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_outer_join_indexer(self):
        typemap = [('int32', algos.algos.outer_join_indexer_int32),
                   ('int64', algos.algos.outer_join_indexer_int64),
                   ('float32', algos.algos.outer_join_indexer_float32),
                   ('float64', algos.algos.outer_join_indexer_float64),
                   ('object', algos.algos.outer_join_indexer_object)]

        for dtype, indexer in typemap:
            left = np.arange(3, dtype = dtype)
            right = np.arange(2,5, dtype = dtype)
            empty = np.array([], dtype = dtype)

            result, lindexer, rindexer = indexer(left, right)
            tm.assertIsInstance(result, np.ndarray)
            tm.assertIsInstance(lindexer, np.ndarray)
            tm.assertIsInstance(rindexer, np.ndarray)
            tm.assert_numpy_array_equal(result, np.arange(5, dtype = dtype))
            tm.assert_numpy_array_equal(lindexer, np.array([0, 1, 2, -1, -1]))
            tm.assert_numpy_array_equal(rindexer, np.array([-1, -1, 0, 1, 2]))

            result, lindexer, rindexer = indexer(empty, right)
            tm.assert_numpy_array_equal(result, right)
            tm.assert_numpy_array_equal(lindexer, np.array([-1, -1, -1]))
            tm.assert_numpy_array_equal(rindexer, np.array([0, 1, 2]))

            result, lindexer, rindexer = indexer(left, empty)
            tm.assert_numpy_array_equal(result, left)
            tm.assert_numpy_array_equal(lindexer, np.array([0, 1, 2]))
            tm.assert_numpy_array_equal(rindexer, np.array([-1, -1, -1]))

class TestUnique(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_ints(self):
        arr = np.random.randint(0, 100, size=50)

        result = algos.unique(arr)
        tm.assertIsInstance(result, np.ndarray)

    def test_objects(self):
        arr = np.random.randint(0, 100, size=50).astype('O')

        result = algos.unique(arr)
        tm.assertIsInstance(result, np.ndarray)

    def test_object_refcount_bug(self):
        lst = ['A', 'B', 'C', 'D', 'E']
        for i in range(1000):
            len(algos.unique(lst))

    def test_on_index_object(self):

        mindex = pd.MultiIndex.from_arrays([np.arange(5).repeat(5),
                                            np.tile(np.arange(5), 5)])
        expected = mindex.values
        expected.sort()

        mindex = mindex.repeat(2)

        result = pd.unique(mindex)
        result.sort()

        tm.assert_almost_equal(result, expected)

    def test_datetime64_dtype_array_returned(self):
        # GH 9431
        expected = np.array(['2015-01-03T00:00:00.000000000+0000',
                             '2015-01-01T00:00:00.000000000+0000'], dtype='M8[ns]')

        dt_index = pd.to_datetime(['2015-01-03T00:00:00.000000000+0000',
                                   '2015-01-01T00:00:00.000000000+0000',
                                   '2015-01-01T00:00:00.000000000+0000'])
        result = algos.unique(dt_index)
        tm.assert_numpy_array_equal(result, expected)
        self.assertEqual(result.dtype, expected.dtype)

        s = pd.Series(dt_index)
        result = algos.unique(s)
        tm.assert_numpy_array_equal(result, expected)
        self.assertEqual(result.dtype, expected.dtype)

        arr = s.values
        result = algos.unique(arr)
        tm.assert_numpy_array_equal(result, expected)
        self.assertEqual(result.dtype, expected.dtype)

    def test_timedelta64_dtype_array_returned(self):
        # GH 9431
        expected = np.array([31200, 45678, 10000], dtype='m8[ns]')

        td_index = pd.to_timedelta([31200, 45678, 31200, 10000, 45678])
        result = algos.unique(td_index)
        tm.assert_numpy_array_equal(result, expected)
        self.assertEqual(result.dtype, expected.dtype)

        s = pd.Series(td_index)
        result = algos.unique(s)
        tm.assert_numpy_array_equal(result, expected)
        self.assertEqual(result.dtype, expected.dtype)

        arr = s.values
        result = algos.unique(arr)
        tm.assert_numpy_array_equal(result, expected)
        self.assertEqual(result.dtype, expected.dtype)

class TestIsin(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_invalid(self):

        self.assertRaises(TypeError, lambda : algos.isin(1,1))
        self.assertRaises(TypeError, lambda : algos.isin(1,[1]))
        self.assertRaises(TypeError, lambda : algos.isin([1],1))

    def test_basic(self):

        result = algos.isin([1,2],[1])
        expected = np.array([True,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(np.array([1,2]),[1])
        expected = np.array([True,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(pd.Series([1,2]),[1])
        expected = np.array([True,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(pd.Series([1,2]),pd.Series([1]))
        expected = np.array([True,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(['a','b'],['a'])
        expected = np.array([True,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(pd.Series(['a','b']),pd.Series(['a']))
        expected = np.array([True,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(['a','b'],[1])
        expected = np.array([False,False])
        tm.assert_numpy_array_equal(result, expected)

        arr = pd.date_range('20130101',periods=3).values
        result = algos.isin(arr,[arr[0]])
        expected = np.array([True,False,False])
        tm.assert_numpy_array_equal(result, expected)

        result = algos.isin(arr,arr[0:2])
        expected = np.array([True,True,False])
        tm.assert_numpy_array_equal(result, expected)

        arr = pd.timedelta_range('1 day',periods=3).values
        result = algos.isin(arr,[arr[0]])
        expected = np.array([True,False,False])
        tm.assert_numpy_array_equal(result, expected)

    def test_large(self):

        s = pd.date_range('20000101',periods=2000000,freq='s').values
        result = algos.isin(s,s[0:2])
        expected = np.zeros(len(s),dtype=bool)
        expected[0] = True
        expected[1] = True
        tm.assert_numpy_array_equal(result, expected)

class TestValueCounts(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_value_counts(self):
        np.random.seed(1234)
        from pandas.tools.tile import cut

        arr = np.random.randn(4)
        factor = cut(arr, 4)

        tm.assertIsInstance(factor, Categorical)
        result = algos.value_counts(factor)
        cats = ['(-1.194, -0.535]',
                '(-0.535, 0.121]',
                '(0.121, 0.777]',
                '(0.777, 1.433]'
        ]
        expected_index = CategoricalIndex(cats, cats, ordered=True)
        expected = Series([1, 1, 1, 1],
                          index=expected_index)
        tm.assert_series_equal(result.sort_index(), expected.sort_index())

    def test_value_counts_bins(self):
        s = [1, 2, 3, 4]
        result = algos.value_counts(s, bins=1)
        self.assertEqual(result.tolist(), [4])
        self.assertEqual(result.index[0], 0.997)

        result = algos.value_counts(s, bins=2, sort=False)
        self.assertEqual(result.tolist(), [2, 2])
        self.assertEqual(result.index[0], 0.997)
        self.assertEqual(result.index[1], 2.5)

    def test_value_counts_dtypes(self):
        result = algos.value_counts([1, 1.])
        self.assertEqual(len(result), 1)

        result = algos.value_counts([1, 1.], bins=1)
        self.assertEqual(len(result), 1)

        result = algos.value_counts(Series([1, 1., '1']))  # object
        self.assertEqual(len(result), 2)

        self.assertRaises(TypeError, lambda s: algos.value_counts(s, bins=1), ['1', 1])

    def test_value_counts_nat(self):
        td = Series([np.timedelta64(10000), pd.NaT], dtype='timedelta64[ns]')
        dt = pd.to_datetime(['NaT', '2014-01-01'])

        for s in [td, dt]:
            vc = algos.value_counts(s)
            vc_with_na = algos.value_counts(s, dropna=False)
            self.assertEqual(len(vc), 1)
            self.assertEqual(len(vc_with_na), 2)

        exp_dt = pd.Series({pd.Timestamp('2014-01-01 00:00:00'): 1})
        tm.assert_series_equal(algos.value_counts(dt), exp_dt)
        # TODO same for (timedelta)

    def test_categorical(self):
        s = Series(pd.Categorical(list('aaabbc')))
        result = s.value_counts()
        expected = pd.Series([3, 2, 1], index=pd.CategoricalIndex(['a', 'b', 'c']))
        tm.assert_series_equal(result, expected, check_index_type=True)

        # preserve order?
        s = s.cat.as_ordered()
        result = s.value_counts()
        expected.index = expected.index.as_ordered()
        tm.assert_series_equal(result, expected, check_index_type=True)

    def test_categorical_nans(self):
        s = Series(pd.Categorical(list('aaaaabbbcc'))) # 4,3,2,1 (nan)
        s.iloc[1] = np.nan
        result = s.value_counts()
        expected = pd.Series([4, 3, 2],
                             index=pd.CategoricalIndex(['a', 'b', 'c'],
                                                       categories=['a', 'b', 'c']))
        tm.assert_series_equal(result, expected, check_index_type=True)
        result = s.value_counts(dropna=False)
        expected = pd.Series([4, 3, 2, 1], index=pd.CategoricalIndex(
            ['a', 'b',  'c', np.nan]))
        tm.assert_series_equal(result, expected, check_index_type=True)

        # out of order
        s = Series(pd.Categorical(list('aaaaabbbcc'),
                                  ordered=True, categories=['b', 'a', 'c']))
        s.iloc[1] = np.nan
        result = s.value_counts()
        expected = pd.Series([4, 3, 2],
                             index=pd.CategoricalIndex(['a', 'b', 'c'],
                                                       categories=['b', 'a', 'c'],
                                                       ordered=True))
        tm.assert_series_equal(result, expected, check_index_type=True)

        result = s.value_counts(dropna=False)
        expected = pd.Series([4, 3, 2, 1], index=pd.CategoricalIndex(
            ['a', 'b',  'c', np.nan], categories=['b', 'a', 'c'], ordered=True))
        tm.assert_series_equal(result, expected, check_index_type=True)

    def test_categorical_zeroes(self):
        # keep the `d` category with 0
        s = Series(pd.Categorical(list('bbbaac'), categories=list('abcd'),
                                  ordered=True))
        result = s.value_counts()
        expected = Series([3, 2, 1, 0], index=pd.Categorical(
            ['b', 'a', 'c', 'd'], categories=list('abcd'), ordered=True))
        tm.assert_series_equal(result, expected, check_index_type=True)

    def test_dropna(self):
        # https://github.com/pydata/pandas/issues/9443#issuecomment-73719328

        tm.assert_series_equal(
            pd.Series([True, True, False]).value_counts(dropna=True),
            pd.Series([2, 1], index=[True, False]))
        tm.assert_series_equal(
            pd.Series([True, True, False]).value_counts(dropna=False),
            pd.Series([2, 1], index=[True, False]))

        tm.assert_series_equal(
            pd.Series([True, True, False, None]).value_counts(dropna=True),
            pd.Series([2, 1], index=[True, False]))
        tm.assert_series_equal(
            pd.Series([True, True, False, None]).value_counts(dropna=False),
            pd.Series([2, 1, 1], index=[True, False, np.nan]))
        tm.assert_series_equal(
            pd.Series([10.3, 5., 5.]).value_counts(dropna=True),
            pd.Series([2, 1], index=[5., 10.3]))
        tm.assert_series_equal(
            pd.Series([10.3, 5., 5.]).value_counts(dropna=False),
            pd.Series([2, 1], index=[5., 10.3]))

        tm.assert_series_equal(
            pd.Series([10.3, 5., 5., None]).value_counts(dropna=True),
            pd.Series([2, 1], index=[5., 10.3]))

        # 32-bit linux has a different ordering
        if not compat.is_platform_32bit():
            tm.assert_series_equal(
                pd.Series([10.3, 5., 5., None]).value_counts(dropna=False),
                pd.Series([2, 1, 1], index=[5., 10.3, np.nan]))

class GroupVarTestMixin(object):

    def test_group_var_generic_1d(self):
        prng = RandomState(1234)

        out = (np.nan * np.ones((5, 1))).astype(self.dtype)
        counts = np.zeros(5, dtype='int64')
        values = 10 * prng.rand(15, 1).astype(self.dtype)
        labels = np.tile(np.arange(5), (3, )).astype('int64')

        expected_out = (np.squeeze(values)
                        .reshape((5, 3), order='F')
                        .std(axis=1, ddof=1) ** 2)[:, np.newaxis]
        expected_counts = counts + 3

        self.algo(out, counts, values, labels)
        np.testing.assert_allclose(out, expected_out, self.rtol)
        tm.assert_numpy_array_equal(counts, expected_counts)

    def test_group_var_generic_1d_flat_labels(self):
        prng = RandomState(1234)

        out = (np.nan * np.ones((1, 1))).astype(self.dtype)
        counts = np.zeros(1, dtype='int64')
        values = 10 * prng.rand(5, 1).astype(self.dtype)
        labels = np.zeros(5, dtype='int64')

        expected_out = np.array([[values.std(ddof=1) ** 2]])
        expected_counts = counts + 5

        self.algo(out, counts, values, labels)

        np.testing.assert_allclose(out, expected_out, self.rtol)
        tm.assert_numpy_array_equal(counts, expected_counts)

    def test_group_var_generic_2d_all_finite(self):
        prng = RandomState(1234)

        out = (np.nan * np.ones((5, 2))).astype(self.dtype)
        counts = np.zeros(5, dtype='int64')
        values = 10 * prng.rand(10, 2).astype(self.dtype)
        labels = np.tile(np.arange(5), (2, )).astype('int64')

        expected_out = np.std(
            values.reshape(2, 5, 2), ddof=1, axis=0) ** 2
        expected_counts = counts + 2

        self.algo(out, counts, values, labels)
        np.testing.assert_allclose(out, expected_out, self.rtol)
        tm.assert_numpy_array_equal(counts, expected_counts)

    def test_group_var_generic_2d_some_nan(self):
        prng = RandomState(1234)

        out = (np.nan * np.ones((5, 2))).astype(self.dtype)
        counts = np.zeros(5, dtype='int64')
        values = 10 * prng.rand(10, 2).astype(self.dtype)
        values[:, 1] = np.nan
        labels = np.tile(np.arange(5), (2, )).astype('int64')

        expected_out = np.vstack([
            values[:, 0].reshape(5, 2, order='F').std(ddof=1, axis=1) ** 2,
            np.nan * np.ones(5)
        ]).T
        expected_counts = counts + 2

        self.algo(out, counts, values, labels)
        np.testing.assert_allclose(out, expected_out, self.rtol)
        tm.assert_numpy_array_equal(counts, expected_counts)

    def test_group_var_constant(self):
        # Regression test from GH 10448.

        out = np.array([[np.nan]], dtype=self.dtype)
        counts = np.array([0],dtype='int64')
        values = 0.832845131556193 * np.ones((3, 1), dtype=self.dtype)
        labels = np.zeros(3, dtype='int64')

        self.algo(out, counts, values, labels)

        self.assertEqual(counts[0], 3)
        self.assertTrue(out[0, 0] >= 0)  # Python 2.6 has no assertGreaterEqual
        tm.assert_almost_equal(out[0, 0], 0.0)

class TestGroupVarFloat64(tm.TestCase, GroupVarTestMixin):
    __test__ = True
    _multiprocess_can_split_ = True

    algo = algos.algos.group_var_float64
    dtype = np.float64
    rtol = 1e-5

    def test_group_var_large_inputs(self):

        prng = RandomState(1234)

        out = np.array([[np.nan]], dtype=self.dtype)
        counts = np.array([0],dtype='int64')
        values = (prng.rand(10 ** 6) + 10 ** 12).astype(self.dtype)
        values.shape = (10 ** 6, 1)
        labels = np.zeros(10 ** 6, dtype='int64')

        self.algo(out, counts, values, labels)

        self.assertEqual(counts[0], 10 ** 6)
        tm.assert_almost_equal(out[0, 0], 1.0 / 12, check_less_precise=True)

class TestGroupVarFloat32(tm.TestCase, GroupVarTestMixin):
    __test__ = True
    _multiprocess_can_split_ = True

    algo = algos.algos.group_var_float32
    dtype = np.float32
    rtol = 1e-2

def test_quantile():
    s = Series(np.random.randn(100))

    result = algos.quantile(s, [0, .25, .5, .75, 1.])
    expected = algos.quantile(s.values, [0, .25, .5, .75, 1.])
    tm.assert_almost_equal(result, expected)

def test_unique_label_indices():
    from pandas.hashtable import unique_label_indices

    a = np.random.randint(1, 1 << 10, 1 << 15).astype('i8')

    left = unique_label_indices(a)
    right = np.unique(a, return_index=True)[1]

    tm.assert_numpy_array_equal(left, right)

    a[np.random.choice(len(a), 10)] = -1
    left= unique_label_indices(a)
    right = np.unique(a, return_index=True)[1][1:]
    tm.assert_numpy_array_equal(left, right)

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)


'''
Copyright (C) 2016 Turi
All rights reserved.

This software may be modified and distributed under the terms
of the BSD license. See the LICENSE file for details.
'''
import unittest
import pandas
import array
from .. import SFrame
from pandas.util.testing import assert_frame_equal
from sys import version_info

class DataFrameTest(unittest.TestCase):
    def test_empty(self):
        expected = pandas.DataFrame()
        assert_frame_equal(SFrame(expected).to_dataframe(), expected)
        expected['int'] = []
        expected['float'] = []
        expected['str'] = []
        assert_frame_equal(SFrame(expected).to_dataframe(), expected)

    def test_simple_dataframe(self):
        expected = pandas.DataFrame()
        expected['int'] = [i for i in range(10)]
        expected['float'] = [float(i) for i in range(10)]
        expected['str'] = [str(i) for i in range(10)]
        if version_info.major == 2:
            expected['unicode'] = [unicode(i) for i in range(10)]
        expected['array'] = [array.array('d', [i]) for i in range(10)]
        expected['ls'] = [[str(i)] for i in range(10)]
        assert_frame_equal(SFrame(expected).to_dataframe(), expected)

    def test_sparse_dataframe(self):
        expected = pandas.DataFrame()
        expected['sparse_int'] = [i if i % 2 == 0 else None for i in range(10)]
        expected['sparse_float'] = [float(i) if i % 2  == 1 else None for i in range(10)]
        expected['sparse_str'] = [str(i) if i % 3 == 0 else None for i in range(10)]
        expected['sparse_array'] = [array.array('d', [i]) if i % 5 == 0 else None for i in range(10)]
        expected['sparse_list'] = [[str(i)] if i % 7 == 0 else None for i in range(10)]
        assert_frame_equal(SFrame(expected).to_dataframe(), expected)


import csv
import matplotlib.pyplot as plt
from numpy import *
import scipy.interpolate
import math
from pylab import *
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
import matplotlib.patches as patches
from matplotlib.path import Path
import os

# ------------------------------------------------------------------------------------------------------
#inputs
for file in os.listdir('.'):
    if file.endswith(".grd"):
    	inputfile = file

for file in os.listdir('.'):
    if file.endswith(".txt"):
    	inputfile2 = file
# ------------------------------------------------------------------------------------------------------
#Patches data

#for the Kewley and Levesque data
verts = [
    (1., 7.97712125471966000000), # left, bottom
    (1., 9.57712125471966000000), # left, top
    (2., 10.57712125471970000000), # right, top
    (2., 8.97712125471966000000), # right, bottom
    (0., 0.), # ignored
    ]
codes = [Path.MOVETO,
         Path.LINETO,
         Path.LINETO,
         Path.LINETO,
         Path.CLOSEPOLY,
         ]

path = Path(verts, codes)
# ------------------------
#for the Kewley 01 data
verts2 = [
    (2.4, 9.243038049), # left, bottom
    (2.4, 11.0211893), # left, top
    (2.6, 11.0211893), # right, top
    (2.6, 9.243038049), # right, bottom
    (0, 0.), # ignored
    ]
path = Path(verts, codes)
path2 = Path(verts2, codes)
# -------------------------
#for the  Moy et al data
verts3 = [
    (1., 6.86712125471966000000), # left, bottom
    (1., 10.18712125471970000000), # left, top
    (3., 12.18712125471970000000), # right, top
    (3., 8.86712125471966000000), # right, bottom
    (0., 0.), # ignored
    ]
path = Path(verts, codes)
path3 = Path(verts3, codes)
# ------------------------------------------------------------------------------------------------------

#the routine to add patches for others peoples' data onto our plots.
def add_patches(ax):
	patch3 = patches.PathPatch(path3, facecolor='yellow', lw=0)
	patch2 = patches.PathPatch(path2, facecolor='green', lw=0)
	patch = patches.PathPatch(path, facecolor='red', lw=0)
	ax1.add_patch(patch3)
	ax1.add_patch(patch2)
	ax1.add_patch(patch)
# ------------------------------------------------------------------------------------------------------

#the subplot routine

def add_sub_plot(sub_num):
	numplots = 16

	plt.subplot(numplots/4.,4,sub_num)

	rbf = scipy.interpolate.Rbf(x, y, z[:,sub_num-1], function='linear')
	zi = rbf(xi, yi)

	contour = plt.contour(xi,yi,zi, levels, colors='c', linestyles = 'dashed')
	contour2 = plt.contour(xi,yi,zi, levels2, colors='k', linewidths=1.5)

	plt.scatter(max_values[line[sub_num-1],2], max_values[line[sub_num-1],3], c ='k',marker = '*')
	plt.annotate(headers[line[sub_num-1]], xy=(8,11),  xytext=(6,8.5), fontsize = 10)
	plt.annotate(max_values[line[sub_num-1],0], xy= (max_values[line[sub_num-1],2], max_values[line[sub_num-1],3]), xytext = (0, -10), textcoords = 'offset points', ha = 'right', va = 'bottom', fontsize=10)

	if sub_num == numplots / 2.:
		print "half the plots are complete"
#axis limits

	yt_min = 8
	yt_max = 23
	xt_min = 0
	xt_max = 12
	plt.ylim(yt_min,yt_max)
	plt.xlim(xt_min,xt_max)
	plt.yticks(arange(yt_min+1,yt_max,1),fontsize=10)
	plt.xticks(arange(xt_min+1,xt_max,1), fontsize = 10)

	if sub_num in [2,3,4,6,7,8,10,11,12,14,15,16]:
		plt.tick_params(labelleft = 'off')
	else:
		plt.tick_params(labelleft = 'on')
		plt.ylabel('Log ($ \phi  _{\mathrm{H}}  $)')

	if sub_num in [1,2,3,4,5,6,7,8,9,10,11,12]:
		plt.tick_params(labelbottom = 'off')
	else:
		plt.tick_params(labelbottom = 'on')
		plt.xlabel('Log($n _{\mathrm{H}}  $)')

	if sub_num == 1:
		plt.yticks(arange(yt_min+1,yt_max+1,1),fontsize=10)
	if sub_num == 13:
		plt.yticks(arange(yt_min,yt_max,1),fontsize=10)
		plt.xticks(arange(xt_min,xt_max,1), fontsize = 10)
	if sub_num == 16 :
		plt.xticks(arange(xt_min+1,xt_max+1,1), fontsize = 10)

# ---------------------------------------------------
#this is where the grid information (phi and hdens) is read in and saved to grid.
grid = [];
with open(inputfile, 'rb') as f:
	csvReader = csv.reader(f,delimiter='\t')
	for row in csvReader:
		grid.append(row);
	grid  = asarray(grid)

#here is where the data for each line is read in and saved to dataEmissionlines
dataEmissionlines = [];
with open(inputfile2, 'rb') as f:
	csvReader = csv.reader(f,delimiter='\t')
	headers = csvReader.next()
	for row in csvReader:
		dataEmissionlines.append(row);
	dataEmissionlines  = asarray(dataEmissionlines)
print "import files complete"
# ---------------------------------------------------
#for grid
phi_values = grid[1:len(dataEmissionlines)+1,6]
hdens_values = grid[1:len(dataEmissionlines)+1,7]

#for lines
headers = headers[1:]
Emissionlines = dataEmissionlines[:, 1:]
concatenated_data = zeros((len(Emissionlines),len(Emissionlines[0])))
max_values = zeros((len(Emissionlines[0]),4))
#select the scaling factor

#for 1215
#incident = Emissionlines[1:,4]

#for 4860
incident = Emissionlines[:,57]

#take the ratio of incident and all the lines and put it all in an array concatenated_data

for i in range(len(Emissionlines)):
	for j in range(len(Emissionlines[0])):
			if math.log(4860.*(float(Emissionlines[i,j])/float(Emissionlines[i,57])), 10) > 0:
				concatenated_data[i,j] = math.log(4860.*(float(Emissionlines[i,j])/float(Emissionlines[i,57])), 10)
			else:
				concatenated_data[i,j] == 0
# for 1215
#for i in range(len(Emissionlines)):
#	for j in range(len(Emissionlines[0])):
#			if math.log(1215.*(float(Emissionlines[i,j])/float(Emissionlines[i,4])), 10) > 0:
#				concatenated_data[i,j] = math.log(1215.*(float(Emissionlines[i,j])/float(Emissionlines[i,4])), 10)
#			else:
#				concatenated_data[i,j] == 0

#find the maxima to plot onto the contour plots
for j in range(len(concatenated_data[0])):
	max_values[j,0] = max(concatenated_data[:,j])
	max_values[j,1] = argmax(concatenated_data[:,j], axis = 0)
	max_values[j,2] = hdens_values[max_values[j,1]]
	max_values[j,3] = phi_values[max_values[j,1]]

#to round off the maxima
max_values[:,0] = [ '%.1f' % elem for elem in max_values[:,0] ]
print "data arranged"

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

#Creating the grid to interpolate with for contours.
gridarray = zeros((len(Emissionlines),2))
gridarray[:,0] = hdens_values
gridarray[:,1] = phi_values

x = gridarray[:,0]
y = gridarray[:,1]

#change desired lines here!

line = [18, #1549
		19, #1640
		20, #1665
		21, #1671
		23, #1750
		24, #1860
		25, #1888
		26, #1907
		27, #2297
		28, #2321
		29, #2471
		30, #2326
		31, #2335
		32, #2665
		33, #2798
		34] #2803
#create z array for this plot
z = concatenated_data[:,line[:]]

# ---------------------------------------------------
# Interpolate
print "starting interpolation"
xi, yi = linspace(x.min(), x.max(), 10), linspace(y.min(), y.max(), 10)
xi, yi = meshgrid(xi, yi)
# ---------------------------------------------------
print "interpolatation complete; now plotting"
#plot
plt.subplots_adjust(wspace=0, hspace=0) #remove space between plots
levels = arange(10**-1,10, .2)
levels2 = arange(10**-2,10**2, 1)
plt.suptitle("UV Lines Continued", fontsize=14)
# ---------------------------------------------------
for i in range(16):
	add_sub_plot(i)
ax1 = plt.subplot(4,4,1)
add_patches(ax1)

print "complete"

plt.savefig('UV_Lines_cntd.pdf')
plt.clf()


from __future__ import division

import numpy as np

from sklearn.utils.linear_assignment_ import linear_assignment
from sklearn.utils.validation import check_consistent_length, check_array

__all__ = ["consensus_score"]

def _check_rows_and_columns(a, b):
    """Unpacks the row and column arrays and checks their shape."""
    check_consistent_length(*a)
    check_consistent_length(*b)
    checks = lambda x: check_array(x, ensure_2d=False)
    a_rows, a_cols = map(checks, a)
    b_rows, b_cols = map(checks, b)
    return a_rows, a_cols, b_rows, b_cols

def _jaccard(a_rows, a_cols, b_rows, b_cols):
    """Jaccard coefficient on the elements of the two biclusters."""
    intersection = ((a_rows * b_rows).sum() *
                    (a_cols * b_cols).sum())

    a_size = a_rows.sum() * a_cols.sum()
    b_size = b_rows.sum() * b_cols.sum()

    return intersection / (a_size + b_size - intersection)

def _pairwise_similarity(a, b, similarity):
    """Computes pairwise similarity matrix.

    result[i, j] is the Jaccard coefficient of a's bicluster i and b's
    bicluster j.

    """
    a_rows, a_cols, b_rows, b_cols = _check_rows_and_columns(a, b)
    n_a = a_rows.shape[0]
    n_b = b_rows.shape[0]
    result = np.array(list(list(similarity(a_rows[i], a_cols[i],
                                           b_rows[j], b_cols[j])
                                for j in range(n_b))
                           for i in range(n_a)))
    return result

def consensus_score(a, b, similarity="jaccard"):
    """The similarity of two sets of biclusters.

    Similarity between individual biclusters is computed. Then the
    best matching between sets is found using the Hungarian algorithm.
    The final score is the sum of similarities divided by the size of
    the larger set.

    Read more in the :ref:`User Guide <biclustering>`.

    Parameters
    ----------
    a : (rows, columns)
        Tuple of row and column indicators for a set of biclusters.

    b : (rows, columns)
        Another set of biclusters like ``a``.

    similarity : string or function, optional, default: "jaccard"
        May be the string "jaccard" to use the Jaccard coefficient, or
        any function that takes four arguments, each of which is a 1d
        indicator vector: (a_rows, a_columns, b_rows, b_columns).

    References
    ----------

    * Hochreiter, Bodenhofer, et. al., 2010. `FABIA: factor analysis
      for bicluster acquisition
      <https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2881408/>`__.

    """
    if similarity == "jaccard":
        similarity = _jaccard
    matrix = _pairwise_similarity(a, b, similarity)
    indices = linear_assignment(1. - matrix)
    n_a = len(a[0])
    n_b = len(b[0])
    return matrix[indices[:, 0], indices[:, 1]].sum() / max(n_a, n_b)


# Copyright 2018 Google LLC
#
# 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.

"""
Used to plot the accuracy of the policy and value networks in
predicting professional game moves and results over the course
of training. Check FLAGS for default values for what models to
load and what sgf files to parse.

Usage:
python training_curve.py

Sample 3 positions from each game
python training_curve.py --num_positions=3

Only grab games after 2005 (default is 2000)
python training_curve.py --min_year=2005
"""
import sys
sys.path.insert(0, '.')

import os.path

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from absl import app, flags
from tqdm import tqdm

import coords
from rl_loop import fsdb
import oneoff_utils

flags.DEFINE_string("sgf_dir", None, "sgf database")

flags.DEFINE_string("plot_dir", "data", "Where to save the plots.")
flags.DEFINE_integer("min_year", "2000",
                     "Only take sgf games with date >= min_year")
flags.DEFINE_string("komi", "7.5", "Only take sgf games with given komi")
flags.DEFINE_integer("idx_start", 150, "Only take models after given idx")
flags.DEFINE_integer("num_positions", 1,
                     "How many positions from each game to sample from.")
flags.DEFINE_integer("eval_every", 5,
                     "Eval every k models to generate the curve")

flags.mark_flag_as_required('sgf_dir')

FLAGS = flags.FLAGS

def batch_run_many(player, positions, batch_size=100):
    """Used to avoid a memory oveflow issue when running the network
    on too many positions. TODO: This should be a member function of
    player.network?"""
    prob_list = []
    value_list = []
    for idx in range(0, len(positions), batch_size):
        probs, values = player.network.run_many(positions[idx:idx + batch_size])
        prob_list.append(probs)
        value_list.append(values)
    return np.concatenate(prob_list, axis=0), np.concatenate(value_list, axis=0)

def eval_player(player, positions, moves, results):
    probs, values = batch_run_many(player, positions)
    policy_moves = [coords.from_flat(c) for c in np.argmax(probs, axis=1)]
    top_move_agree = [moves[idx] == policy_moves[idx]
                      for idx in range(len(moves))]
    square_err = (values - results) ** 2 / 4
    return top_move_agree, square_err

def sample_positions_from_games(sgf_files, num_positions=1):
    pos_data = []
    move_data = []
    result_data = []
    move_idxs = []

    fail_count = 0
    for path in tqdm(sgf_files, desc="loading sgfs", unit="games"):
        try:
            positions, moves, results = oneoff_utils.parse_sgf_to_examples(path)
        except KeyboardInterrupt:
            raise
        except Exception as e:
            print("Parse exception:", e)
            fail_count += 1
            continue

        # add entire game
        if num_positions == -1:
            pos_data.extend(positions)
            move_data.extend(moves)
            move_idxs.extend(range(len(positions)))
            result_data.extend(results)
        else:
            for idx in np.random.choice(len(positions), num_positions):
                pos_data.append(positions[idx])
                move_data.append(moves[idx])
                result_data.append(results[idx])
                move_idxs.append(idx)
    print("Sampled {} positions, failed to parse {} files".format(
        len(pos_data), fail_count))
    return pos_data, move_data, result_data, move_idxs

def get_training_curve_data(
        model_dir, pos_data, move_data, result_data, idx_start, eval_every):
    model_paths = oneoff_utils.get_model_paths(model_dir)
    df = pd.DataFrame()
    player = None

    print("Evaluating models {}-{}, eval_every={}".format(
        idx_start, len(model_paths), eval_every))
    for idx in tqdm(range(idx_start, len(model_paths), eval_every)):
        if player:
            oneoff_utils.restore_params(model_paths[idx], player)
        else:
            player = oneoff_utils.load_player(model_paths[idx])

        correct, squared_errors = eval_player(
            player=player, positions=pos_data,
            moves=move_data, results=result_data)

        avg_acc = np.mean(correct)
        avg_mse = np.mean(squared_errors)
        print("Model: {}, acc: {:.4f}, mse: {:.4f}".format(
            model_paths[idx], avg_acc, avg_mse))
        df = df.append({"num": idx, "acc": avg_acc,
                        "mse": avg_mse}, ignore_index=True)
    return df

def save_plots(data_dir, df):
    plt.plot(df["num"], df["acc"])
    plt.xlabel("Model idx")
    plt.ylabel("Accuracy")
    plt.title("Accuracy in Predicting Professional Moves")
    plot_path = os.path.join(data_dir, "move_acc.pdf")
    plt.savefig(plot_path)

    plt.figure()

    plt.plot(df["num"], df["mse"])
    plt.xlabel("Model idx")
    plt.ylabel("MSE/4")
    plt.title("MSE in predicting outcome")
    plot_path = os.path.join(data_dir, "value_mse.pdf")
    plt.savefig(plot_path)

def main(unusedargv):
    sgf_files = oneoff_utils.find_and_filter_sgf_files(
        FLAGS.sgf_dir, FLAGS.min_year, FLAGS.komi)
    pos_data, move_data, result_data, move_idxs = sample_positions_from_games(
        sgf_files=sgf_files, num_positions=FLAGS.num_positions)
    df = get_training_curve_data(fsdb.models_dir(), pos_data, move_data,
                                 result_data, FLAGS.idx_start, FLAGS.eval_every)
    save_plots(FLAGS.plot_dir, df)

if __name__ == "__main__":
    app.run(main)


"""
================
Confusion matrix
================

Example of confusion matrix usage to evaluate the quality
of the output of a classifier on the iris data set. The
diagonal elements represent the number of points for which
the predicted label is equal to the true label, while
off-diagonal elements are those that are mislabeled by the
classifier. The higher the diagonal values of the confusion
matrix the better, indicating many correct predictions.

The figures show the confusion matrix with and without
normalization by class support size (number of elements
in each class). This kind of normalization can be
interesting in case of class imbalance to have a more
visual interpretation of which class is being misclassified.

Here the results are not as good as they could be as our
choice for the regularization parameter C was not the best.
In real life applications this parameter is usually chosen
using :ref:`grid_search`.

"""

print(__doc__)

import numpy as np
import matplotlib.pyplot as plt

from sklearn import svm, datasets
from sklearn.cross_validation import train_test_split
from sklearn.metrics import confusion_matrix

# import some data to play with
iris = datasets.load_iris()
X = iris.data
y = iris.target

# Split the data into a training set and a test set
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

# Run classifier, using a model that is too regularized (C too low) to see
# the impact on the results
classifier = svm.SVC(kernel='linear', C=0.01)
y_pred = classifier.fit(X_train, y_train).predict(X_test)

def plot_confusion_matrix(cm, title='Confusion matrix', cmap=plt.cm.Blues):
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(iris.target_names))
    plt.xticks(tick_marks, iris.target_names, rotation=45)
    plt.yticks(tick_marks, iris.target_names)
    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')

# Compute confusion matrix
cm = confusion_matrix(y_test, y_pred)
np.set_printoptions(precision=2)
print('Confusion matrix, without normalization')
print(cm)
plt.figure()
plot_confusion_matrix(cm)

# Normalize the confusion matrix by row (i.e by the number of samples
# in each class)
cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
print('Normalized confusion matrix')
print(cm_normalized)
plt.figure()
plot_confusion_matrix(cm_normalized, title='Normalized confusion matrix')

plt.show()


import pandas as pd
import numpy as np
import json

from unidecode import unidecode

def extract_domain(url):
    # extract domains
    domain = url.lower().split('/')[2]
    domain_parts = domain.split('.')

    # e.g. co.uk
    if domain_parts[-2] not in ['com', 'co']:
        return '.'.join(domain_parts[-2:])
    else:
        return '.'.join(domain_parts[-3:])

def load_csv(filename):
    return pd.read_table(filename)

def parse_data(df):
    data = []
    columns = df.columns

    for key, row in df.iterrows():
        item = {}

        for column in columns:
            item[column] = row[column]

        # parse url

        item['real_url'] = row['url'].lower()
        item['domain'] = extract_domain(row['url'])
        item['tld'] = item['domain'].split('.')[-1]

        # parse boilerplate
        boilerplate = json.loads(row['boilerplate'])

        for f in ['title', 'url', 'body']:
            item[f] = boilerplate[f] if f in boilerplate else u''
            item[f] = unidecode(item[f]) if item[f] else ''

        if 'label' in row:
            item['label'] = row['label']
        else:
            item['label'] = np.nan

        data.append(item)

    return data

def get_train():
    train = load_csv('../data/raw/train.tsv')

    return (parse_data(train))

def get_test():
    test = load_csv('../data/raw/test.tsv')

    return (parse_data(test))


#-----------------------------------------------------------------------------
# Copyright (c) 2012 - 2019, Anaconda, Inc., and Bokeh Contributors.
# All rights reserved.
#
# The full license is in the file LICENSE.txt, distributed with this software.
#-----------------------------------------------------------------------------
''' Functions useful for dealing with hexagonal tilings.

For more information on the concepts employed here, see this informative page

    https://www.redblobgames.com/grids/hexagons/

'''

#-----------------------------------------------------------------------------
# Boilerplate
#-----------------------------------------------------------------------------
from __future__ import absolute_import, division, print_function, unicode_literals

import logging
log = logging.getLogger(__name__)

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

# Standard library imports

# External imports
import numpy as np

# Bokeh imports
from .dependencies import import_required

#-----------------------------------------------------------------------------
# Globals and constants
#-----------------------------------------------------------------------------

__all__ = (
    'axial_to_cartesian',
    'cartesian_to_axial',
    'hexbin',
)

#-----------------------------------------------------------------------------
# General API
#-----------------------------------------------------------------------------

def axial_to_cartesian(q, r, size, orientation, aspect_scale=1):
    ''' Map axial *(q,r)* coordinates to cartesian *(x,y)* coordinates of
    tiles centers.

    This function can be useful for positioning other Bokeh glyphs with
    cartesian coordinates in relation to a hex tiling.

    This function was adapted from:

        https://www.redblobgames.com/grids/hexagons/#hex-to-pixel

    Args:
        q (array[float]) :
            A NumPy array of q-coordinates for binning

        r (array[float]) :
            A NumPy array of r-coordinates for binning

        size (float) :
            The size of the hexagonal tiling.

            The size is defined as the distance from the center of a hexagon
            to the top corner for "pointytop" orientation, or from the center
            to a side corner for "flattop" orientation.

        orientation (str) :
            Whether the hex tile orientation should be "pointytop" or
            "flattop".

        aspect_scale (float, optional) :
            Scale the hexagons in the "cross" dimension.

            For "pointytop" orientations, hexagons are scaled in the horizontal
            direction. For "flattop", they are scaled in vertical direction.

            When working with a plot with ``aspect_scale != 1``, it may be
            useful to set this value to match the plot.

    Returns:
        (array[int], array[int])

    '''
    if orientation == "pointytop":
        x = size * np.sqrt(3) * (q + r/2.0) / aspect_scale
        y = -size * 3/2.0 * r
    else:
        x = size * 3/2.0 * q
        y = -size * np.sqrt(3) * (r + q/2.0) * aspect_scale

    return (x, y)

def cartesian_to_axial(x, y, size, orientation, aspect_scale=1):
    ''' Map Cartesion *(x,y)* points to axial *(q,r)* coordinates of enclosing
    tiles.

    This function was adapted from:

        https://www.redblobgames.com/grids/hexagons/#pixel-to-hex

    Args:
        x (array[float]) :
            A NumPy array of x-coordinates to convert

        y (array[float]) :
            A NumPy array of y-coordinates to convert

        size (float) :
            The size of the hexagonal tiling.

            The size is defined as the distance from the center of a hexagon
            to the top corner for "pointytop" orientation, or from the center
            to a side corner for "flattop" orientation.

        orientation (str) :
            Whether the hex tile orientation should be "pointytop" or
            "flattop".

        aspect_scale (float, optional) :
            Scale the hexagons in the "cross" dimension.

            For "pointytop" orientations, hexagons are scaled in the horizontal
            direction. For "flattop", they are scaled in vertical direction.

            When working with a plot with ``aspect_scale != 1``, it may be
            useful to set this value to match the plot.

    Returns:
        (array[int], array[int])

    '''
    HEX_FLAT = [2.0/3.0, 0.0, -1.0/3.0, np.sqrt(3.0)/3.0]
    HEX_POINTY = [np.sqrt(3.0)/3.0, -1.0/3.0, 0.0, 2.0/3.0]

    coords = HEX_FLAT if orientation == 'flattop' else HEX_POINTY

    x =  x / size * (aspect_scale if orientation == "pointytop" else 1)
    y = -y / size / (aspect_scale if orientation == "flattop" else 1)

    q = coords[0] * x + coords[1] * y
    r = coords[2] * x + coords[3] * y

    return _round_hex(q, r)

def hexbin(x, y, size, orientation="pointytop", aspect_scale=1):
    ''' Perform an equal-weight binning of data points into hexagonal tiles.

    For more sophisticated use cases, e.g. weighted binning or scaling
    individual tiles proportional to some other quantity, consider using
    HoloViews.

    Args:
        x (array[float]) :
            A NumPy array of x-coordinates for binning

        y (array[float]) :
            A NumPy array of y-coordinates for binning

        size (float) :
            The size of the hexagonal tiling.

            The size is defined as the distance from the center of a hexagon
            to the top corner for "pointytop" orientation, or from the center
            to a side corner for "flattop" orientation.

        orientation (str, optional) :
            Whether the hex tile orientation should be "pointytop" or
            "flattop". (default: "pointytop")

        aspect_scale (float, optional) :
            Match a plot's aspect ratio scaling.

            When working with a plot with ``aspect_scale != 1``, this
            parameter can be set to match the plot, in order to draw
            regular hexagons (instead of "stretched" ones).

            This is roughly equivalent to binning in "screen space", and
            it may be better to use axis-aligned rectangular bins when
            plot aspect scales are not one.

    Returns:
        DataFrame

        The resulting DataFrame will have columns *q* and *r* that specify
        hexagon tile locations in axial coordinates, and a column *counts* that
        provides the count for each tile.

    .. warning::
        Hex binning only functions on linear scales, i.e. not on log plots.

    '''
    pd = import_required('pandas','hexbin requires pandas to be installed')

    q, r = cartesian_to_axial(x, y, size, orientation, aspect_scale=aspect_scale)

    df = pd.DataFrame(dict(r=r, q=q))
    return df.groupby(['q', 'r']).size().reset_index(name='counts')

#-----------------------------------------------------------------------------
# Dev API
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Private API
#-----------------------------------------------------------------------------

def _round_hex(q, r):
    ''' Round floating point axial hex coordinates to integer *(q,r)*
    coordinates.

    This code was adapted from:

        https://www.redblobgames.com/grids/hexagons/#rounding

    Args:
        q (array[float]) :
            NumPy array of Floating point axial *q* coordinates to round

        r (array[float]) :
            NumPy array of Floating point axial *q* coordinates to round

    Returns:
        (array[int], array[int])

    '''
    x = q
    z = r
    y = -x-z

    rx = np.round(x)
    ry = np.round(y)
    rz = np.round(z)

    dx = np.abs(rx - x)
    dy = np.abs(ry - y)
    dz = np.abs(rz - z)

    cond = (dx > dy) & (dx > dz)
    q = np.where(cond              , -(ry + rz), rx)
    r = np.where(~cond & ~(dy > dz), -(rx + ry), rz)

    return q.astype(int), r.astype(int)

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------