# -*- coding: utf-8 -*-
"""
===================================
Demo of DBSCAN clustering algorithm
===================================

Finds core samples of high density and expands clusters from them.

"""
print(__doc__)

import numpy as np

from sklearn.cluster import DBSCAN
from sklearn import metrics
from sklearn.datasets.samples_generator import make_blobs
from sklearn.preprocessing import StandardScaler

##############################################################################
# Generate sample data
centers = [[1, 1], [-1, -1], [1, -1]]
X, labels_true = make_blobs(n_samples=750, centers=centers, cluster_std=0.4,
                            random_state=0)

X = StandardScaler().fit_transform(X)

##############################################################################
# Compute DBSCAN
db = DBSCAN(eps=0.3, min_samples=10).fit(X)
core_samples_mask = np.zeros_like(db.labels_, dtype=bool)
core_samples_mask[db.core_sample_indices_] = True
labels = db.labels_

# Number of clusters in labels, ignoring noise if present.
n_clusters_ = len(set(labels)) - (1 if -1 in labels else 0)

print('Estimated number of clusters: %d' % n_clusters_)
print("Homogeneity: %0.3f" % metrics.homogeneity_score(labels_true, labels))
print("Completeness: %0.3f" % metrics.completeness_score(labels_true, labels))
print("V-measure: %0.3f" % metrics.v_measure_score(labels_true, labels))
print("Adjusted Rand Index: %0.3f"
      % metrics.adjusted_rand_score(labels_true, labels))
print("Adjusted Mutual Information: %0.3f"
      % metrics.adjusted_mutual_info_score(labels_true, labels))
print("Silhouette Coefficient: %0.3f"
      % metrics.silhouette_score(X, labels))

##############################################################################
# Plot result
import matplotlib.pyplot as plt

# Black removed and is used for noise instead.
unique_labels = set(labels)
colors = plt.cm.Spectral(np.linspace(0, 1, len(unique_labels)))
for k, col in zip(unique_labels, colors):
    if k == -1:
        # Black used for noise.
        col = 'k'

    class_member_mask = (labels == k)

    xy = X[class_member_mask & core_samples_mask]
    plt.plot(xy[:, 0], xy[:, 1], 'o', markerfacecolor=col,
             markeredgecolor='k', markersize=14)

    xy = X[class_member_mask & ~core_samples_mask]
    plt.plot(xy[:, 0], xy[:, 1], 'o', markerfacecolor=col,
             markeredgecolor='k', markersize=6)

plt.title('Estimated number of clusters: %d' % n_clusters_)
plt.show()


import os
import sys
import pickle
import pandas as pd
import numpy as np
import hashlib
import os.path

from daisylu_config import *
from daisylu_vectors import *
from sentences import *
from daisylu_output import *
""

def addWikificationToDFrames(sents, sTypes, sentenceAttr):
    # need to split this up into manageable file sizes, the wikifier dies with out of memory error currently

    maxPartitionLen=200000
    resultsDir = getSystemPath('daisyluPython') + '/wikificationData'

    # if the working directory does not exist, create it.
    # './wikificationData/input/test0.txt'
    if not os.path.exists(resultsDir):
        os.makedirs(resultsDir)
    if not os.path.exists(resultsDir+'/input'):
        os.makedirs(resultsDir+'/input')
    if not os.path.exists(resultsDir+'/output'):
        os.makedirs(resultsDir+'/output')

    for sType in sTypes:
        partitions = [   {'textString':'', 'charMapper':{}}  ]
        for sentIX in range(len(sents[sType])):
            if len(partitions[-1]['textString']) > maxPartitionLen:
                partitions.append( {'textString':'', 'charMapper':{}})
                print 'THERE ARE NOW %d PARTITIONS' % len(partitions)
                print '====================================================================================================================='
                print '====================================================================================================================='
                print

            sentence = sents[sType][sentIX]
            if not sentIX % 100: print 'addWikificationToDFrames', sType, sentIX
            if not hasattr(sentence, sentenceAttr):
                continue
            sdf = getattr(sentence, sentenceAttr)
            if sdf.empty:
                continue
            sdf['NERForm']    = ''
            sdf['NERLabel']   = 'O'
            sdf['WKCategory'] = ''
            sdf['WKLink']     = ''
            sdf['WKLinker']   = np.nan
            sdf['WKRanker']   = np.nan

            df = sdf[['wordIX','words','txFunc','txBIOES','nameCategory','wiki']].copy()
            df['type'] = sType
            df['sentIX'] = sentIX
            df['allTStart'] = -1
            df['allTEnd'] = -1
            for i,t in enumerate(sentence.tokens):
                startOffset = len(partitions[-1]['textString'])
                partitions[-1]['textString'] += t
                endOffset = len(partitions[-1]['textString'])
                if (any(df.wordIX == i)):
                    df['allTStart']=startOffset
                    df['allTEnd']=endOffset
                partitions[-1]['charMapper'][startOffset] = (sentIX, i, t)
                partitions[-1]['textString'] += ' '
            partitions[-1]['textString'] += '\n\n'

        allText = ''
        for x in partitions:
            allText += x['textString']
        m = hashlib.md5()
        m.update(allText)
        md5 = m.hexdigest()
        print md5
        cacheFn = 'wikificationData/' + md5 + '.pcl'

        if not os.path.isfile(cacheFn):   # calculate and archive the info, use it later if the same set of sentences is called for
            wconfigs = []
            wconfigs.append({
                                'config'   : 'configs/STAND_ALONE_NO_INFERENCE.xml',
                                'inputFn'  : resultsDir + '/input/test%d.txt',
                                'outputDn' : '%s/output/' % resultsDir,
                             })
            info = { 'NER':{}, 'wiki':{} }

            # partitions = pickle.load( open( 'wikiPartitions.pcl' ) )
            """
            If you're using this system, please cite the paper.

            Relational Inference for Wikification
            Xiao Cheng and Dan Roth
            EMNLP 2013
            """
            for p, partition in enumerate(partitions):
                for wtype in wconfigs:
                    tfile = open(wtype['inputFn'] % p, 'wb')
                    tfile.write(partition['textString'])
                    tfile.close()

                    direc     = getSystemPath('Wikifier2013')
                    config    = wtype['config']
                    inputFn   = wtype['inputFn'] % p
                    outputDn  = wtype['outputDn']
                    stencil   = '/usr/bin/java -Xmx10G -jar dist/wikifier-3.0-jar-with-dependencies.jar -annotateData %s %s false %s'
                    cmd = stencil % (inputFn, outputDn, config)
                    print cmd
                    errorCode = os.system('cd %s; %s' % (direc, cmd)   )
                    if errorCode:
                        raise ValueError('ERROR!\n    non zero error code %d' % errorCode)
                        exit(1)

            import xmltodict
            from bs4 import BeautifulSoup

            for p, partition in enumerate(partitions):
                print 'Partition %d' % p
                charMapper = partitions[p]['charMapper']
                html = open('%s/output/' % resultsDir + '/test%d.txt.NER.tagged' % p).read()
                parsed_html = BeautifulSoup(html, "lxml")
                ner={ 'start':[], 'end':[], 'form':[], 'label':[]}
                for item in parsed_html.find_all('start'):
                    ner['start'].append(int(item.text))
                for item in parsed_html.find_all('end'):
                    ner['end'].append(int(item.text))
                for item in parsed_html.find_all('form'):
                    ner['form'].append(item.text)
                for item in parsed_html.find_all('label' ):
                    ner['label'].append(item.text)

                for i in range(len(ner['start'])):
                    if not i % 100: print 'ner', i
                    tset = set()
                    for z in range(ner['start'][i],ner['end'][i]):
                        if z in charMapper:
                            tset.add(charMapper[z])
                    for trip in list(tset):
                        (six, wix, _) = trip
                        if not six in info['NER']:
                            info['NER'][six] = { 'NERForm':{}, 'NERLabel':{} }
                        info['NER'][six]['NERForm'][wix]       =  ner['form'][i]
                        info['NER'][six]['NERLabel'][wix]      =  ner['label'][i]

                with open('%s/output/' % resultsDir + '/test%d.txt.wikification.tagged.full.xml' % p) as fd:
                    obj = xmltodict.parse(fd.read())
                    if obj['WikifierOutput']['WikifiedEntities']:
                        entities = obj['WikifierOutput']['WikifiedEntities']['Entity']
                        for entity in entities:
                            #entityText = entity['EntitySurfaceForm']
                            entityStartOffset = int(entity['EntityTextStart'])
                            entityEndOffset   = int(entity['EntityTextEnd'])
                            linkerScore       = float(entity['LinkerScore'])
                            rankerScore       = float(entity['TopDisambiguation']['RankerScore'])
                            wikiTitle         = entity['TopDisambiguation']['WikiTitle']
                            attributes        = entity['TopDisambiguation']['Attributes']

                            #print   entityText, entityStartOffset, entityEndOffset, textString[entityStartOffset:entityEndOffset]
                            tset = set()
                            for z in range(entityStartOffset,entityEndOffset+1):
                                if z in charMapper:
                                    tset.add(charMapper[z])

                            for trip in list(tset):
                                (six, wix, _) = trip

                                if not six in info['wiki']:
                                    info['wiki'][six] = { 'WKCategory':{}, 'WKLink':{}, 'WKLinker':{}, 'WKRanker':{} }
                                info['wiki'][six]['WKCategory'][wix]       = attributes
                                info['wiki'][six]['WKLink'][wix]           = wikiTitle
                                info['wiki'][six]['WKLinker'][wix]         = linkerScore
                                info['wiki'][six]['WKRanker'][wix]         = rankerScore

            pickle.dump( info, open( cacheFn, "wb" ) )

        else:
            info = pickle.load( open( cacheFn, "rb" ) )

        for six in info['NER']:
            sentence = sents[sType][six]
            if not hasattr(sentence, sentenceAttr):
                continue
            sdf = getattr(sentence, sentenceAttr)
            for wix in info['NER'][six]['NERForm']:
                sdf.loc[  (sdf.wordIX == wix), 'NERForm']     = info['NER'][six]['NERForm'][wix]
                sdf.loc[  (sdf.wordIX == wix), 'NERLabel']    = info['NER'][six]['NERLabel'][wix]
        for six in info['wiki']:
            sentence = sents[sType][six]
            if not hasattr(sentence, sentenceAttr):
                continue
            sdf = getattr(sentence, sentenceAttr)
            for wix in info['wiki'][six]['WKCategory']:
                sdf.loc[  (sdf.wordIX == wix), 'WKCategory']  = info['wiki'][six]['WKCategory'][wix]
                sdf.loc[  (sdf.wordIX == wix), 'WKLink']      = info['wiki'][six]['WKLink'][wix]
                sdf.loc[  (sdf.wordIX == wix), 'WKLinker']    = info['wiki'][six]['WKLinker'][wix]
                sdf.loc[  (sdf.wordIX == wix), 'WKRanker']    = info['wiki'][six]['WKRanker'][wix]

def initializePredictionDataFrames(sents, ixList=None, NEWPrediction=False):
    if not ixList:
        ixList = range(len(sents['test']))

    for sentIX in ixList:
        sentence = sents['test'][sentIX]
        tagList = getSentenceDFTagList()
        if NEWPrediction:
            tagList += ['NEWPrediction']
        df = pd.DataFrame( columns=tagList )
        if not (sentIX %1000):
            print 'initializing pred frome ', sentIX
        df['wordIX']  = range(len(sentence.tokens))
        df['sentIX']  = sentIX
        df['words']   = sentence.tokens
        df['txBIOES'] = 'O'
        sentence.predictedDFrame = df

    addWikificationToDFrames(sents, ['test'], 'predictedDFrame')
    print 'CLIPPING ALL SENTENCES TO LENGTH 100'
    for sentIX in ixList:
        sentence = sents['test'][sentIX]
        sentence.predictedDFrame = sentence.predictedDFrame[sentence.predictedDFrame['wordIX'] < 100]

def createVectorsFromDataFrames(sents, sentenceAttr, dbf, dbtf, systemName, keepSense=True,  L0OnlyFromFeaturesDB=False, useDistSG=False ):
    wordDF = []
    dbfn = getSystemPath('daisylu') + 'data/%s' % dbf
    dbTestFn = getSystemPath('daisylu') + 'data/%s' %  dbtf

    merged = mergeSentenceDataFrames(None, ['test'], None,  sents=sents, sentenceAttr=sentenceAttr)
    if systemName== 'AMRL0NoNER':
        createAMRL0Vectors(None, dbTestFn,  100.0, keepSense, sTypes=['test'], vectors=merged, featuresDB=dbfn, maxSents=None, useNER=False)
    elif systemName== 'AMRL0':
        wordDF = createAMRL0Vectors(None, dbTestFn,  100.0, keepSense, sTypes=['test'], vectors=merged, featuresDB=dbfn, maxSents=None )
    elif systemName== 'AMRL0Args':
        createAMRL0ArgVectors(None, dbTestFn,  100.0, keepSense, sTypes=['test'], vectors=merged, featuresDB=dbfn, maxSents=None,
                              L0OnlyFromFeaturesDB=L0OnlyFromFeaturesDB, useDistSG=useDistSG )
    elif systemName== 'AMRL0Nargs':
        createAMRL0NargVectors(None, dbTestFn,  100.0, keepSense, sTypes=['test'], vectors=merged, featuresDB=dbfn, maxSents=None,
                               L0OnlyFromFeaturesDB=L0OnlyFromFeaturesDB, useDistSG=useDistSG )
    elif systemName== 'AMRL0Attr':
        createAMRL0AttrVectors(None, dbTestFn,  100.0, keepSense, sTypes=['test'], vectors=merged, featuresDB=dbfn, maxSents=None,
                               L0OnlyFromFeaturesDB=L0OnlyFromFeaturesDB, useDistSG=useDistSG )
    elif systemName== 'AMRL0Ncat':
        createAMRL0NcatVectors(None, dbTestFn,  100.0, keepSense, sTypes=['test'], vectors=merged, featuresDB=dbfn, maxSents=None,
                               L0OnlyFromFeaturesDB=L0OnlyFromFeaturesDB, useDistSG=useDistSG )
    else:
        assert('error, invalid system name')
    return wordDF

def runKerasNetwork(networkType, vectorDBFn, modelFn, resultsFn, sType='test'):
    direc = getSystemPath( 'NNModels' )
    mm, ww = modelFn.split('@')
    cmd = getSystemPath( 'python' )
    cmd   = cmd +' AMR_NN_Forward.py  -v %s -m %s -w %s -r %s -s %s' % ('../data/' + vectorDBFn, mm, ww, '../results/' + resultsFn, sType)
    print direc, cmd
    print direc, cmd
    print direc, cmd
    errorCode = os.system('cd %s; %s' % (direc, cmd)   )
    if errorCode:
        raise ValueError('ERROR!\n    non zero error code %d' % errorCode)
        exit(1)
    if errorCode:
        print vectorDBFn, modelFn, resultsFn, sType
        raise ValueError('ERROR!\n    non zero error code %d' % errorCode)
        exit(1)

def runNetwork(networkType, vectorDBFn, modelFn, resultsFn, sType='test'):
    if '@' in modelFn:
        runKerasNetwork(networkType, vectorDBFn, modelFn, resultsFn, sType)
    else:
        assert('error, Torch networks are no longer supported')


from __future__ import print_function, division
import matplotlib
import logging
from sys import stdout
matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab!
from neuralnilm import (Net, RealApplianceSource,
                        BLSTMLayer, DimshuffleLayer,
                        BidirectionalRecurrentLayer)
from neuralnilm.source import standardise, discretize, fdiff, power_and_fdiff
from neuralnilm.experiment import run_experiment, init_experiment
from neuralnilm.net import TrainingError
from neuralnilm.layers import MixtureDensityLayer
from neuralnilm.objectives import (scaled_cost, mdn_nll,
                                   scaled_cost_ignore_inactive, ignore_inactive,
                                   scaled_cost3)
from neuralnilm.plot import MDNPlotter, CentralOutputPlotter

from lasagne.nonlinearities import sigmoid, rectify, tanh
from lasagne.objectives import mse, binary_crossentropy
from lasagne.init import Uniform, Normal
from lasagne.layers import (LSTMLayer, DenseLayer, Conv1DLayer,
                            ReshapeLayer, FeaturePoolLayer, RecurrentLayer)
from lasagne.updates import nesterov_momentum, momentum
from functools import partial
import os
import __main__
from copy import deepcopy
from math import sqrt
import numpy as np
import theano.tensor as T

NAME = os.path.splitext(os.path.split(__main__.__file__)[1])[0]
PATH = "/homes/dk3810/workspace/python/neuralnilm/figures"
SAVE_PLOT_INTERVAL = 5000
GRADIENT_STEPS = 100

source_dict = dict(
    filename='/data/dk3810/ukdale.h5',
    appliances=[
        ['fridge freezer', 'fridge', 'freezer'],
        'hair straighteners',
        'television',
        'dish washer',
        ['washer dryer', 'washing machine']
    ],
    max_appliance_powers=[300, 500, 200, 2500, 2400],
    on_power_thresholds=[5] * 5,
    max_input_power=5900,
    min_on_durations=[60, 60, 60, 1800, 1800],
    min_off_durations=[12, 12, 12, 1800, 600],
    window=("2013-06-01", "2014-07-01"),
    seq_length=512,
    output_one_appliance=False,
    boolean_targets=False,
    train_buildings=[1],
    validation_buildings=[1],
    skip_probability=0.7,
    one_target_per_seq=False,
    n_seq_per_batch=16,
#    subsample_target=2,
    include_diff=False,
    clip_appliance_power=True,
    target_is_prediction=False,
#    independently_center_inputs = True,
    standardise_input=True,
    unit_variance_targets=True,
#    input_padding=8,
    lag=0,
    classification=True
#    reshape_target_to_2D=True
    # input_stats={'mean': np.array([ 0.05526326], dtype=np.float32),
    #              'std': np.array([ 0.12636775], dtype=np.float32)},
    # target_stats={
    #     'mean': np.array([ 0.04066789,  0.01881946,
    #                        0.24639061,  0.17608672,  0.10273963],
    #                      dtype=np.float32),
    #     'std': np.array([ 0.11449792,  0.07338708,
    #                    0.26608968,  0.33463112,  0.21250485],
    #                  dtype=np.float32)}
)

N = 50
net_dict = dict(
    save_plot_interval=SAVE_PLOT_INTERVAL,
#    loss_function=partial(ignore_inactive, loss_func=mdn_nll, seq_length=SEQ_LENGTH),
#    loss_function=lambda x, t: mdn_nll(x, t).mean(),
#    loss_function=lambda x, t: mse(x, t).mean(),
    loss_function=lambda x, t: binary_crossentropy(x, t).mean(),
#    loss_function=partial(scaled_cost, loss_func=mse),
#    loss_function=ignore_inactive,
#    loss_function=partial(scaled_cost3, ignore_inactive=False),
    updates_func=momentum,
    learning_rate=1e-4,
    learning_rate_changes_by_iteration={
        # 200: 1e-2,
        # 400: 1e-3,
        # 800: 1e-4
#        500: 1e-3
       #  4000: 1e-03,
       # 6000: 5e-06,
       # 7000: 1e-06
       # 2000: 5e-06
        # 3000: 1e-05
        # 7000: 5e-06,
        # 10000: 1e-06,
        # 15000: 5e-07,
        # 50000: 1e-07
    },
    do_save_activations=True,
    auto_reshape=False,
    plotter=CentralOutputPlotter
#    plotter=MDNPlotter
)

"""
          ||||||||||
        ||||||||||
      ||||||||||
    ||||||||||
  ||||||||||
||||||||||
12345678901234567890

"""

def exp_a(name):
    global source
    # source_dict_copy = deepcopy(source_dict)
    # source = RealApplianceSource(**source_dict_copy)
    net_dict_copy = deepcopy(net_dict)
    net_dict_copy.update(dict(
        experiment_name=name,
        source=source
    ))
    N = 512
    output_shape = source.output_shape_after_processing()
    net_dict_copy['layers_config'] = [
        {
            'type': DimshuffleLayer,
            'pattern': (0, 2, 1)  # (batch, features, time)
        },
        {
            'type': Conv1DLayer, # convolve over the time axis
            'num_filters': 16,
            'filter_length': 4,
            'stride': 1,
            'nonlinearity': rectify,
            'border_mode': 'same'
        },
        {
            'type': FeaturePoolLayer,
            'ds': 4, # number of feature maps to be pooled together
            'axis': 2, # pool over the time axis
            'pool_function': T.max
        },
        {
            'type': Conv1DLayer, # convolve over the time axis
            'num_filters': 16,
            'filter_length': 4,
            'stride': 1,
            'nonlinearity': rectify,
            'border_mode': 'same'
        },
        {
            'type': FeaturePoolLayer,
            'ds': 4, # number of feature maps to be pooled together
            'axis': 2, # pool over the time axis
            'pool_function': T.max
        },
        {
            'type': DimshuffleLayer,
            'pattern': (0, 2, 1) # back to (batch, time, features)
        },
        {
            'type': DenseLayer,
            'num_units': N,
            'nonlinearity': rectify
        },
        {
            'type': DenseLayer,
            'num_units': N // 2,
            'nonlinearity': rectify
        },
        {
            'type': DenseLayer,
            'num_units': N // 4,
            'nonlinearity': rectify
        },
        {
            'type': DenseLayer,
            'num_units': output_shape[1] * output_shape[2],
            'nonlinearity': sigmoid
        }
    ]
    net = Net(**net_dict_copy)
    return net

def main():
    #     EXPERIMENTS = list('abcdefghijklmnopqrstuvwxyz')
    EXPERIMENTS = list('a')
    for experiment in EXPERIMENTS:
        full_exp_name = NAME + experiment
        func_call = init_experiment(PATH, experiment, full_exp_name)
        logger = logging.getLogger(full_exp_name)
        try:
            net = eval(func_call)
            run_experiment(net, epochs=None)
        except KeyboardInterrupt:
            logger.info("KeyboardInterrupt")
            break
        except Exception as exception:
            logger.exception("Exception")
            raise
        finally:
            logging.shutdown()

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-
"""Algorithms for spectral clustering"""

# Author: Gael Varoquaux gael.varoquaux@normalesup.org
#         Brian Cheung
#         Wei LI <kuantkid@gmail.com>
# License: BSD 3 clause
import warnings

import numpy as np

from ..base import BaseEstimator, ClusterMixin
from ..utils import check_random_state, as_float_array
from ..utils.validation import check_array
from ..utils.extmath import norm
from ..metrics.pairwise import pairwise_kernels
from ..neighbors import kneighbors_graph
from ..manifold import spectral_embedding
from .k_means_ import k_means

def discretize(vectors, copy=True, max_svd_restarts=30, n_iter_max=20,
               random_state=None):
    """Search for a partition matrix (clustering) which is closest to the
    eigenvector embedding.

    Parameters
    ----------
    vectors : array-like, shape: (n_samples, n_clusters)
        The embedding space of the samples.

    copy : boolean, optional, default: True
        Whether to copy vectors, or perform in-place normalization.

    max_svd_restarts : int, optional, default: 30
        Maximum number of attempts to restart SVD if convergence fails

    n_iter_max : int, optional, default: 30
        Maximum number of iterations to attempt in rotation and partition
        matrix search if machine precision convergence is not reached

    random_state: int seed, RandomState instance, or None (default)
        A pseudo random number generator used for the initialization of the
        of the rotation matrix

    Returns
    -------
    labels : array of integers, shape: n_samples
        The labels of the clusters.

    References
    ----------

    - Multiclass spectral clustering, 2003
      Stella X. Yu, Jianbo Shi
      http://www1.icsi.berkeley.edu/~stellayu/publication/doc/2003kwayICCV.pdf

    Notes
    -----

    The eigenvector embedding is used to iteratively search for the
    closest discrete partition.  First, the eigenvector embedding is
    normalized to the space of partition matrices. An optimal discrete
    partition matrix closest to this normalized embedding multiplied by
    an initial rotation is calculated.  Fixing this discrete partition
    matrix, an optimal rotation matrix is calculated.  These two
    calculations are performed until convergence.  The discrete partition
    matrix is returned as the clustering solution.  Used in spectral
    clustering, this method tends to be faster and more robust to random
    initialization than k-means.

    """

    from scipy.sparse import csc_matrix
    from scipy.linalg import LinAlgError

    random_state = check_random_state(random_state)

    vectors = as_float_array(vectors, copy=copy)

    eps = np.finfo(float).eps
    n_samples, n_components = vectors.shape

    # Normalize the eigenvectors to an equal length of a vector of ones.
    # Reorient the eigenvectors to point in the negative direction with respect
    # to the first element.  This may have to do with constraining the
    # eigenvectors to lie in a specific quadrant to make the discretization
    # search easier.
    norm_ones = np.sqrt(n_samples)
    for i in range(vectors.shape[1]):
        vectors[:, i] = (vectors[:, i] / norm(vectors[:, i])) \
            * norm_ones
        if vectors[0, i] != 0:
            vectors[:, i] = -1 * vectors[:, i] * np.sign(vectors[0, i])

    # Normalize the rows of the eigenvectors.  Samples should lie on the unit
    # hypersphere centered at the origin.  This transforms the samples in the
    # embedding space to the space of partition matrices.
    vectors = vectors / np.sqrt((vectors ** 2).sum(axis=1))[:, np.newaxis]

    svd_restarts = 0
    has_converged = False

    # If there is an exception we try to randomize and rerun SVD again
    # do this max_svd_restarts times.
    while (svd_restarts < max_svd_restarts) and not has_converged:

        # Initialize first column of rotation matrix with a row of the
        # eigenvectors
        rotation = np.zeros((n_components, n_components))
        rotation[:, 0] = vectors[random_state.randint(n_samples), :].T

        # To initialize the rest of the rotation matrix, find the rows
        # of the eigenvectors that are as orthogonal to each other as
        # possible
        c = np.zeros(n_samples)
        for j in range(1, n_components):
            # Accumulate c to ensure row is as orthogonal as possible to
            # previous picks as well as current one
            c += np.abs(np.dot(vectors, rotation[:, j - 1]))
            rotation[:, j] = vectors[c.argmin(), :].T

        last_objective_value = 0.0
        n_iter = 0

        while not has_converged:
            n_iter += 1

            t_discrete = np.dot(vectors, rotation)

            labels = t_discrete.argmax(axis=1)
            vectors_discrete = csc_matrix(
                (np.ones(len(labels)), (np.arange(0, n_samples), labels)),
                shape=(n_samples, n_components))

            t_svd = vectors_discrete.T * vectors

            try:
                U, S, Vh = np.linalg.svd(t_svd)
                svd_restarts += 1
            except LinAlgError:
                print("SVD did not converge, randomizing and trying again")
                break

            ncut_value = 2.0 * (n_samples - S.sum())
            if ((abs(ncut_value - last_objective_value) < eps) or
                    (n_iter > n_iter_max)):
                has_converged = True
            else:
                # otherwise calculate rotation and continue
                last_objective_value = ncut_value
                rotation = np.dot(Vh.T, U.T)

    if not has_converged:
        raise LinAlgError('SVD did not converge')
    return labels

def spectral_clustering(affinity, n_clusters=8, n_components=None,
                        eigen_solver=None, random_state=None, n_init=10,
                        eigen_tol=0.0, assign_labels='kmeans'):
    """Apply clustering to a projection to the normalized laplacian.

    In practice Spectral Clustering is very useful when the structure of
    the individual clusters is highly non-convex or more generally when
    a measure of the center and spread of the cluster is not a suitable
    description of the complete cluster. For instance when clusters are
    nested circles on the 2D plan.

    If affinity is the adjacency matrix of a graph, this method can be
    used to find normalized graph cuts.

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

    Parameters
    -----------
    affinity : array-like or sparse matrix, shape: (n_samples, n_samples)
        The affinity matrix describing the relationship of the samples to
        embed. **Must be symmetric**.

        Possible examples:
          - adjacency matrix of a graph,
          - heat kernel of the pairwise distance matrix of the samples,
          - symmetric k-nearest neighbours connectivity matrix of the samples.

    n_clusters : integer, optional
        Number of clusters to extract.

    n_components : integer, optional, default is n_clusters
        Number of eigen vectors to use for the spectral embedding

    eigen_solver : {None, 'arpack', 'lobpcg', or 'amg'}
        The eigenvalue decomposition strategy to use. AMG requires pyamg
        to be installed. It can be faster on very large, sparse problems,
        but may also lead to instabilities

    random_state : int seed, RandomState instance, or None (default)
        A pseudo random number generator used for the initialization
        of the lobpcg eigen vectors decomposition when eigen_solver == 'amg'
        and by the K-Means initialization.

    n_init : int, optional, default: 10
        Number of time the k-means algorithm will be run with different
        centroid seeds. The final results will be the best output of
        n_init consecutive runs in terms of inertia.

    eigen_tol : float, optional, default: 0.0
        Stopping criterion for eigendecomposition of the Laplacian matrix
        when using arpack eigen_solver.

    assign_labels : {'kmeans', 'discretize'}, default: 'kmeans'
        The strategy to use to assign labels in the embedding
        space.  There are two ways to assign labels after the laplacian
        embedding.  k-means can be applied and is a popular choice. But it can
        also be sensitive to initialization. Discretization is another
        approach which is less sensitive to random initialization. See
        the 'Multiclass spectral clustering' paper referenced below for
        more details on the discretization approach.

    Returns
    -------
    labels : array of integers, shape: n_samples
        The labels of the clusters.

    References
    ----------

    - Normalized cuts and image segmentation, 2000
      Jianbo Shi, Jitendra Malik
      http://citeseer.ist.psu.edu/viewdoc/summary?doi=10.1.1.160.2324

    - A Tutorial on Spectral Clustering, 2007
      Ulrike von Luxburg
      http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.165.9323

    - Multiclass spectral clustering, 2003
      Stella X. Yu, Jianbo Shi
      http://www1.icsi.berkeley.edu/~stellayu/publication/doc/2003kwayICCV.pdf

    Notes
    ------
    The graph should contain only one connect component, elsewhere
    the results make little sense.

    This algorithm solves the normalized cut for k=2: it is a
    normalized spectral clustering.
    """
    if assign_labels not in ('kmeans', 'discretize'):
        raise ValueError("The 'assign_labels' parameter should be "
                         "'kmeans' or 'discretize', but '%s' was given"
                         % assign_labels)

    random_state = check_random_state(random_state)
    n_components = n_clusters if n_components is None else n_components
    maps = spectral_embedding(affinity, n_components=n_components,
                              eigen_solver=eigen_solver,
                              random_state=random_state,
                              eigen_tol=eigen_tol, drop_first=False)

    if assign_labels == 'kmeans':
        _, labels, _ = k_means(maps, n_clusters, random_state=random_state,
                               n_init=n_init)
    else:
        labels = discretize(maps, random_state=random_state)

    return labels

class SpectralClustering(BaseEstimator, ClusterMixin):
    """Apply clustering to a projection to the normalized laplacian.

    In practice Spectral Clustering is very useful when the structure of
    the individual clusters is highly non-convex or more generally when
    a measure of the center and spread of the cluster is not a suitable
    description of the complete cluster. For instance when clusters are
    nested circles on the 2D plan.

    If affinity is the adjacency matrix of a graph, this method can be
    used to find normalized graph cuts.

    When calling ``fit``, an affinity matrix is constructed using either
    kernel function such the Gaussian (aka RBF) kernel of the euclidean
    distanced ``d(X, X)``::

            np.exp(-gamma * d(X,X) ** 2)

    or a k-nearest neighbors connectivity matrix.

    Alternatively, using ``precomputed``, a user-provided affinity
    matrix can be used.

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

    Parameters
    -----------
    n_clusters : integer, optional
        The dimension of the projection subspace.

    affinity : string, array-like or callable, default 'rbf'
        If a string, this may be one of 'nearest_neighbors', 'precomputed',
        'rbf' or one of the kernels supported by
        `sklearn.metrics.pairwise_kernels`.

        Only kernels that produce similarity scores (non-negative values that
        increase with similarity) should be used. This property is not checked
        by the clustering algorithm.

    gamma : float
        Scaling factor of RBF, polynomial, exponential chi^2 and
        sigmoid affinity kernel. Ignored for
        ``affinity='nearest_neighbors'``.

    degree : float, default=3
        Degree of the polynomial kernel. Ignored by other kernels.

    coef0 : float, default=1
        Zero coefficient for polynomial and sigmoid kernels.
        Ignored by other kernels.

    n_neighbors : integer
        Number of neighbors to use when constructing the affinity matrix using
        the nearest neighbors method. Ignored for ``affinity='rbf'``.

    eigen_solver : {None, 'arpack', 'lobpcg', or 'amg'}
        The eigenvalue decomposition strategy to use. AMG requires pyamg
        to be installed. It can be faster on very large, sparse problems,
        but may also lead to instabilities

    random_state : int seed, RandomState instance, or None (default)
        A pseudo random number generator used for the initialization
        of the lobpcg eigen vectors decomposition when eigen_solver == 'amg'
        and by the K-Means initialization.

    n_init : int, optional, default: 10
        Number of time the k-means algorithm will be run with different
        centroid seeds. The final results will be the best output of
        n_init consecutive runs in terms of inertia.

    eigen_tol : float, optional, default: 0.0
        Stopping criterion for eigendecomposition of the Laplacian matrix
        when using arpack eigen_solver.

    assign_labels : {'kmeans', 'discretize'}, default: 'kmeans'
        The strategy to use to assign labels in the embedding
        space. There are two ways to assign labels after the laplacian
        embedding. k-means can be applied and is a popular choice. But it can
        also be sensitive to initialization. Discretization is another approach
        which is less sensitive to random initialization.

    kernel_params : dictionary of string to any, optional
        Parameters (keyword arguments) and values for kernel passed as
        callable object. Ignored by other kernels.

    Attributes
    ----------
    affinity_matrix_ : array-like, shape (n_samples, n_samples)
        Affinity matrix used for clustering. Available only if after calling
        ``fit``.

    labels_ :
        Labels of each point

    Notes
    -----
    If you have an affinity matrix, such as a distance matrix,
    for which 0 means identical elements, and high values means
    very dissimilar elements, it can be transformed in a
    similarity matrix that is well suited for the algorithm by
    applying the Gaussian (RBF, heat) kernel::

        np.exp(- X ** 2 / (2. * delta ** 2))

    Another alternative is to take a symmetric version of the k
    nearest neighbors connectivity matrix of the points.

    If the pyamg package is installed, it is used: this greatly
    speeds up computation.

    References
    ----------

    - Normalized cuts and image segmentation, 2000
      Jianbo Shi, Jitendra Malik
      http://citeseer.ist.psu.edu/viewdoc/summary?doi=10.1.1.160.2324

    - A Tutorial on Spectral Clustering, 2007
      Ulrike von Luxburg
      http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.165.9323

    - Multiclass spectral clustering, 2003
      Stella X. Yu, Jianbo Shi
      http://www1.icsi.berkeley.edu/~stellayu/publication/doc/2003kwayICCV.pdf
    """

    def __init__(self, n_clusters=8, eigen_solver=None, random_state=None,
                 n_init=10, gamma=1., affinity='rbf', n_neighbors=10,
                 eigen_tol=0.0, assign_labels='kmeans', degree=3, coef0=1,
                 kernel_params=None):
        self.n_clusters = n_clusters
        self.eigen_solver = eigen_solver
        self.random_state = random_state
        self.n_init = n_init
        self.gamma = gamma
        self.affinity = affinity
        self.n_neighbors = n_neighbors
        self.eigen_tol = eigen_tol
        self.assign_labels = assign_labels
        self.degree = degree
        self.coef0 = coef0
        self.kernel_params = kernel_params

    def fit(self, X, y=None):
        """Creates an affinity matrix for X using the selected affinity,
        then applies spectral clustering to this affinity matrix.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            OR, if affinity==`precomputed`, a precomputed affinity
            matrix of shape (n_samples, n_samples)
        """
        X = check_array(X, accept_sparse=['csr', 'csc', 'coo'],
                        dtype=np.float64)
        if X.shape[0] == X.shape[1] and self.affinity != "precomputed":
            warnings.warn("The spectral clustering API has changed. ``fit``"
                          "now constructs an affinity matrix from data. To use"
                          " a custom affinity matrix, "
                          "set ``affinity=precomputed``.")

        if self.affinity == 'nearest_neighbors':
            connectivity = kneighbors_graph(X, n_neighbors=self.n_neighbors, include_self=True)
            self.affinity_matrix_ = 0.5 * (connectivity + connectivity.T)
        elif self.affinity == 'precomputed':
            self.affinity_matrix_ = X
        else:
            params = self.kernel_params
            if params is None:
                params = {}
            if not callable(self.affinity):
                params['gamma'] = self.gamma
                params['degree'] = self.degree
                params['coef0'] = self.coef0
            self.affinity_matrix_ = pairwise_kernels(X, metric=self.affinity,
                                                     filter_params=True,
                                                     **params)

        random_state = check_random_state(self.random_state)
        self.labels_ = spectral_clustering(self.affinity_matrix_,
                                           n_clusters=self.n_clusters,
                                           eigen_solver=self.eigen_solver,
                                           random_state=random_state,
                                           n_init=self.n_init,
                                           eigen_tol=self.eigen_tol,
                                           assign_labels=self.assign_labels)
        return self

    @property
    def _pairwise(self):
        return self.affinity == "precomputed"


#
# Copyright (C) 2019 Kevin Thornton <krthornt@uci.edu>
#
# This file is part of fwdpy11.
#
# fwdpy11 is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# fwdpy11 is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with fwdpy11.  If not, see <http://www.gnu.org/licenses/>.
#
"""
Local adaptation of a quantitative trait to differing optima.
"""

import argparse
import math
import sys
from collections import namedtuple

import numpy as np
import pandas as pd

import fwdpy11

# Simulations with tree sequence recording need
# to know the max position in a genome.  Here,
# we use a length of 1.0. Thus, all mutation
# and recombination events will be uniform
# random variables on the continuous interval
# [0, GENOME_LENGTH).
GENOME_LENGTH = 1.0

# When recording quant-genetic statistics during a simulation,
# we will use this type. Named tuples are extremely efficient,
# and they are easily converted into Pandas DataFrame objects,
# which is very convenient for analysis and output.
Datum = namedtuple("Data", ["generation", "deme", "gbar", "vg", "wbar"])

def make_parser():
    """
    Create a command-line interface to the script.
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    required = parser.add_argument_group("Required arguments")
    required.add_argument("--popsize", "-N", type=int, help="Diploid population size")
    required.add_argument(
        "--mu", "-m", type=float, help="Mutation rate (per gamete, per generation)"
    )
    required.add_argument(
        "--sigma",
        "-s",
        type=float,
        help="Standard deviation of Gaussian" "distribution of mutational effects",
    )

    optional = parser.add_argument_group("Optional arguments")
    optional.add_argument(
        "--rho", type=float, default=1000.0, help="Scaled recombination rate, rho=4Nr"
    )
    optional.add_argument(
        "--VS",
        type=float,
        default=10.0,
        help="Inverse strength of stabilizing selection",
    )
    optional.add_argument(
        "--opt", type=float, default=1.0, help="Value of new phenotypic optimum"
    )
    optional.add_argument(
        "--migrates",
        type=float,
        nargs=2,
        default=None,
        help="Migration rates from 0 to 1 and 1 to 0, respectively.",
    )
    optional.add_argument(
        "--time",
        type=float,
        default=0.1,
        help="Amount of time to simulate past" "optimum shift, in units of N",
    )
    optional.add_argument(
        "--plotfile", type=str, default=None, help="File name for plot"
    )
    optional.add_argument("--seed", type=int, default=42, help="Random number seed.")

    return parser

def validate_arguments(args):
    """
    Validate input arguments.
    Note: this is likely incomplete.
    """
    if args.popsize is None:
        raise ValueError("popsize cannot be None")
    if args.mu < 0:
        raise ValueError("mu must be non-negative")
    if args.mu is None:
        raise ValueError("mu cannot be None")
    if args.mu < 0 or math.isfinite(args.mu) is False:
        raise ValueError("Mutation rate must be non-negative and finite")
    if args.sigma is None:
        raise ValueError("sigma cannot be none")
    if args.sigma < 0 or math.isfinite(args.sigma) is False:
        raise ValueError(
            "Std. dev. of distribution of effect sizes"
            "must be non-negative and finite"
        )

    if args.migrates is not None:
        for m in args.migrates:
            if m < 0 or m > 1:
                raise ValueError("migration rates must be 0 <= m <= 1")

def make_migmatrix(migrates):
    if migrates is None:
        return None
    mm = np.zeros(4).reshape(2, 2)
    mm[0, 1] = migrates[1]
    mm[1, 0] = migrates[0]
    rs = np.sum(mm, axis=1)
    np.fill_diagonal(mm, 1.0 - rs)
    return fwdpy11.MigrationMatrix(mm)

class Recorder(object):
    """
    fwdpy11 allows you to define objects that record data
    from populations during simulation.  Such objects must
    be callable, and the easiest way to do things is to
    create a class with a __call__ function.
    """

    def __init__(self, start):
        self.data = []
        self.start = start

    def __call__(self, pop, recorder):
        if pop.generation >= self.start:
            # Record mean trait value each generation.
            md = np.array(pop.diploid_metadata, copy=False)
            demes = np.unique(md["deme"])
            for d in demes:
                w = np.where(md["deme"] == d)[0]
                gbar = md["g"][w].mean()
                vg = md["g"][w].var()
                wbar = md["w"][w].mean()
                self.data.append(Datum(pop.generation, d, gbar, vg, wbar))

def plot_output(data, filename):
    import matplotlib.pyplot as plt
    import matplotlib.gridspec as gridspec

    fig = plt.figure(figsize=(9, 3))
    gs = gridspec.GridSpec(ncols=3, nrows=1, figure=fig)
    ax_gbar = fig.add_subplot(gs[0, 0])
    ax_vg = fig.add_subplot(gs[0, 1])
    ax_wbar = fig.add_subplot(gs[0, 2])

    df = pd.DataFrame(data, columns=Datum._fields)
    g = df.groupby(["deme"])
    for n, gi in g:
        ax_gbar.plot(gi["generation"], gi["gbar"], label="Deme {}".format(n))
        ax_vg.plot(gi["generation"], gi["vg"], label="Deme {}".format(n))
        ax_wbar.plot(gi["generation"], gi["wbar"], label="Deme {}".format(n))

    for ax in [ax_gbar, ax_vg, ax_wbar]:
        ax.set_xlabel("Generation")

    ax_gbar.set_ylabel(r"$\bar{g}$")
    ax_vg.set_ylabel(r"$V(G)$")
    ax_wbar.set_ylabel(r"$\bar{w}$")
    ax_gbar.legend()
    plt.tight_layout()
    plt.savefig(filename)

def runsim(args):
    """
    Run the simulation.
    """
    pop = fwdpy11.DiploidPopulation(2 * args.popsize, GENOME_LENGTH)

    np.random.seed(args.seed)
    rng = fwdpy11.GSLrng(args.seed)

    GSSmo0 = fwdpy11.GSSmo(
        [
            fwdpy11.Optimum(when=0, optimum=0.0, VS=args.VS),
            fwdpy11.Optimum(when=10 * args.popsize, optimum=args.opt, VS=args.VS),
        ]
    )
    GSSmo1 = fwdpy11.GSSmo(
        [
            fwdpy11.Optimum(when=0, optimum=0.0, VS=args.VS),
            fwdpy11.Optimum(
                when=10 * args.popsize, optimum=-1.0 * args.opt, VS=args.VS
            ),
        ]
    )

    mm = make_migmatrix(args.migrates)
    dd = fwdpy11.DiscreteDemography(
        mass_migrations=[fwdpy11.move_individuals(0, 0, 1, 0.5)], migmatrix=mm
    )

    p = {
        "nregions": [],  # No neutral mutations -- add them later!
        "gvalue": [fwdpy11.Additive(2.0, GSSmo0), fwdpy11.Additive(2.0, GSSmo1)],
        "sregions": [fwdpy11.GaussianS(0, GENOME_LENGTH, 1, args.sigma)],
        "recregions": [fwdpy11.Region(0, GENOME_LENGTH, 1)],
        "rates": (0.0, args.mu, args.rho / float(4 * args.popsize)),
        # Keep mutations at frequency 1 in the pop if they affect fitness.
        "prune_selected": False,
        "demography": dd,
        "simlen": 10 * args.popsize + int(args.popsize * args.time),
    }
    params = fwdpy11.ModelParams(**p)

    r = Recorder(10 * args.popsize)
    fwdpy11.evolvets(rng, pop, params, 100, r, suppress_table_indexing=True)
    if args.plotfile is not None:
        plot_output(r.data, args.plotfile)

if __name__ == "__main__":
    parser = make_parser()
    args = parser.parse_args(sys.argv[1:])
    validate_arguments(args)
    runsim(args)


import matplotlib.pyplot as plt

with open("test_pen", "r") as f:
    data = f.read()
data = data.split("\n")[:-1]
data = [data[i].split(" ") for i in range(0, len(data))]
pen = [float(data[i][0]) for i in range(len(data))]
u = [float(data[i][1]) for i in range(len(data))]
d = [float(data[i][2]) for i in range(len(data))]
gain = [((d[i-1] - d[i])) / (u[i] - u[i - 1]) for i in range(1, len(data))]
gain = [gain[0]] + gain
print(u, d, gain)
fig, ax1 = plt.subplots()
pu, = ax1.plot(pen, u, color="r", label="Uncertainty rate")
ax1.scatter(pen, u, color="k")
#ax1.axhline(9000, color="r", linestyle="--")
#ax1.set_title("Cost evolution according to the number of iterations")
ax1.set_xlabel("Penalization coefficient")
ax1.set_ylabel("Uncertainty rate")
ax2 = ax1.twinx()
pd, = ax2.plot(pen, d, color="b", linestyle="--", label="Distance")
ax2.scatter(pen, d, color="k")
ax2.set_ylabel("Distance")
#ax2.axhline(0.99, color="b", linestyle="--")
#plt.axvline(4000000, color="k",linestyle = ":")
plt.legend(handles=[pu, pd], loc=7)
plt.show()


#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""Routines for writing TeX for tables."""
from __future__ import absolute_import
import os
import sys
import string
import numpy

from . import toolsstats
from pdb import set_trace

#GLOBAL VARIABLES DEFINITION
alphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
#conversion of matplotlib elements to LaTeX
latex_marker_map = {'o': r'$\circ$',
              'd': r'$\diamondsuit$',
              's': r'$\Box$',
              'v': r'$\triangledown$',
              '*': r'$\star$',
              'h': r'$\varhexagon$', # need \usepackage{wasysymb}
              '^': r'$\triangle$',
              'p': r'$\pentagon$', # need \usepackage{wasysymb}
              'H': r'$\hexagon$', # need \usepackage{wasysymb}
              '<': r'$\triangleleft$',
              'D': r'$\Diamond$',
              '>': r'$\triangleright$',
              '1': r'$\downY$', # need \usepackage{MnSymbol}
              '2': r'$\upY$', # need \usepackage{MnSymbol}
              '3': r'$\rightY$', # need \usepackage{MnSymbol}
              '4': r'$\leftY$'} # need \usepackage{MnSymbol}
html_marker_map = {
              'o': r'○',
              'd': r'♢',
              's': r'◻',
              'v': r'▽',
              '*': r'☆',
              'h': r'varhexagon',
              '^': r'△',
              'p': r'pentagon',
              'H': r'hexagon',
              '<': r'◁',
              'D': r'◇',
              '>': r'▷',
              '1': r'downY',
              '2': r'upY',
              '3': r'rightY',
              '4': r'leftY'}
latex_color_map_old = {
             'g': 'green!45!black',
             'r': 'red',
             'c': 'cyan',
             'm': 'magenta',
             'y': 'yellow',
             'k': 'black',
             'b': 'blue'}
latex_color_map = {
             'c': 'cyan',
             'm': 'magenta',
             'y': 'yellow',
             'b': 'blue',
             'g': 'green',
             '#000080': 'NavyBlue',
             'r': 'red',
             '#ffd700': 'Goldenrod',
             '#d02090': 'VioletRed',
             'k': 'Black',
             '#6495ed': 'CornflowerBlue',
             '#ff4500': 'OrangeRed',
             '#ffff00': 'Yellow',
             '#ff00ff': 'Magenta',
             '#bebebe': 'Gray',
             '#87ceeb': 'SkyBlue',
             '#ffa500': 'Orange',
             '#ffc0cb': 'Lavender',
             '#4169e1': 'RoyalBlue',
             '#228b22': 'ForestGreen',
             '#32cd32': 'LimeGreen',
             '#9acd32': 'YellowGreen',
             '#adff2f': 'GreenYellow'}

#CLASS DEFINITION
class Error(Exception):
    """ Base class for errors. """
    pass

class WrongInputSizeError(Error):
    """Error if an array has the wrong size for the following operation.

    :returns: message containing the size of the array and the required
              size.

    """
    def __init__(self,arrName, arrSize, reqSize):
        self.arrName = arrName
        self.arrSize = arrSize
        self.reqSize = reqSize

    def __str__(self):
        message = 'The size of %s is %s. One dimension must be of length %s!' %\
                  (self.arrName,str(self.arrSize), str(self.reqSize))
        return repr(message)

#TOP LEVEL METHODS
def color_to_latex(color):
    try:
        res = '\color{%s}' % latex_color_map[color]
    except KeyError, err:
        try:
            float(color)
            res = '\color[gray]{%s}' % color
        except ValueError:
            raise err
    return res

def marker_to_latex(marker):
    return latex_marker_map[marker]

def marker_to_html(marker):
    return html_marker_map[marker]

def numtotext(n):
    """Returns a text from a positive integer.

    Is to be used for generating command names: they cannot include number
    characters.

    WARNING: n should not be larger than (53*52)-1 = 2755 for the moment

    """
    if n < 52:
        str = alphabet[n]
    elif n < 53*52:
        str = alphabet[(n-52)//52] + alphabet[n-n//52*52]
    else:
        raise Exception('Cannot handle a number of algorithms that large.')

    return str

def writeLabels(label):
    """Format text to be output by LaTeX."""
    return label.replace('_', r'\_')

def writeFEvals(fevals, precision='.2'):
    """Returns string representation of a number of function evaluations."""

    if numpy.isinf(fevals):
        return r'$\infty$'

    tmp = (('%' + precision + 'g') % fevals)
    res = tmp.split('e')
    if len(res) > 1:
        res[1] = '%d' % int(res[1])
        res = '%s' % 'e'.join(res)
        pr2 = str(float(precision) + .2)
        #res2 = (('%' + pr2 + 'g') % fevals)
        res2 = (('%' + pr2 + 'g') % float(tmp))
        # To have the same number of significant digits.
        if len(res) >= len(res2):
            res = res2
    else:
        res = res[0]
    return res

def writeFEvals2(fevals, precision=2, maxdigits=None, isscientific=False):
    """Returns string representation of a number of function evaluations.

    This method is supposed to be used for filling up a LaTeX tabular.

    To address the eventual need to keep their string representation
    short, the method here proposes the shortest representation between
    the full representation and a modified scientific representation.

    :param float fevals:
    :param int precision: number of significant digits
    :param int maxdigits:
    :param bool isscientific:

    Examples:

    ======   =========   =====================
    Number   Precision   Output Representation
    ======   =========   =====================
    102345   2 digits    1.0e5
    ======   =========   =====================

    """

    #Printf:
    # %[flags][width][.precision][length]specifier

    assert not numpy.isnan(fevals)

    if numpy.isinf(fevals):
        return r'$\infty$'

    if maxdigits is None:
        precision = int(precision)

        #repr1 is the alternative scientific notation
        #repr2 is the full notation but with a number of significant digits given
        #by the variable precision.

        res = (('%.' + str(precision-1) + 'e') % fevals)
        repr1 = res
        tmp = repr1.split('e')
        tmp[1] = '%d' % int(tmp[1]) # Drop the eventual plus sign and trailing zero
        repr1 = 'e'.join(tmp)

        repr2 = (('%.' + str(precision+1) + 'f') % float(res)).rstrip('0').rstrip('.')
        #set_trace()
        if len(repr1) > len(repr2) and not isscientific:
            return repr2

        return repr1

    else:
        # takes precedence, in this case we expect a positive integer
        if not isinstance(fevals, int):
            return '%d' % fevals

        repr2 = '%.0f' % fevals
        if len(repr2) > maxdigits:
            precision = maxdigits - 4
            # 1) one symbol for the most significant digit
            # 2) one for the dot, 3) one for the e, 4) one for the exponent
            if numpy.log10(fevals) > 10:
                precision -= 1
            if precision < 0:
                precision = 0
            repr1 = (('%.' + str(precision) + 'e') % fevals).split('e')
            repr1[1] = '%d' % int(repr1[1]) # drop the sign and trailing zero
            repr1 = 'e'.join(repr1)
            return repr1

        return repr2

def writeFEvalsMaxSymbols(fevals, maxsymbols, isscientific=False):
    """Return the smallest string representation of a number.

    This method is only concerned with the maximum number of significant
    digits.

    Two alternatives:

    1) modified scientific notation (without the trailing + and zero in
       the exponent)
    2) float notation

    :returns: string representation of a number of function evaluations
              or ERT.

    """

    #Compared to writeFEvals2?
    #Printf:
    # %[flags][width][.precision][length]specifier

    assert not numpy.isnan(fevals)

    if numpy.isinf(fevals):
        return r'$\infty$'

    #repr1 is the alternative scientific notation
    #repr2 is the full notation but with a number of significant digits given
    #by the variable precision.

    # modified scientific notation:
    #smallest representation of the decimal part
    #drop + and starting zeros of the exponent part
    repr1 = (('%.' + str(maxsymbols) + 'e') % fevals)
    size1 = len(repr1)
    tmp = repr1.split('e', 1)
    tmp2 = tmp[-1].lstrip('+-0')
    if float(tmp[-1]) < 0:
        tmp2 = '-' + tmp2
    tmp[-1] = tmp2
    remainingsymbols = max(maxsymbols - len(tmp2) - 2, 0)
    tmp[0] = (('%.' + str(remainingsymbols) + 'f') % float(tmp[0]))
    repr1 = 'e'.join(tmp)
    #len(repr1) <= maxsymbols is not always the case but should be most usual

    tmp = '%.0f' % fevals
    remainingsymbols = max(maxsymbols - len(tmp), 0)
    repr2 = (('%.' + str(remainingsymbols) + 'f') % fevals)
    tmp = repr2.split('.', 1)
    if len(tmp) > 1:
        tmp[-1] = tmp[-1].rstrip('0')
    repr2 = '.'.join(tmp)
    repr2 = repr2.rstrip('.')
    #set_trace()

    if len(repr1)-repr1.count('.') < len(repr2)-repr2.count('.') or isscientific:
        return repr1

    #tmp1 = '%4.0f' % bestalgdata[-1]
    #tmp2 = ('%2.2g' % bestalgdata[-1]).split('e', 1)
    #if len(tmp2) > 1:
    #    tmp2[-1] = tmp2[-1].lstrip('+0')
    #    tmp2 = 'e'.join(tmp2)
    #    tmp = tmp1
    #    if len(tmp1) >= len(tmp2):
    #        tmp = tmp2
    #    curline.append(r'\multicolumn{2}{c|}{%s}' % tmp)

    return repr2

def writeFEvalsMaxPrec(entry, SIG, maxfloatrepr=1e5):
    """Return a string representation of a number.

    Two alternatives:

    1) float notation with a precision smaller or equal to SIG (if the
       entry is one, then the result is 1).
    2) if the number is larger or equal to maxfloatrepr, a modified
       scientific notation (without the trailing + and zero in the
       exponent)

    :returns: string representation of a number of function evaluations
              or ERT.

    """
    #CAVE: what if entry is smaller than 10**(-SIG)?
    #Printf:
    # %[flags][width][.precision][length]specifier

    assert not numpy.isnan(entry)

    if numpy.isinf(entry):
        return r'$\infty$'

    if entry == 1.:
        res = '1'
    elif entry < maxfloatrepr:
        # the full notation but with given maximum precision
        corr = 1 if abs(entry) < 1 else 0
        tmp = '%.0f' % entry
        remainingsymbols = max(SIG - len(tmp) + corr, 0)
        res = (('%.' + str(remainingsymbols) + 'f') % entry)
    else:
        # modified scientific notation:
        #smallest representation of the decimal part
        #drop + and starting zeros of the exponent part
        res = (('%.' + str(max([0, SIG - 1])) + 'e') % entry)
        size1 = len(res)
        tmp = res.split('e', 1)
        tmp2 = tmp[-1].lstrip('+-0')
        if float(tmp[-1]) < 0:
            tmp2 = '-' + tmp2
        tmp[-1] = tmp2
        if len(tmp) > 1 and tmp[-1]:
            res = 'e'.join(tmp)
        else:
            res = tmp[0]

    return res

def tableLaTeX(table, spec, extraeol=()):
    """Generates a tabular from a sequence of sequence of strings.

    :param seq table: sequence of sequence of strings
    :param string spec: string for table specification, see
                        http://en.wikibooks.org/wiki/LaTeX/Tables#The_tabular_environment
    :param seq extraeol: sequence of string the same length as the table
                         (same number of lines) which are added at the
                         end of each line.
    :returns: sequence of strings of a LaTeX tabular.

    """

    if not extraeol:
        extraeol = len(table) * ['']

    # TODO: check that spec and extraeol have the right format?

    res = [r'\begin{tabular}{%s}' % spec]
    for i, line in enumerate(table[:-1]):
        curline = ' & '.join(line) + r'\\' + extraeol[i]
#        curline = ' & '.join(line) + r'\\\hline' + extraeol[i]
        res.append(curline)
    res.append(' & '.join(table[-1]) + extraeol[-1])

    res.append(r'\end{tabular}')
    res = '\n'.join(res)
    return res

def tableXLaTeX(table, spec, extraeol=()):
    """Generates a tabular from a sequence of sequence of strings.

    :param seq table: sequence of sequence of strings
    :param string spec: string for table specification, see
                        http://en.wikibooks.org/wiki/LaTeX/Tables#The_tabular_environment
    :param seq extraeol: sequence of string the same length as the table
                         (same number of lines) which are added at the
                         end of each line.
    :returns: sequence of strings of a LaTeX tabular.

    """

    if not extraeol:
        extraeol = len(table) * ['']

    # TODO: check that spec and extraeol have the right format?
    if 1 < 3:
        res = [r'\begin{tabularx}{1.0\textwidth}{%s}' % spec]
        for i, line in enumerate(table[:-1]):
            curline = ' & '.join(line) + r'\\' + extraeol[i]
            res.append(curline)
    else: # format with hline, when is it needed, for non-paper tables?
        res = [r'\begin{tabularx}{1.3\textwidth}{%s}' % spec]
        for i, line in enumerate(table[:-1]):
            curline = ' & '.join(line) + r'\\\hline' + extraeol[i]
            res.append(curline)

    res.append(' & '.join(table[-1]) + extraeol[-1])

    res.append(r'\end{tabularx}')
    res = '\n'.join(res)
    return res

def tableLaTeXStar(table, width, spec, extraeol=()):
    """Generates a tabular\* from a sequence of sequence of strings

    :param seq table: sequence of sequence of strings
    :param string width: string for the width of the table
    :param strin spec: string for table specification, see
                       http://en.wikibooks.org/wiki/LaTeX/Tables#The_tabular_environment
    :param seq extraeol: sequence of string the same length as the table
                         (same number of lines) which are added at the
                         end of each line.

    """
    if not extraeol:
        extraeol = len(table) * ['']

    # TODO: check that spec and extraeol have the right format?

    res = [r'\begin{tabular*}{%s}{%s}' % (width, spec)]
    for i, line in enumerate(table[:-1]):
        curline = ' & '.join(line) + r'\\' + extraeol[i]
        res.append(curline)
    res.append(' & '.join(table[-1]) + extraeol[-1])

    res.append(r'\end{tabular*}')
    res = '\n'.join(res)
    return res

class DataTable(list):
    pass


"""Calibration of predicted probabilities."""

# Author: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#         Balazs Kegl <balazs.kegl@gmail.com>
#         Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
#         Mathieu Blondel <mathieu@mblondel.org>
#
# License: BSD 3 clause

from __future__ import division
import inspect
import warnings

from math import log
import numpy as np

from scipy.optimize import fmin_bfgs

from .base import BaseEstimator, ClassifierMixin, RegressorMixin, clone
from .preprocessing import LabelBinarizer
from .utils import check_X_y, check_array, indexable, column_or_1d
from .utils.validation import check_is_fitted
from .isotonic import IsotonicRegression
from .svm import LinearSVC
from .cross_validation import check_cv
from .metrics.classification import _check_binary_probabilistic_predictions

class CalibratedClassifierCV(BaseEstimator, ClassifierMixin):
    """Probability calibration with isotonic regression or sigmoid.

    With this class, the base_estimator is fit on the train set of the
    cross-validation generator and the test set is used for calibration.
    The probabilities for each of the folds are then averaged
    for prediction. In case that cv="prefit" is passed to __init__,
    it is it is assumed that base_estimator has been
    fitted already and all data is used for calibration. Note that
    data for fitting the classifier and for calibrating it must be disjpint.

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

    Parameters
    ----------
    base_estimator : instance BaseEstimator
        The classifier whose output decision function needs to be calibrated
        to offer more accurate predict_proba outputs. If cv=prefit, the
        classifier must have been fit already on data.

    method : 'sigmoid' | 'isotonic'
        The method to use for calibration. Can be 'sigmoid' which
        corresponds to Platt's method or 'isotonic' which is a
        non-parameteric approach. It is not advised to use isotonic calibration
        with too few calibration samples (<<1000) since it tends to overfit.
        Use sigmoids (Platt's calibration) in this case.

    cv : integer or cross-validation generator or "prefit", optional
        If an integer is passed, it is the number of folds (default 3).
        Specific cross-validation objects can be passed, see
        sklearn.cross_validation module for the list of possible objects.
        If "prefit" is passed, it is assumed that base_estimator has been
        fitted already and all data is used for calibration.

    Attributes
    ----------
    classes_ : array, shape (n_classes)
        The class labels.

    calibrated_classifiers_: list (len() equal to cv or 1 if cv == "prefit")
        The list of calibrated classifiers, one for each crossvalidation fold,
        which has been fitted on all but the validation fold and calibrated
        on the validation fold.

    References
    ----------
    .. [1] Obtaining calibrated probability estimates from decision trees
           and naive Bayesian classifiers, B. Zadrozny & C. Elkan, ICML 2001

    .. [2] Transforming Classifier Scores into Accurate Multiclass
           Probability Estimates, B. Zadrozny & C. Elkan, (KDD 2002)

    .. [3] Probabilistic Outputs for Support Vector Machines and Comparisons to
           Regularized Likelihood Methods, J. Platt, (1999)

    .. [4] Predicting Good Probabilities with Supervised Learning,
           A. Niculescu-Mizil & R. Caruana, ICML 2005
    """
    def __init__(self, base_estimator=None, method='sigmoid', cv=3):
        self.base_estimator = base_estimator
        self.method = method
        self.cv = cv

    def fit(self, X, y, sample_weight=None):
        """Fit the calibrated model

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        y : array-like, shape (n_samples,)
            Target values.

        sample_weight : array-like, shape = [n_samples] or None
            Sample weights. If None, then samples are equally weighted.

        Returns
        -------
        self : object
            Returns an instance of self.
        """
        X, y = check_X_y(X, y, accept_sparse=['csc', 'csr', 'coo'],
                         force_all_finite=False)
        X, y = indexable(X, y)
        lb = LabelBinarizer().fit(y)
        self.classes_ = lb.classes_

        # Check that we each cross-validation fold can have at least one
        # example per class
        n_folds = self.cv if isinstance(self.cv, int) \
            else self.cv.n_folds if hasattr(self.cv, "n_folds") else None
        if n_folds and \
           np.any([np.sum(y == class_) < n_folds for class_ in self.classes_]):
            raise ValueError("Requesting %d-fold cross-validation but provided"
                             " less than %d examples for at least one class."
                             % (n_folds, n_folds))

        self.calibrated_classifiers_ = []
        if self.base_estimator is None:
            # we want all classifiers that don't expose a random_state
            # to be deterministic (and we don't want to expose this one).
            base_estimator = LinearSVC(random_state=0)
        else:
            base_estimator = self.base_estimator

        if self.cv == "prefit":
            calibrated_classifier = _CalibratedClassifier(
                base_estimator, method=self.method)
            if sample_weight is not None:
                calibrated_classifier.fit(X, y, sample_weight)
            else:
                calibrated_classifier.fit(X, y)
            self.calibrated_classifiers_.append(calibrated_classifier)
        else:
            cv = check_cv(self.cv, X, y, classifier=True)
            arg_names = inspect.getargspec(base_estimator.fit)[0]
            estimator_name = type(base_estimator).__name__
            if (sample_weight is not None
                    and "sample_weight" not in arg_names):
                warnings.warn("%s does not support sample_weight. Samples"
                              " weights are only used for the calibration"
                              " itself." % estimator_name)
                base_estimator_sample_weight = None
            else:
                base_estimator_sample_weight = sample_weight
            for train, test in cv:
                this_estimator = clone(base_estimator)
                if base_estimator_sample_weight is not None:
                    this_estimator.fit(
                        X[train], y[train],
                        sample_weight=base_estimator_sample_weight[train])
                else:
                    this_estimator.fit(X[train], y[train])

                calibrated_classifier = _CalibratedClassifier(
                    this_estimator, method=self.method)
                if sample_weight is not None:
                    calibrated_classifier.fit(X[test], y[test],
                                              sample_weight[test])
                else:
                    calibrated_classifier.fit(X[test], y[test])
                self.calibrated_classifiers_.append(calibrated_classifier)

        return self

    def predict_proba(self, X):
        """Posterior probabilities of classification

        This function returns posterior probabilities of classification
        according to each class on an array of test vectors X.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            The samples.

        Returns
        -------
        C : array, shape (n_samples, n_classes)
            The predicted probas.
        """
        check_is_fitted(self, ["classes_", "calibrated_classifiers_"])
        X = check_array(X, accept_sparse=['csc', 'csr', 'coo'],
                        force_all_finite=False)
        # Compute the arithmetic mean of the predictions of the calibrated
        # classfiers
        mean_proba = np.zeros((X.shape[0], len(self.classes_)))
        for calibrated_classifier in self.calibrated_classifiers_:
            proba = calibrated_classifier.predict_proba(X)
            mean_proba += proba

        mean_proba /= len(self.calibrated_classifiers_)

        return mean_proba

    def predict(self, X):
        """Predict the target of new samples. Can be different from the
        prediction of the uncalibrated classifier.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            The samples.

        Returns
        -------
        C : array, shape (n_samples,)
            The predicted class.
        """
        check_is_fitted(self, ["classes_", "calibrated_classifiers_"])
        return self.classes_[np.argmax(self.predict_proba(X), axis=1)]

class _CalibratedClassifier(object):
    """Probability calibration with isotonic regression or sigmoid.

    It assumes that base_estimator has already been fit, and trains the
    calibration on the input set of the fit function. Note that this class
    should not be used as an estimator directly. Use CalibratedClassifierCV
    with cv="prefit" instead.

    Parameters
    ----------
    base_estimator : instance BaseEstimator
        The classifier whose output decision function needs to be calibrated
        to offer more accurate predict_proba outputs. No default value since
        it has to be an already fitted estimator.

    method : 'sigmoid' | 'isotonic'
        The method to use for calibration. Can be 'sigmoid' which
        corresponds to Platt's method or 'isotonic' which is a
        non-parameteric approach based on isotonic regression.

    References
    ----------
    .. [1] Obtaining calibrated probability estimates from decision trees
           and naive Bayesian classifiers, B. Zadrozny & C. Elkan, ICML 2001

    .. [2] Transforming Classifier Scores into Accurate Multiclass
           Probability Estimates, B. Zadrozny & C. Elkan, (KDD 2002)

    .. [3] Probabilistic Outputs for Support Vector Machines and Comparisons to
           Regularized Likelihood Methods, J. Platt, (1999)

    .. [4] Predicting Good Probabilities with Supervised Learning,
           A. Niculescu-Mizil & R. Caruana, ICML 2005
    """
    def __init__(self, base_estimator, method='sigmoid'):
        self.base_estimator = base_estimator
        self.method = method

    def _preproc(self, X):
        n_classes = len(self.classes_)
        if hasattr(self.base_estimator, "decision_function"):
            df = self.base_estimator.decision_function(X)
            if df.ndim == 1:
                df = df[:, np.newaxis]
        elif hasattr(self.base_estimator, "predict_proba"):
            df = self.base_estimator.predict_proba(X)
            if n_classes == 2:
                df = df[:, 1:]
        else:
            raise RuntimeError('classifier has no decision_function or '
                               'predict_proba method.')

        idx_pos_class = np.arange(df.shape[1])

        return df, idx_pos_class

    def fit(self, X, y, sample_weight=None):
        """Calibrate the fitted model

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data.

        y : array-like, shape (n_samples,)
            Target values.

        sample_weight : array-like, shape = [n_samples] or None
            Sample weights. If None, then samples are equally weighted.

        Returns
        -------
        self : object
            Returns an instance of self.
        """
        lb = LabelBinarizer()
        Y = lb.fit_transform(y)
        self.classes_ = lb.classes_

        df, idx_pos_class = self._preproc(X)
        self.calibrators_ = []

        for k, this_df in zip(idx_pos_class, df.T):
            if self.method == 'isotonic':
                calibrator = IsotonicRegression(out_of_bounds='clip')
            elif self.method == 'sigmoid':
                calibrator = _SigmoidCalibration()
            else:
                raise ValueError('method should be "sigmoid" or '
                                 '"isotonic". Got %s.' % self.method)
            calibrator.fit(this_df, Y[:, k], sample_weight)
            self.calibrators_.append(calibrator)

        return self

    def predict_proba(self, X):
        """Posterior probabilities of classification

        This function returns posterior probabilities of classification
        according to each class on an array of test vectors X.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            The samples.

        Returns
        -------
        C : array, shape (n_samples, n_classes)
            The predicted probas. Can be exact zeros.
        """
        n_classes = len(self.classes_)
        proba = np.zeros((X.shape[0], n_classes))

        df, idx_pos_class = self._preproc(X)

        for k, this_df, calibrator in \
                zip(idx_pos_class, df.T, self.calibrators_):
            if n_classes == 2:
                k += 1
            proba[:, k] = calibrator.predict(this_df)

        # Normalize the probabilities
        if n_classes == 2:
            proba[:, 0] = 1. - proba[:, 1]
        else:
            proba /= np.sum(proba, axis=1)[:, np.newaxis]

        # XXX : for some reason all probas can be 0
        proba[np.isnan(proba)] = 1. / n_classes

        # Deal with cases where the predicted probability minimally exceeds 1.0
        proba[(1.0 < proba) & (proba <= 1.0 + 1e-5)] = 1.0

        return proba

def _sigmoid_calibration(df, y, sample_weight=None):
    """Probability Calibration with sigmoid method (Platt 2000)

    Parameters
    ----------
    df : ndarray, shape (n_samples,)
        The decision function or predict proba for the samples.

    y : ndarray, shape (n_samples,)
        The targets.

    sample_weight : array-like, shape = [n_samples] or None
        Sample weights. If None, then samples are equally weighted.

    Returns
    -------
    a : float
        The slope.

    b : float
        The intercept.

    References
    ----------
    Platt, "Probabilistic Outputs for Support Vector Machines"
    """
    df = column_or_1d(df)
    y = column_or_1d(y)

    F = df  # F follows Platt's notations
    tiny = np.finfo(np.float).tiny  # to avoid division by 0 warning

    # Bayesian priors (see Platt end of section 2.2)
    prior0 = float(np.sum(y <= 0))
    prior1 = y.shape[0] - prior0
    T = np.zeros(y.shape)
    T[y > 0] = (prior1 + 1.) / (prior1 + 2.)
    T[y <= 0] = 1. / (prior0 + 2.)
    T1 = 1. - T

    def objective(AB):
        # From Platt (beginning of Section 2.2)
        E = np.exp(AB[0] * F + AB[1])
        P = 1. / (1. + E)
        l = -(T * np.log(P + tiny) + T1 * np.log(1. - P + tiny))
        if sample_weight is not None:
            return (sample_weight * l).sum()
        else:
            return l.sum()

    def grad(AB):
        # gradient of the objective function
        E = np.exp(AB[0] * F + AB[1])
        P = 1. / (1. + E)
        TEP_minus_T1P = P * (T * E - T1)
        if sample_weight is not None:
            TEP_minus_T1P *= sample_weight
        dA = np.dot(TEP_minus_T1P, F)
        dB = np.sum(TEP_minus_T1P)
        return np.array([dA, dB])

    AB0 = np.array([0., log((prior0 + 1.) / (prior1 + 1.))])
    AB_ = fmin_bfgs(objective, AB0, fprime=grad, disp=False)
    return AB_[0], AB_[1]

class _SigmoidCalibration(BaseEstimator, RegressorMixin):
    """Sigmoid regression model.

    Attributes
    ----------
    a_ : float
        The slope.

    b_ : float
        The intercept.
    """
    def fit(self, X, y, sample_weight=None):
        """Fit the model using X, y as training data.

        Parameters
        ----------
        X : array-like, shape (n_samples,)
            Training data.

        y : array-like, shape (n_samples,)
            Training target.

        sample_weight : array-like, shape = [n_samples] or None
            Sample weights. If None, then samples are equally weighted.

        Returns
        -------
        self : object
            Returns an instance of self.
        """
        X = column_or_1d(X)
        y = column_or_1d(y)
        X, y = indexable(X, y)

        self.a_, self.b_ = _sigmoid_calibration(X, y, sample_weight)
        return self

    def predict(self, T):
        """Predict new data by linear interpolation.

        Parameters
        ----------
        T : array-like, shape (n_samples,)
            Data to predict from.

        Returns
        -------
        T_ : array, shape (n_samples,)
            The predicted data.
        """
        T = column_or_1d(T)
        return 1. / (1. + np.exp(self.a_ * T + self.b_))

def calibration_curve(y_true, y_prob, normalize=False, n_bins=5):
    """Compute true and predicted probabilities for a calibration curve.

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

    Parameters
    ----------
    y_true : array, shape (n_samples,)
        True targets.

    y_prob : array, shape (n_samples,)
        Probabilities of the positive class.

    normalize : bool, optional, default=False
        Whether y_prob needs to be normalized into the bin [0, 1], i.e. is not
        a proper probability. If True, the smallest value in y_prob is mapped
        onto 0 and the largest one onto 1.

    n_bins : int
        Number of bins. A bigger number requires more data.

    Returns
    -------
    prob_true : array, shape (n_bins,)
        The true probability in each bin (fraction of positives).

    prob_pred : array, shape (n_bins,)
        The mean predicted probability in each bin.

    References
    ----------
    Alexandru Niculescu-Mizil and Rich Caruana (2005) Predicting Good
    Probabilities With Supervised Learning, in Proceedings of the 22nd
    International Conference on Machine Learning (ICML).
    See section 4 (Qualitative Analysis of Predictions).
    """
    y_true = column_or_1d(y_true)
    y_prob = column_or_1d(y_prob)

    if normalize:  # Normalize predicted values into interval [0, 1]
        y_prob = (y_prob - y_prob.min()) / (y_prob.max() - y_prob.min())
    elif y_prob.min() < 0 or y_prob.max() > 1:
        raise ValueError("y_prob has values outside [0, 1] and normalize is "
                         "set to False.")

    y_true = _check_binary_probabilistic_predictions(y_true, y_prob)

    bins = np.linspace(0., 1. + 1e-8, n_bins + 1)
    binids = np.digitize(y_prob, bins) - 1

    bin_sums = np.bincount(binids, weights=y_prob, minlength=len(bins))
    bin_true = np.bincount(binids, weights=y_true, minlength=len(bins))
    bin_total = np.bincount(binids, minlength=len(bins))

    nonzero = bin_total != 0
    prob_true = (bin_true[nonzero] / bin_total[nonzero])
    prob_pred = (bin_sums[nonzero] / bin_total[nonzero])

    return prob_true, prob_pred