"""
The rcsetup module contains the default values and the validation code for
customization using matplotlib's rc settings.

Each rc setting is assigned a default value and a function used to validate any
attempted changes to that setting. The default values and validation functions
are defined in the rcsetup module, and are used to construct the rcParams global
object which stores the settings and is referenced throughout matplotlib.

These default values should be consistent with the default matplotlibrc file
that actually reflects the values given here. Any additions or deletions to the
parameter set listed here should also be visited to the
:file:`matplotlibrc.template` in matplotlib's root source directory.
"""

import os
import warnings
from matplotlib.fontconfig_pattern import parse_fontconfig_pattern
from matplotlib.colors import is_color_like

#interactive_bk = ['gtk', 'gtkagg', 'gtkcairo', 'fltkagg', 'qtagg', 'qt4agg',
#                  'tkagg', 'wx', 'wxagg', 'cocoaagg']
# The capitalized forms are needed for ipython at present; this may
# change for later versions.

interactive_bk = ['GTK', 'GTKAgg', 'GTKCairo', 'FltkAgg', 'MacOSX',
                  'QtAgg', 'Qt4Agg', 'TkAgg', 'WX', 'WXAgg', 'CocoaAgg']

non_interactive_bk = ['agg', 'cairo', 'emf', 'gdk',
                      'pdf', 'ps', 'svg', 'template']
all_backends = interactive_bk + non_interactive_bk

class ValidateInStrings:
    def __init__(self, key, valid, ignorecase=False):
        'valid is a list of legal strings'
        self.key = key
        self.ignorecase = ignorecase
        def func(s):
            if ignorecase: return s.lower()
            else: return s
        self.valid = dict([(func(k),k) for k in valid])

    def __call__(self, s):
        if self.ignorecase: s = s.lower()
        if s in self.valid: return self.valid[s]
        raise ValueError('Unrecognized %s string "%s": valid strings are %s'
                         % (self.key, s, self.valid.values()))

def validate_path_exists(s):
    'If s is a path, return s, else False'
    if os.path.exists(s): return s
    else:
        raise RuntimeError('"%s" should be a path but it does not exist'%s)

def validate_bool(b):
    'Convert b to a boolean or raise'
    if type(b) is str:
        b = b.lower()
    if b in ('t', 'y', 'yes', 'on', 'true', '1', 1, True): return True
    elif b in ('f', 'n', 'no', 'off', 'false', '0', 0, False): return False
    else:
        raise ValueError('Could not convert "%s" to boolean' % b)

def validate_bool_maybe_none(b):
    'Convert b to a boolean or raise'
    if type(b) is str:
        b = b.lower()
    if b=='none': return None
    if b in ('t', 'y', 'yes', 'on', 'true', '1', 1, True): return True
    elif b in ('f', 'n', 'no', 'off', 'false', '0', 0, False): return False
    else:
        raise ValueError('Could not convert "%s" to boolean' % b)

def validate_float(s):
    'convert s to float or raise'
    try: return float(s)
    except ValueError:
        raise ValueError('Could not convert "%s" to float' % s)

def validate_int(s):
    'convert s to int or raise'
    try: return int(s)
    except ValueError:
        raise ValueError('Could not convert "%s" to int' % s)

def validate_fonttype(s):
    'confirm that this is a Postscript of PDF font type that we know how to convert to'
    fonttypes = { 'type3':    3,
                  'truetype': 42 }
    try:
        fonttype = validate_int(s)
    except ValueError:
        if s.lower() in fonttypes.keys():
            return fonttypes[s.lower()]
        raise ValueError('Supported Postscript/PDF font types are %s' % fonttypes.keys())
    else:
        if fonttype not in fonttypes.values():
            raise ValueError('Supported Postscript/PDF font types are %s' % fonttypes.values())
        return fonttype

#validate_backend = ValidateInStrings('backend', all_backends, ignorecase=True)
_validate_standard_backends = ValidateInStrings('backend', all_backends, ignorecase=True)
def validate_backend(s):
    if s.startswith('module://'): return s
    else: return _validate_standard_backends(s)

validate_numerix = ValidateInStrings('numerix',[
    'Numeric','numarray','numpy',
    ], ignorecase=True)

validate_toolbar = ValidateInStrings('toolbar',[
    'None','classic','toolbar2',
    ], ignorecase=True)

def validate_autolayout(v):
    if v:
        warnings.warn("figure.autolayout is not currently supported")

class validate_nseq_float:
    def __init__(self, n):
        self.n = n
    def __call__(self, s):
        'return a seq of n floats or raise'
        if type(s) is str:
            ss = s.split(',')
            if len(ss) != self.n:
                raise ValueError('You must supply exactly %d comma separated values'%self.n)
            try:
                return [float(val) for val in ss]
            except ValueError:
                raise ValueError('Could not convert all entries to floats')
        else:
            assert type(s) in (list,tuple)
            if len(s) != self.n:
                raise ValueError('You must supply exactly %d values'%self.n)
            return [float(val) for val in s]

class validate_nseq_int:
    def __init__(self, n):
        self.n = n
    def __call__(self, s):
        'return a seq of n ints or raise'
        if type(s) is str:
            ss = s.split(',')
            if len(ss) != self.n:
                raise ValueError('You must supply exactly %d comma separated values'%self.n)
            try:
                return [int(val) for val in ss]
            except ValueError:
                raise ValueError('Could not convert all entries to ints')
        else:
            assert type(s) in (list,tuple)
            if len(s) != self.n:
                raise ValueError('You must supply exactly %d values'%self.n)
            return [int(val) for val in s]

def validate_color(s):
    'return a valid color arg'
    if s.lower() == 'none':
        return 'None'
    if is_color_like(s):
        return s
    stmp = '#' + s
    if is_color_like(stmp):
        return stmp
    # If it is still valid, it must be a tuple.
    colorarg = s
    msg = ''
    if s.find(',')>=0:
        # get rid of grouping symbols
        stmp = ''.join([ c for c in s if c.isdigit() or c=='.' or c==','])
        vals = stmp.split(',')
        if len(vals)!=3:
            msg = '\nColor tuples must be length 3'
        else:
            try:
                colorarg = [float(val) for val in vals]
            except ValueError:
                msg = '\nCould not convert all entries to floats'

    if not msg and is_color_like(colorarg):
        return colorarg

    raise ValueError('%s does not look like a color arg%s'%(s, msg))

def validate_stringlist(s):
    'return a list'
    if type(s) is str:
        return [ v.strip() for v in s.split(',') ]
    else:
        assert type(s) in [list,tuple]
        return [ str(v) for v in s ]

validate_orientation = ValidateInStrings('orientation',[
    'landscape', 'portrait',
    ])

def validate_aspect(s):
    if s in ('auto', 'equal'):
        return s
    try:
        return float(s)
    except ValueError:
        raise ValueError('not a valid aspect specification')

def validate_fontsize(s):
    if type(s) is str:
        s = s.lower()
    if s in ['xx-small', 'x-small', 'small', 'medium', 'large', 'x-large',
             'xx-large', 'smaller', 'larger']:
        return s
    try:
        return float(s)
    except ValueError:
        raise ValueError('not a valid font size')

def validate_font_properties(s):
    parse_fontconfig_pattern(s)
    return s

validate_fontset = ValidateInStrings('fontset', ['cm', 'stix', 'stixsans', 'custom'])

validate_verbose = ValidateInStrings('verbose',[
    'silent', 'helpful', 'debug', 'debug-annoying',
    ])

validate_cairo_format = ValidateInStrings('cairo_format',
                            ['png', 'ps', 'pdf', 'svg'],
                            ignorecase=True)

validate_ps_papersize = ValidateInStrings('ps_papersize',[
    'auto', 'letter', 'legal', 'ledger',
    'a0', 'a1', 'a2','a3', 'a4', 'a5', 'a6', 'a7', 'a8', 'a9', 'a10',
    'b0', 'b1', 'b2', 'b3', 'b4', 'b5', 'b6', 'b7', 'b8', 'b9', 'b10',
    ], ignorecase=True)

def validate_ps_distiller(s):
    if type(s) is str:
        s = s.lower()

    if s in ('none',None):
        return None
    elif s in ('false', False):
        return False
    elif s in ('ghostscript', 'xpdf'):
        return s
    else:
        raise ValueError('matplotlibrc ps.usedistiller must either be none, ghostscript or xpdf')

validate_joinstyle = ValidateInStrings('joinstyle',['miter', 'round', 'bevel'], ignorecase=True)

validate_capstyle = ValidateInStrings('capstyle',['butt', 'round', 'projecting'], ignorecase=True)

validate_negative_linestyle = ValidateInStrings('negative_linestyle',['solid', 'dashed'], ignorecase=True)

def validate_negative_linestyle_legacy(s):
    try:
        res = validate_negative_linestyle(s)
        return res
    except ValueError:
        dashes = validate_nseq_float(2)(s)
        warnings.warn("Deprecated negative_linestyle specification; use 'solid' or 'dashed'")
        return (0, dashes)  # (offset, (solid, blank))

validate_legend_loc = ValidateInStrings('legend_loc',[
  'best',
  'upper right',
  'upper left',
  'lower left',
  'lower right',
  'right',
  'center left',
  'center right',
  'lower center',
  'upper center',
  'center',
], ignorecase=True)

class ValidateInterval:
    """
    Value must be in interval
    """
    def __init__(self, vmin, vmax, closedmin=True, closedmax=True):
        self.vmin = vmin
        self.vmax = vmax
        self.cmin = closedmin
        self.cmax = closedmax

    def __call__(self, s):
        try: s = float(s)
        except: raise RuntimeError('Value must be a float; found "%s"'%s)

        if self.cmin and s<self.vmin:
            raise RuntimeError('Value must be >= %f; found "%f"'%(self.vmin, s))
        elif not self.cmin and s<=self.vmin:
            raise RuntimeError('Value must be > %f; found "%f"'%(self.vmin, s))

        if self.cmax and s>self.vmax:
            raise RuntimeError('Value must be <= %f; found "%f"'%(self.vmax, s))
        elif not self.cmax and s>=self.vmax:
            raise RuntimeError('Value must be < %f; found "%f"'%(self.vmax, s))
        return s

# a map from key -> value, converter
defaultParams = {
    'backend'           : ['Agg', validate_backend], # agg is certainly present
    'backend_fallback'  : [True, validate_bool], # agg is certainly present
    'numerix'           : ['numpy', validate_numerix],
    'maskedarray'       : [False, validate_bool],
    'toolbar'           : ['toolbar2', validate_toolbar],
    'datapath'          : [None, validate_path_exists],   # handled by _get_data_path_cached
    'units'             : [False, validate_bool],
    'interactive'       : [False, validate_bool],
    'timezone'          : ['UTC', str],

    # the verbosity setting
    'verbose.level'     : ['silent', validate_verbose],
    'verbose.fileo'     : ['sys.stdout', str],

    # line props
    'lines.linewidth'       : [1.0, validate_float],     # line width in points
    'lines.linestyle'       : ['-', str],                # solid line
    'lines.color'           : ['b', validate_color],     # blue
    'lines.marker'          : ['None', str],     # black
    'lines.markeredgewidth' : [0.5, validate_float],
    'lines.markersize'      : [6, validate_float],       # markersize, in points
    'lines.antialiased'     : [True, validate_bool],     # antialised (no jaggies)
    'lines.dash_joinstyle'  : ['miter', validate_joinstyle],
    'lines.solid_joinstyle' : ['miter', validate_joinstyle],
    'lines.dash_capstyle'   : ['butt', validate_capstyle],
    'lines.solid_capstyle'  : ['projecting', validate_capstyle],

    # patch props
    'patch.linewidth'   : [1.0, validate_float], # line width in points
    'patch.edgecolor'   : ['k', validate_color], # black
    'patch.facecolor'   : ['b', validate_color], # blue
    'patch.antialiased' : [True, validate_bool], # antialised (no jaggies)

    # font props
    'font.family'       : ['sans-serif', str],       # used by text object
    'font.style'        : ['normal', str],           #
    'font.variant'      : ['normal', str],           #
    'font.stretch'      : ['normal', str],           #
    'font.weight'       : ['normal', str],           #
    'font.size'         : [12.0, validate_float], #
    'font.serif'        : [['Bitstream Vera Serif', 'DejaVu Serif',
                            'New Century Schoolbook', 'Century Schoolbook L',
                            'Utopia', 'ITC Bookman', 'Bookman',
                            'Nimbus Roman No9 L','Times New Roman',
                            'Times','Palatino','Charter','serif'],
                           validate_stringlist],
    'font.sans-serif'   : [['Bitstream Vera Sans', 'DejaVu Sans',
                            'Lucida Grande', 'Verdana', 'Geneva', 'Lucid',
                            'Arial', 'Helvetica', 'Avant Garde', 'sans-serif'],
                           validate_stringlist],
    'font.cursive'      : [['Apple Chancery','Textile','Zapf Chancery',
                           'Sand','cursive'], validate_stringlist],
    'font.fantasy'      : [['Comic Sans MS','Chicago','Charcoal','Impact'
                           'Western','fantasy'], validate_stringlist],
    'font.monospace'    : [['Bitstream Vera Sans Mono', 'DejaVu Sans Mono',
                            'Andale Mono', 'Nimbus Mono L', 'Courier New',
                            'Courier','Fixed', 'Terminal','monospace'],
                           validate_stringlist],

    # text props
    'text.color'          : ['k', validate_color],     # black
    'text.usetex'         : [False, validate_bool],
    'text.latex.unicode'  : [False, validate_bool],
    'text.latex.preamble' : [[''], validate_stringlist],
    'text.dvipnghack'     : [None, validate_bool_maybe_none],
    'text.fontstyle'      : ['normal', str],
    'text.fontangle'      : ['normal', str],
    'text.fontvariant'    : ['normal', str],
    'text.fontweight'     : ['normal', str],
    'text.fontsize'       : ['medium', validate_fontsize],

    'mathtext.cal'        : ['cursive', validate_font_properties],
    'mathtext.rm'         : ['serif', validate_font_properties],
    'mathtext.tt'         : ['monospace', validate_font_properties],
    'mathtext.it'         : ['serif:italic', validate_font_properties],
    'mathtext.bf'         : ['serif:bold', validate_font_properties],
    'mathtext.sf'         : ['sans\-serif', validate_font_properties],
    'mathtext.fontset'    : ['cm', validate_fontset],
    'mathtext.fallback_to_cm' : [True, validate_bool],

    'image.aspect'        : ['equal', validate_aspect],  # equal, auto, a number
    'image.interpolation' : ['bilinear', str],
    'image.cmap'          : ['jet', str],        # one of gray, jet, etc
    'image.lut'           : [256, validate_int],  # lookup table
    'image.origin'        : ['upper', str],  # lookup table
    'image.resample'      : [False, validate_bool],

    'contour.negative_linestyle' : ['dashed', validate_negative_linestyle_legacy],

    # axes props
    'axes.axisbelow'        : [False, validate_bool],
    'axes.hold'             : [True, validate_bool],
    'axes.facecolor'        : ['w', validate_color],    # background color; white
    'axes.edgecolor'        : ['k', validate_color],    # edge color; black
    'axes.linewidth'        : [1.0, validate_float],    # edge linewidth
    'axes.titlesize'        : ['large', validate_fontsize], # fontsize of the axes title
    'axes.grid'             : [False, validate_bool],   # display grid or not
    'axes.labelsize'        : ['medium', validate_fontsize], # fontsize of the x any y labels
    'axes.labelcolor'       : ['k', validate_color],    # color of axis label
    'axes.formatter.limits' : [[-7, 7], validate_nseq_int(2)],
                               # use scientific notation if log10
                               # of the axis range is smaller than the
                               # first or larger than the second
    'axes.unicode_minus'        : [True, validate_bool],

    'polaraxes.grid'        : [True, validate_bool],   # display polar grid or not

    #legend properties
    'legend.fancybox'         : [False,validate_bool],
    'legend.loc'         : ['upper right',validate_legend_loc], # at some point, this should be changed to 'best'
    'legend.isaxes'      : [True,validate_bool],  # this option is internally ignored - it never served any useful purpose
    'legend.numpoints'   : [2, validate_int],     # the number of points in the legend line
    'legend.fontsize'    : ['large', validate_fontsize],
    'legend.pad'         : [0,   validate_float], # was 0.2, deprecated; the fractional whitespace inside the legend border
    'legend.borderpad'   : [0.4, validate_float], # units are fontsize
    'legend.markerscale' : [1.0, validate_float], # the relative size of legend markers vs. original

    # the following dimensions are in axes coords
    'legend.labelsep'      : [0.010, validate_float], # the vertical space between the legend entries
    'legend.handlelen'     : [0.05, validate_float], # the length of the legend lines
    'legend.handletextsep' : [0.02, validate_float], # the space between the legend line and legend text
    'legend.axespad'       : [0.02, validate_float], # the border between the axes and legend edge
    'legend.shadow'        : [False, validate_bool],

    'legend.labelspacing'      : [0.5, validate_float], # the vertical space between the legend entries
    'legend.handlelength'     : [2., validate_float], # the length of the legend lines
    'legend.handletextpad' : [.8, validate_float], # the space between the legend line and legend text
    'legend.borderaxespad'       : [0.5, validate_float], # the border between the axes and legend edge
    'legend.columnspacing'       : [2., validate_float], # the border between the axes and legend edge

    'legend.markerscale' : [1.0, validate_float], # the relative size of legend markers vs. original

    # the following dimensions are in axes coords
    'legend.labelsep'      : [0.010, validate_float], # the vertical space between the legend entries
    'legend.handlelen'     : [0.05, validate_float], # the length of the legend lines
    'legend.handletextsep' : [0.02, validate_float], # the space between the legend line and legend text
    'legend.axespad'       : [0.5, validate_float], # the border between the axes and legend edge
    'legend.shadow'        : [False, validate_bool],

    # tick properties
    'xtick.major.size' : [4, validate_float],      # major xtick size in points
    'xtick.minor.size' : [2, validate_float],      # minor xtick size in points
    'xtick.major.pad'  : [4, validate_float],      # distance to label in points
    'xtick.minor.pad'  : [4, validate_float],      # distance to label in points
    'xtick.color'      : ['k', validate_color],    # color of the xtick labels
    'xtick.labelsize'  : ['medium', validate_fontsize], # fontsize of the xtick labels
    'xtick.direction'  : ['in', str],            # direction of xticks

    'ytick.major.size' : [4, validate_float],      # major ytick size in points
    'ytick.minor.size' : [2, validate_float],      # minor ytick size in points
    'ytick.major.pad'  : [4, validate_float],      # distance to label in points
    'ytick.minor.pad'  : [4, validate_float],      # distance to label in points
    'ytick.color'      : ['k', validate_color],    # color of the ytick labels
    'ytick.labelsize'  : ['medium', validate_fontsize], # fontsize of the ytick labels
    'ytick.direction'  : ['in', str],            # direction of yticks

    'grid.color'       : ['k', validate_color],       # grid color
    'grid.linestyle'   : [':', str],       # dotted
    'grid.linewidth'   : [0.5, validate_float],     # in points

    # figure props
    # figure size in inches: width by height
    'figure.figsize'    : [ [8.0,6.0], validate_nseq_float(2)],
    'figure.dpi'        : [ 80, validate_float],   # DPI
    'figure.facecolor'  : [ '0.75', validate_color], # facecolor; scalar gray
    'figure.edgecolor'  : [ 'w', validate_color],  # edgecolor; white
    'figure.autolayout' : [ False, validate_autolayout],

    'figure.subplot.left'   : [0.125, ValidateInterval(0, 1, closedmin=True, closedmax=True)],
    'figure.subplot.right'  : [0.9, ValidateInterval(0, 1, closedmin=True, closedmax=True)],
    'figure.subplot.bottom' : [0.1, ValidateInterval(0, 1, closedmin=True, closedmax=True)],
    'figure.subplot.top'    : [0.9, ValidateInterval(0, 1, closedmin=True, closedmax=True)],
    'figure.subplot.wspace' : [0.2, ValidateInterval(0, 1, closedmin=True, closedmax=False)],
    'figure.subplot.hspace' : [0.2, ValidateInterval(0, 1, closedmin=True, closedmax=False)],

    'savefig.dpi'         : [100, validate_float],   # DPI
    'savefig.facecolor'   : ['w', validate_color],  # facecolor; white
    'savefig.edgecolor'   : ['w', validate_color],  # edgecolor; white
    'savefig.orientation' : ['portrait', validate_orientation],  # edgecolor; white

    'cairo.format'       : ['png', validate_cairo_format],
    'tk.window_focus'    : [False, validate_bool],  # Maintain shell focus for TkAgg
    'tk.pythoninspect'   : [False, validate_bool],  # Set PYTHONINSPECT
    'ps.papersize'       : ['letter', validate_ps_papersize], # Set the papersize/type
    'ps.useafm'          : [False, validate_bool],  # Set PYTHONINSPECT
    'ps.usedistiller'    : [False, validate_ps_distiller], # use ghostscript or xpdf to distill ps output
    'ps.distiller.res'   : [6000, validate_int],     # dpi
    'ps.fonttype'        : [3, validate_fonttype], # 3 (Type3) or 42 (Truetype)
    'pdf.compression'    : [6, validate_int],        # compression level from 0 to 9; 0 to disable
    'pdf.inheritcolor'   : [False, validate_bool],   # ignore any color-setting commands from the frontend
    'pdf.use14corefonts' : [False, validate_bool],  # use only the 14 PDF core fonts
                                                    # embedded in every PDF viewing application
    'pdf.fonttype'      : [3, validate_fonttype],  # 3 (Type3) or 42 (Truetype)
    'svg.image_inline'  : [True, validate_bool],    # write raster image data directly into the svg file
    'svg.image_noscale' : [False, validate_bool],  # suppress scaling of raster data embedded in SVG
    'svg.embed_char_paths' : [True, validate_bool],  # True to save all characters as paths in the SVG

    'docstring.hardcopy' : [False, validate_bool],  # set this when you want to generate hardcopy docstring
    'plugins.directory' : ['.matplotlib_plugins', str], # where plugin directory is locate

    'path.simplify' : [False, validate_bool],
    'agg.path.chunksize' : [0, validate_int]       # 0 to disable chunking;
                                                   # recommend about 20000 to
                                                   # enable. Experimental.
}

if __name__ == '__main__':
    rc = defaultParams
    rc['datapath'][0] = '/'
    for key in rc:
        if not rc[key][1](rc[key][0]) == rc[key][0]:
            print "%s: %s != %s"%(key, rc[key][1](rc[key][0]), rc[key][0])


# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
# Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# Using BLAST, create a CSV file that lists the % identity of the specified sequence to all sequences from the specified germline

__author__ = 'William Lees'
__docformat__ = "restructuredtext en"

import os.path
import sys
import argparse
import csv
import re
import subprocess
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio import pairwise2
from Bio.Alphabet import generic_nucleotide
from Bio import SeqIO
from Bio import Phylo
from itertools import izip

def main(argv):
    parser = argparse.ArgumentParser(description='Create an Identity/Divergence plot.')
    parser.add_argument('repertoire', help='file containing repertoire sequence identities (CSV)')
    parser.add_argument('-a', '--adjust', help='Adjust labels to prevent overlap (requires package adjustText)', action='store_true')
    parser.add_argument('-b', '--bar', help='Plot a colour bar', action='store_true')
    parser.add_argument('-c', '--colourmap', help='colourmap')
    parser.add_argument('-g', '--background', help='Set the contour colourwhere the density is zero')
    parser.add_argument('-mx', '--maxx', help='max divergence value to show')
    parser.add_argument('-my', '--miny', help='min identity value to show')
    parser.add_argument('-p', '--points', help='comma-seperated list of identity files and formats')
    parser.add_argument('-s', '--save', help='Save output to file (as opposed to interactive display)')
    args = parser.parse_args()

    if args.adjust:
        from adjustText import adjust_text

    colourmap = args.colourmap if args.colourmap else 'hot_r'

    plist = args.points.split(',')
    points = []

    repertoire = read_file(args.repertoire)

    def pairwise(iterable):
        a = iter(iterable)
        return izip(a, a)

    if len(plist) > 0:
        try:
            for file, format in pairwise(plist):
                points.append((read_file(file), format))
        except IOError:
            print 'file "%s" cannot be opened.' % file
        except:
            print '"points" must consist of pairs of files and formats.'
            quit()

    max_divergence = int(args.maxx) if args.maxx else None
    min_identity = int(args.miny) if args.miny else None
    savefile = args.save if args.save else None

    if not max_divergence:
        max_divergence = max(repertoire['GermlineDist'])
        for point in points:
            max_divergence = max(max_divergence, max(point[0]['GermlineDist']))
        max_divergence = int(max_divergence) + 1

    if not min_identity:
        min_identity = min(repertoire['TargetDist'])
        for point in points:
            min_identity = min(min_identity, min(point[0]['TargetDist']))
        min_identity = int(min_identity)

    H, yedges, xedges = np.histogram2d(repertoire['TargetDist'], repertoire['GermlineDist'], bins=[101-min_identity, max_divergence+1], range=[[min_identity, 101], [-1, max_divergence]], normed=False)

    # For alternative interpolations and plots, see http://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram2d.html
    # For colour maps, see http://matplotlib.org/examples/color/colormaps_reference.html

    fig = plt.figure()

    cm = plt.cm.get_cmap(colourmap)

    if args.background:
        cm.set_under(color=args.background)

    ax = fig.add_subplot(1,1,1)
    im = plt.imshow(H, interpolation='bilinear', origin='low', extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]], vmin=0.0000001, cmap=cm)
    ax.set_xlim(xedges[0], xedges[-1])
    ax.set_ylim(yedges[0], yedges[-1])

    if args.bar:
        cb = plt.colorbar(im, shrink=0.8, extend='neither')
        cb.ax.set_ylabel('sequences', rotation=90)

    texts = []

    for point in points:
        df, format = point
        markersize = 5
        label = False
        labelcolour = 'black'
        fmt = format.split('/')
        format = fmt[0]
        for f in fmt[1:]:
            if f[0] == 'm':
                markersize = int(f[1:])
            elif f[0] == 'l':
                label = True
                if len(f) > 1:
                    labelcolour = f[1:]
            else:
                print 'Unrecognised format string: %s' % format
        for index, row in df.iterrows():
            if label:
                if args.adjust:
                    texts.append(plt.text(row['GermlineDist'], row['TargetDist'], row['SequenceId'], bbox={'pad':0, 'alpha':0}, fontdict={ 'color': labelcolour}))
                else:
                    texts.append(plt.text(row['GermlineDist'] + 0.2, row['TargetDist'] - 0.2, row['SequenceId'], bbox={'pad':0, 'alpha':0}, fontdict={ 'color': labelcolour}))
            ax.plot(row['GermlineDist'], row['TargetDist'], format, markersize=markersize)

    if args.adjust:
        adjust_text(texts)

    plt.xlabel('Germline Divergence (%)')
    plt.ylabel('Target Ab Identity (%)')

    if savefile:
        plt.savefig(savefile)
    else:
        plt.show()

def read_file(file):
    df = pd.read_csv(file, converters={'SequenceId': lambda x: x})
    for key in ("SequenceId", "TargetDist", "GermlineDist"):
        if key not in df.keys():
            print 'File %s does not contain a column "%s"' % (file, key)
            quit()

    for index, row in df.iterrows():
        try:
            x = row[1] * row[2]     # check they behave like numbers
        except:
            print 'Error in file %s: malformed row at %s.' % (file, row[0])

    if len(df) < 1:
        print '%s: empty file.' % file
        quit()

    return df

if __name__=="__main__":
    main(sys.argv)


__all__ = [
    'TensorMeshReader',
    'TensorMeshAppender',
    'TopoMeshAppender',
]

__displayname__ = 'Tensor Mesh'

import os
import sys

import numpy as np
import pandas as pd
import vtk

from .. import _helpers, interface
from ..base import AlgorithmBase
from .two_file_base import ModelAppenderBase, ubcMeshReaderBase

if sys.version_info < (3,):
    from StringIO import StringIO
else:
    from io import StringIO

class TensorMeshReader(ubcMeshReaderBase):
    """UBC Mesh 2D/3D models are defined using a 2-file format. The "mesh" file
    describes how the data is discretized. The "model" file lists the physical
    property values for all cells in a mesh. A model file is meaningless without
    an associated mesh file. The reader will automatically detect if the mesh is
    2D or 3D and read the remainder of the data with that dimensionality
    assumption. If the mesh file is 2D, then then model file must also be in the
    2D format (same for 3D).

    Note:
        Model File is optional. Reader will still construct
        ``vtkRectilinearGrid`` safely.
    """

    __displayname__ = 'UBC Tensor Mesh Reader'
    __category__ = 'reader'
    description = 'PVGeo: UBC Mesh 2D/3D Two-File Format'

    def __init__(self, nOutputPorts=1, outputType='vtkRectilinearGrid', **kwargs):
        ubcMeshReaderBase.__init__(
            self, nOutputPorts=nOutputPorts, outputType=outputType, **kwargs
        )

        self.__mesh = vtk.vtkRectilinearGrid()
        self.__models = []

    @staticmethod
    def place_model_on_mesh(mesh, model, data_name='Data'):
        """Places model data onto a mesh. This is for the UBC Grid data reaers
        to associate model data with the mesh grid.

        Args:
            mesh (vtkRectilinearGrid): The ``vtkRectilinearGrid`` that is the
                mesh to place the model data upon.
            model (np.array): A NumPy float array that holds all of the data to
                place inside of the mesh's cells.
            data_name (str) : The name of the model data array once placed on the
                ``vtkRectilinearGrid``.

        Return:
            vtkRectilinearGrid :
                Returns the input ``vtkRectilinearGrid`` with model data appended.
        """
        if isinstance(model, dict):
            for key in model.keys():
                TensorMeshReader.place_model_on_mesh(mesh, model[key], data_name=key)
            return mesh

        # model.GetNumberOfValues() if model is vtkDataArray
        # Make sure this model file fits the dimensions of the mesh
        ext = mesh.GetExtent()
        n1, n2, n3 = ext[1], ext[3], ext[5]
        if n1 * n2 * n3 < len(model):
            raise _helpers.PVGeoError(
                'Model `%s` has more data than the given mesh has cells to hold.'
                % data_name
            )
        elif n1 * n2 * n3 > len(model):
            raise _helpers.PVGeoError(
                'Model `%s` does not have enough data to fill the given mesh\'s cells.'
                % data_name
            )

        # Swap axes because VTK structures the coordinates a bit differently
        # -  This is absolutely crucial!
        # -  Do not play with unless you know what you are doing!
        if model.ndim > 1 and model.ndim < 3:
            ncomp = model.shape[1]
            model = np.reshape(model, (n1, n2, n3, ncomp))
            model = np.swapaxes(model, 0, 1)
            model = np.swapaxes(model, 0, 2)
            # Now reverse Z axis
            model = model[::-1, :, :, :]  # Note it is in Fortran ordering
            model = np.reshape(model, (n1 * n2 * n3, ncomp))
        else:
            model = np.reshape(model, (n1, n2, n3))
            model = np.swapaxes(model, 0, 1)
            model = np.swapaxes(model, 0, 2)
            # Now reverse Z axis
            model = model[::-1, :, :]  # Note it is in Fortran ordering
            model = model.flatten()

        # Convert data to VTK data structure and append to output
        c = interface.convert_array(model, name=data_name, deep=True)
        # THIS IS CELL DATA! Add the model data to CELL data:
        mesh.GetCellData().AddArray(c)
        return mesh

    # ------------------------------------------------------------------#
    # ----------------------     UBC MESH 2D    ------------------------#
    # ------------------------------------------------------------------#

    @staticmethod
    def ubc_mesh_2d(FileName, output):
        """This method reads a UBC 2D Mesh file and builds an empty
        ``vtkRectilinearGrid`` for data to be inserted into. `Format Specs`_.

        .. _Format Specs: http://giftoolscookbook.readthedocs.io/en/latest/content/fileFormats/mesh2Dfile.html

        Args:
            FileName (str) : The mesh filename as an absolute path for the input
                mesh file in UBC 3D Mesh Format.
            output (vtkRectilinearGrid) : The output data object

        Return:
            vtkRectilinearGrid :
                a ``vtkRectilinearGrid`` generated from the UBC 3D Mesh grid.
                Mesh is defined by the input mesh file.
                No data attributes here, simply an empty mesh. Use the
                ``place_model_on_mesh()`` method to associate with model data.
        """
        # Read in data from file
        xpts, xdisc, zpts, zdisc = ubcMeshReaderBase._ubc_mesh_2d_part(FileName)

        nx = np.sum(np.array(xdisc, dtype=int)) + 1
        nz = np.sum(np.array(zdisc, dtype=int)) + 1

        # Now generate the vtkRectilinear Grid
        def _genCoords(pts, disc, z=False):
            c = [float(pts[0])]
            for i in range(len(pts) - 1):
                start = float(pts[i])
                stop = float(pts[i + 1])
                num = int(disc[i])
                w = (stop - start) / num

                for j in range(1, num):
                    c.append(start + (j) * w)
                c.append(stop)
            c = np.array(c, dtype=float)
            if z:
                c = -c[::-1]
            return interface.convert_array(c, deep=True)

        xcoords = _genCoords(xpts, xdisc)
        zcoords = _genCoords(zpts, zdisc, z=True)
        ycoords = interface.convert_array(np.zeros(1), deep=True)

        output.SetDimensions(nx, 2, nz)  # note this subtracts 1
        output.SetXCoordinates(xcoords)
        output.SetYCoordinates(ycoords)
        output.SetZCoordinates(zcoords)

        return output

    @staticmethod
    def ubc_model_2d(FileName):
        """Reads a 2D model file and returns a 1D NumPy float array. Use the
        ``place_model_on_mesh()`` method to associate with a grid.

        Note:
            Only supports single component data

        Args:
            FileName (str) : The model filename as an absolute path for the
                input model file in UBCMesh Model Format. Also accepts a list of
                string file names.

        Return:
            np.array :
                a NumPy float array that holds the model data read from
                the file. Use the ``place_model_on_mesh()`` method to associate
                with a grid. If a list of file names is given then it will
                return a dictionary of NumPy float array with keys as the
                basenames of the files.
        """
        if isinstance(FileName, (list, tuple)):
            out = {}
            for f in FileName:
                out[os.path.basename(f)] = TensorMeshReader.ubc_model_2d(f)
            return out

        dim = np.genfromtxt(
            FileName, dtype=int, delimiter=None, comments='!', max_rows=1
        )
        names = ['col%d' % i for i in range(dim[0])]
        df = pd.read_csv(
            FileName, names=names, delim_whitespace=True, skiprows=1, comment='!'
        )
        data = df.values
        if np.shape(data)[0] != dim[1] and np.shape(data)[1] != dim[0]:
            raise _helpers.PVGeoError('Mode file `%s` improperly formatted.' % FileName)
        return data.flatten(order='F')

    def __ubc_mesh_data_2d(self, filename_mesh, filename_models, output):
        """Helper method to read a 2D mesh"""
        # Construct/read the mesh
        if self.need_to_readMesh():
            TensorMeshReader.ubc_mesh_2d(filename_mesh, self.__mesh)
            self.need_to_readMesh(flag=False)
        output.DeepCopy(self.__mesh)
        if self.need_to_readModels() and self.this_has_models():
            self.__models = []
            for f in filename_models:
                # Read the model data
                self.__models.append(TensorMeshReader.ubc_model_2d(f))
            self.need_to_readModels(flag=False)
        return output

    # ------------------------------------------------------------------#
    # ----------------------     UBC MESH 3D    ------------------------#
    # ------------------------------------------------------------------#

    @staticmethod
    def ubc_mesh_3d(FileName, output):
        """This method reads a UBC 3D Mesh file and builds an empty
        ``vtkRectilinearGrid`` for data to be inserted into.

        Args:
            FileName (str) : The mesh filename as an absolute path for the input
                mesh file in UBC 3D Mesh Format.
            output (vtkRectilinearGrid) : The output data object

        Return:
            vtkRectilinearGrid :
                a ``vtkRectilinearGrid`` generated from the UBC 3D Mesh grid.
                Mesh is defined by the input mesh file.
                No data attributes here, simply an empty mesh. Use the
                ``place_model_on_mesh()`` method to associate with model data.
        """

        # --- Read in the mesh ---#
        fileLines = np.genfromtxt(FileName, dtype=str, delimiter='\n', comments='!')

        # Get mesh dimensions
        dim = np.array(fileLines[0].split('!')[0].split(), dtype=int)
        dim = (dim[0] + 1, dim[1] + 1, dim[2] + 1)

        # The origin corner (Southwest-top)
        # - Remember UBC format specifies down as the positive Z
        # - Easting, Northing, Altitude
        oo = np.array(fileLines[1].split('!')[0].split(), dtype=float)
        ox, oy, oz = oo[0], oo[1], oo[2]

        # Read cell sizes for each line in the UBC mesh files
        def _readCellLine(line):
            line_list = []
            for seg in line.split():
                if '*' in seg:
                    sp = seg.split('*')
                    seg_arr = np.ones((int(sp[0]),), dtype=float) * float(sp[1])
                else:
                    seg_arr = np.array([float(seg)], dtype=float)
                line_list.append(seg_arr)
            return np.concatenate(line_list)

        # Read the cell sizes
        cx = _readCellLine(fileLines[2].split('!')[0])
        cy = _readCellLine(fileLines[3].split('!')[0])
        cz = _readCellLine(fileLines[4].split('!')[0])
        # Invert the indexing of the vector to start from the bottom.
        cz = cz[::-1]
        # Adjust the reference point to the bottom south west corner
        oz = oz - np.sum(cz)

        # Now generate the coordinates for from cell width and origin
        cox = ox + np.cumsum(cx)
        cox = np.insert(cox, 0, ox)
        coy = oy + np.cumsum(cy)
        coy = np.insert(coy, 0, oy)
        coz = oz + np.cumsum(cz)
        coz = np.insert(coz, 0, oz)

        # Set the dims and coordinates for the output
        output.SetDimensions(dim[0], dim[1], dim[2])
        # Convert to VTK array for setting coordinates
        output.SetXCoordinates(interface.convert_array(cox, deep=True))
        output.SetYCoordinates(interface.convert_array(coy, deep=True))
        output.SetZCoordinates(interface.convert_array(coz, deep=True))

        return output

    def __ubc_mesh_data_3d(self, filename_mesh, filename_models, output):
        """Helper method to read a 3D mesh"""
        # Construct/read the mesh
        if self.need_to_readMesh():
            TensorMeshReader.ubc_mesh_3d(filename_mesh, self.__mesh)
            self.need_to_readMesh(flag=False)
        output.DeepCopy(self.__mesh)
        if self.need_to_readModels() and self.this_has_models():
            self.__models = []
            for f in filename_models:
                # Read the model data
                self.__models.append(TensorMeshReader.ubc_model_3d(f))
            self.need_to_readModels(flag=False)
        return output

    def __ubc_tensor_mesh(self, filename_mesh, filename_models, output):
        """Wrapper to Read UBC GIF 2D and 3D meshes. UBC Mesh 2D/3D models are
        defined using a 2-file format. The "mesh" file describes how the data is
        descritized. The "model" file lists the physical property values for all
        cells in a mesh. A model file is meaningless without an associated mesh
        file. If the mesh file is 2D, then then model file must also be in the
        2D format (same for 3D).

        Args:
            filename_mesh (str) : The mesh filename as an absolute path for the
                input mesh file in UBC 2D/3D Mesh Format
            filename_models (str or list(str)) : The model filename(s) as an
                absolute path for the input model file in UBC 2D/3D Model Format.
            output (vtkRectilinearGrid) : The output data object

        Return:
            vtkRectilinearGrid :
                a ``vtkRectilinearGrid`` generated from the UBC 2D/3D Mesh grid.
                Mesh is defined by the input mesh file.
                Cell data is defined by the input model file.

        """
        # Check if the mesh is a UBC 2D mesh
        if self.is_2d():
            self.__ubc_mesh_data_2d(filename_mesh, filename_models, output)
        # Check if the mesh is a UBC 3D mesh
        elif self.is_3d():
            self.__ubc_mesh_data_3d(filename_mesh, filename_models, output)
        else:
            raise _helpers.PVGeoError('File format not recognized')
        return output

    def RequestData(self, request, inInfo, outInfo):
        """Handles data request by the pipeline."""
        # Get output:
        output = self.GetOutputData(outInfo, 0)
        # Get requested time index
        i = _helpers.get_requested_time(self, outInfo)
        self.__ubc_tensor_mesh(
            self.get_mesh_filename(), self.get_model_filenames(), output
        )
        # Place the model data for given timestep onto the mesh
        if len(self.__models) > i:
            TensorMeshReader.place_model_on_mesh(
                output, self.__models[i], self.get_data_name()
            )
        return 1

    def RequestInformation(self, request, inInfo, outInfo):
        """Handles info request by pipeline about timesteps and grid extents."""
        # Call parent to handle time stuff
        ubcMeshReaderBase.RequestInformation(self, request, inInfo, outInfo)
        # Now set whole output extent
        if self.need_to_readMesh():
            ext = self._read_extent()
            info = outInfo.GetInformationObject(0)
            # Set WHOLE_EXTENT: This is absolutely necessary
            info.Set(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(), ext, 6)
        return 1

    def clear_mesh(self):
        """Use to clean/rebuild the mesh"""
        self.__mesh = vtk.vtkRectilinearGrid()
        ubcMeshReaderBase.clear_models(self)

    def clear_models(self):
        """Use to clean the models and reread"""
        self.__models = []
        ubcMeshReaderBase.clear_models(self)

###############################################################################

class TensorMeshAppender(ModelAppenderBase):
    """This filter reads a timeseries of models and appends it to an input
    ``vtkRectilinearGrid``
    """

    __displayname__ = 'UBC Tensor Mesh Appender'
    __category__ = 'filter'

    def __init__(self, **kwargs):
        ModelAppenderBase.__init__(
            self,
            inputType='vtkRectilinearGrid',
            outputType='vtkRectilinearGrid',
            **kwargs
        )

    def _read_up_front(self):
        """Internal helepr to read data at start"""
        reader = ubcMeshReaderBase.ubc_model_3d
        if not self._is_3D:
            # Note how in UBC format, 2D grids are specified on an XZ plane (no Y component)
            # This will only work prior to rotations to account for real spatial reference
            reader = TensorMeshReader.ubc_model_2d
        self._models = []
        for f in self._model_filenames:
            # Read the model data
            self._models.append(reader(f))
        self.need_to_read(flag=False)
        return

    def _place_on_mesh(self, output, idx=0):
        """Internal helepr to place a model on the mesh for a given index"""
        TensorMeshReader.place_model_on_mesh(
            output, self._models[idx], self.get_data_name()
        )
        return

###############################################################################

class TopoMeshAppender(AlgorithmBase):
    """This filter reads a single discrete topography file and appends it as a
    boolean data array.
    """

    __displayname__ = 'Append UBC Discrete Topography'
    __category__ = 'filter'

    def __init__(
        self, inputType='vtkRectilinearGrid', outputType='vtkRectilinearGrid', **kwargs
    ):
        AlgorithmBase.__init__(
            self,
            nInputPorts=1,
            inputType=inputType,
            nOutputPorts=1,
            outputType=outputType,
        )
        self._topoFileName = kwargs.get('filename', None)
        self.__indices = None
        self.__need_to_read = True
        self.__ne, self.__nn = None, None

    def need_to_read(self, flag=None):
        """Ask self if the reader needs to read the files again

        Args:
            flag (bool): if the flag is set then this method will set the read
                status

        Return:
            bool:
                The status of the reader aspect of the filter.
        """
        if flag is not None and isinstance(flag, (bool, int)):
            self.__need_to_read = flag
        return self.__need_to_read

    def Modified(self, read_again=True):
        """Call modified if the files needs to be read again again."""
        if read_again:
            self.__need_to_read = read_again
        AlgorithmBase.Modified(self)

    def modified(self, read_again=True):
        """Call modified if the files needs to be read again again."""
        return self.Modified(read_again=read_again)

    def _read_up_front(self):
        """Internal helepr to read data at start"""
        # Read the file
        content = np.genfromtxt(
            self._topoFileName, dtype=str, delimiter='\n', comments='!'
        )
        dim = content[0].split()
        self.__ne, self.__nn = int(dim[0]), int(dim[1])
        self.__indices = pd.read_csv(
            StringIO("\n".join(content[1::])),
            names=['i', 'j', 'k'],
            delim_whitespace=True,
        )
        # NOTE: K indices are inverted
        self.need_to_read(flag=False)
        return

    def _place_on_mesh(self, output):
        """Internal helepr to place an active cells model on the mesh"""
        # Check mesh extents to math topography
        nx, ny, nz = output.GetDimensions()
        nx, ny, nz = nx - 1, ny - 1, nz - 1  # because GetDimensions counts the nodes
        topz = np.max(self.__indices['k']) + 1
        if nx != self.__nn or ny != self.__ne or topz > nz:
            raise _helpers.PVGeoError(
                'Dimension mismatch between input grid and topo file.'
            )
        # # Adjust the k indices to be in caarteian system
        # self.__indices['k'] = nz - self.__indices['k']
        # Fill out the topo and add it as model as it will be in UBC format
        # Create a 3D array of 1s and zeros (1 means beneath topo or active)
        topo = np.empty((ny, nx, nz), dtype=float)
        topo[:] = np.nan
        for row in self.__indices.values:
            i, j, k = row
            topo[i, j, k + 1 :] = 0
            topo[i, j, : k + 1] = 1
        # Add as model... ``place_model_on_mesh`` handles the rest
        TensorMeshReader.place_model_on_mesh(
            output, topo.flatten(), 'Active Topography'
        )
        return

    def RequestData(self, request, inInfo, outInfo):
        """Used by pipeline to generate output"""
        # Get input/output of Proxy
        pdi = self.GetInputData(inInfo, 0, 0)
        output = self.GetOutputData(outInfo, 0)
        output.DeepCopy(pdi)  # ShallowCopy if you want changes to propagate upstream
        # Perfrom task:
        if self.__need_to_read:
            self._read_up_front()
        # Place the model data for given timestep onto the mesh
        self._place_on_mesh(output)
        return 1

    #### Setters and Getters ####

    def clear_topo_file(self):
        """Use to clear data file name."""
        self._topoFileName = None
        self.Modified(read_again=True)

    def set_topo_filename(self, filename):
        """Use to set the file names for the reader. Handles single strings only"""
        if filename is None:
            return  # do nothing if None is passed by a constructor on accident
        elif isinstance(filename, str) and self._topoFileName != filename:
            self._topoFileName = filename
            self.Modified()
        return 1

###############################################################################

#
# import numpy as np
# indices = np.array([[0,0,1],
#                     [0,1,1],
#                     [0,2,1],
#                     [1,0,1],
#                     [1,1,1],
#                     [1,2,1],
#                     [2,0,1],
#                     [2,1,1],
#                     [2,2,1],
#                     ])
#
# topo = np.empty((3,3,3), dtype=float)
# topo[:] = np.nan
#
# for row in indices:
#     i, j, k = row
#     topo[i, j, k:] = 0
#     topo[i, j, :k] = 1
# topo


# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Vincent Michel <vincent.michel@inria.fr>
#          Gilles Louppe <g.louppe@gmail.com>
#
# License: BSD 3 clause

"""Recursive feature elimination for feature ranking"""

import numpy as np
from ..utils import check_X_y, safe_sqr
from ..utils.metaestimators import if_delegate_has_method
from ..base import BaseEstimator
from ..base import MetaEstimatorMixin
from ..base import clone
from ..base import is_classifier
from ..model_selection import check_cv
from ..model_selection._validation import _safe_split, _score
from ..metrics.scorer import check_scoring
from .base import SelectorMixin

class RFE(BaseEstimator, MetaEstimatorMixin, SelectorMixin):
    """Feature ranking with recursive feature elimination.

    Given an external estimator that assigns weights to features (e.g., the
    coefficients of a linear model), the goal of recursive feature elimination
    (RFE) is to select features by recursively considering smaller and smaller
    sets of features. First, the estimator is trained on the initial set of
    features and weights are assigned to each one of them. Then, features whose
    absolute weights are the smallest are pruned from the current set features.
    That procedure is recursively repeated on the pruned set until the desired
    number of features to select is eventually reached.

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

    Parameters
    ----------
    estimator : object
        A supervised learning estimator with a `fit` method that updates a
        `coef_` attribute that holds the fitted parameters. Important features
        must correspond to high absolute values in the `coef_` array.

        For instance, this is the case for most supervised learning
        algorithms such as Support Vector Classifiers and Generalized
        Linear Models from the `svm` and `linear_model` modules.

    n_features_to_select : int or None (default=None)
        The number of features to select. If `None`, half of the features
        are selected.

    step : int or float, optional (default=1)
        If greater than or equal to 1, then `step` corresponds to the (integer)
        number of features to remove at each iteration.
        If within (0.0, 1.0), then `step` corresponds to the percentage
        (rounded down) of features to remove at each iteration.

    verbose : int, default=0
        Controls verbosity of output.

    Attributes
    ----------
    n_features_ : int
        The number of selected features.

    support_ : array of shape [n_features]
        The mask of selected features.

    ranking_ : array of shape [n_features]
        The feature ranking, such that ``ranking_[i]`` corresponds to the
        ranking position of the i-th feature. Selected (i.e., estimated
        best) features are assigned rank 1.

    estimator_ : object
        The external estimator fit on the reduced dataset.

    Examples
    --------
    The following example shows how to retrieve the 5 right informative
    features in the Friedman #1 dataset.

    >>> from sklearn.datasets import make_friedman1
    >>> from sklearn.feature_selection import RFE
    >>> from sklearn.svm import SVR
    >>> X, y = make_friedman1(n_samples=50, n_features=10, random_state=0)
    >>> estimator = SVR(kernel="linear")
    >>> selector = RFE(estimator, 5, step=1)
    >>> selector = selector.fit(X, y)
    >>> selector.support_ # doctest: +NORMALIZE_WHITESPACE
    array([ True,  True,  True,  True,  True,
            False, False, False, False, False], dtype=bool)
    >>> selector.ranking_
    array([1, 1, 1, 1, 1, 6, 4, 3, 2, 5])

    References
    ----------

    .. [1] Guyon, I., Weston, J., Barnhill, S., & Vapnik, V., "Gene selection
           for cancer classification using support vector machines",
           Mach. Learn., 46(1-3), 389--422, 2002.
    """
    def __init__(self, estimator, n_features_to_select=None, step=1,
                 verbose=0):
        self.estimator = estimator
        self.n_features_to_select = n_features_to_select
        self.step = step
        self.verbose = verbose

    @property
    def _estimator_type(self):
        return self.estimator._estimator_type

    def fit(self, X, y):
        """Fit the RFE model and then the underlying estimator on the selected
           features.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples]
            The target values.
        """
        return self._fit(X, y)

    def _fit(self, X, y, step_score=None):
        X, y = check_X_y(X, y, "csc")
        # Initialization
        n_features = X.shape[1]
        if self.n_features_to_select is None:
            n_features_to_select = n_features // 2
        else:
            n_features_to_select = self.n_features_to_select

        if 0.0 < self.step < 1.0:
            step = int(max(1, self.step * n_features))
        else:
            step = int(self.step)
        if step <= 0:
            raise ValueError("Step must be >0")

        support_ = np.ones(n_features, dtype=np.bool)
        ranking_ = np.ones(n_features, dtype=np.int)

        if step_score:
            self.scores_ = []

        # Elimination
        while np.sum(support_) > n_features_to_select:
            # Remaining features
            features = np.arange(n_features)[support_]

            # Rank the remaining features
            estimator = clone(self.estimator)
            if self.verbose > 0:
                print("Fitting estimator with %d features." % np.sum(support_))

            estimator.fit(X[:, features], y)

            # Get coefs
            if hasattr(estimator, 'coef_'):
                coefs = estimator.coef_
            elif hasattr(estimator, 'feature_importances_'):
                coefs = estimator.feature_importances_
            else:
                raise RuntimeError('The classifier does not expose '
                                   '"coef_" or "feature_importances_" '
                                   'attributes')

            # Get ranks
            if coefs.ndim > 1:
                ranks = np.argsort(safe_sqr(coefs).sum(axis=0))
            else:
                ranks = np.argsort(safe_sqr(coefs))

            # for sparse case ranks is matrix
            ranks = np.ravel(ranks)

            # Eliminate the worse features
            threshold = min(step, np.sum(support_) - n_features_to_select)

            # Compute step score on the previous selection iteration
            # because 'estimator' must use features
            # that have not been eliminated yet
            if step_score:
                self.scores_.append(step_score(estimator, features))
            support_[features[ranks][:threshold]] = False
            ranking_[np.logical_not(support_)] += 1

        # Set final attributes
        features = np.arange(n_features)[support_]
        self.estimator_ = clone(self.estimator)
        self.estimator_.fit(X[:, features], y)

        # Compute step score when only n_features_to_select features left
        if step_score:
            self.scores_.append(step_score(self.estimator_, features))
        self.n_features_ = support_.sum()
        self.support_ = support_
        self.ranking_ = ranking_

        return self

    @if_delegate_has_method(delegate='estimator')
    def predict(self, X):
        """Reduce X to the selected features and then predict using the
           underlying estimator.

        Parameters
        ----------
        X : array of shape [n_samples, n_features]
            The input samples.

        Returns
        -------
        y : array of shape [n_samples]
            The predicted target values.
        """
        return self.estimator_.predict(self.transform(X))

    @if_delegate_has_method(delegate='estimator')
    def score(self, X, y):
        """Reduce X to the selected features and then return the score of the
           underlying estimator.

        Parameters
        ----------
        X : array of shape [n_samples, n_features]
            The input samples.

        y : array of shape [n_samples]
            The target values.
        """
        return self.estimator_.score(self.transform(X), y)

    def _get_support_mask(self):
        return self.support_

    @if_delegate_has_method(delegate='estimator')
    def decision_function(self, X):
        return self.estimator_.decision_function(self.transform(X))

    @if_delegate_has_method(delegate='estimator')
    def predict_proba(self, X):
        return self.estimator_.predict_proba(self.transform(X))

    @if_delegate_has_method(delegate='estimator')
    def predict_log_proba(self, X):
        return self.estimator_.predict_log_proba(self.transform(X))

class RFECV(RFE, MetaEstimatorMixin):
    """Feature ranking with recursive feature elimination and cross-validated
    selection of the best number of features.

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

    Parameters
    ----------
    estimator : object
        A supervised learning estimator with a `fit` method that updates a
        `coef_` attribute that holds the fitted parameters. Important features
        must correspond to high absolute values in the `coef_` array.

        For instance, this is the case for most supervised learning
        algorithms such as Support Vector Classifiers and Generalized
        Linear Models from the `svm` and `linear_model` modules.

    step : int or float, optional (default=1)
        If greater than or equal to 1, then `step` corresponds to the (integer)
        number of features to remove at each iteration.
        If within (0.0, 1.0), then `step` corresponds to the percentage
        (rounded down) of features to remove at each iteration.

    cv : int, cross-validation generator or an iterable, optional
        Determines the cross-validation splitting strategy.
        Possible inputs for cv are:

        - None, to use the default 3-fold cross-validation,
        - integer, to specify the number of folds.
        - An object to be used as a cross-validation generator.
        - An iterable yielding train/test splits.

        For integer/None inputs, if ``y`` is binary or multiclass,
        :class:`StratifiedKFold` used. If the estimator is a classifier
        or if ``y`` is neither binary nor multiclass, :class:`KFold` is used.

        Refer :ref:`User Guide <cross_validation>` for the various
        cross-validation strategies that can be used here.

    scoring : string, callable or None, optional, default: None
        A string (see model evaluation documentation) or
        a scorer callable object / function with signature
        ``scorer(estimator, X, y)``.

    verbose : int, default=0
        Controls verbosity of output.

    Attributes
    ----------
    n_features_ : int
        The number of selected features with cross-validation.

    support_ : array of shape [n_features]
        The mask of selected features.

    ranking_ : array of shape [n_features]
        The feature ranking, such that `ranking_[i]`
        corresponds to the ranking
        position of the i-th feature.
        Selected (i.e., estimated best)
        features are assigned rank 1.

    grid_scores_ : array of shape [n_subsets_of_features]
        The cross-validation scores such that
        ``grid_scores_[i]`` corresponds to
        the CV score of the i-th subset of features.

    estimator_ : object
        The external estimator fit on the reduced dataset.

    Notes
    -----
    The size of ``grid_scores_`` is equal to ceil((n_features - 1) / step) + 1,
    where step is the number of features removed at each iteration.

    Examples
    --------
    The following example shows how to retrieve the a-priori not known 5
    informative features in the Friedman #1 dataset.

    >>> from sklearn.datasets import make_friedman1
    >>> from sklearn.feature_selection import RFECV
    >>> from sklearn.svm import SVR
    >>> X, y = make_friedman1(n_samples=50, n_features=10, random_state=0)
    >>> estimator = SVR(kernel="linear")
    >>> selector = RFECV(estimator, step=1, cv=5)
    >>> selector = selector.fit(X, y)
    >>> selector.support_ # doctest: +NORMALIZE_WHITESPACE
    array([ True,  True,  True,  True,  True,
            False, False, False, False, False], dtype=bool)
    >>> selector.ranking_
    array([1, 1, 1, 1, 1, 6, 4, 3, 2, 5])

    References
    ----------

    .. [1] Guyon, I., Weston, J., Barnhill, S., & Vapnik, V., "Gene selection
           for cancer classification using support vector machines",
           Mach. Learn., 46(1-3), 389--422, 2002.
    """
    def __init__(self, estimator, step=1, cv=None, scoring=None, verbose=0):
        self.estimator = estimator
        self.step = step
        self.cv = cv
        self.scoring = scoring
        self.verbose = verbose

    def fit(self, X, y):
        """Fit the RFE model and automatically tune the number of selected
           features.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training vector, where `n_samples` is the number of samples and
            `n_features` is the total number of features.

        y : array-like, shape = [n_samples]
            Target values (integers for classification, real numbers for
            regression).
        """
        X, y = check_X_y(X, y, "csr")

        # Initialization
        cv = check_cv(self.cv, y, is_classifier(self.estimator))
        scorer = check_scoring(self.estimator, scoring=self.scoring)
        n_features = X.shape[1]
        n_features_to_select = 1

        # Determine the number of subsets of features
        scores = []

        # Cross-validation
        for n, (train, test) in enumerate(cv.split(X, y)):
            X_train, y_train = _safe_split(self.estimator, X, y, train)
            X_test, y_test = _safe_split(self.estimator, X, y, test, train)

            rfe = RFE(estimator=self.estimator,
                      n_features_to_select=n_features_to_select,
                      step=self.step, verbose=self.verbose - 1)

            rfe._fit(X_train, y_train, lambda estimator, features:
                     _score(estimator, X_test[:, features], y_test, scorer))
            scores.append(np.array(rfe.scores_[::-1]).reshape(1, -1))
        scores = np.sum(np.concatenate(scores, 0), 0)
        # The index in 'scores' when 'n_features' features are selected
        n_feature_index = np.ceil((n_features - n_features_to_select) /
                                  float(self.step))
        n_features_to_select = max(n_features_to_select,
                                   n_features - ((n_feature_index -
                                                 np.argmax(scores)) *
                                                 self.step))
        # Re-execute an elimination with best_k over the whole set
        rfe = RFE(estimator=self.estimator,
                  n_features_to_select=n_features_to_select, step=self.step)

        rfe.fit(X, y)

        # Set final attributes
        self.support_ = rfe.support_
        self.n_features_ = rfe.n_features_
        self.ranking_ = rfe.ranking_
        self.estimator_ = clone(self.estimator)
        self.estimator_.fit(self.transform(X), y)

        # Fixing a normalization error, n is equal to get_n_splits(X, y) - 1
        # here, the scores are normalized by get_n_splits(X, y)
        self.grid_scores_ = scores / cv.get_n_splits(X, y)
        return self


__author__ = 'saeedamen'  # Saeed Amen

#
# Copyright 2016 Cuemacro
#
# 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.
#

import pandas

# support Quandl 3.x.x
try:
    import quandl as Quandl
except:
    # if import fails use Quandl 2.x.x
    import Quandl

from chartpy import Chart, Style

# get your own free bQuandl API key from https://www.quandl.com/
try:
    from chartpy.chartcred import ChartCred

    cred = ChartCred()
    quandl_api_key = cred.quandl_api_key
except:
    quandl_api_key = "x"

# choose run_example = 0 for everything
# run_example = 1 - plot US GDP QoQ (real) and nominal with Plotly/Bokeh/Matplotlib with subplots for each line
# run_example = 2 - plot US GDP QoQ (real + nominal) in two double plots (passing an array of dataframes)
run_example = 0

if run_example == 1 or run_example == 0:
    df = Quandl.get(["FRED/A191RL1Q225SBEA", "FRED/A191RP1Q027SBEA"], authtoken=quandl_api_key)

    df.columns = ["Real QoQ", "Nominal QoQ"]

    # set the style of the plot
    style = Style(title="US GDP", source="Quandl/Fred", subplots=True)

    # Chart object is initialised with the dataframe and our chart style
    chart = Chart(df=df, chart_type='line', style=style)

    chart.plot(engine='matplotlib')
    chart.plot(engine='bokeh')
    chart.plot(engine='plotly')

if run_example == 2 or run_example == 0:
    df = Quandl.get(["FRED/A191RL1Q225SBEA", "FRED/A191RP1Q027SBEA"], authtoken=quandl_api_key)

    df.columns = ["Real QoQ", "Nominal QoQ"]

    df = [df, df]

    # set the style of the plot
    style = Style(title="US GDP double plot", source="Quandl/Fred", subplots=True)

    # Chart object is initialised with the dataframe and our chart style
    chart = Chart(df=df, chart_type='line', style=style)

    chart.plot(engine='bokeh')
    chart.plot(engine='matplotlib')
    chart.plot(engine='plotly') # TODO fix legends though


# -*- coding: utf-8 -*-
"""

Created on Sun May 26 13:23:40 2013

Author: Josef Perktold, based on Enrico Giampieri's multiOLS
"""

#import numpy as np
import pandas as pd

import statsmodels.api as sm
from statsmodels.sandbox.multilinear import multiOLS, multigroup

data = sm.datasets.longley.load_pandas()
df = data.exog
df['TOTEMP'] = data.endog

#This will perform the specified linear model on all the
#other columns of the dataframe
res0 = multiOLS('GNP + 1', df)

#This select only a certain subset of the columns
res = multiOLS('GNP + 0', df, ['GNPDEFL', 'TOTEMP', 'POP'])
print(res.to_string())

url = "http://vincentarelbundock.github.com/"
url = url + "Rdatasets/csv/HistData/Guerry.csv"
df = pd.read_csv(url, index_col=1) #'dept')

#evaluate the relationship between the various parameters whith the Wealth
pvals = multiOLS('Wealth', df)['adj_pvals', '_f_test']

#define the groups
groups = {}
groups['crime'] = ['Crime_prop', 'Infanticide',
                   'Crime_parents', 'Desertion', 'Crime_pers']
groups['religion'] = ['Donation_clergy', 'Clergy', 'Donations']
groups['wealth'] = ['Commerce', 'Lottery', 'Instruction', 'Literacy']

#do the analysis of the significance
res3 = multigroup(pvals < 0.05, groups)
print(res3)


import numpy as np
from ..qt import QtWidgets, has_qt, FigureManagerQT, FigureCanvasQTAgg
from ..._shared.utils import warn
import matplotlib as mpl
from matplotlib.figure import Figure
from matplotlib import _pylab_helpers
from matplotlib.colors import LinearSegmentedColormap

if has_qt and 'agg' not in mpl.get_backend().lower():
    warn("Recommended matplotlib backend is `Agg` for full "
         "skimage.viewer functionality.")

__all__ = ['init_qtapp', 'start_qtapp', 'RequiredAttr', 'figimage',
           'LinearColormap', 'ClearColormap', 'FigureCanvas', 'new_plot',
           'update_axes_image']

QApp = None

def init_qtapp():
    """Initialize QAppliction.

    The QApplication needs to be initialized before creating any QWidgets
    """
    global QApp
    QApp = QtWidgets.QApplication.instance()
    if QApp is None:
        QApp = QtWidgets.QApplication([])
    return QApp

def is_event_loop_running(app=None):
    """Return True if event loop is running."""
    if app is None:
        app = init_qtapp()
    if hasattr(app, '_in_event_loop'):
        return app._in_event_loop
    else:
        return False

def start_qtapp(app=None):
    """Start Qt mainloop"""
    if app is None:
        app = init_qtapp()
    if not is_event_loop_running(app):
        app._in_event_loop = True
        app.exec_()
        app._in_event_loop = False
    else:
        app._in_event_loop = True

class RequiredAttr(object):
    """A class attribute that must be set before use."""

    instances = dict()

    def __init__(self, init_val=None):
        self.instances[self, None] = init_val

    def __get__(self, obj, objtype):
        value = self.instances[self, obj]
        if value is None:
            raise AttributeError('Required attribute not set')
        return value

    def __set__(self, obj, value):
        self.instances[self, obj] = value

class LinearColormap(LinearSegmentedColormap):
    """LinearSegmentedColormap in which color varies smoothly.

    This class is a simplification of LinearSegmentedColormap, which doesn't
    support jumps in color intensities.

    Parameters
    ----------
    name : str
        Name of colormap.

    segmented_data : dict
        Dictionary of 'red', 'green', 'blue', and (optionally) 'alpha' values.
        Each color key contains a list of `x`, `y` tuples. `x` must increase
        monotonically from 0 to 1 and corresponds to input values for a
        mappable object (e.g. an image). `y` corresponds to the color
        intensity.

    """
    def __init__(self, name, segmented_data, **kwargs):
        segmented_data = dict((key, [(x, y, y) for x, y in value])
                              for key, value in segmented_data.items())
        LinearSegmentedColormap.__init__(self, name, segmented_data, **kwargs)

class ClearColormap(LinearColormap):
    """Color map that varies linearly from alpha = 0 to 1
    """
    def __init__(self, rgb, max_alpha=1, name='clear_color'):
        r, g, b = rgb
        cg_speq = {'blue':  [(0.0, b), (1.0, b)],
                   'green': [(0.0, g), (1.0, g)],
                   'red':   [(0.0, r), (1.0, r)],
                   'alpha': [(0.0, 0.0), (1.0, max_alpha)]}
        LinearColormap.__init__(self, name, cg_speq)

class FigureCanvas(FigureCanvasQTAgg):
    """Canvas for displaying images."""
    def __init__(self, figure, **kwargs):
        self.fig = figure
        FigureCanvasQTAgg.__init__(self, self.fig)
        FigureCanvasQTAgg.setSizePolicy(self,
                                        QtWidgets.QSizePolicy.Expanding,
                                        QtWidgets.QSizePolicy.Expanding)
        FigureCanvasQTAgg.updateGeometry(self)

    def resizeEvent(self, event):
        FigureCanvasQTAgg.resizeEvent(self, event)
        # Call to `resize_event` missing in FigureManagerQT.
        # See https://github.com/matplotlib/matplotlib/pull/1585
        self.resize_event()

def new_canvas(*args, **kwargs):
    """Return a new figure canvas."""
    allnums = _pylab_helpers.Gcf.figs.keys()
    num = max(allnums) + 1 if allnums else 1

    FigureClass = kwargs.pop('FigureClass', Figure)
    figure = FigureClass(*args, **kwargs)
    canvas = FigureCanvas(figure)
    fig_manager = FigureManagerQT(canvas, num)
    return fig_manager.canvas

def new_plot(parent=None, subplot_kw=None, **fig_kw):
    """Return new figure and axes.

    Parameters
    ----------
    parent : QtWidget
        Qt widget that displays the plot objects. If None, you must manually
        call ``canvas.setParent`` and pass the parent widget.
    subplot_kw : dict
        Keyword arguments passed ``matplotlib.figure.Figure.add_subplot``.
    fig_kw : dict
        Keyword arguments passed ``matplotlib.figure.Figure``.
    """
    if subplot_kw is None:
        subplot_kw = {}
    canvas = new_canvas(**fig_kw)
    canvas.setParent(parent)

    fig = canvas.figure
    ax = fig.add_subplot(1, 1, 1, **subplot_kw)
    return fig, ax

def figimage(image, scale=1, dpi=None, **kwargs):
    """Return figure and axes with figure tightly surrounding image.

    Unlike pyplot.figimage, this actually plots onto an axes object, which
    fills the figure. Plotting the image onto an axes allows for subsequent
    overlays of axes artists.

    Parameters
    ----------
    image : array
        image to plot
    scale : float
        If scale is 1, the figure and axes have the same dimension as the
        image.  Smaller values of `scale` will shrink the figure.
    dpi : int
        Dots per inch for figure. If None, use the default rcParam.
    """
    dpi = dpi if dpi is not None else mpl.rcParams['figure.dpi']
    kwargs.setdefault('interpolation', 'nearest')
    kwargs.setdefault('cmap', 'gray')

    h, w, d = np.atleast_3d(image).shape
    figsize = np.array((w, h), dtype=float) / dpi * scale

    fig, ax = new_plot(figsize=figsize, dpi=dpi)
    fig.subplots_adjust(left=0, bottom=0, right=1, top=1)

    ax.set_axis_off()
    ax.imshow(image, **kwargs)
    ax.figure.canvas.draw()
    return fig, ax

def update_axes_image(image_axes, image):
    """Update the image displayed by an image plot.

    This sets the image plot's array and updates its shape appropriately

    Parameters
    ----------
    image_axes : `matplotlib.image.AxesImage`
        Image axes to update.
    image : array
        Image array.
    """
    image_axes.set_array(image)

    # Adjust size if new image shape doesn't match the original
    h, w = image.shape[:2]
    image_axes.set_extent((0, w, h, 0))


"""
=========================================================
Using FunctionTransformer to select columns
=========================================================

Shows how to use a function transformer in a pipeline. If you know your
dataset's first principle component is irrelevant for a classification task,
you can use the FunctionTransformer to select all but the first column of the
PCA transformed data.
"""
import matplotlib.pyplot as plt
import numpy as np

from sklearn.cross_validation import train_test_split
from sklearn.decomposition import PCA
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import FunctionTransformer

def _generate_vector(shift=0.5, noise=15):
    return np.arange(1000) + (np.random.rand(1000) - shift) * noise

def generate_dataset():
    """
    This dataset is two lines with a slope ~ 1, where one has
    a y offset of ~100
    """
    return np.vstack((
        np.vstack((
            _generate_vector(),
            _generate_vector() + 100,
        )).T,
        np.vstack((
            _generate_vector(),
            _generate_vector(),
        )).T,
    )), np.hstack((np.zeros(1000), np.ones(1000)))

def all_but_first_column(X):
    return X[:, 1:]

def drop_first_component(X, y):
    """
    Create a pipeline with PCA and the column selector and use it to
    transform the dataset.
    """
    pipeline = make_pipeline(
        PCA(), FunctionTransformer(all_but_first_column),
    )
    X_train, X_test, y_train, y_test = train_test_split(X, y)
    pipeline.fit(X_train, y_train)
    return pipeline.transform(X_test), y_test

if __name__ == '__main__':
    X, y = generate_dataset()
    plt.scatter(X[:, 0], X[:, 1], c=y, s=50)
    plt.show()
    X_transformed, y_transformed = drop_first_component(*generate_dataset())
    plt.scatter(
        X_transformed[:, 0],
        np.zeros(len(X_transformed)),
        c=y_transformed,
        s=50,
    )
    plt.show()


"""
=======================================================================================
Topic extraction with Non-negative Matrix Factorization and Latent Dirichlet Allocation
=======================================================================================

This is an example of applying Non-negative Matrix Factorization
and Latent Dirichlet Allocation on a corpus of documents and
extract additive models of the topic structure of the corpus.
The output is a list of topics, each represented as a list of terms
(weights are not shown).

The default parameters (n_samples / n_features / n_topics) should make
the example runnable in a couple of tens of seconds. You can try to
increase the dimensions of the problem, but be aware that the time
complexity is polynomial in NMF. In LDA, the time complexity is
proportional to (n_samples * iterations).
"""

# Author: Olivier Grisel <olivier.grisel@ensta.org>
#         Lars Buitinck <L.J.Buitinck@uva.nl>
#         Chyi-Kwei Yau <chyikwei.yau@gmail.com>
# License: BSD 3 clause

from __future__ import print_function
from time import time

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.decomposition import NMF, LatentDirichletAllocation
from sklearn.datasets import fetch_20newsgroups

n_samples = 2000
n_features = 1000
n_topics = 10
n_top_words = 20

def print_top_words(model, feature_names, n_top_words):
    for topic_idx, topic in enumerate(model.components_):
        print("Topic #%d:" % topic_idx)
        print(" ".join([feature_names[i]
                        for i in topic.argsort()[:-n_top_words - 1:-1]]))
    print()

# Load the 20 newsgroups dataset and vectorize it. We use a few heuristics
# to filter out useless terms early on: the posts are stripped of headers,
# footers and quoted replies, and common English words, words occurring in
# only one document or in at least 95% of the documents are removed.

print("Loading dataset...")
t0 = time()
dataset = fetch_20newsgroups(shuffle=True, random_state=1,
                             remove=('headers', 'footers', 'quotes'))
data_samples = dataset.data
print("done in %0.3fs." % (time() - t0))

# Use tf-idf features for NMF.
print("Extracting tf-idf features for NMF...")
tfidf_vectorizer = TfidfVectorizer(max_df=0.95, min_df=2, #max_features=n_features,
                                   stop_words='english')
t0 = time()
tfidf = tfidf_vectorizer.fit_transform(data_samples)
print("done in %0.3fs." % (time() - t0))

# Use tf (raw term count) features for LDA.
print("Extracting tf features for LDA...")
tf_vectorizer = CountVectorizer(max_df=0.95, min_df=2, max_features=n_features,
                                stop_words='english')
t0 = time()
tf = tf_vectorizer.fit_transform(data_samples)
print("done in %0.3fs." % (time() - t0))

# Fit the NMF model
print("Fitting the NMF model with tf-idf features,"
      "n_samples=%d and n_features=%d..."
      % (n_samples, n_features))
t0 = time()
nmf = NMF(n_components=n_topics, random_state=1, alpha=.1, l1_ratio=.5).fit(tfidf)
exit()
print("done in %0.3fs." % (time() - t0))

print("\nTopics in NMF model:")
tfidf_feature_names = tfidf_vectorizer.get_feature_names()
print_top_words(nmf, tfidf_feature_names, n_top_words)

print("Fitting LDA models with tf features, n_samples=%d and n_features=%d..."
      % (n_samples, n_features))
lda = LatentDirichletAllocation(n_topics=n_topics, max_iter=5,
                                learning_method='online', learning_offset=50.,
                                random_state=0)
t0 = time()
lda.fit(tf)
print("done in %0.3fs." % (time() - t0))

print("\nTopics in LDA model:")
tf_feature_names = tf_vectorizer.get_feature_names()
print_top_words(lda, tf_feature_names, n_top_words)