#!/usr/bin/env python
# Description = Visualizes different output data from APRICOT analysis

from collections import defaultdict
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import os
import sys

try:
    import subprocess
except ImportError:
    print('Python package subprocess is missing. Please install/update.\n')
    sys.exit(0)
try:
    import shutil
except ImportError:
    print('Python package shutil is missing. Please install/update.\n')
    sys.exit(0)

class VizApricotAnalysis(object):
    def __init__(self, annotation_scoring_data,
                 domain_file,
                 additional_annotation,
                 outpath):
        self._annotation_scoring_data = annotation_scoring_data
        self._domain_file = domain_file
        self._additional_annotation = additional_annotation
        self._outpath = outpath
        self._sec_str = self._outpath+'/secondary_structure'
        self._dom_highlight = self._outpath+'/domain_highlighting'
        self._pdb_msa = self._outpath+'/homologous_pdb_msa'
        self._overview = self._outpath+'/overview_and_statistics'
        self._localize = self._outpath+'/subcellular_localization'
        self._annotation_data = []
        self._filter_viz_dict = {}
        self._highlight_dict = {}
        self._uid_key_dict = {}
        self._dom_rank = {}
        self._fasta_dict = {}
        self._secstr_dict = {}
        self._dom_annotation = {}
        self._location_dict = defaultdict(lambda: defaultdict(lambda: []))
        self._sec_str_color = {'H': '#FF6666', 'E': '#33CCCC', 'C': '#FFFFCC'}
        self._localization_dict = defaultdict(
            lambda: defaultdict(lambda: float))
        self._color_list = (
            "Blue", "Green", "Teal", "Lime", "SeaGreen", "MediumTurquoise",
            "Pink", "DarkOliveGreen", "Indigo", "Orange", "SlateBlue",
            "LawnGreen", "Brown", "LightSkyBlue", "LightGreen", "DarkOrchid",
            "GoldenRod", "MidnightBlue", "LightPink", "Gold")

    def viz_all_the_visualization_files(self):
        self.viz_domain_data()
        self.domain_highlight()
        self.viz_annotation_scoring()
        self.viz_secondary_structure()
        self.viz_subcellular_localization()
        self.viz_homologous_pdb_msa()

    def viz_domain_data(self):
        with open(self._domain_file, 'r') as in_fh:
            for entry in in_fh:
                if not entry.startswith('Entry'):
                    domain_info = DomainDataColumns(
                        entry.strip().split('\t'))
                    prot_name = domain_info.entry_name
                    prot_end = int(domain_info.length)-1
                    prot_key = '\n'.join(
                        ["\tstart: 0,", "\tend: %s,"
                         % prot_end, '\tname: "%s",' % prot_name,
                         '\thref: "http://www.uniprot.org/uniprot/%s"'
                         % domain_info.uid])
                    self._uid_key_dict[domain_info.uid] = prot_key
                    self._location_dict[
                        domain_info.uid][domain_info.domain_id].append(
                            '\t{start: %s, end: %s}' % (
                                domain_info.start, domain_info.stop))
                    self._dom_annotation[
                        domain_info.domain_id] = domain_info.full_name
                    src = domain_info.resource
                    if src == 'CDD':
                        self._dom_rank.setdefault(
                            domain_info.uid+':CDD', []).append(
                            domain_info.domain_id)
                        self._highlight_dict.setdefault(
                            prot_key, []).append('\n'.join(
                                ['\t\tstart: %s,' % domain_info.start,
                                 '\t\tend: %s,' % domain_info.stop,
                                 '\t\tdomain: {', '\t\t\tname: "%s",'
                                 % domain_info.domain_id,
                                 '\t\t\tid: %s,' % len(
                                     self._dom_rank[domain_info.uid+':CDD']),
                                 '\t\t\tdescription: "%s"},' %
                                 domain_info.short_name,
                                 '\t\tsource: {', '\t\t\tname: "CDD",',
                                 '\t\t\thref: null,', '\t\t\tid: 1}']))
                    else:
                        self._dom_rank.setdefault(
                            domain_info.uid+':IPR', []).append(
                            domain_info.domain_id)
                        self._highlight_dict.setdefault(
                            prot_key, []).append('\n'.join(
                                ['start: %s,' % domain_info.start,
                                 'end: %s,' % domain_info.stop,
                                 'domain: {', '\t\tname: "%s",' %
                                 domain_info.domain_id,
                                 '\t\tid: %s,' % len(
                                     self._dom_rank[domain_info.uid+':IPR']),
                                 '\t\tdescription: "%s"},' % domain_info.short_name,
                                 'source: {', '\t\tname: "InterPro",',
                                 '\t\thref: null,', '\t\tid: 2}']))
            return self._uid_key_dict, self._location_dict, self._dom_annotation, self._dom_highlight, self._highlight_dict

    def domain_highlight(self):
        for uid in self._uid_key_dict.keys():
            header = '\n'.join(['<meta charset="UTF-8">'
            '<link type="text/css" rel="stylesheet" href="http://parce.li/bundle/biojs-vis-protein-viewer@0.1.4">',
            '',
            '
', '
', ''])
            with open(self._dom_highlight+'/%s.html' % uid, 'w') as out_fh:
                out_fh.write('\n'.join([header, body, panel, footer]))

    def viz_annotation_scoring(self):
        if os.path.exists(self._annotation_scoring_data):
            with open(self._annotation_scoring_data, 'r') as in_fh:
                for entry in in_fh:
                    if not entry.startswith('Entry'):
                        self._filter_viz_dict.setdefault('filter1_list', []).append(
                            float(entry.strip().split('\t')[-5]))
                        self._filter_viz_dict.setdefault('filter2_list', []).append(
                            float(entry.strip().split('\t')[-4]))
                        self._filter_viz_dict.setdefault('filter3_list', []).append(
                            float(entry.strip().split('\t')[-3]))
                        self._filter_viz_dict.setdefault('filter4_list', []).append(
                            float(entry.strip().split('\t')[-2]))
                        self._filter_viz_dict.setdefault('bayscore_list', []).append(
                            float(entry.strip().split('\t')[-1]))
                try:
                    label_list = range(0, len(self._filter_viz_dict['bayscore_list']))
                    plt.plot(sorted(self._filter_viz_dict['filter1_list']), 'ro', label='Filter-1 Score')
                    plt.plot(sorted(self._filter_viz_dict['filter2_list']), 'ys', label='Filter-2 Score')
                    plt.plot(sorted(self._filter_viz_dict['filter3_list']), 'g8', label='Filter-3 Score')
                    plt.plot(sorted(self._filter_viz_dict['filter4_list']), 'mp', label='Filter-4 Score')
                    plt.plot(sorted(self._filter_viz_dict['bayscore_list']), 'b^', label='Bayesian Score')
                    plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
                       ncol=3, mode="expand", borderaxespad=0.)
                    plt.xticks(label_list)
                    plt.xlabel('Annotation scores of selected proteins')
                    plt.ylabel('Filter/Bayesian score')
                    plt.savefig(os.path.join(self._overview, 'viz_annotation_scoring.png'))
                except KeyError:
                    print("!!! The annotation scoring file seems to be empty."
                          " Please reanalyse annotation score using the subcommand 'annoscore' !!!")
        else:
            print('The data for annotation scores do not exist,'
                  'please calculate the annotation score using the subcommand'
                  '"annoscore", the flag "-nd" can be used to specify the absolute path for needle.')

    def viz_secondary_structure(self):
        for uid in self._uid_key_dict.keys():
            if uid+'.horiz' in os.listdir(
            self._additional_annotation+'/protein_secondary_structure/'):
                files = uid+'.horiz'
            elif uid+'.plain' in os.listdir(
                self._additional_annotation+'/protein_secondary_structure/'):
                files = uid+'.plain'
                print("\nRaptorX secondary structure files are unavailable.")
                print("Visualizing secondary structure using literature based analysis.\n")
            else:
                print("\nRaptorX/literature-based secondary structure files are unavailable.")
                print("Exiting the current analysis.")
                print("Please re-run the secondary structure prediction by RaptorX\n")
                return
            secstr_list = []
            uid_secstr_dict = {}
            sec_data_sites = []
            with open(self._additional_annotation+
                '/protein_secondary_structure/'+files, 'r') as in_fh:
                for entry in in_fh:
                    if 'AA: ' in entry:
                        self._fasta_dict.setdefault(uid,
                        []).append(entry.strip().split('AA: ')[1])
                    if 'Pred: ' in entry:
                        try:
                            secstr_list.append(entry.strip().split('Pred: ')[1])
                        except IndexError:
                            print("\nRaptorX output file is incomplete. Exiting the current analysis.")
                            print("Please re-run the secondary structure prediction by RaptorX\n")
                            return
            for i, pred_data in enumerate(''.join(secstr_list)):
                    uid_secstr_dict[i] = pred_data
            for j in range(len(uid_secstr_dict)-1):
                if j == 0:
                    sec_data_sites.append(j)
                if not uid_secstr_dict[j] == uid_secstr_dict[j+1]:
                    sec_data_sites.append(j+1)
                    self._secstr_dict.setdefault(uid, []).append(
                        'mySequence.addHighlight({start:%s, end:%s, color:"Black", background:"%s"});'
                        %(int(sec_data_sites[-2])+1, int(j)+1,
                          self._sec_str_color[uid_secstr_dict[j]]))
            self._secstr_dict.setdefault(uid, []).append(
                'mySequence.addHighlight({start:%s, end:%s, color:"Black", background:"%s"});'
                %(int(sec_data_sites[-1])+1, int(list(uid_secstr_dict.keys())[-1])+1,
                  self._sec_str_color[uid_secstr_dict[j]]))
        self.sec_str_script()

    def sec_str_script(self):
        for uid in self._fasta_dict.keys():
            header = '\n'.join(['<meta charset="UTF-8">',
            '',
            '',
            '',
            '

', ''])
            body1 = '\n'.join(['var theSequence = "%s";' %
            ''.join(self._fasta_dict[uid]), 'yourDiv.textContent = "";',
            'window.onload = function() {', 'var mySequence = new Seq({',
            '\tsequence : theSequence,', '\ttarget : yourDiv.id,',
            '\tformat : "CODATA",', '\tformatOptions : {',
            '\ttitle:false,', '\tfooter:false', '\t},', '\tid : "%s"' % uid, '});'])
            body2 = '\n'.join(self._secstr_dict[uid])
            dom_list = sorted(list(self._location_dict[uid].keys()))
            annotation_list = []
            for dom_id in dom_list:
                dom_idx = dom_list.index(dom_id)
                annotation_list.append('\n'.join([
                    'mySequence.addAnnotation({', 'name:"Domain-%s",' % str(int(dom_idx)+1),
                    'html:"
%s
%s",' % (dom_id,
                    self._dom_annotation[dom_id]), 'color:"%s",' % self._color_list[dom_idx],
                    'regions: [', ',\n'.join(self._location_dict[uid][dom_id]), ']});']))
            with open(self._sec_str+'/'+uid+'.html', 'w') as out_fh:
                out_fh.write('\n'.join([header, body1, '\n'.join(annotation_list),
                                        body2, footer]))

    def viz_subcellular_localization(self):
        ''''''
        if 'psortb_data_summary.csv' in os.listdir(
            self._additional_annotation+'/protein_localization'):
            total_loc = set()
            with open(
            self._additional_annotation+'/protein_localization/psortb_data_summary.csv',
            'r') as in_fh:
                for entry in in_fh:
                    if not 'Localization' in entry:
                        protein = entry.strip().split('\t')[0]
                        localization = entry.strip().split('\t')[1]
                        if not localization.lower() == 'unknown':
                            score = float(entry.strip().split('\t')[2])
                            self._localization_dict[protein][localization] = score
                            total_loc.add(localization)
            with open(self._localize+'/localization_table.csv', 'w') as out_fh:
                out_fh.write('Proteins\t%s\n' % '\t'.join(sorted(list(total_loc))))
                for each_prot in self._localization_dict.keys():
                    for localization in self._localization_dict[each_prot]:
                        entry_list = list('0'*len(total_loc))
                        loc_idx = sorted(list(total_loc)).index(localization)
                        entry_list[loc_idx] = self._localization_dict[each_prot][localization]
                        out_fh.write("%s\t%s\n" % (each_prot, '\t'.join(map(str, entry_list))))
            self._create_localization_heatmap()
        else:
            print("\nPsortB-based localization prediction files are unavailable.")
            print("Exiting the current analysis.")
            print("Please re-run the localization prediction by PsortB\n")
            return

    def _create_localization_heatmap(self):
        ''''''
        plot_file = self._localize+'/localization_heatmap.pdf'
        infile = self._localize+'/localization_table.csv'
        with open(self._localize+'/localization_heatmap.R', 'w') as r_fh:
            r_fh.write('\n'.join(['library(gplots)', 'library(RColorBrewer)', 'display.brewer.all()',
        'data <- read.csv("%s", header=T, sep = "\\t")' % infile,
        'rnames <- data[,1]',  'data_matrix <- data.matrix(data[,2:ncol(data)])',
        'data_matrix[is.na(data_matrix)] <- 0', 'data_matrix[is.nan(data_matrix)] <- 0',
        'data_matrix[is.infinite(data_matrix)] <- max(data_matrix)',
        'rownames(data_matrix) <- rnames', 'pdf(file="%s")' % plot_file,
        'out_map <- heatmap.2(data_matrix, dendrogram = "none", Rowv = FALSE, \
        Colv = FALSE, col=brewer.pal(9,"YlGn"), margins=c(5,8), \
        cexCol=0.8, cexRow=0.8, key.title="PsortB Pred-value", key.xlab="", key.ylab="")',
        'dev.off()']))
        subprocess.Popen(['Rscript %s/localization_heatmap.R' %
                          self._localize], shell=True).wait()

    def viz_homologous_pdb_msa(self):
        header = '\n'.join(['<meta charset="UTF-8">',
        '<link type="text/css" rel="stylesheet" href="http://parce.li/bundle/msa@0.4.8">',
        '',
        '',
        '',
        '',
        '',
        '

', ''])
        body = '//EDIT PATH\n'.join([
        'var url = "https://github.com/malvikasharan/APRICOT/blob/master/Biojs_dependencies/data/biojs_msa_tab.clustal";'
        'var opts = {', '\tel: msaDiv', '};'])
        with open(self._pdb_msa+'/Biojs_pdb_msa_tab.html', 'w') as out_fh:
            out_fh.write('\n'.join([header, body, footer]))

        for files in os.listdir(self._additional_annotation+'/pdb_sequence_prediction/'):
            if '_top5.fasta' in files:
                shutil.copyfile(
                self._additional_annotation+'/pdb_sequence_prediction/'+files,
                self._pdb_msa+'/'+files)
                subprocess.Popen(['bin/reference_db_files/clustal/clustalw2 %s' %
                    self._pdb_msa+'/'+files], shell=True).wait()

        print("\nPlease open the BioJS MSA tab generated in Biojs_pdb_msa_tab.html.")
        print("Import MSA files (.aln) in the BioJS MSA tab to visualize the alignment.\n")

class AnnotationScoringColumns(object):
    '''Column information of annotation scoring file'''
    def __init__(self, row):
        self.uid = row[0]
        self.entry_name = row[1]
        self.prot_name = row[2]
        self.species = row[3]
        self.length = row[4]
        self.resource = row[5]
        self.resource_id = row[6]
        self.domain_id = row[7]
        self.short_name = row[8]
        self.full_name = row[9]
        self.domain_length = row[10]
        self.start = row[11]
        self.stop = row[12]
        self.ref_seq = row[13]
        self.q_seq = row[14]
        self.ref_ss = row[15]
        self.q_ss = row[16]
        self.mol_mass = row[17]
        self.iso_pt = row[18]
        self.solub = row[19]
        self.vdw = row[20]
        self.coverage = row[21]
        self.cov_by_dom = row[22]
        self.identity = row[23]
        self.iden_by_cov = row[24]
        self.similarity = row[25]
        self.sim_by_cov = row[26]
        self.gap = row[27]
        self.gap_by_cov = row[28]
        self.AA_RO = row[29]
        self.SS_RO = row[30]
        self.PC_RO = row[31]
        self.AAC_ED = row[32]
        self.PCC_ED = row[33]
        self.DPC_ED = row[34]
        self.TPC_ED = row[35]

class DomainDataColumns(object):
    '''Column information of domain annotation file'''
    def __init__(self, row):
        self.uid = row[0]
        self.entry_name = row[1]
        self.prot_name = row[2]
        self.species = row[3]
        self.length = row[4]
        self.gene_name = row[5]
        self.locus_tag = row[6]
        self.existance = row[7]
        self.go = row[8]
        self.embl_id = row[9]
        self.pdb_id = row[10]
        self.kegg_id = row[11]
        self.interpro_id = row[12]
        self.pfam_id = row[13]
        self.pubmed_id = row[14]
        self.resource = row[15]
        self.resource_id = row[16]
        self.domain_id = row[17]
        self.short_name = row[18]
        self.full_name = row[19]
        self.dom_kw = row[20]
        self.dom_go = row[21]
        self.members = row[22]
        self.dom_len = row[23]
        self.start = row[24]
        self.stop = row[25]
        self.evalue = row[26]
        self.bitscore = row[27]
        self.bits = row[28]
        self.cover_len = row[29]
        self.cov_prcnt = row[30]
        self.identity = row[31]
        self.iden_prcnt = row[32]
        self.similarity = row[33]
        self.sim_prcnt = row[34]
        self.gaps = row[35]
        self.gap_prcnt = row[36]
        self.filter_tag = row[37]


"""
Base IO code for all datasets
"""

# Copyright (c) 2007 David Cournapeau <cournape@gmail.com>
#               2010 Fabian Pedregosa <fabian.pedregosa@inria.fr>
#               2010 Olivier Grisel <olivier.grisel@ensta.org>
# License: BSD 3 clause

import os
import csv
import shutil
from os import environ
from os.path import dirname
from os.path import join
from os.path import exists
from os.path import expanduser
from os.path import isdir
from os import listdir
from os import makedirs

import numpy as np

from ..utils import check_random_state

class Bunch(dict):
    """Container object for datasets

    Dictionary-like object that exposes its keys as attributes.

    >>> b = Bunch(a=1, b=2)
    >>> b['b']
    2
    >>> b.b
    2
    >>> b.a = 3
    >>> b['a']
    3
    >>> b.c = 6
    >>> b['c']
    6

    """

    def __init__(self, **kwargs):
        dict.__init__(self, kwargs)

    def __setattr__(self, key, value):
        self[key] = value

    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError:
            raise AttributeError(key)

    def __getstate__(self):
        return self.__dict__

def get_data_home(data_home=None):
    """Return the path of the scikit-learn data dir.

    This folder is used by some large dataset loaders to avoid
    downloading the data several times.

    By default the data dir is set to a folder named 'scikit_learn_data'
    in the user home folder.

    Alternatively, it can be set by the 'SCIKIT_LEARN_DATA' environment
    variable or programmatically by giving an explicit folder path. The
    '~' symbol is expanded to the user home folder.

    If the folder does not already exist, it is automatically created.
    """
    if data_home is None:
        data_home = environ.get('SCIKIT_LEARN_DATA',
                                join('~', 'scikit_learn_data'))
    data_home = expanduser(data_home)
    if not exists(data_home):
        makedirs(data_home)
    return data_home

def clear_data_home(data_home=None):
    """Delete all the content of the data home cache."""
    data_home = get_data_home(data_home)
    shutil.rmtree(data_home)

def load_files(container_path, description=None, categories=None,
               load_content=True, shuffle=True, encoding=None,
               decode_error='strict', random_state=0):
    """Load text files with categories as subfolder names.

    Individual samples are assumed to be files stored a two levels folder
    structure such as the following:

        container_folder/
            category_1_folder/
                file_1.txt
                file_2.txt
                ...
                file_42.txt
            category_2_folder/
                file_43.txt
                file_44.txt
                ...

    The folder names are used as supervised signal label names. The
    individual file names are not important.

    This function does not try to extract features into a numpy array or
    scipy sparse matrix. In addition, if load_content is false it
    does not try to load the files in memory.

    To use text files in a scikit-learn classification or clustering
    algorithm, you will need to use the `sklearn.feature_extraction.text`
    module to build a feature extraction transformer that suits your
    problem.

    If you set load_content=True, you should also specify the encoding of
    the text using the 'encoding' parameter. For many modern text files,
    'utf-8' will be the correct encoding. If you leave encoding equal to None,
    then the content will be made of bytes instead of Unicode, and you will
    not be able to use most functions in `sklearn.feature_extraction.text`.

    Similar feature extractors should be built for other kind of unstructured
    data input such as images, audio, video, ...

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

    Parameters
    ----------
    container_path : string or unicode
        Path to the main folder holding one subfolder per category

    description: string or unicode, optional (default=None)
        A paragraph describing the characteristic of the dataset: its source,
        reference, etc.

    categories : A collection of strings or None, optional (default=None)
        If None (default), load all the categories.
        If not None, list of category names to load (other categories ignored).

    load_content : boolean, optional (default=True)
        Whether to load or not the content of the different files. If
        true a 'data' attribute containing the text information is present
        in the data structure returned. If not, a filenames attribute
        gives the path to the files.

    encoding : string or None (default is None)
        If None, do not try to decode the content of the files (e.g. for
        images or other non-text content).
        If not None, encoding to use to decode text files to Unicode if
        load_content is True.

    decode_error: {'strict', 'ignore', 'replace'}, optional
        Instruction on what to do if a byte sequence is given to analyze that
        contains characters not of the given `encoding`. Passed as keyword
        argument 'errors' to bytes.decode.

    shuffle : bool, optional (default=True)
        Whether or not to shuffle the data: might be important for models that
        make the assumption that the samples are independent and identically
        distributed (i.i.d.), such as stochastic gradient descent.

    random_state : int, RandomState instance or None, optional (default=0)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    Returns
    -------
    data : Bunch
        Dictionary-like object, the interesting attributes are: either
        data, the raw text data to learn, or 'filenames', the files
        holding it, 'target', the classification labels (integer index),
        'target_names', the meaning of the labels, and 'DESCR', the full
        description of the dataset.
    """
    target = []
    target_names = []
    filenames = []

    folders = [f for f in sorted(listdir(container_path))
               if isdir(join(container_path, f))]

    if categories is not None:
        folders = [f for f in folders if f in categories]

    for label, folder in enumerate(folders):
        target_names.append(folder)
        folder_path = join(container_path, folder)
        documents = [join(folder_path, d)
                     for d in sorted(listdir(folder_path))]
        target.extend(len(documents) * [label])
        filenames.extend(documents)

    # convert to array for fancy indexing
    filenames = np.array(filenames)
    target = np.array(target)

    if shuffle:
        random_state = check_random_state(random_state)
        indices = np.arange(filenames.shape[0])
        random_state.shuffle(indices)
        filenames = filenames[indices]
        target = target[indices]

    if load_content:
        data = []
        for filename in filenames:
            with open(filename, 'rb') as f:
                data.append(f.read())
        if encoding is not None:
            data = [d.decode(encoding, decode_error) for d in data]
        return Bunch(data=data,
                     filenames=filenames,
                     target_names=target_names,
                     target=target,
                     DESCR=description)

    return Bunch(filenames=filenames,
                 target_names=target_names,
                 target=target,
                 DESCR=description)

def load_iris():
    """Load and return the iris dataset (classification).

    The iris dataset is a classic and very easy multi-class classification
    dataset.

    =================   ==============
    Classes                          3
    Samples per class               50
    Samples total                  150
    Dimensionality                   4
    Features            real, positive
    =================   ==============

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

    Returns
    -------
    data : Bunch
        Dictionary-like object, the interesting attributes are:
        'data', the data to learn, 'target', the classification labels,
        'target_names', the meaning of the labels, 'feature_names', the
        meaning of the features, and 'DESCR', the
        full description of the dataset.

    Examples
    --------
    Let's say you are interested in the samples 10, 25, and 50, and want to
    know their class name.

    >>> from sklearn.datasets import load_iris
    >>> data = load_iris()
    >>> data.target[[10, 25, 50]]
    array([0, 0, 1])
    >>> list(data.target_names)
    ['setosa', 'versicolor', 'virginica']
    """
    module_path = dirname(__file__)
    with open(join(module_path, 'data', 'iris.csv')) as csv_file:
        data_file = csv.reader(csv_file)
        temp = next(data_file)
        n_samples = int(temp[0])
        n_features = int(temp[1])
        target_names = np.array(temp[2:])
        data = np.empty((n_samples, n_features))
        target = np.empty((n_samples,), dtype=np.int)

        for i, ir in enumerate(data_file):
            data[i] = np.asarray(ir[:-1], dtype=np.float)
            target[i] = np.asarray(ir[-1], dtype=np.int)

    with open(join(module_path, 'descr', 'iris.rst')) as rst_file:
        fdescr = rst_file.read()

    return Bunch(data=data, target=target,
                 target_names=target_names,
                 DESCR=fdescr,
                 feature_names=['sepal length (cm)', 'sepal width (cm)',
                                'petal length (cm)', 'petal width (cm)'])

def load_digits(n_class=10):
    """Load and return the digits dataset (classification).

    Each datapoint is a 8x8 image of a digit.

    =================   ==============
    Classes                         10
    Samples per class             ~180
    Samples total                 1797
    Dimensionality                  64
    Features             integers 0-16
    =================   ==============

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

    Parameters
    ----------
    n_class : integer, between 0 and 10, optional (default=10)
        The number of classes to return.

    Returns
    -------
    data : Bunch
        Dictionary-like object, the interesting attributes are:
        'data', the data to learn, 'images', the images corresponding
        to each sample, 'target', the classification labels for each
        sample, 'target_names', the meaning of the labels, and 'DESCR',
        the full description of the dataset.

    Examples
    --------
    To load the data and visualize the images::

        >>> from sklearn.datasets import load_digits
        >>> digits = load_digits()
        >>> print(digits.data.shape)
        (1797, 64)
        >>> import pylab as pl #doctest: +SKIP
        >>> pl.gray() #doctest: +SKIP
        >>> pl.matshow(digits.images[0]) #doctest: +SKIP
        >>> pl.show() #doctest: +SKIP
    """
    module_path = dirname(__file__)
    data = np.loadtxt(join(module_path, 'data', 'digits.csv.gz'),
                      delimiter=',')
    with open(join(module_path, 'descr', 'digits.rst')) as f:
        descr = f.read()
    target = data[:, -1]
    flat_data = data[:, :-1]
    images = flat_data.view()
    images.shape = (-1, 8, 8)

    if n_class < 10:
        idx = target < n_class
        flat_data, target = flat_data[idx], target[idx]
        images = images[idx]

    return Bunch(data=flat_data,
                 target=target.astype(np.int),
                 target_names=np.arange(10),
                 images=images,
                 DESCR=descr)

def load_diabetes():
    """Load and return the diabetes dataset (regression).

    ==============      ==================
    Samples total       442
    Dimensionality      10
    Features            real, -.2 < x < .2
    Targets             integer 25 - 346
    ==============      ==================

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

    Returns
    -------
    data : Bunch
        Dictionary-like object, the interesting attributes are:
        'data', the data to learn and 'target', the regression target for each
        sample.
    """
    base_dir = join(dirname(__file__), 'data')
    data = np.loadtxt(join(base_dir, 'diabetes_data.csv.gz'))
    target = np.loadtxt(join(base_dir, 'diabetes_target.csv.gz'))
    return Bunch(data=data, target=target)

def load_linnerud():
    """Load and return the linnerud dataset (multivariate regression).

    Samples total: 20
    Dimensionality: 3 for both data and targets
    Features: integer
    Targets: integer

    Returns
    -------
    data : Bunch
        Dictionary-like object, the interesting attributes are: 'data' and
        'targets', the two multivariate datasets, with 'data' corresponding to
        the exercise and 'targets' corresponding to the physiological
        measurements, as well as 'feature_names' and 'target_names'.
    """
    base_dir = join(dirname(__file__), 'data/')
    # Read data
    data_exercise = np.loadtxt(base_dir + 'linnerud_exercise.csv', skiprows=1)
    data_physiological = np.loadtxt(base_dir + 'linnerud_physiological.csv',
                                    skiprows=1)
    # Read header
    with open(base_dir + 'linnerud_exercise.csv') as f:
        header_exercise = f.readline().split()
    with open(base_dir + 'linnerud_physiological.csv') as f:
        header_physiological = f.readline().split()
    with open(dirname(__file__) + '/descr/linnerud.rst') as f:
        descr = f.read()

    return Bunch(data=data_exercise, feature_names=header_exercise,
                 target=data_physiological,
                 target_names=header_physiological,
                 DESCR=descr)

def load_boston():
    """Load and return the boston house-prices dataset (regression).

    ==============     ==============
    Samples total                 506
    Dimensionality                 13
    Features           real, positive
    Targets             real 5. - 50.
    ==============     ==============

    Returns
    -------
    data : Bunch
        Dictionary-like object, the interesting attributes are:
        'data', the data to learn, 'target', the regression targets,
        and 'DESCR', the full description of the dataset.

    Examples
    --------
    >>> from sklearn.datasets import load_boston
    >>> boston = load_boston()
    >>> print(boston.data.shape)
    (506, 13)
    """
    module_path = dirname(__file__)

    fdescr_name = join(module_path, 'descr', 'boston_house_prices.rst')
    with open(fdescr_name) as f:
        descr_text = f.read()

    data_file_name = join(module_path, 'data', 'boston_house_prices.csv')
    with open(data_file_name) as f:
        data_file = csv.reader(f)
        temp = next(data_file)
        n_samples = int(temp[0])
        n_features = int(temp[1])
        data = np.empty((n_samples, n_features))
        target = np.empty((n_samples,))
        temp = next(data_file)  # names of features
        feature_names = np.array(temp)

        for i, d in enumerate(data_file):
            data[i] = np.asarray(d[:-1], dtype=np.float)
            target[i] = np.asarray(d[-1], dtype=np.float)

    return Bunch(data=data,
                 target=target,
                 # last column is target value
                 feature_names=feature_names[:-1],
                 DESCR=descr_text)

def load_sample_images():
    """Load sample images for image manipulation.
    Loads both, ``china`` and ``flower``.

    Returns
    -------
    data : Bunch
        Dictionary-like object with the following attributes :
        'images', the two sample images, 'filenames', the file
        names for the images, and 'DESCR'
        the full description of the dataset.

    Examples
    --------
    To load the data and visualize the images:

    >>> from sklearn.datasets import load_sample_images
    >>> dataset = load_sample_images()     #doctest: +SKIP
    >>> len(dataset.images)                #doctest: +SKIP
    2
    >>> first_img_data = dataset.images[0] #doctest: +SKIP
    >>> first_img_data.shape               #doctest: +SKIP
    (427, 640, 3)
    >>> first_img_data.dtype               #doctest: +SKIP
    dtype('uint8')
    """
    # Try to import imread from scipy. We do this lazily here to prevent
    # this module from depending on PIL.
    try:
        try:
            from scipy.misc import imread
        except ImportError:
            from scipy.misc.pilutil import imread
    except ImportError:
        raise ImportError("The Python Imaging Library (PIL) "
                          "is required to load data from jpeg files")
    module_path = join(dirname(__file__), "images")
    with open(join(module_path, 'README.txt')) as f:
        descr = f.read()
    filenames = [join(module_path, filename)
                 for filename in os.listdir(module_path)
                 if filename.endswith(".jpg")]
    # Load image data for each image in the source folder.
    images = [imread(filename) for filename in filenames]

    return Bunch(images=images,
                 filenames=filenames,
                 DESCR=descr)

def load_sample_image(image_name):
    """Load the numpy array of a single sample image

    Parameters
    -----------
    image_name: {`china.jpg`, `flower.jpg`}
        The name of the sample image loaded

    Returns
    -------
    img: 3D array
        The image as a numpy array: height x width x color

    Examples
    ---------

    >>> from sklearn.datasets import load_sample_image
    >>> china = load_sample_image('china.jpg')   # doctest: +SKIP
    >>> china.dtype                              # doctest: +SKIP
    dtype('uint8')
    >>> china.shape                              # doctest: +SKIP
    (427, 640, 3)
    >>> flower = load_sample_image('flower.jpg') # doctest: +SKIP
    >>> flower.dtype                             # doctest: +SKIP
    dtype('uint8')
    >>> flower.shape                             # doctest: +SKIP
    (427, 640, 3)
    """
    images = load_sample_images()
    index = None
    for i, filename in enumerate(images.filenames):
        if filename.endswith(image_name):
            index = i
            break
    if index is None:
        raise AttributeError("Cannot find sample image: %s" % image_name)
    return images.images[index]


"""
========================================
Plot multi-class SGD on the iris dataset
========================================

Plot decision surface of multi-class SGD on iris dataset.
The hyperplanes corresponding to the three one-versus-all (OVA) classifiers
are represented by the dashed lines.

"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.linear_model import SGDClassifier

# import some data to play with
iris = datasets.load_iris()
X = iris.data[:, :2]  # we only take the first two features. We could
                      # avoid this ugly slicing by using a two-dim dataset
y = iris.target
colors = "bry"

# shuffle
idx = np.arange(X.shape[0])
np.random.seed(13)
np.random.shuffle(idx)
X = X[idx]
y = y[idx]

# standardize
mean = X.mean(axis=0)
std = X.std(axis=0)
X = (X - mean) / std

h = .02  # step size in the mesh

clf = SGDClassifier(alpha=0.001, n_iter=100).fit(X, y)

# create a mesh to plot in
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                     np.arange(y_min, y_max, h))

# Plot the decision boundary. For that, we will assign a color to each
# point in the mesh [x_min, x_max]x[y_min, y_max].
Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
# Put the result into a color plot
Z = Z.reshape(xx.shape)
cs = plt.contourf(xx, yy, Z, cmap=plt.cm.Paired)
plt.axis('tight')

# Plot also the training points
for i, color in zip(clf.classes_, colors):
    idx = np.where(y == i)
    plt.scatter(X[idx, 0], X[idx, 1], c=color, label=iris.target_names[i],
                cmap=plt.cm.Paired)
plt.title("Decision surface of multi-class SGD")
plt.axis('tight')

# Plot the three one-against-all classifiers
xmin, xmax = plt.xlim()
ymin, ymax = plt.ylim()
coef = clf.coef_
intercept = clf.intercept_

def plot_hyperplane(c, color):
    def line(x0):
        return (-(x0 * coef[c, 0]) - intercept[c]) / coef[c, 1]

    plt.plot([xmin, xmax], [line(xmin), line(xmax)],
             ls="--", color=color)

for i, color in zip(clf.classes_, colors):
    plot_hyperplane(i, color)
plt.legend()
plt.show()