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

import numpy as np
import os
import pandas as pd
import pkg_resources

import build_coicop_nomenclature

legislation_directory = os.path.join(
    pkg_resources.get_distribution('openfisca_france_indirect_taxation').location,
    'openfisca_france_indirect_taxation',
    'assets',
    'legislation',
    )

sub_levels = ['divisions', 'groupes', 'classes', 'sous_classes', 'postes']
divisions = ['0{}'.format(i) for i in range(1, 10)] + ['11', '12']  # TODO: fix this
taxe_by_categorie_fiscale_number = {
    0: '',
    1: 'tva_taux_super_reduit',
    2: 'tva_taux_reduit',
    3: 'tva_taux_plein',
    4: 'tva_taux_intermediaire',
    7: 'cigarettes',
    8: 'cigares',
    9: 'tabac_a_rouler',
    10: 'alcools_forts',
    11: 'tva_taux_plein',
    12: 'vin',
    13: 'biere',
    14: 'ticpe',
    15: 'assurance_transport',
    16: 'assurance_sante',
    17: 'autres_assurances'
    }

def extract_informations_from_coicop_to_categorie_fiscale():

    def format_exceptions(exceptions):
        grouped = exceptions.groupby(
            by = [exceptions.annee - np.arange(exceptions.shape[0]), 'posteCOICOP', 'categoriefiscale']
            )
        for k, g in grouped:
            print g.posteCOICOP.unique()[0], g.description.unique()[0], g.annee.min(), g.annee.max(), \
                taxe_by_categorie_fiscale_number[int(g.categoriefiscale.unique())]

    def get_dominant_and_exceptions(division):
        assert division in divisions
        parametres_fiscalite_file_path = os.path.join(legislation_directory, 'coicop_to_categorie_fiscale.csv')
        parametres_fiscalite_data_frame = pd.read_csv(
            parametres_fiscalite_file_path,
            converters = {'posteCOICOP': unicode}
            )
        parametres_fiscalite_data_frame['division'] = parametres_fiscalite_data_frame['posteCOICOP'].str[:2].copy()
        division_dataframe = parametres_fiscalite_data_frame.query('division == @division')
        dominant_fiscal_category = division_dataframe.categoriefiscale.value_counts().argmax()
        exceptions = division_dataframe.query('categoriefiscale != @dominant_fiscal_category')
        return dominant_fiscal_category, exceptions

    for coicop_division in divisions:
        dominant, exceptions = get_dominant_and_exceptions(coicop_division)
        print u'\nDivision: {}.\nCatégorie fiscale dominante: {}.\nExceptions:'.format(
            coicop_division,
            taxe_by_categorie_fiscale_number[dominant]
            )
        format_exceptions(exceptions)

def extract_infra_labels_from_coicop_code(coicop_nomenclature = None, coicop_code = None, label = None):
    assert coicop_nomenclature  is not None
    assert coicop_code is not None
    assert label is not None
    known_levels = sub_levels[:len(coicop_code.split('.')) - 1]
    coicop_sub_code = coicop_code[:len(coicop_code) - 2]
    assert known_levels, 'No know levels for COICOP {}'.format(coicop_code)
    labels_by_sub_level = coicop_nomenclature.loc[
        coicop_nomenclature.code_coicop.str[:len(coicop_sub_code)] == coicop_sub_code,
        ['label_{}'.format(level[:-1]) for level in known_levels]
        ].drop_duplicates().dropna().to_dict(orient = 'records')[0]
    for key, value in labels_by_sub_level.iteritems():
        labels_by_sub_level[key] = value
    modified_level = sub_levels[len(coicop_code.split('.')) - 1][:-1]
    labels_by_sub_level['label_{}'.format(modified_level)] = label
    return labels_by_sub_level

def apply_modification(coicop_nomenclature = None, value = None, categorie_fiscale = None,
        origin = None, start = 1994, stop = 2014, label = ''):
    assert coicop_nomenclature is not None
    assert categorie_fiscale in taxe_by_categorie_fiscale_number.values()
    assert 1994 <= start < stop <= 2014, "Invalid start={} and/or stop={}".format(start, stop)

    if isinstance(value, int):
        value_str = '0' + str(value) if value < 10 else str(value)
        selection = coicop_nomenclature.code_coicop.str[:2] == value_str
    elif isinstance(value, str):
        selection = coicop_nomenclature.code_coicop.str[:len(value)] == value
    elif isinstance(value, list):
        selection = coicop_nomenclature.code_coicop.isin(value)

    if selection.any():  # la coicop existe
        filled_start_stop = (
            (coicop_nomenclature.loc[selection, 'start'].unique() != 0).any() or
            (coicop_nomenclature.loc[selection, 'stop'].unique() != 0).any()
            )
        if not filled_start_stop:
            coicop_nomenclature.loc[selection, 'start'] = 1994
            coicop_nomenclature.loc[selection, 'stop'] = 2014
            coicop_nomenclature.loc[selection, 'categorie_fiscale'] = categorie_fiscale

        else:
            equal_start = coicop_nomenclature.start == start
            equal_stop = coicop_nomenclature.stop == stop

            selection_bis = selection & (coicop_nomenclature.start <= start) & (coicop_nomenclature.stop >= stop)

            if (selection_bis & equal_start & equal_stop).any():  # meme intervalle
                coicop_nomenclature.loc[
                    selection_bis & equal_start & equal_stop, 'categorie_fiscale'
                    ] = categorie_fiscale

            elif (selection_bis & equal_start).any():  # recouvrement au debut
                coicop_nomenclature.loc[selection_bis & equal_start, 'start'] = stop + 1
                coicop_copy = coicop_nomenclature.loc[selection_bis & equal_start].copy()
                coicop_copy['categorie_fiscale'] = categorie_fiscale
                coicop_copy['start'] = start
                coicop_copy['stop'] = stop
                coicop_nomenclature = coicop_nomenclature.append(coicop_copy)
                coicop_nomenclature.reset_index(inplace = True, drop = True)
                coicop_nomenclature.sort_values(by = 'code_coicop', inplace = True)

            elif (selection_bis & equal_stop).any():  # recouvrement a la fin
                coicop_nomenclature.loc[selection_bis & equal_stop, 'stop'] = start - 1
                coicop_copy = coicop_nomenclature.loc[selection_bis & equal_stop].copy()
                coicop_copy['categorie_fiscale'] = categorie_fiscale
                coicop_copy['start'] = start
                coicop_copy['stop'] = stop
                coicop_nomenclature = coicop_nomenclature.append(coicop_copy)
                coicop_nomenclature.reset_index(inplace = True, drop = True)
                coicop_nomenclature.sort_values(by = 'code_coicop', inplace = True)

            else:  # recouvrement au milieu sans affecter les extermités
                coicop_copy_inf = coicop_nomenclature.loc[selection_bis].copy()
                coicop_copy_inf['stop'] = start - 1
                coicop_copy_sup = coicop_nomenclature.loc[selection_bis].copy()
                coicop_copy_sup['start'] = stop + 1

                coicop_nomenclature.loc[selection_bis, 'categorie_fiscale'] = categorie_fiscale
                coicop_nomenclature.loc[selection_bis, 'start'] = start
                coicop_nomenclature.loc[selection_bis, 'stop'] = stop
                coicop_nomenclature = coicop_nomenclature.append(coicop_copy_inf)
                coicop_nomenclature = coicop_nomenclature.append(coicop_copy_sup)
                coicop_nomenclature.reset_index(inplace = True, drop = True)
                coicop_nomenclature.sort_values(by = 'code_coicop', inplace = True)

    else:
        assert origin is not None
        assert label is not None
        infra_labels = extract_infra_labels_from_coicop_code(coicop_nomenclature, str(value), label)
        additional_row = pd.DataFrame(columns = coicop_nomenclature.columns)
        additional_dict = {
            'code_coicop': str(value),
            'categorie_fiscale': categorie_fiscale,
            'start': start,
            'stop': stop,
            'origin': origin
            }
        additional_dict.update(infra_labels)
        for item, val in additional_dict.iteritems():
            additional_row[item] = [val]

        coicop_nomenclature = coicop_nomenclature.append(additional_row)
        coicop_nomenclature.reset_index(inplace = True, drop = True)
        coicop_nomenclature.sort_values(by = 'code_coicop', inplace = True)

    # print coicop_nomenclature.categorie_fiscale.value_counts()
    return coicop_nomenclature

def build_coicop_nomenclature_with_fiscal_categories(to_csv = False):
    coicop_nomenclature = build_coicop_nomenclature.build_complete_coicop_nomenclature()
    # On  ajoute des colonnes
    # période d'effet de la législation
    coicop_nomenclature['start'] = 0
    coicop_nomenclature['stop'] = 0
    # origine de du poste
    coicop_nomenclature['origin'] = 'COICOP INSEE'
    # 01 Produits alimentaires et boissons non alcoolisées
    # ils sont tous à taux réduit
    alimentation = dict(
        value = 1,
        categorie_fiscale = 'tva_taux_reduit'
        )
    # sauf la margarine à taux plein
    margarine = dict(
        value = '01.1.5.2.2',
        categorie_fiscale = 'tva_taux_plein',
        )
    saindoux = dict(
        value = '01.1.5.2.3',
        categorie_fiscale = 'tva_taux_reduit',
        label = "Saindoux autres graisses d'origine animale",
        origin = 'TAXIPP',
        )
    # et les confiseries et le chocolat à taux plein http://bofip.impots.gouv.fr/bofip/1438-PGP.html
    confiserie = dict(
        value = ['01.1.8.1.3', '01.1.8.2.1', '01.1.8.2.2'],
        categorie_fiscale = 'tva_taux_plein'
        )
    # 02 Boissons alcoolisées et tabac
    # alccols forts
    alcools = dict(
        value = '02.1.1',
        categorie_fiscale = 'alcools_forts',
        )
    # vins et boissons fermentées
    vin = dict(
        value = '02.1.2',
        categorie_fiscale = 'vin',
        )
    # bière
    biere = dict(
        value = '02.1.3',
        categorie_fiscale = 'biere',
        )
    # tabac
    cigares = dict(
        value = '02.2.1',
        categorie_fiscale = 'cigares',
        label = 'Cigares et cigarillos',
        origin = 'TAXIPP',
        )
    cigarettes = dict(
        value = '02.2.2',
        categorie_fiscale = 'cigarettes',
        label = 'Cigarettes',
        origin = 'TAXIPP',
        )
    tabac_a_rouler = dict(
        value = '02.2.3',
        categorie_fiscale = 'tabac_a_rouler',
        label = 'Tabac a rouler',  # TODO je n'arrive aps à mettre des accents
        origin = 'TAXIPP',
        )
    stupefiants = dict(
        value = '02.3',
        categorie_fiscale = '',
        label = 'Stupefiants',
        origin = 'COICOP UN',
        )
    # 03 Habillement et chaussures
    habillement = dict(
        value = 3,
        categorie_fiscale = 'tva_taux_plein'
        )

    # 04 Logement, eau, gaz, électricité et autres combustibles
    logement = dict(
        value = 4,
        categorie_fiscale = 'tva_taux_plein',
        )
    # sauf distribution d'eau, enlèvement des ordures ménagères, assainissement, autres services liés au logement n.d.a.
    # qui sont au taux réduit de 1994 à 2011
    eau_ordures_assainissement = dict(
        value = ['04.4.1.1.1', '04.4.1.2.1', '04.4.1.3.1'],
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    # avant de passer au taux intermédiaire
    eau_ordures_assainissement_reforme_2012 = dict(
        value = ['04.4.1.1.1', '04.4.1.2.1', '04.4.1.3.1'],
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # et pas de taxation des loyers
    loyers = dict(
        value = ['04.1.1.1.1', '04.1.1.2.1'],
        categorie_fiscale = '',
        )
    # TODO ajouter loyers fictifs
    # 05 Ameublement, équipement ménager et entretien courant de la maison
    ameublement = dict(
        value = 5,
        categorie_fiscale = 'tva_taux_plein',
        )
    # sauf Services domestiques et autres services pour l'habitation
    services_domestiques = dict(
        value = '05.6.2',
        categorie_fiscale = 'tva_taux_reduit',
        )
    # 06 Santé pas taxée
    sante = dict(
        value = 6,
        categorie_fiscale = '',
        )
    # sauf pharmacie
    pharmacie = dict(
        value = '06.1.1.1',
        categorie_fiscale = 'tva_taux_super_reduit',
        )
    # parapharmacie
    parapharmacie = dict(
        value = '06.1.1.2',
        categorie_fiscale = 'tva_taux_plein',
        )
    # materiel therapeutique
    materiel_therapeutique = dict(
        value = '06.1.1.3',
        categorie_fiscale = 'tva_taux_reduit',
        )
    # 07 Transports
    transports = dict(
        value = 7,
        categorie_fiscale = 'tva_taux_plein',
        )
    # Transport combine de passagers change en 2012 #
    transport_combine_passagers = dict(
        value = '07.3.5',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    transport_combine_passagers_reforme_2012 = dict(
        value = '07.3.5',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Transport maritime et fluvial de passagers change en 2011  Attention 07.3.4 dans enquête BDF
    transport_maritime = dict(
        value = '07.3.6',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    transport_maritime_reforme_2012 = dict(
        value = '07.3.6',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Transport aérien de passagers change en 2012
    transport_aerien = dict(
        value = '07.3.3',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    transport_aerien_reforme_2012 = dict(
        value = '07.3.3',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Transport routier de passagers 1994 2011 change en 2012
    transport_routier = dict(
        value = '07.3.2',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    transport_routier_reforme_2012 = dict(
        value = '07.3.2',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Transport ferroviaire de passagers change en 2012
    transport_ferroviaire = dict(
        value = '07.3.1',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    transport_ferroviaire_reforme_2012 = dict(
        value = '07.3.1',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Carburants et lubrifiants pour véhicules de tourisme 1994 2014
    carburants_lubrifiants = dict(
        value = '07.2.2',
        categorie_fiscale = 'ticpe',
        )
    # 08 Communications
    communications = dict(
        value = 8,
        categorie_fiscale = 'tva_taux_plein',
        )
    services_postaux = dict(
        value = '08.1.1.1',
        categorie_fiscale = '',
        )
    # 09 Loisirs et cutures
    loisirs_cuture = dict(
        value = 9,
        categorie_fiscale = 'tva_taux_plein',
        )
    # Journaux et publications périodiques']
    journaux_periodiques = dict(
        value = '09.5.2',
        categorie_fiscale = 'tva_taux_super_reduit',
        )
    # Livre'
    livre = dict(
        value = '09.5.1',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    livre_reforme_2012 = dict(
        value = '09.5.1',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Jeux de hasard
    jeux_hasard = dict(
        value = '09.4.3',
        categorie_fiscale = '',
        label = 'Jeux de hasard',
        origin = 'COICOP UN',
        )
    # Services culturels
    services_culturels = dict(
        value = '09.4.2',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    # Services culturels
    services_culturels_reforme_2012 = dict(
        value = '09.4.2',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Services récréatifs et sportifs
    services_recreatifs_sportifs = dict(
        value = '09.4.1',
        categorie_fiscale = 'tva_taux_reduit',
        stop = 2011,
        )
    services_recreatifs_sportifs_reforme_2012 = dict(
        value = '09.4.1',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # 10 Education
    education = dict(
        value = 10,
        categorie_fiscale = '',
        )
    # 11 Hotellerie restauration
    hotellerie_restauration = dict(
        value = 11,
        categorie_fiscale = 'tva_taux_reduit',
        )
    # Consommation de boissons alcoolisées
    consommation_boissons_alcoolisees = dict(
        value = ['11.1.1.2.2', '11.1.1.2.3', '11.1.1.2.4'],
        categorie_fiscale = 'tva_taux_plein', # TODO sauf en corse à 10%
        )
    # Restauration sur place 1994 2009
    restauration_sur_place = dict(
        value = '11.1.1.1.1',
        categorie_fiscale = 'tva_taux_plein',
        stop = 2009,
        )
    restauration_sur_place_reforme_2010 = dict(
        value = '11.1.1.1.1',
        categorie_fiscale = 'tva_taux_reduit',
        start = 2010,
        stop = 2011,
        )
    # Restauration sur place 2012 2014
    restauration_sur_place_reforme_2012 = dict(
        value = '11.1.1.1.1',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Restauration à emporter 1994 1997
    restauration_a_emporter = dict(
        value = '11.1.1.1.2',
        categorie_fiscale = 'tva_taux_plein',
        stop = 1997,
        )
    # Restauration à emporter 2010 2011
    restauration_a_emporter_reforme_2010 = dict(
        value = '11.1.1.1.2',
        categorie_fiscale = 'tva_taux_reduit',
        start = 1998,
        stop = 2011,
        )
    # Restauration à emporter 2012 2014
    restauration_a_emporter_reforme_2012 = dict(
        value = '11.1.1.1.2',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Cantines'] 1994
    cantines = dict(
        value = '11.1.2',
        categorie_fiscale = '',
        )
    # Services d'hébergement 2012 2014
    service_hebergement = dict(
        value = '11.2.1',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # 12 Autres biens et services
    autres_biens_et_services = dict(
        value = 12,
        categorie_fiscale = 'tva_taux_plein',
        )
    # Prostitution
    prostitution = dict(
        value = '12.2',
        categorie_fiscale = '',
        label = 'Prostitution',
        origin = 'COICOP UN',
        )
    # Protection sociale TODO: check tva_taux_plein avant 2000
    protection_sociale_reforme_2000 = dict(
        value = '12.4',
        categorie_fiscale = 'tva_taux_reduit',
        start = 2000,
        stop = 2011,
        )
    # Protection sociale
    protection_sociale_reforme_2012 = dict(
        value = '12.4',
        categorie_fiscale = 'tva_taux_intermediaire',
        start = 2012,
        )
    # Autres assurances
    autres_assurances = dict(
        value = '12.5.5',
        categorie_fiscale = 'autres_assurances',
        label = 'Autres assurances',
        origin = 'COICOP UN',
        )
    # Assurance_transports
    assurance_transports = dict(
        value = '12.5.4',
        categorie_fiscale = 'assurance_transport',
        )
    # Assurance maladie
    assurance_maladie = dict(
        value = '12.5.3',
        categorie_fiscale = 'assurance_sante',
        )
    # Assurance habitation
    assurance_habitation = dict(
        value = '12.5.2',
        categorie_fiscale = 'autres_assurances',
        )
    # Assurance vie
    assurance_vie = dict(
        value = '12.5.1',
        categorie_fiscale = 'autres_assurances',
        label = 'Assurance vie',
        origin = 'COICOP UN',
        )
    # Couts des services d'intermédiation financière indirectement mesurés
    intermediation_financiere = dict(
        value = '12.6',
        categorie_fiscale = '',
        )

    for member in [
        # 01
        alimentation,
        margarine, saindoux, confiserie,
        # 02
        alcools, vin, biere,
        cigares, cigarettes, tabac_a_rouler, stupefiants,
        # 03
        habillement,
        # 04
        logement, eau_ordures_assainissement, eau_ordures_assainissement_reforme_2012, loyers,
        # 05
        ameublement, services_domestiques,
        # 06
        sante, pharmacie, parapharmacie, materiel_therapeutique,
        # 07
        transports,
        transport_combine_passagers, transport_combine_passagers_reforme_2012,
        transport_maritime, transport_maritime_reforme_2012,
        transport_aerien, transport_aerien_reforme_2012,
        transport_routier, transport_routier_reforme_2012,
        transport_ferroviaire, transport_ferroviaire_reforme_2012,
        carburants_lubrifiants,
        # 08
        communications, services_postaux,
        # 09
        loisirs_cuture, journaux_periodiques, livre, livre_reforme_2012, jeux_hasard,
        services_culturels, services_culturels_reforme_2012,
        services_recreatifs_sportifs, services_recreatifs_sportifs_reforme_2012,
        # 10 Education
        education,
        # 11 Hotellerie restauration
        hotellerie_restauration, cantines, service_hebergement,
        consommation_boissons_alcoolisees,
        restauration_sur_place, restauration_sur_place_reforme_2010, restauration_sur_place_reforme_2012,
        restauration_a_emporter, restauration_a_emporter_reforme_2010, restauration_a_emporter_reforme_2012,
        # 12
        autres_biens_et_services,
        protection_sociale_reforme_2000, protection_sociale_reforme_2012,
        prostitution,
        intermediation_financiere,
        autres_assurances, assurance_transports, assurance_vie, assurance_maladie, assurance_habitation,
        ]:
        coicop_nomenclature = apply_modification(coicop_nomenclature, **member)

    coicop_legislation = coicop_nomenclature.copy()
    if to_csv:
        coicop_legislation.to_csv(
            os.path.join(legislation_directory, 'coicop_legislation.csv'),
            )
    return coicop_legislation.copy()

def get_categorie_fiscale(value, year = None, assertion_error = True):
    coicop_nomenclature = pd.read_csv(
        os.path.join(legislation_directory, 'coicop_legislation.csv'),
        converters = {'posteCOICOP': unicode}
        )
    build_coicop_nomenclature_with_fiscal_categories()
    if isinstance(value, int):
        value_str = '0' + str(value) if value < 10 else str(value)
        selection = coicop_nomenclature.code_coicop.str[:2] == value_str
    elif isinstance(value, str):
        selection = coicop_nomenclature.code_coicop.str[:len(value)] == value
    elif isinstance(value, list):
        selection = coicop_nomenclature.code_coicop.isin(value)

    if year is not None:
        selection = selection & (coicop_nomenclature.start <= year) & (year <= coicop_nomenclature.stop)

    categorie_fiscale = coicop_nomenclature.loc[selection, 'categorie_fiscale'].unique()
    if assertion_error:
        assert len(categorie_fiscale) == 1, 'Ther categorie fiscale is not unique. Candidates are: {}'.format(
            categorie_fiscale)
        return categorie_fiscale[0]
    else:
        return categorie_fiscale

def test_coicop_legislation():
    coicop_nomenclature = build_coicop_nomenclature_with_fiscal_categories(to_csv = True)
    if coicop_nomenclature.categorie_fiscale.isnull().any():
        return coicop_nomenclature.loc[coicop_nomenclature.categorie_fiscale.isnull()]

if __name__ == "__main__":
    # extract_informations_from_coicop_to_categorie_fiscale()
    coicop_nomenclature = build_coicop_nomenclature_with_fiscal_categories(to_csv = True)
    # print test_coicop_legislation(coicop_nomenclature)
    # TODO créer des sous-catégories pour tabac

    # print get_categorie_fiscale('11.1.1.1.1', year = 2010)


"""
=================================
Gaussian Mixture Model Ellipsoids
=================================

Plot the confidence ellipsoids of a mixture of two Gaussians
obtained with Expectation Maximisation (``GaussianMixture`` class) and
Variational Inference (``BayesianGaussianMixture`` class models with
a Dirichlet process prior).

Both models have access to five components with which to fit the data. Note
that the Expectation Maximisation model will necessarily use all five
components while the Variational Inference model will effectively only use as
many as are needed for a good fit. Here we can see that the Expectation
Maximisation model splits some components arbitrarily, because it is trying to
fit too many components, while the Dirichlet Process model adapts it number of
state automatically.

This example doesn't show it, as we're in a low-dimensional space, but
another advantage of the Dirichlet process model is that it can fit
full covariance matrices effectively even when there are less examples
per cluster than there are dimensions in the data, due to
regularization properties of the inference algorithm.
"""

import itertools

import numpy as np
from scipy import linalg
import matplotlib.pyplot as plt
import matplotlib as mpl

from sklearn import mixture

color_iter = itertools.cycle(['navy', 'c', 'cornflowerblue', 'gold',
                              'darkorange'])

def plot_results(X, Y_, means, covariances, index, title):
    splot = plt.subplot(2, 1, 1 + index)
    for i, (mean, covar, color) in enumerate(zip(
            means, covariances, color_iter)):
        v, w = linalg.eigh(covar)
        v = 2. * np.sqrt(2.) * np.sqrt(v)
        u = w[0] / linalg.norm(w[0])
        # as the DP will not use every component it has access to
        # unless it needs it, we shouldn't plot the redundant
        # components.
        if not np.any(Y_ == i):
            continue
        plt.scatter(X[Y_ == i, 0], X[Y_ == i, 1], .8, color=color)

        # Plot an ellipse to show the Gaussian component
        angle = np.arctan(u[1] / u[0])
        angle = 180. * angle / np.pi  # convert to degrees
        ell = mpl.patches.Ellipse(mean, v[0], v[1], 180. + angle, color=color)
        ell.set_clip_box(splot.bbox)
        ell.set_alpha(0.5)
        splot.add_artist(ell)

    plt.xlim(-9., 5.)
    plt.ylim(-3., 6.)
    plt.xticks(())
    plt.yticks(())
    plt.title(title)

# Number of samples per component
n_samples = 500

# Generate random sample, two components
np.random.seed(0)
C = np.array([[0., -0.1], [1.7, .4]])
X = np.r_[np.dot(np.random.randn(n_samples, 2), C),
          .7 * np.random.randn(n_samples, 2) + np.array([-6, 3])]

# Fit a Gaussian mixture with EM using five components
gmm = mixture.GaussianMixture(n_components=5, covariance_type='full').fit(X)
plot_results(X, gmm.predict(X), gmm.means_, gmm.covariances_, 0,
             'Gaussian Mixture')

# Fit a Dirichlet process Gaussian mixture using five components
dpgmm = mixture.BayesianGaussianMixture(n_components=5,
                                        covariance_type='full').fit(X)
plot_results(X, dpgmm.predict(X), dpgmm.means_, dpgmm.covariances_, 1,
             'Bayesian Gaussian Mixture with a Dirichlet process prior')

plt.show()


import numpy as np
import matplotlib.pyplot as plt
from grid_world import standard_grid, negative_grid
from iterative_policy_evaluation import print_values, print_policy

# NOTE: this is only policy evaluation, not optimization

# we'll try to obtain the same result as our other MC script
from monte_carlo_random import random_action, play_game, SMALL_ENOUGH, GAMMA, ALL_POSSIBLE_ACTIONS

LEARNING_RATE = 0.001

if __name__ == '__main__':
  # use the standard grid again (0 for every step) so that we can compare
  # to iterative policy evaluation
  grid = standard_grid()

  # print rewards
  print("rewards:")
  print_values(grid.rewards, grid)

  # state -> action
  # found by policy_iteration_random on standard_grid
  # MC method won't get exactly this, but should be close
  # values:
  # ---------------------------
  #  0.43|  0.56|  0.72|  0.00|
  # ---------------------------
  #  0.33|  0.00|  0.21|  0.00|
  # ---------------------------
  #  0.25|  0.18|  0.11| -0.17|
  # policy:
  # ---------------------------
  #   R  |   R  |   R  |      |
  # ---------------------------
  #   U  |      |   U  |      |
  # ---------------------------
  #   U  |   L  |   U  |   L  |
  policy = {
    (2, 0): 'U',
    (1, 0): 'U',
    (0, 0): 'R',
    (0, 1): 'R',
    (0, 2): 'R',
    (1, 2): 'U',
    (2, 1): 'L',
    (2, 2): 'U',
    (2, 3): 'L',
  }

  # initialize theta
  # our model is V_hat = theta.dot(x)
  # where x = [row, col, row*col, 1] - 1 for bias term
  theta = np.random.randn(4) / 2
  def s2x(s):
    return np.array([s[0] - 1, s[1] - 1.5, s[0]*s[1] - 3, 1])

  # repeat until convergence
  deltas = []
  t = 1.0
  for it in range(20000):
    if it % 100 == 0:
      t += 0.01
    alpha = LEARNING_RATE/t
    # generate an episode using pi
    biggest_change = 0
    states_and_returns = play_game(grid, policy)
    seen_states = set()
    for s, G in states_and_returns:
      # check if we have already seen s
      # called "first-visit" MC policy evaluation
      if s not in seen_states:
        old_theta = theta.copy()
        x = s2x(s)
        V_hat = theta.dot(x)
        # grad(V_hat) wrt theta = x
        theta += alpha*(G - V_hat)*x
        biggest_change = max(biggest_change, np.abs(old_theta - theta).sum())
        seen_states.add(s)
    deltas.append(biggest_change)

  plt.plot(deltas)
  plt.show()

  # obtain predicted values
  V = {}
  states = grid.all_states()
  for s in states:
    if s in grid.actions:
      V[s] = theta.dot(s2x(s))
    else:
      # terminal state or state we can't otherwise get to
      V[s] = 0

  print("values:")
  print_values(V, grid)
  print("policy:")
  print_policy(policy, grid)


#!/usr/bin/env python

"""Top level ``eval`` module.
"""

import warnings
import tokenize
from pandas.io.formats.printing import pprint_thing
from pandas.core.computation import _NUMEXPR_INSTALLED
from pandas.core.computation.expr import Expr, _parsers, tokenize_string
from pandas.core.computation.scope import _ensure_scope
from pandas.compat import string_types
from pandas.core.computation.engines import _engines
from pandas.util._validators import validate_bool_kwarg

def _check_engine(engine):
    """Make sure a valid engine is passed.

    Parameters
    ----------
    engine : str

    Raises
    ------
    KeyError
      * If an invalid engine is passed
    ImportError
      * If numexpr was requested but doesn't exist

    Returns
    -------
    string engine

    """

    if engine is None:
        if _NUMEXPR_INSTALLED:
            engine = 'numexpr'
        else:
            engine = 'python'

    if engine not in _engines:
        raise KeyError('Invalid engine {0!r} passed, valid engines are'
                       ' {1}'.format(engine, list(_engines.keys())))

    # TODO: validate this in a more general way (thinking of future engines
    # that won't necessarily be import-able)
    # Could potentially be done on engine instantiation
    if engine == 'numexpr':
        if not _NUMEXPR_INSTALLED:
            raise ImportError("'numexpr' is not installed or an "
                              "unsupported version. Cannot use "
                              "engine='numexpr' for query/eval "
                              "if 'numexpr' is not installed")

    return engine

def _check_parser(parser):
    """Make sure a valid parser is passed.

    Parameters
    ----------
    parser : str

    Raises
    ------
    KeyError
      * If an invalid parser is passed
    """
    if parser not in _parsers:
        raise KeyError('Invalid parser {0!r} passed, valid parsers are'
                       ' {1}'.format(parser, _parsers.keys()))

def _check_resolvers(resolvers):
    if resolvers is not None:
        for resolver in resolvers:
            if not hasattr(resolver, '__getitem__'):
                name = type(resolver).__name__
                raise TypeError('Resolver of type %r does not implement '
                                'the __getitem__ method' % name)

def _check_expression(expr):
    """Make sure an expression is not an empty string

    Parameters
    ----------
    expr : object
        An object that can be converted to a string

    Raises
    ------
    ValueError
      * If expr is an empty string
    """
    if not expr:
        raise ValueError("expr cannot be an empty string")

def _convert_expression(expr):
    """Convert an object to an expression.

    Thus function converts an object to an expression (a unicode string) and
    checks to make sure it isn't empty after conversion. This is used to
    convert operators to their string representation for recursive calls to
    :func:`~pandas.eval`.

    Parameters
    ----------
    expr : object
        The object to be converted to a string.

    Returns
    -------
    s : unicode
        The string representation of an object.

    Raises
    ------
    ValueError
      * If the expression is empty.
    """
    s = pprint_thing(expr)
    _check_expression(s)
    return s

def _check_for_locals(expr, stack_level, parser):
    at_top_of_stack = stack_level == 0
    not_pandas_parser = parser != 'pandas'

    if not_pandas_parser:
        msg = "The '@' prefix is only supported by the pandas parser"
    elif at_top_of_stack:
        msg = ("The '@' prefix is not allowed in "
               "top-level eval calls, \nplease refer to "
               "your variables by name without the '@' "
               "prefix")

    if at_top_of_stack or not_pandas_parser:
        for toknum, tokval in tokenize_string(expr):
            if toknum == tokenize.OP and tokval == '@':
                raise SyntaxError(msg)

def eval(expr, parser='pandas', engine=None, truediv=True,
         local_dict=None, global_dict=None, resolvers=(), level=0,
         target=None, inplace=None):
    """Evaluate a Python expression as a string using various backends.

    The following arithmetic operations are supported: ``+``, ``-``, ``*``,
    ``/``, ``**``, ``%``, ``//`` (python engine only) along with the following
    boolean operations: ``|`` (or), ``&`` (and), and ``~`` (not).
    Additionally, the ``'pandas'`` parser allows the use of :keyword:`and`,
    :keyword:`or`, and :keyword:`not` with the same semantics as the
    corresponding bitwise operators.  :class:`~pandas.Series` and
    :class:`~pandas.DataFrame` objects are supported and behave as they would
    with plain ol' Python evaluation.

    Parameters
    ----------
    expr : str or unicode
        The expression to evaluate. This string cannot contain any Python
        `statements
        <http://docs.python.org/2/reference/simple_stmts.html#simple-statements>`__,
        only Python `expressions
        <http://docs.python.org/2/reference/simple_stmts.html#expression-statements>`__.
    parser : string, default 'pandas', {'pandas', 'python'}
        The parser to use to construct the syntax tree from the expression. The
        default of ``'pandas'`` parses code slightly different than standard
        Python. Alternatively, you can parse an expression using the
        ``'python'`` parser to retain strict Python semantics.  See the
        :ref:`enhancing performance <enhancingperf.eval>` documentation for
        more details.
    engine : string or None, default 'numexpr', {'python', 'numexpr'}

        The engine used to evaluate the expression. Supported engines are

        - None         : tries to use ``numexpr``, falls back to ``python``
        - ``'numexpr'``: This default engine evaluates pandas objects using
                         numexpr for large speed ups in complex expressions
                         with large frames.
        - ``'python'``: Performs operations as if you had ``eval``'d in top
                        level python. This engine is generally not that useful.

        More backends may be available in the future.

    truediv : bool, optional
        Whether to use true division, like in Python >= 3
    local_dict : dict or None, optional
        A dictionary of local variables, taken from locals() by default.
    global_dict : dict or None, optional
        A dictionary of global variables, taken from globals() by default.
    resolvers : list of dict-like or None, optional
        A list of objects implementing the ``__getitem__`` special method that
        you can use to inject an additional collection of namespaces to use for
        variable lookup. For example, this is used in the
        :meth:`~pandas.DataFrame.query` method to inject the
        :attr:`~pandas.DataFrame.index` and :attr:`~pandas.DataFrame.columns`
        variables that refer to their respective :class:`~pandas.DataFrame`
        instance attributes.
    level : int, optional
        The number of prior stack frames to traverse and add to the current
        scope. Most users will **not** need to change this parameter.
    target : a target object for assignment, optional, default is None
        essentially this is a passed in resolver
    inplace : bool, default True
        If expression mutates, whether to modify object inplace or return
        copy with mutation.

        WARNING: inplace=None currently falls back to to True, but
        in a future version, will default to False.  Use inplace=True
        explicitly rather than relying on the default.

    Returns
    -------
    ndarray, numeric scalar, DataFrame, Series

    Notes
    -----
    The ``dtype`` of any objects involved in an arithmetic ``%`` operation are
    recursively cast to ``float64``.

    See the :ref:`enhancing performance <enhancingperf.eval>` documentation for
    more details.

    See Also
    --------
    pandas.DataFrame.query
    pandas.DataFrame.eval
    """
    inplace = validate_bool_kwarg(inplace, 'inplace')
    first_expr = True
    if isinstance(expr, string_types):
        _check_expression(expr)
        exprs = [e.strip() for e in expr.splitlines() if e.strip() != '']
    else:
        exprs = [expr]
    multi_line = len(exprs) > 1

    if multi_line and target is None:
        raise ValueError("multi-line expressions are only valid in the "
                         "context of data, use DataFrame.eval")

    first_expr = True
    for expr in exprs:
        expr = _convert_expression(expr)
        engine = _check_engine(engine)
        _check_parser(parser)
        _check_resolvers(resolvers)
        _check_for_locals(expr, level, parser)

        # get our (possibly passed-in) scope
        env = _ensure_scope(level + 1, global_dict=global_dict,
                            local_dict=local_dict, resolvers=resolvers,
                            target=target)

        parsed_expr = Expr(expr, engine=engine, parser=parser, env=env,
                           truediv=truediv)

        # construct the engine and evaluate the parsed expression
        eng = _engines[engine]
        eng_inst = eng(parsed_expr)
        ret = eng_inst.evaluate()

        if parsed_expr.assigner is None and multi_line:
            raise ValueError("Multi-line expressions are only valid"
                             " if all expressions contain an assignment")

        # assign if needed
        if env.target is not None and parsed_expr.assigner is not None:
            if inplace is None:
                warnings.warn(
                    "eval expressions containing an assignment currently"
                    "default to operating inplace.\nThis will change in "
                    "a future version of pandas, use inplace=True to "
                    "avoid this warning.",
                    FutureWarning, stacklevel=3)
                inplace = True

            # if returning a copy, copy only on the first assignment
            if not inplace and first_expr:
                target = env.target.copy()
            else:
                target = env.target

            target[parsed_expr.assigner] = ret

            if not resolvers:
                resolvers = ({parsed_expr.assigner: ret},)
            else:
                # existing resolver needs updated to handle
                # case of mutating existing column in copy
                for resolver in resolvers:
                    if parsed_expr.assigner in resolver:
                        resolver[parsed_expr.assigner] = ret
                        break
                else:
                    resolvers += ({parsed_expr.assigner: ret},)

            ret = None
            first_expr = False

    if not inplace and inplace is not None:
        return target

    return ret


import numpy as np
import pandas as pd
import pandas.util.testing as tm

class TestSparseArrayArithmetics(object):

    _base = np.array
    _klass = pd.SparseArray

    def _assert(self, a, b):
        tm.assert_numpy_array_equal(a, b)

    def _check_numeric_ops(self, a, b, a_dense, b_dense):
        with np.errstate(invalid='ignore', divide='ignore'):
            # Unfortunately, trying to wrap the computation of each expected
            # value is with np.errstate() is too tedious.

            # sparse & sparse
            self._assert((a + b).to_dense(), a_dense + b_dense)
            self._assert((b + a).to_dense(), b_dense + a_dense)

            self._assert((a - b).to_dense(), a_dense - b_dense)
            self._assert((b - a).to_dense(), b_dense - a_dense)

            self._assert((a * b).to_dense(), a_dense * b_dense)
            self._assert((b * a).to_dense(), b_dense * a_dense)

            # pandas uses future division
            self._assert((a / b).to_dense(), a_dense * 1.0 / b_dense)
            self._assert((b / a).to_dense(), b_dense * 1.0 / a_dense)

            # ToDo: FIXME in GH 13843
            if not (self._base == pd.Series and a.dtype == 'int64'):
                self._assert((a // b).to_dense(), a_dense // b_dense)
                self._assert((b // a).to_dense(), b_dense // a_dense)

            self._assert((a % b).to_dense(), a_dense % b_dense)
            self._assert((b % a).to_dense(), b_dense % a_dense)

            self._assert((a ** b).to_dense(), a_dense ** b_dense)
            self._assert((b ** a).to_dense(), b_dense ** a_dense)

            # sparse & dense
            self._assert((a + b_dense).to_dense(), a_dense + b_dense)
            self._assert((b_dense + a).to_dense(), b_dense + a_dense)

            self._assert((a - b_dense).to_dense(), a_dense - b_dense)
            self._assert((b_dense - a).to_dense(), b_dense - a_dense)

            self._assert((a * b_dense).to_dense(), a_dense * b_dense)
            self._assert((b_dense * a).to_dense(), b_dense * a_dense)

            # pandas uses future division
            self._assert((a / b_dense).to_dense(), a_dense * 1.0 / b_dense)
            self._assert((b_dense / a).to_dense(), b_dense * 1.0 / a_dense)

            # ToDo: FIXME in GH 13843
            if not (self._base == pd.Series and a.dtype == 'int64'):
                self._assert((a // b_dense).to_dense(), a_dense // b_dense)
                self._assert((b_dense // a).to_dense(), b_dense // a_dense)

            self._assert((a % b_dense).to_dense(), a_dense % b_dense)
            self._assert((b_dense % a).to_dense(), b_dense % a_dense)

            self._assert((a ** b_dense).to_dense(), a_dense ** b_dense)
            self._assert((b_dense ** a).to_dense(), b_dense ** a_dense)

    def _check_bool_result(self, res):
        assert isinstance(res, self._klass)
        assert res.dtype == np.bool
        assert isinstance(res.fill_value, bool)

    def _check_comparison_ops(self, a, b, a_dense, b_dense):
        with np.errstate(invalid='ignore'):
            # Unfortunately, trying to wrap the computation of each expected
            # value is with np.errstate() is too tedious.
            #
            # sparse & sparse
            self._check_bool_result(a == b)
            self._assert((a == b).to_dense(), a_dense == b_dense)

            self._check_bool_result(a != b)
            self._assert((a != b).to_dense(), a_dense != b_dense)

            self._check_bool_result(a >= b)
            self._assert((a >= b).to_dense(), a_dense >= b_dense)

            self._check_bool_result(a <= b)
            self._assert((a <= b).to_dense(), a_dense <= b_dense)

            self._check_bool_result(a > b)
            self._assert((a > b).to_dense(), a_dense > b_dense)

            self._check_bool_result(a < b)
            self._assert((a < b).to_dense(), a_dense < b_dense)

            # sparse & dense
            self._check_bool_result(a == b_dense)
            self._assert((a == b_dense).to_dense(), a_dense == b_dense)

            self._check_bool_result(a != b_dense)
            self._assert((a != b_dense).to_dense(), a_dense != b_dense)

            self._check_bool_result(a >= b_dense)
            self._assert((a >= b_dense).to_dense(), a_dense >= b_dense)

            self._check_bool_result(a <= b_dense)
            self._assert((a <= b_dense).to_dense(), a_dense <= b_dense)

            self._check_bool_result(a > b_dense)
            self._assert((a > b_dense).to_dense(), a_dense > b_dense)

            self._check_bool_result(a < b_dense)
            self._assert((a < b_dense).to_dense(), a_dense < b_dense)

    def _check_logical_ops(self, a, b, a_dense, b_dense):
        # sparse & sparse
        self._check_bool_result(a & b)
        self._assert((a & b).to_dense(), a_dense & b_dense)

        self._check_bool_result(a | b)
        self._assert((a | b).to_dense(), a_dense | b_dense)
        # sparse & dense
        self._check_bool_result(a & b_dense)
        self._assert((a & b_dense).to_dense(), a_dense & b_dense)

        self._check_bool_result(a | b_dense)
        self._assert((a | b_dense).to_dense(), a_dense | b_dense)

    def test_float_scalar(self):
        values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])

        for kind in ['integer', 'block']:
            a = self._klass(values, kind=kind)
            self._check_numeric_ops(a, 1, values, 1)
            self._check_numeric_ops(a, 0, values, 0)
            self._check_numeric_ops(a, 3, values, 3)

            a = self._klass(values, kind=kind, fill_value=0)
            self._check_numeric_ops(a, 1, values, 1)
            self._check_numeric_ops(a, 0, values, 0)
            self._check_numeric_ops(a, 3, values, 3)

            a = self._klass(values, kind=kind, fill_value=2)
            self._check_numeric_ops(a, 1, values, 1)
            self._check_numeric_ops(a, 0, values, 0)
            self._check_numeric_ops(a, 3, values, 3)

    def test_float_scalar_comparison(self):
        values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])

        for kind in ['integer', 'block']:
            a = self._klass(values, kind=kind)
            self._check_comparison_ops(a, 1, values, 1)
            self._check_comparison_ops(a, 0, values, 0)
            self._check_comparison_ops(a, 3, values, 3)

            a = self._klass(values, kind=kind, fill_value=0)
            self._check_comparison_ops(a, 1, values, 1)
            self._check_comparison_ops(a, 0, values, 0)
            self._check_comparison_ops(a, 3, values, 3)

            a = self._klass(values, kind=kind, fill_value=2)
            self._check_comparison_ops(a, 1, values, 1)
            self._check_comparison_ops(a, 0, values, 0)
            self._check_comparison_ops(a, 3, values, 3)

    def test_float_same_index(self):
        # when sp_index are the same
        for kind in ['integer', 'block']:
            values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
            rvalues = self._base([np.nan, 2, 3, 4, np.nan, 0, 1, 3, 2, np.nan])

            a = self._klass(values, kind=kind)
            b = self._klass(rvalues, kind=kind)
            self._check_numeric_ops(a, b, values, rvalues)

            values = self._base([0., 1., 2., 6., 0., 0., 1., 2., 1., 0.])
            rvalues = self._base([0., 2., 3., 4., 0., 0., 1., 3., 2., 0.])

            a = self._klass(values, kind=kind, fill_value=0)
            b = self._klass(rvalues, kind=kind, fill_value=0)
            self._check_numeric_ops(a, b, values, rvalues)

    def test_float_same_index_comparison(self):
        # when sp_index are the same
        for kind in ['integer', 'block']:
            values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
            rvalues = self._base([np.nan, 2, 3, 4, np.nan, 0, 1, 3, 2, np.nan])

            a = self._klass(values, kind=kind)
            b = self._klass(rvalues, kind=kind)
            self._check_comparison_ops(a, b, values, rvalues)

            values = self._base([0., 1., 2., 6., 0., 0., 1., 2., 1., 0.])
            rvalues = self._base([0., 2., 3., 4., 0., 0., 1., 3., 2., 0.])

            a = self._klass(values, kind=kind, fill_value=0)
            b = self._klass(rvalues, kind=kind, fill_value=0)
            self._check_comparison_ops(a, b, values, rvalues)

    def test_float_array(self):
        values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
        rvalues = self._base([2, np.nan, 2, 3, np.nan, 0, 1, 5, 2, np.nan])

        for kind in ['integer', 'block']:
            a = self._klass(values, kind=kind)
            b = self._klass(rvalues, kind=kind)
            self._check_numeric_ops(a, b, values, rvalues)
            self._check_numeric_ops(a, b * 0, values, rvalues * 0)

            a = self._klass(values, kind=kind, fill_value=0)
            b = self._klass(rvalues, kind=kind)
            self._check_numeric_ops(a, b, values, rvalues)

            a = self._klass(values, kind=kind, fill_value=0)
            b = self._klass(rvalues, kind=kind, fill_value=0)
            self._check_numeric_ops(a, b, values, rvalues)

            a = self._klass(values, kind=kind, fill_value=1)
            b = self._klass(rvalues, kind=kind, fill_value=2)
            self._check_numeric_ops(a, b, values, rvalues)

    def test_float_array_different_kind(self):
        values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
        rvalues = self._base([2, np.nan, 2, 3, np.nan, 0, 1, 5, 2, np.nan])

        a = self._klass(values, kind='integer')
        b = self._klass(rvalues, kind='block')
        self._check_numeric_ops(a, b, values, rvalues)
        self._check_numeric_ops(a, b * 0, values, rvalues * 0)

        a = self._klass(values, kind='integer', fill_value=0)
        b = self._klass(rvalues, kind='block')
        self._check_numeric_ops(a, b, values, rvalues)

        a = self._klass(values, kind='integer', fill_value=0)
        b = self._klass(rvalues, kind='block', fill_value=0)
        self._check_numeric_ops(a, b, values, rvalues)

        a = self._klass(values, kind='integer', fill_value=1)
        b = self._klass(rvalues, kind='block', fill_value=2)
        self._check_numeric_ops(a, b, values, rvalues)

    def test_float_array_comparison(self):
        values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
        rvalues = self._base([2, np.nan, 2, 3, np.nan, 0, 1, 5, 2, np.nan])

        for kind in ['integer', 'block']:
            a = self._klass(values, kind=kind)
            b = self._klass(rvalues, kind=kind)
            self._check_comparison_ops(a, b, values, rvalues)
            self._check_comparison_ops(a, b * 0, values, rvalues * 0)

            a = self._klass(values, kind=kind, fill_value=0)
            b = self._klass(rvalues, kind=kind)
            self._check_comparison_ops(a, b, values, rvalues)

            a = self._klass(values, kind=kind, fill_value=0)
            b = self._klass(rvalues, kind=kind, fill_value=0)
            self._check_comparison_ops(a, b, values, rvalues)

            a = self._klass(values, kind=kind, fill_value=1)
            b = self._klass(rvalues, kind=kind, fill_value=2)
            self._check_comparison_ops(a, b, values, rvalues)

    def test_int_array(self):
        # have to specify dtype explicitly until fixing GH 667
        dtype = np.int64

        values = self._base([0, 1, 2, 0, 0, 0, 1, 2, 1, 0], dtype=dtype)
        rvalues = self._base([2, 0, 2, 3, 0, 0, 1, 5, 2, 0], dtype=dtype)

        for kind in ['integer', 'block']:
            a = self._klass(values, dtype=dtype, kind=kind)
            assert a.dtype == dtype
            b = self._klass(rvalues, dtype=dtype, kind=kind)
            assert b.dtype == dtype

            self._check_numeric_ops(a, b, values, rvalues)
            self._check_numeric_ops(a, b * 0, values, rvalues * 0)

            a = self._klass(values, fill_value=0, dtype=dtype, kind=kind)
            assert a.dtype == dtype
            b = self._klass(rvalues, dtype=dtype, kind=kind)
            assert b.dtype == dtype

            self._check_numeric_ops(a, b, values, rvalues)

            a = self._klass(values, fill_value=0, dtype=dtype, kind=kind)
            assert a.dtype == dtype
            b = self._klass(rvalues, fill_value=0, dtype=dtype, kind=kind)
            assert b.dtype == dtype
            self._check_numeric_ops(a, b, values, rvalues)

            a = self._klass(values, fill_value=1, dtype=dtype, kind=kind)
            assert a.dtype == dtype
            b = self._klass(rvalues, fill_value=2, dtype=dtype, kind=kind)
            assert b.dtype == dtype
            self._check_numeric_ops(a, b, values, rvalues)

    def test_int_array_comparison(self):

        # int32 NI ATM
        for dtype in ['int64']:
            values = self._base([0, 1, 2, 0, 0, 0, 1, 2, 1, 0], dtype=dtype)
            rvalues = self._base([2, 0, 2, 3, 0, 0, 1, 5, 2, 0], dtype=dtype)

            for kind in ['integer', 'block']:
                a = self._klass(values, dtype=dtype, kind=kind)
                b = self._klass(rvalues, dtype=dtype, kind=kind)
                self._check_comparison_ops(a, b, values, rvalues)
                self._check_comparison_ops(a, b * 0, values, rvalues * 0)

                a = self._klass(values, dtype=dtype, kind=kind, fill_value=0)
                b = self._klass(rvalues, dtype=dtype, kind=kind)
                self._check_comparison_ops(a, b, values, rvalues)

                a = self._klass(values, dtype=dtype, kind=kind, fill_value=0)
                b = self._klass(rvalues, dtype=dtype, kind=kind, fill_value=0)
                self._check_comparison_ops(a, b, values, rvalues)

                a = self._klass(values, dtype=dtype, kind=kind, fill_value=1)
                b = self._klass(rvalues, dtype=dtype, kind=kind, fill_value=2)
                self._check_comparison_ops(a, b, values, rvalues)

    def test_bool_same_index(self):
        # GH 14000
        # when sp_index are the same
        for kind in ['integer', 'block']:
            values = self._base([True, False, True, True], dtype=np.bool)
            rvalues = self._base([True, False, True, True], dtype=np.bool)

            for fill_value in [True, False, np.nan]:
                a = self._klass(values, kind=kind, dtype=np.bool,
                                fill_value=fill_value)
                b = self._klass(rvalues, kind=kind, dtype=np.bool,
                                fill_value=fill_value)
                self._check_logical_ops(a, b, values, rvalues)

    def test_bool_array_logical(self):
        # GH 14000
        # when sp_index are the same
        for kind in ['integer', 'block']:
            values = self._base([True, False, True, False, True, True],
                                dtype=np.bool)
            rvalues = self._base([True, False, False, True, False, True],
                                 dtype=np.bool)

            for fill_value in [True, False, np.nan]:
                a = self._klass(values, kind=kind, dtype=np.bool,
                                fill_value=fill_value)
                b = self._klass(rvalues, kind=kind, dtype=np.bool,
                                fill_value=fill_value)
                self._check_logical_ops(a, b, values, rvalues)

    def test_mixed_array_float_int(self):

        for rdtype in ['int64']:
            values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
            rvalues = self._base([2, 0, 2, 3, 0, 0, 1, 5, 2, 0], dtype=rdtype)

            for kind in ['integer', 'block']:
                a = self._klass(values, kind=kind)
                b = self._klass(rvalues, kind=kind)
                assert b.dtype == rdtype

                self._check_numeric_ops(a, b, values, rvalues)
                self._check_numeric_ops(a, b * 0, values, rvalues * 0)

                a = self._klass(values, kind=kind, fill_value=0)
                b = self._klass(rvalues, kind=kind)
                assert b.dtype == rdtype
                self._check_numeric_ops(a, b, values, rvalues)

                a = self._klass(values, kind=kind, fill_value=0)
                b = self._klass(rvalues, kind=kind, fill_value=0)
                assert b.dtype == rdtype
                self._check_numeric_ops(a, b, values, rvalues)

                a = self._klass(values, kind=kind, fill_value=1)
                b = self._klass(rvalues, kind=kind, fill_value=2)
                assert b.dtype == rdtype
                self._check_numeric_ops(a, b, values, rvalues)

    def test_mixed_array_comparison(self):

        # int32 NI ATM
        for rdtype in ['int64']:
            values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
            rvalues = self._base([2, 0, 2, 3, 0, 0, 1, 5, 2, 0], dtype=rdtype)

            for kind in ['integer', 'block']:
                a = self._klass(values, kind=kind)
                b = self._klass(rvalues, kind=kind)
                assert b.dtype == rdtype

                self._check_comparison_ops(a, b, values, rvalues)
                self._check_comparison_ops(a, b * 0, values, rvalues * 0)

                a = self._klass(values, kind=kind, fill_value=0)
                b = self._klass(rvalues, kind=kind)
                assert b.dtype == rdtype
                self._check_comparison_ops(a, b, values, rvalues)

                a = self._klass(values, kind=kind, fill_value=0)
                b = self._klass(rvalues, kind=kind, fill_value=0)
                assert b.dtype == rdtype
                self._check_comparison_ops(a, b, values, rvalues)

                a = self._klass(values, kind=kind, fill_value=1)
                b = self._klass(rvalues, kind=kind, fill_value=2)
                assert b.dtype == rdtype
                self._check_comparison_ops(a, b, values, rvalues)

class TestSparseSeriesArithmetic(TestSparseArrayArithmetics):

    _base = pd.Series
    _klass = pd.SparseSeries

    def _assert(self, a, b):
        tm.assert_series_equal(a, b)

    def test_alignment(self):
        da = pd.Series(np.arange(4))
        db = pd.Series(np.arange(4), index=[1, 2, 3, 4])

        sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=0)
        sb = pd.SparseSeries(np.arange(4), index=[1, 2, 3, 4],
                             dtype=np.int64, fill_value=0)
        self._check_numeric_ops(sa, sb, da, db)

        sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=np.nan)
        sb = pd.SparseSeries(np.arange(4), index=[1, 2, 3, 4],
                             dtype=np.int64, fill_value=np.nan)
        self._check_numeric_ops(sa, sb, da, db)

        da = pd.Series(np.arange(4))
        db = pd.Series(np.arange(4), index=[10, 11, 12, 13])

        sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=0)
        sb = pd.SparseSeries(np.arange(4), index=[10, 11, 12, 13],
                             dtype=np.int64, fill_value=0)
        self._check_numeric_ops(sa, sb, da, db)

        sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=np.nan)
        sb = pd.SparseSeries(np.arange(4), index=[10, 11, 12, 13],
                             dtype=np.int64, fill_value=np.nan)
        self._check_numeric_ops(sa, sb, da, db)


"""
=================================
Compare BIRCH and MiniBatchKMeans
=================================

This example compares the timing of Birch (with and without the global
clustering step) and MiniBatchKMeans on a synthetic dataset having
100,000 samples and 2 features generated using make_blobs.

If ``n_clusters`` is set to None, the data is reduced from 100,000
samples to a set of 158 clusters. This can be viewed as a preprocessing
step before the final (global) clustering step that further reduces these
158 clusters to 100 clusters.
"""

# Authors: Manoj Kumar <manojkumarsivaraj334@gmail.com
#          Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
# License: BSD 3 clause

print(__doc__)

from itertools import cycle
from time import time
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors

from sklearn.preprocessing import StandardScaler
from sklearn.cluster import Birch, MiniBatchKMeans
from sklearn.datasets.samples_generator import make_blobs

# Generate centers for the blobs so that it forms a 10 X 10 grid.
xx = np.linspace(-22, 22, 10)
yy = np.linspace(-22, 22, 10)
xx, yy = np.meshgrid(xx, yy)
n_centres = np.hstack((np.ravel(xx)[:, np.newaxis],
                       np.ravel(yy)[:, np.newaxis]))

# Generate blobs to do a comparison between MiniBatchKMeans and Birch.
X, y = make_blobs(n_samples=100000, centers=n_centres, random_state=0)

# Use all colors that matplotlib provides by default.
colors_ = cycle(colors.cnames.keys())

fig = plt.figure(figsize=(12, 4))
fig.subplots_adjust(left=0.04, right=0.98, bottom=0.1, top=0.9)

# Compute clustering with Birch with and without the final clustering step
# and plot.
birch_models = [Birch(threshold=1.7, n_clusters=None),
                Birch(threshold=1.7, n_clusters=100)]
final_step = ['without global clustering', 'with global clustering']

for ind, (birch_model, info) in enumerate(zip(birch_models, final_step)):
    t = time()
    birch_model.fit(X)
    time_ = time() - t
    print("Birch %s as the final step took %0.2f seconds" % (
          info, (time() - t)))

    # Plot result
    labels = birch_model.labels_
    centroids = birch_model.subcluster_centers_
    n_clusters = np.unique(labels).size
    print("n_clusters : %d" % n_clusters)

    ax = fig.add_subplot(1, 3, ind + 1)
    for this_centroid, k, col in zip(centroids, range(n_clusters), colors_):
        mask = labels == k
        ax.plot(X[mask, 0], X[mask, 1], 'w',
                markerfacecolor=col, marker='.')
        if birch_model.n_clusters is None:
            ax.plot(this_centroid[0], this_centroid[1], '+', markerfacecolor=col,
                    markeredgecolor='k', markersize=5)
    ax.set_ylim([-25, 25])
    ax.set_xlim([-25, 25])
    ax.set_autoscaley_on(False)
    ax.set_title('Birch %s' % info)

# Compute clustering with MiniBatchKMeans.
mbk = MiniBatchKMeans(init='k-means++', n_clusters=100, batch_size=100,
                      n_init=10, max_no_improvement=10, verbose=0,
                      random_state=0)
t0 = time()
mbk.fit(X)
t_mini_batch = time() - t0
print("Time taken to run MiniBatchKMeans %0.2f seconds" % t_mini_batch)
mbk_means_labels_unique = np.unique(mbk.labels_)

ax = fig.add_subplot(1, 3, 3)
for this_centroid, k, col in zip(mbk.cluster_centers_,
                                 range(n_clusters), colors_):
    mask = mbk.labels_ == k
    ax.plot(X[mask, 0], X[mask, 1], 'w', markerfacecolor=col, marker='.')
    ax.plot(this_centroid[0], this_centroid[1], '+', markeredgecolor='k',
            markersize=5)
ax.set_xlim([-25, 25])
ax.set_ylim([-25, 25])
ax.set_title("MiniBatchKMeans")
ax.set_autoscaley_on(False)
plt.show()


# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD 3 clause

import numpy as np
from scipy import sparse

from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_raises

from sklearn.linear_model.randomized_l1 import (lasso_stability_path,
                                                RandomizedLasso,
                                                RandomizedLogisticRegression)
from sklearn.datasets import load_diabetes, load_iris
from sklearn.feature_selection import f_regression, f_classif
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model.base import center_data

diabetes = load_diabetes()
X = diabetes.data
y = diabetes.target
X = StandardScaler().fit_transform(X)
X = X[:, [2, 3, 6, 7, 8]]

# test that the feature score of the best features
F, _ = f_regression(X, y)

def test_lasso_stability_path():
    # Check lasso stability path
    # Load diabetes data and add noisy features
    scaling = 0.3
    coef_grid, scores_path = lasso_stability_path(X, y, scaling=scaling,
                                                  random_state=42,
                                                  n_resampling=30)

    assert_array_equal(np.argsort(F)[-3:],
                       np.argsort(np.sum(scores_path, axis=1))[-3:])

def test_randomized_lasso():
    # Check randomized lasso
    scaling = 0.3
    selection_threshold = 0.5

    # or with 1 alpha
    clf = RandomizedLasso(verbose=False, alpha=1, random_state=42,
                          scaling=scaling,
                          selection_threshold=selection_threshold)
    feature_scores = clf.fit(X, y).scores_
    assert_array_equal(np.argsort(F)[-3:], np.argsort(feature_scores)[-3:])

    # or with many alphas
    clf = RandomizedLasso(verbose=False, alpha=[1, 0.8], random_state=42,
                          scaling=scaling,
                          selection_threshold=selection_threshold)
    feature_scores = clf.fit(X, y).scores_
    assert_equal(clf.all_scores_.shape, (X.shape[1], 2))
    assert_array_equal(np.argsort(F)[-3:], np.argsort(feature_scores)[-3:])

    X_r = clf.transform(X)
    X_full = clf.inverse_transform(X_r)
    assert_equal(X_r.shape[1], np.sum(feature_scores > selection_threshold))
    assert_equal(X_full.shape, X.shape)

    clf = RandomizedLasso(verbose=False, alpha='aic', random_state=42,
                          scaling=scaling)
    feature_scores = clf.fit(X, y).scores_
    assert_array_equal(feature_scores, X.shape[1] * [1.])

    clf = RandomizedLasso(verbose=False, scaling=-0.1)
    assert_raises(ValueError, clf.fit, X, y)

    clf = RandomizedLasso(verbose=False, scaling=1.1)
    assert_raises(ValueError, clf.fit, X, y)

def test_randomized_logistic():
    # Check randomized sparse logistic regression
    iris = load_iris()
    X = iris.data[:, [0, 2]]
    y = iris.target
    X = X[y != 2]
    y = y[y != 2]

    F, _ = f_classif(X, y)

    scaling = 0.3
    clf = RandomizedLogisticRegression(verbose=False, C=1., random_state=42,
                                       scaling=scaling, n_resampling=50,
                                       tol=1e-3)
    X_orig = X.copy()
    feature_scores = clf.fit(X, y).scores_
    assert_array_equal(X, X_orig)   # fit does not modify X
    assert_array_equal(np.argsort(F), np.argsort(feature_scores))

    clf = RandomizedLogisticRegression(verbose=False, C=[1., 0.5],
                                       random_state=42, scaling=scaling,
                                       n_resampling=50, tol=1e-3)
    feature_scores = clf.fit(X, y).scores_
    assert_array_equal(np.argsort(F), np.argsort(feature_scores))

def test_randomized_logistic_sparse():
    # Check randomized sparse logistic regression on sparse data
    iris = load_iris()
    X = iris.data[:, [0, 2]]
    y = iris.target
    X = X[y != 2]
    y = y[y != 2]

    # center here because sparse matrices are usually not centered
    X, y, _, _, _ = center_data(X, y, True, True)

    X_sp = sparse.csr_matrix(X)

    F, _ = f_classif(X, y)

    scaling = 0.3
    clf = RandomizedLogisticRegression(verbose=False, C=1., random_state=42,
                                       scaling=scaling, n_resampling=50,
                                       tol=1e-3)
    feature_scores = clf.fit(X, y).scores_
    clf = RandomizedLogisticRegression(verbose=False, C=1., random_state=42,
                                       scaling=scaling, n_resampling=50,
                                       tol=1e-3)
    feature_scores_sp = clf.fit(X_sp, y).scores_
    assert_array_equal(feature_scores, feature_scores_sp)


from __future__ import print_function
from show3d_balls import *
import argparse
import os
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from datasets import PartDataset
from pointnet import PointGen, PointGenR
import torch.nn.functional as F
import matplotlib.pyplot as plt

#showpoints(np.random.randn(2500,3), c1 = np.random.uniform(0,1,size = (2500)))

parser = argparse.ArgumentParser()

parser.add_argument('--model', type=str, default = '',  help='model path')

opt = parser.parse_args()
print (opt)

gen = PointGenR()
gen.load_state_dict(torch.load(opt.model))

#sim_noise = Variable(torch.randn(5, 2, 20))
#
#sim_noises = Variable(torch.zeros(5, 15, 20))
#
#for i in range(15):
#    x = i/15.0
#    sim_noises[:,i,:] = sim_noise[:,0,:] * x + sim_noise[:,1,:] * (1-x)
#
#points = gen(sim_noises)
#point_np = points.transpose(2,1).data.numpy()

sim_noise = Variable(torch.randn(5, 6, 20))
sim_noises = Variable(torch.zeros(5, 30 * 5,20))

for j in range(5):
    for i in range(30):
        x = (1-i/30.0)
        sim_noises[:,i + 30 * j,:] = sim_noise[:,j,:] * x + sim_noise[:,(j+1) % 5,:] * (1-x)

points = gen(sim_noises)
point_np = points.transpose(2,1).data.numpy()
print(point_np.shape)

for i in range(150):
    print(i)
    frame = showpoints_frame(point_np[i])
    plt.imshow(frame)
    plt.axis('off')
    plt.savefig('%s/%04d.png' %('out_rgan', i), bbox_inches='tight')
    plt.clf()

#showpoints(point_np)

#sim_noise = Variable(torch.randn(5, 1000, 20))
#points = gen(sim_noise)
#point_np = points.transpose(2,1).data.numpy()
#print(point_np.shape)
#choice = np.random.choice(2500, 2048, replace=False)
#print(point_np[:, choice, :].shape)

#showpoints(point_np)

#np.savez('rgan.npz', points = point_np[:, choice, :])


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

"""
=========================================================
Logistic Regression 3-class Classifier
=========================================================

Show below is a logistic-regression classifiers decision boundaries on the
`iris <http://en.wikipedia.org/wiki/Iris_flower_data_set>`_ dataset. The
datapoints are colored according to their labels.

"""
print(__doc__)

# Code source: Gaël Varoquaux
# Modified for documentation by Jaques Grobler
# License: BSD 3 clause

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

# import some data to play with
iris = datasets.load_iris()
X = iris.data[:, :2]  # we only take the first two features.
Y = iris.target

h = .02  # step size in the mesh

logreg = linear_model.LogisticRegression(C=1e5)

# we create an instance of Neighbours Classifier and fit the data.
logreg.fit(X, Y)

# Plot the decision boundary. For that, we will assign a color to each
# point in the mesh [x_min, m_max]x[y_min, y_max].
x_min, x_max = X[:, 0].min() - .5, X[:, 0].max() + .5
y_min, y_max = X[:, 1].min() - .5, X[:, 1].max() + .5
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
Z = logreg.predict(np.c_[xx.ravel(), yy.ravel()])

# Put the result into a color plot
Z = Z.reshape(xx.shape)
plt.figure(1, figsize=(4, 3))
plt.pcolormesh(xx, yy, Z, cmap=plt.cm.Paired)

# Plot also the training points
plt.scatter(X[:, 0], X[:, 1], c=Y, edgecolors='k', cmap=plt.cm.Paired)
plt.xlabel('Sepal length')
plt.ylabel('Sepal width')

plt.xlim(xx.min(), xx.max())
plt.ylim(yy.min(), yy.max())
plt.xticks(())
plt.yticks(())

plt.show()


"""
Segmentation creation and prediction
"""

import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn import metrics
from sklearn.cross_validation import train_test_split
from crankshaft.analysis_data_provider import AnalysisDataProvider

# NOTE: added optional param here

class Segmentation(object):
    """
        Add docstring
    """

    def __init__(self, data_provider=None):
        if data_provider is None:
            self.data_provider = AnalysisDataProvider()
        else:
            self.data_provider = data_provider

    def create_and_predict_segment_agg(self, target, features, target_features,
                                       target_ids, model_parameters):
        """
        Version of create_and_predict_segment that works on arrays that come
            straight form the SQL calling the function.

            Input:
                @param target: The 1D array of length NSamples containing the
                target variable we want the model to predict
                @param features: The 2D array of size NSamples * NFeatures that
                    form the input to the model
                @param target_ids: A 1D array of target_ids that will be used
                to associate the results of the prediction with the rows which
                    they come from
                @param model_parameters: A dictionary containing parameters for
                the model.
        """
        clean_target, _ = replace_nan_with_mean(target)
        clean_features, _ = replace_nan_with_mean(features)
        target_features,  _ = replace_nan_with_mean(target_features)

        model, accuracy = train_model(clean_target, clean_features,
                                      model_parameters, 0.2)
        prediction = model.predict(target_features)
        accuracy_array = [accuracy] * prediction.shape[0]
        return zip(target_ids, prediction, accuracy_array)

    def create_and_predict_segment(self, query, variable, feature_columns,
                                   target_query, model_params,
                                   id_col='cartodb_id'):
        """
        generate a segment with machine learning
        Stuart Lynn
                @param query: subquery that data is pulled from for packaging
                @param variable: name of the target variable
                @param feature_columns: list of column names
                @target_query: The query to run to obtain the data to predict
                @param model_params: A dictionary of model parameters, the full
                        specification can be found on the
                        scikit learn page for [GradientBoostingRegressor]
                        (http://scikit-learn.org/stable/modules/generated/sklearn.ensemble.GradientBoostingRegressor.html)
        """
        params = {"subquery": target_query,
                  "id_col": id_col}

        (target, features, target_mean,
            feature_means) = self.clean_data(query, variable, feature_columns)

        model, accuracy = train_model(target, features, model_params, 0.2)
        result = self.predict_segment(model, feature_columns, target_query,
                                      feature_means)
        accuracy_array = [accuracy] * result.shape[0]

        rowid = self.data_provider.get_segmentation_data(params)
        '''
        rowid = [{'ids': [2.9, 4.9, 4, 5, 6]}]
        '''
        return zip(rowid[0]['ids'], result, accuracy_array)

    def predict_segment(self, model, feature_columns, target_query,
                        feature_means):
        """
        Use the provided model to predict the values for the new feature set
            Input:
                @param model: The pretrained model
                @features_col: A list of features to use in the
                    model prediction (list of column names)
                @target_query: The query to run to obtain the data to predict
                    on and the cartodb_ids associated with it.
        """

        batch_size = 1000
        params = {"subquery": target_query,
                  "feature_columns": feature_columns}

        results = []
        cursors = self.data_provider.get_segmentation_predict_data(params)
        '''
         cursors = [{'features': [[m1[0],m2[0],m3[0]],[m1[1],m2[1],m3[1]],
                                  [m1[2],m2[2],m3[2]]]}]
        '''

        while True:
            rows = cursors.fetch(batch_size)
            if not rows:
                break
            batch = np.row_stack([np.array(row['features'])
                                  for row in rows]).astype(float)

            batch = replace_nan_with_mean(batch, feature_means)[0]
            prediction = model.predict(batch)
            results.append(prediction)

        # NOTE: we removed the cartodb_ids calculation in here
        return np.concatenate(results)

    def clean_data(self, query, variable, feature_columns):
        """
            Add docstring
        """
        params = {"subquery": query,
                  "target": variable,
                  "features": feature_columns}

        data = self.data_provider.get_segmentation_model_data(params)

        '''
        data = [{'target': [2.9, 4.9, 4, 5, 6],
        'feature1': [1,2,3,4], 'feature2' : [2,3,4,5]}]
        '''

        # extract target data from data_provider object
        target = np.array(data[0]['target'], dtype=float)

        # put n feature data arrays into an n x m array of arrays
        features = np.column_stack([np.array(data[0][col])
                                    for col in feature_columns]).astype(float)

        features, feature_means = replace_nan_with_mean(features)
        target, target_mean = replace_nan_with_mean(target)
        return target, features, target_mean, feature_means

def replace_nan_with_mean(array, means=None):
    """
        Input:
            @param array: an array of floats which may have null-valued
                          entries
        Output:
            array with nans filled in with the mean of the dataset
    """

    # returns an array of rows and column indices
    nanvals = np.isnan(array)
    indices = np.where(nanvals)

    def loops(array, axis):
        try:
            return np.shape(array)[axis]
        except IndexError:
            return 1
    ran = loops(array, 1)

    if means is None:
        means = {}

        if ran == 1:
            array = np.array(array)
            means[0] = np.mean(array[~np.isnan(array)])
            for row in zip(*indices):
                array[row] = means[0]
        else:
            for col in range(ran):
                means[col] = np.mean(array[~np.isnan(array[:, col]), col])
            for row, col in zip(*indices):
                array[row, col] = means[col]
    else:
        if ran == 1:
            for row in zip(*indices):
                array[row] = means[0]
        else:
            for row, col in zip(*indices):
                array[row, col] = means[col]

    return array, means

def train_model(target, features, model_params, test_split):
    """
        Train the Gradient Boosting model on the provided data to calculate
        the accuracy of the model
        Input:
            @param target: 1D Array of the variable that the model is to be
                trained to predict
            @param features: 2D Array NSamples *NFeatures to use in training
                the model
            @param model_params: A dictionary of model parameters, the full
                specification can be found on the
                scikit learn page for [GradientBoostingRegressor]
                (http://scikit-learn.org/stable/modules/generated/sklearn.ensemble.GradientBoostingRegressor.html)
            @parma test_split: The fraction of the data to be withheld for
                testing the model / calculating the accuray
    """
    features_train, features_test, \
        target_train, target_test = train_test_split(features, target,
                                                     test_size=test_split)
    model = GradientBoostingRegressor(**model_params)
    model.fit(features_train, target_train)
    accuracy = calculate_model_accuracy(model, features_test, target_test)
    return model, accuracy

def calculate_model_accuracy(model, features_test, target_test):
    """
        Calculate the mean squared error of the model prediction
        Input:
            @param model: model trained from input features
            @param features_test: test features set to make prediction from
            @param target_test: test target set to compare predictions to
        Output:
            mean squared error of the model prection compared target_test
    """
    prediction = model.predict(features_test)
    return metrics.mean_squared_error(prediction, target_test)


#!/usr/bin/env python

##
# wrapping: A program making it easy to use hyperparameter
# optimization software.
# Copyright (C) 2013 Katharina Eggensperger and Matthias Feurer
#
# This program 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.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

from argparse import ArgumentParser
import cPickle
import itertools
import sys

from matplotlib.pyplot import tight_layout, figure, subplots_adjust, subplot, savefig, show
import matplotlib.gridspec
import numpy as np

from HPOlib.Plotting import plot_util

__authors__ = ["Katharina Eggensperger", "Matthias Feurer"]
__contact__ = "automl.org"

def plot_optimization_trace_cv(trial_list, name_list, optimum=0, title="",
                               log=True, save="", y_max=0, y_min=0):
    markers =plot_util.get_plot_markers()
    colors = plot_util.get_plot_colors()
    linestyles = itertools.cycle(['-'])
    size = 1

    ratio = 5
    gs = matplotlib.gridspec.GridSpec(ratio, 1)
    fig = figure(1, dpi=100)
    fig.suptitle(title, fontsize=16)
    ax1 = subplot(gs[0:ratio, :])
    ax1.grid(True, linestyle='-', which='major', color='lightgrey', alpha=0.5)
    min_val = sys.maxint
    max_val = -sys.maxint
    max_trials = 0

    fig.suptitle(title, fontsize=16)

    # Plot the average error and std
    for i in range(len(name_list)):
        m = markers.next()
        c = colors.next()
        l = linestyles.next()
        leg = False
        for tr in trial_list[i]:
            if log:
                tr = np.log10(tr)
            x = range(1, len(tr)+1)
            y = tr
            if not leg:
                ax1.plot(x, y, color=c, linewidth=size, linestyle=l, label=name_list[i][0])
                leg = True
            ax1.plot(x, y, color=c, linewidth=size, linestyle=l)
            min_val = min(min_val, min(tr))
            max_val = max(max_val, max(tr))
            max_trials = max(max_trials, len(tr))

    # Maybe plot on logscale
    ylabel = ""

    if log:
        ax1.set_ylabel("log10(Minfunction value)" + ylabel)
    else:
        ax1.set_ylabel("Minfunction value" + ylabel)

    # Descript and label the stuff
    leg = ax1.legend(loc='best', fancybox=True)
    leg.get_frame().set_alpha(0.5)
    ax1.set_xlabel("#Function evaluations")

    if y_max == y_min:
         # Set axes limits
        ax1.set_ylim([min_val-0.1*abs((max_val-min_val)), max_val+0.1*abs((max_val-min_val))])
    else:
        ax1.set_ylim([y_min, y_max])
    ax1.set_xlim([0, max_trials + 1])

    tight_layout()
    subplots_adjust(top=0.85)
    if save != "":
        savefig(save, dpi=100, facecolor='w', edgecolor='w',
                orientation='portrait', papertype=None, format=None,
                transparent=False, bbox_inches="tight", pad_inches=0.1)
    else:
        show()

def main(pkl_list, name_list, autofill, optimum=0, save="", title="", log=False,
         y_min=0, y_max=0):

    trial_list = list()
    for i in range(len(pkl_list)):
        tmp_trial_list = list()
        max_len = -sys.maxint
        for pkl in pkl_list[i]:
            fh = open(pkl, "r")
            trials = cPickle.load(fh)
            fh.close()

            trace = plot_util.get_Trace_cv(trials)
            tmp_trial_list.append(trace)
            max_len = max(max_len, len(trace))
        trial_list.append(list())
        for tr in tmp_trial_list:
        #    if len(tr) < max_len:
        #        tr.extend([tr[-1] for idx in range(abs(max_len - len(tr)))])
            trial_list[-1].append(np.array(tr))

    plot_optimization_trace_cv(trial_list, name_list, optimum, title=title, log=log,
                            save=save, y_min=y_min, y_max=y_max)

    if save != "":
        sys.stdout.write("Saved plot to " + save + "\n")
    else:
        sys.stdout.write("..Done\n")

if __name__ == "__main__":
    prog = "python plotTraceWithStd.py WhatIsThis <oneOrMorePickles> [WhatIsThis <oneOrMorePickles>]"
    description = "Plot a Trace with std for multiple experiments"

    parser = ArgumentParser(description=description, prog=prog)

    # Options for specific benchmarks
    parser.add_argument("-o", "--optimum", type=float, dest="optimum",
                        default=0, help="If not set, the optimum is supposed to be zero")

    # Options which are available only for this plot
    parser.add_argument("-a", "--autofill", action="store_true", dest="autofill",
                        default=False, help="Fill trace automatically")

    # General Options
    parser.add_argument("-l", "--log", action="store_true", dest="log",
                        default=False, help="Plot on log scale")
    parser.add_argument("--max", dest="max", type=float,
                        default=0, help="Maximum of the plot")
    parser.add_argument("--min", dest="min", type=float,
                        default=0, help="Minimum of the plot")
    parser.add_argument("-s", "--save", dest="save",
                        default="", help="Where to save plot instead of showing it?")
    parser.add_argument("-t", "--title", dest="title",
                        default="", help="Optional supertitle for plot")

    args, unknown = parser.parse_known_args()

    sys.stdout.write("\nFound " + str(len(unknown)) + " arguments\n")

    pkl_list_main, name_list_main = plot_util.get_pkl_and_name_list(unknown)

    main(pkl_list=pkl_list_main, name_list=name_list_main, autofill=args.autofill, optimum=args.optimum,
         save=args.save, title=args.title, log=args.log, y_min=args.min, y_max=args.max)


#!/usr/bin/python

import sys
import os
import matplotlib

import numpy
import matplotlib
matplotlib.use('AGG')
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.ticker import FormatStrFormatter

def getArg(param, default=""):
    if (sys.argv.count(param) == 0): return default
    i = sys.argv.index(param)
    return sys.argv[i + 1]

lastsecs = int(getArg("lastsecs", 240))

fname = sys.argv[1]
try:
	tdata = numpy.loadtxt(fname, delimiter=" ")
except:
	exit(0)

if len(tdata.shape) < 2 or tdata.shape[0] < 2 or tdata.shape[1] < 2:
    print "Too small data - do not try to plot yet."
    exit(0)

times = tdata[:, 0]
values = tdata[:, 1]

lastt = max(times)

#majorFormatter = FormatStrFormatter('%.2f')

fig = plt.figure(figsize=(3.5, 2.0))
plt.plot(times[times > lastt - lastsecs], values[times > lastt - lastsecs])
plt.gca().xaxis.set_major_locator( MaxNLocator(nbins = 7, prune = 'lower') )
plt.xlim([max(0, lastt - lastsecs), lastt])
#plt.ylim([lastt - lastsecs, lastt])

plt.gca().yaxis.set_major_locator( MaxNLocator(nbins = 7, prune = 'lower') )

#plt.gca().yaxis.set_major_formatter(majorFormatter)
plt.savefig(fname.replace(".dat", ".png"), format="png", bbox_inches='tight')