# -*- coding: utf-8 -*-
#
# test_plotting.py
#
# This file is part of NEST.
#
# Copyright (C) 2004 The NEST Initiative
#
# NEST 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 2 of the License, or
# (at your option) any later version.
#
# NEST 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 NEST.  If not, see <http://www.gnu.org/licenses/>.

"""
Tests for basic spatial plotting functions.
"""

import unittest
import nest
import numpy as np

try:
    import matplotlib.pyplot as plt

    tmp_fig = plt.figure()  # make sure we can open a window; DISPLAY may not be set
    plt.close(tmp_fig)
    PLOTTING_POSSIBLE = True
except:
    PLOTTING_POSSIBLE = False

@unittest.skipIf(not PLOTTING_POSSIBLE,
                 'Plotting impossible because matplotlib or display missing')
class PlottingTestCase(unittest.TestCase):
    def test_PlotLayer(self):
        """Test plotting layer."""
        nest.ResetKernel()
        l = nest.Create('iaf_psc_alpha',
                        positions=nest.spatial.grid(shape=[3, 3],
                                                    extent=[2., 2.],
                                                    edge_wrap=True))
        nest.PlotLayer(l)

        plotted_datapoints = plt.gca().collections[-1].get_offsets().data
        reference_datapoints = nest.GetPosition(l)
        self.assertTrue(np.allclose(plotted_datapoints, reference_datapoints))

    def test_PlotTargets(self):
        """Test plotting targets."""
        delta = 0.05
        mask = {'rectangular': {'lower_left': [-delta, -2/3 - delta], 'upper_right': [2/3 + delta, delta]}}
        cdict = {'rule': 'pairwise_bernoulli', 'p': 1.,
                 'mask': mask}
        sdict = {'synapse_model': 'stdp_synapse'}
        nest.ResetKernel()
        l = nest.Create('iaf_psc_alpha',
                        positions=nest.spatial.grid(shape=[3, 3],
                                                    extent=[2., 2.],
                                                    edge_wrap=True))

        # connect l -> l
        nest.Connect(l, l, cdict, sdict)

        ctr = nest.FindCenterElement(l)
        fig = nest.PlotTargets(ctr, l)
        fig.gca().set_title('Plain call')

        plotted_datapoints = plt.gca().collections[0].get_offsets().data
        eps = 0.01
        pos = np.array(nest.GetPosition(l))
        pos_xmask = pos[np.where(pos[:, 0] > -eps)]
        reference_datapoints = pos_xmask[np.where(pos_xmask[:, 1] < eps)][::-1]
        self.assertTrue(np.array_equal(np.sort(plotted_datapoints, axis=0), np.sort(reference_datapoints, axis=0)))

        fig = nest.PlotTargets(ctr, l, mask=mask)
        ax = fig.gca()
        ax.set_title('Call with mask')
        self.assertGreaterEqual(len(ax.patches), 1)

    def test_plot_probability_kernel(self):
        """Plot parameter probability"""
        nest.ResetKernel()
        plot_shape = [10, 10]
        plot_edges = [-0.5, 0.5, -0.5, 0.5]

        def probability_calculation(distance):
            return 1 - 1.5*distance

        l = nest.Create('iaf_psc_alpha', positions=nest.spatial.grid([10, 10], edge_wrap=False))
        source = l[25]
        source_pos = np.array(nest.GetPosition(source))
        source_x, source_y = source_pos

        # Calculate reference values
        ref_probability = np.zeros(plot_shape[::-1])
        for i, x in enumerate(np.linspace(plot_edges[0], plot_edges[1], plot_shape[0])):
            positions = np.array([[x, y] for y in np.linspace(plot_edges[2], plot_edges[3], plot_shape[1])])
            ref_distances = np.sqrt((positions[:, 0] - source_x)**2 + (positions[:, 1] - source_y)**2)
            values = probability_calculation(ref_distances)
            ref_probability[:, i] = np.maximum(np.minimum(np.array(values), 1.0), 0.0)

        # Create the parameter
        parameter = probability_calculation(nest.spatial.distance)

        fig, ax = plt.subplots()
        nest.PlotProbabilityParameter(source, parameter, ax=ax, shape=plot_shape, edges=plot_edges)

        self.assertEqual(len(ax.images), 1)
        img = ax.images[0]
        img_data = img.get_array().data
        self.assertTrue(np.array_equal(img_data, ref_probability))

    def test_plot_probability_kernel_with_mask(self):
        """Plot parameter probability with mask"""
        nest.ResetKernel()
        plot_shape = [10, 10]
        plot_edges = [-0.5, 0.5, -0.5, 0.5]

        l = nest.Create('iaf_psc_alpha', positions=nest.spatial.grid([10, 10], edge_wrap=False))
        parameter = 1 - 1.5*nest.spatial.distance

        source = l[25]
        masks = [{'circular': {'radius': 0.4}},
                 {'doughnut': {'inner_radius': 0.2, 'outer_radius': 0.45}},
                 {'rectangular': {'lower_left': [-.3, -.3], 'upper_right': [0.3, 0.3]}},
                 {'elliptical': {'major_axis': 0.8, 'minor_axis': 0.4}}]
        fig, axs = plt.subplots(2, 2)
        for mask, ax in zip(masks, axs.flatten()):
            nest.PlotProbabilityParameter(source, parameter, mask=mask, ax=ax, shape=plot_shape, edges=plot_edges)
            self.assertEqual(len(ax.images), 1)
            self.assertGreaterEqual(len(ax.patches), 1)

def suite():
    suite = unittest.makeSuite(PlottingTestCase, 'test')
    return suite

if __name__ == "__main__":
    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite())

    plt.show()


"""
Machine learning module for Python
==================================

sklearn is a Python module integrating classical machine
learning algorithms in the tightly-knit world of scientific Python
packages (numpy, scipy, matplotlib).

It aims to provide simple and efficient solutions to learning problems
that are accessible to everybody and reusable in various contexts:
machine-learning as a versatile tool for science and engineering.

See http://scikit-learn.org for complete documentation.
"""
import sys
import re
import warnings

# Make sure that DeprecationWarning within this package always gets printed
warnings.filterwarnings('always', category=DeprecationWarning,
                        module='^{0}\.'.format(re.escape(__name__)))

# PEP0440 compatible formatted version, see:
# https://www.python.org/dev/peps/pep-0440/
#
# Generic release markers:
#   X.Y
#   X.Y.Z   # For bugfix releases
#
# Admissible pre-release markers:
#   X.YaN   # Alpha release
#   X.YbN   # Beta release
#   X.YrcN  # Release Candidate
#   X.Y     # Final release
#
# Dev branch marker is: 'X.Y.dev' or 'X.Y.devN' where N is an integer.
# 'X.Y.dev0' is the canonical version of 'X.Y.dev'
#
__version__ = '0.18.dev0'

try:
    # This variable is injected in the __builtins__ by the build
    # process. It used to enable importing subpackages of sklearn when
    # the binaries are not built
    __SKLEARN_SETUP__
except NameError:
    __SKLEARN_SETUP__ = False

if __SKLEARN_SETUP__:
    sys.stderr.write('Partial import of sklearn during the build process.\n')
    # We are not importing the rest of the scikit during the build
    # process, as it may not be compiled yet
else:
    from . import __check_build
    from .base import clone
    __check_build  # avoid flakes unused variable error

    __all__ = ['calibration', 'cluster', 'covariance', 'cross_decomposition',
               'cross_validation', 'datasets', 'decomposition', 'dummy',
               'ensemble', 'exceptions', 'externals', 'feature_extraction',
               'feature_selection', 'gaussian_process', 'grid_search',
               'isotonic', 'kernel_approximation', 'kernel_ridge',
               'lda', 'learning_curve', 'linear_model', 'manifold', 'metrics',
               'mixture', 'model_selection', 'multiclass', 'multioutput',
               'naive_bayes', 'neighbors', 'neural_network', 'pipeline',
               'preprocessing', 'qda', 'random_projection', 'semi_supervised',
               'svm', 'tree', 'discriminant_analysis',
               # Non-modules:
               'clone']

def setup_module(module):
    """Fixture for the tests to assure globally controllable seeding of RNGs"""
    import os
    import numpy as np
    import random

    # It could have been provided in the environment
    _random_seed = os.environ.get('SKLEARN_SEED', None)
    if _random_seed is None:
        _random_seed = np.random.uniform() * (2 ** 31 - 1)
    _random_seed = int(_random_seed)
    print("I: Seeding RNGs with %r" % _random_seed)
    np.random.seed(_random_seed)
    random.seed(_random_seed)


# -*- coding: utf-8 -*-
"""
Created on Sun Aug 20 11:25:06 2017

@author: kalvi
"""

#required imports
import pandas as pd
import numpy as np

import csv
import time
import calendar

def get_test_data(test_transformed_path, test_path, earlyWeatherDataPath, weatherData1, weatherData2):
    x_data = pd.read_csv(test_transformed_path, header=0)

    ########## Adding the date back into the data
    dataCSV = open(test_path, 'rt')
    csvData = list(csv.reader(dataCSV))
    csvFields = csvData[0] #['Dates', 'Category', 'Descript', 'DayOfWeek', 'PdDistrict', 'Resolution', 'Address', 'X', 'Y']
    allData = csvData[1:]
    dataCSV.close()

    df = pd.DataFrame(allData)
    df.columns = csvFields
    dates = df['Dates']
    dates = dates.apply(time.strptime, args=("%Y-%m-%d %H:%M:%S",))
    dates = dates.apply(calendar.timegm)

    x_data['secondsFromEpoch'] = dates
    colnames = x_data.columns.tolist()
    colnames = colnames[-1:] + colnames[:-1]
    x_data = x_data[colnames]
    ##########

    #functions for processing the sunrise and sunset times of each day
    def get_hour_and_minute(milTime):
        hour = int(milTime[:-2])
        minute = int(milTime[-2:])
        return [hour, minute]

    def get_date_only(date):
        return time.struct_time(tuple([date[0], date[1], date[2], 0, 0, 0, date[6], date[7], date[8]]))

    def structure_sun_time(timeSeries, dateSeries):
        sunTimes = timeSeries.copy()
        for index in range(len(dateSeries)):
            sunTimes[index] = time.struct_time(tuple([dateSeries[index][0], dateSeries[index][1], dateSeries[index][2], timeSeries[index][0], timeSeries[index][1], dateSeries[index][5], dateSeries[index][6], dateSeries[index][7], dateSeries[index][8]]))
        return sunTimes

    def get_weather_data(data_path):
        dataCSV = open(data_path, 'rt')
        csv_data = list(csv.reader(dataCSV))
        csv_fields = csv_data[0] #['Dates', 'Category', 'Descript', 'DayOfWeek', 'PdDistrict', 'Resolution', 'Address', 'X', 'Y']
        weather_data = csv_data[1:]
        dataCSV.close()

        weather_df = pd.DataFrame(weather_data)
        weather_df.columns = csv_fields
        dates = weather_df['DATE']
        sunrise = weather_df['DAILYSunrise']
        sunset = weather_df['DAILYSunset']

        dates = dates.apply(time.strptime, args=("%Y-%m-%d %H:%M",))
        sunrise = sunrise.apply(get_hour_and_minute)
        sunrise = structure_sun_time(sunrise, dates)
        sunrise = sunrise.apply(calendar.timegm)

        sunset = sunset.apply(get_hour_and_minute)
        sunset = structure_sun_time(sunset, dates)
        sunset = sunset.apply(calendar.timegm)
        dates = dates.apply(calendar.timegm)

        weather_df['DATE'] = dates
        weather_df['DAILYSunrise'] = sunrise
        weather_df['DAILYSunset'] = sunset

        return weather_df

    ########## Adding the weather data into the original crime data

    earlyWeatherDF = get_weather_data(earlyWeatherDataPath)
    weatherDF1 = get_weather_data(weatherData1)
    weatherDF2 = get_weather_data(weatherData2)

    weatherDF = pd.concat([earlyWeatherDF[450:975],weatherDF1,weatherDF2[32:]],ignore_index=True)

    # weather feature selection
    weatherMetrics = weatherDF[['DATE','HOURLYDRYBULBTEMPF','HOURLYRelativeHumidity', 'HOURLYWindSpeed', \
                                'HOURLYSeaLevelPressure', 'HOURLYVISIBILITY', 'DAILYSunrise', 'DAILYSunset']]
    weatherMetrics = weatherMetrics.convert_objects(convert_numeric=True)
    weatherDates = weatherMetrics['DATE']
    #'DATE','HOURLYDRYBULBTEMPF','HOURLYRelativeHumidity', 'HOURLYWindSpeed',
    #'HOURLYSeaLevelPressure', 'HOURLYVISIBILITY'
    timeWindow = 10800 #3 hours
    hourlyDryBulbTemp = []
    hourlyRelativeHumidity = []
    hourlyWindSpeed = []
    hourlySeaLevelPressure = []
    hourlyVisibility = []
    dailySunrise = []
    dailySunset = []
    daylight = []
    test = 0
    for timePoint in dates:#dates is the epoch time from the kaggle data
        relevantWeather = weatherMetrics[(weatherDates <= timePoint) & (weatherDates > timePoint - timeWindow)]
        hourlyDryBulbTemp.append(relevantWeather['HOURLYDRYBULBTEMPF'].mean())
        hourlyRelativeHumidity.append(relevantWeather['HOURLYRelativeHumidity'].mean())
        hourlyWindSpeed.append(relevantWeather['HOURLYWindSpeed'].mean())
        hourlySeaLevelPressure.append(relevantWeather['HOURLYSeaLevelPressure'].mean())
        hourlyVisibility.append(relevantWeather['HOURLYVISIBILITY'].mean())
        dailySunrise.append(relevantWeather['DAILYSunrise'].iloc[-1])
        dailySunset.append(relevantWeather['DAILYSunset'].iloc[-1])
        daylight.append(1.0*((timePoint >= relevantWeather['DAILYSunrise'].iloc[-1]) and (timePoint < relevantWeather['DAILYSunset'].iloc[-1])))

        if test%100000 == 0:
            print(relevantWeather)
        test += 1

    hourlyDryBulbTemp = pd.Series.from_array(np.array(hourlyDryBulbTemp))
    hourlyRelativeHumidity = pd.Series.from_array(np.array(hourlyRelativeHumidity))
    hourlyWindSpeed = pd.Series.from_array(np.array(hourlyWindSpeed))
    hourlySeaLevelPressure = pd.Series.from_array(np.array(hourlySeaLevelPressure))
    hourlyVisibility = pd.Series.from_array(np.array(hourlyVisibility))
    dailySunrise = pd.Series.from_array(np.array(dailySunrise))
    dailySunset = pd.Series.from_array(np.array(dailySunset))
    daylight = pd.Series.from_array(np.array(daylight))

    x_data['HOURLYDRYBULBTEMPF'] = hourlyDryBulbTemp
    x_data['HOURLYRelativeHumidity'] = hourlyRelativeHumidity
    x_data['HOURLYWindSpeed'] = hourlyWindSpeed
    x_data['HOURLYSeaLevelPressure'] = hourlySeaLevelPressure
    x_data['HOURLYVISIBILITY'] = hourlyVisibility
    #x_data['DAILYSunrise'] = dailySunrise
    #x_data['DAILYSunset'] = dailySunset
    x_data['Daylight'] = daylight
    x_data = x_data.drop('secondsFromEpoch', 1)
    x_data = x_data.drop('pd_bayview_binary', 1)

    return x_data

test_transformed_path = "./data/test_transformed.csv"
test_path = "./data/test.csv"
earlyWeatherDataPath = "./data/1049158.csv"
weatherData1 = "./data/1027175.csv"
weatherData2 = "./data/1027176.csv"

write_path = "C:/MIDS/W207 final project/test_data_with_weather.csv"

x_data = get_test_data(test_transformed_path, test_path, earlyWeatherDataPath, weatherData1, weatherData2)
x_data.to_csv(path_or_buf=write_path,index=0)


"""
===========================================================
A demo of K-Means clustering on the handwritten digits data
===========================================================

In this example we compare the various initialization strategies for
K-means in terms of runtime and quality of the results.

As the ground truth is known here, we also apply different cluster
quality metrics to judge the goodness of fit of the cluster labels to the
ground truth.

Cluster quality metrics evaluated (see :ref:`clustering_evaluation` for
definitions and discussions of the metrics):

=========== ========================================================
Shorthand    full name
=========== ========================================================
homo         homogeneity score
compl        completeness score
v-meas       V measure
ARI          adjusted Rand index
AMI          adjusted mutual information
silhouette   silhouette coefficient
=========== ========================================================

"""
print(__doc__)

from time import time
import numpy as np
import matplotlib.pyplot as plt

from sklearn import metrics
from sklearn.cluster import KMeans
from sklearn.datasets import load_digits
from sklearn.decomposition import PCA
from sklearn.preprocessing import scale

np.random.seed(42)

digits = load_digits()
data = scale(digits.data)

n_samples, n_features = data.shape
n_digits = len(np.unique(digits.target))
labels = digits.target

sample_size = 300

print("n_digits: %d, \t n_samples %d, \t n_features %d"
      % (n_digits, n_samples, n_features))

print(79 * '_')
print('% 9s' % 'init'
      '    time  inertia    homo   compl  v-meas     ARI AMI  silhouette')

def bench_k_means(estimator, name, data):
    t0 = time()
    estimator.fit(data)
    print('% 9s   %.2fs    %i   %.3f   %.3f   %.3f   %.3f   %.3f    %.3f'
          % (name, (time() - t0), estimator.inertia_,
             metrics.homogeneity_score(labels, estimator.labels_),
             metrics.completeness_score(labels, estimator.labels_),
             metrics.v_measure_score(labels, estimator.labels_),
             metrics.adjusted_rand_score(labels, estimator.labels_),
             metrics.adjusted_mutual_info_score(labels,  estimator.labels_),
             metrics.silhouette_score(data, estimator.labels_,
                                      metric='euclidean',
                                      sample_size=sample_size)))

bench_k_means(KMeans(init='k-means++', n_clusters=n_digits, n_init=10),
              name="k-means++", data=data)

bench_k_means(KMeans(init='random', n_clusters=n_digits, n_init=10),
              name="random", data=data)

# in this case the seeding of the centers is deterministic, hence we run the
# kmeans algorithm only once with n_init=1
pca = PCA(n_components=n_digits).fit(data)
bench_k_means(KMeans(init=pca.components_, n_clusters=n_digits, n_init=1),
              name="PCA-based",
              data=data)
print(79 * '_')

###############################################################################
# Visualize the results on PCA-reduced data

reduced_data = PCA(n_components=2).fit_transform(data)
kmeans = KMeans(init='k-means++', n_clusters=n_digits, n_init=10)
kmeans.fit(reduced_data)

# Step size of the mesh. Decrease to increase the quality of the VQ.
h = .02     # point in the mesh [x_min, m_max]x[y_min, y_max].

# Plot the decision boundary. For that, we will assign a color to each
x_min, x_max = reduced_data[:, 0].min() - 1, reduced_data[:, 0].max() + 1
y_min, y_max = reduced_data[:, 1].min() - 1, reduced_data[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))

# Obtain labels for each point in mesh. Use last trained model.
Z = kmeans.predict(np.c_[xx.ravel(), yy.ravel()])

# Put the result into a color plot
Z = Z.reshape(xx.shape)
plt.figure(1)
plt.clf()
plt.imshow(Z, interpolation='nearest',
           extent=(xx.min(), xx.max(), yy.min(), yy.max()),
           cmap=plt.cm.Paired,
           aspect='auto', origin='lower')

plt.plot(reduced_data[:, 0], reduced_data[:, 1], 'k.', markersize=2)
# Plot the centroids as a white X
centroids = kmeans.cluster_centers_
plt.scatter(centroids[:, 0], centroids[:, 1],
            marker='x', s=169, linewidths=3,
            color='w', zorder=10)
plt.title('K-means clustering on the digits dataset (PCA-reduced data)\n'
          'Centroids are marked with white cross')
plt.xlim(x_min, x_max)
plt.ylim(y_min, y_max)
plt.xticks(())
plt.yticks(())
plt.show()


import pytest

import numpy as np
import pandas as pd
from pandas.util.testing import assert_frame_equal

import fishact

def test_sniffer():
    n_header, delimiter, line = fishact.parse._sniff_file_info(
                                                'tests/single_gtype.txt')
    assert n_header == 2
    assert delimiter is None
    assert line == '1\n'

    n_header, delimiter, line = fishact.parse._sniff_file_info(
                                                'tests/multiple_gtype.txt')
    assert n_header == 2
    assert delimiter == '\t'
    assert line == '1\t5\t2\n'

def test_gtype_loader():
    df = fishact.parse.load_gtype('tests/single_gtype.txt', quiet=True,
                                  rstrip=True)
    assert all(df.columns == ['genotype', 'location'])
    assert all(df['genotype'] == 'all fish')
    assert all(df['location'] == [1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 86, 87, 88, 89, 90, 91, 92, 93, 94, 96])

def test_resample_array():
    x  = np.arange(10, dtype=float)
    assert np.isclose(fishact.parse._resample_array(x, 10),
                      np.array([45.])).all()
    assert np.isclose(fishact.parse._resample_array(x, 20),
                      np.array([90.])).all()
    assert np.isclose(fishact.parse._resample_array(x, 5),
                      np.array([10., 35.])).all()
    assert np.isclose(fishact.parse._resample_array(x, 3),
                      np.array([3., 12., 21., 27.])).all()

    x[-3] = np.nan
    assert np.isclose(fishact.parse._resample_array(x, 5)[0], 10.0) \
                and np.isnan(fishact.parse._resample_array(x, 5)[1])

def test_resample_segment():
    df = pd.DataFrame({'a': np.arange(10),
                       'b': np.arange(10, 20),
                       'c': np.arange(10, dtype=float),
                       'd': np.arange(10, 20, dtype=float)})

    re_df = fishact.parse._resample_segment(df, 5, ['c'])
    re_df = re_df.reindex_axis(sorted(re_df.columns), axis=1)
    correct_df = pd.DataFrame({'a': [0, 5],
                               'b': [10, 15],
                               'c': [10., 35.],
                               'd': [10., 15.]})
    assert_frame_equal(re_df, correct_df)

    re_df = fishact.parse._resample_segment(df, 5, ['c', 'd'])
    re_df = re_df.reindex_axis(sorted(re_df.columns), axis=1)
    correct_df = pd.DataFrame({'a': [0, 5],
                               'b': [10, 15],
                               'c': [10., 35.],
                               'd': [60., 85.]})
    assert_frame_equal(re_df, correct_df)

    re_df = fishact.parse._resample_segment(df, 3, ['c'])
    re_df = re_df.reindex_axis(sorted(re_df.columns), axis=1)
    correct_df = pd.DataFrame({'a': [0, 3, 6, 9],
                               'b': [10, 13, 16, 19],
                               'c': [3., 12., 21., 27.],
                               'd': [10., 13., 16., 19.]})
    assert_frame_equal(re_df, correct_df)

    re_df = fishact.parse._resample_segment(df, 3, ['c', 'd'])
    re_df = re_df.reindex_axis(sorted(re_df.columns), axis=1)
    correct_df = pd.DataFrame({'a': [0, 3, 6, 9],
                               'b': [10, 13, 16, 19],
                               'c': [3., 12., 21., 27.],
                               'd': [33., 42., 51., 57.]})
    assert_frame_equal(re_df, correct_df)

def test_resample():
    df = pd.DataFrame(
        {'location': np.concatenate((np.ones(10), 2*np.ones(10))).astype(int),
         'exp_time': np.concatenate((np.arange(10),
                                     np.arange(10))).astype(float),
         'exp_ind': np.concatenate((np.arange(10), np.arange(10))).astype(int),
         'zeit': np.concatenate((np.arange(10),
                                 np.arange(10))).astype(float),
         'zeit_ind': np.concatenate((np.arange(10),
                                     np.arange(10))).astype(int),
         'activity': np.concatenate((np.arange(10),
                                     np.arange(10, 20))).astype(float),
         'sleep': np.ones(20, dtype=float),
         'light': [True]*5 + [False]*5 + [True]*5 + [False]*5,
         'day': [5]*10 + [6]*10,
         'genotype': ['wt']*20,
         'acquisition': np.ones(20, dtype=int),
         'instrument':  np.ones(20, dtype=int),
         'trial':  np.ones(20, dtype=int),
         'time': pd.to_datetime(['2017-03-30 14:00:00',
                                 '2017-03-30 14:01:00',
                                 '2017-03-30 14:02:00',
                                 '2017-03-30 14:03:00',
                                 '2017-03-30 14:04:00',
                                 '2017-03-30 14:05:00',
                                 '2017-03-30 14:06:00',
                                 '2017-03-30 14:07:00',
                                 '2017-03-30 14:08:00',
                                 '2017-03-30 14:09:00']*2)})

    re_df = fishact.parse.resample(df, 5, signal=['activity', 'sleep'],
                                   quiet=True)
    re_df = re_df.reindex_axis(sorted(re_df.columns), axis=1)
    correct_df = pd.DataFrame(
        {'activity': np.array([10., 35., 60., 85.]),
         'day': [5, 5, 6, 6],
         'location': np.array([1, 1, 2, 2], dtype=int),
         'genotype': ['wt']*4,
         'light': [True, False, True, False],
         'sleep': np.array([5., 5., 5., 5.]),
         'time': pd.to_datetime(['2017-03-30 14:00:00',
                                 '2017-03-30 14:05:00']*2),
         'exp_time': np.array([0., 5., 0., 5.]),
         'exp_ind': np.array([0, 5, 0, 5], dtype=int),
         'zeit': np.array([0., 5., 0., 5.]),
         'zeit_ind': np.array([0, 5, 0, 5], dtype=int),
         'acquisition': np.ones(4, dtype=int),
         'instrument': np.ones(4, dtype=int),
         'trial': np.ones(4, dtype=int)})
    assert_frame_equal(re_df, correct_df)

    re_df = fishact.parse.resample(df, 3, quiet=True)
    re_df = re_df.reindex_axis(sorted(re_df.columns), axis=1)
    correct_df = pd.DataFrame(
        {'activity': np.array([3., 10.5, 18., 25.5, 33., 40.5, 48., 55.5]),
         'day': [5, 5, 5, 5, 6, 6, 6, 6],
         'location': np.array([1, 1, 1, 1, 2, 2, 2, 2], dtype=int),
         'genotype': ['wt']*8,
         'light': [True, True, False, False, True, True, False, False],
         'sleep': np.array([3., 3., 3., 3., 3., 3., 3., 3.]),
         'time': pd.to_datetime(['2017-03-30 14:00:00',
                                 '2017-03-30 14:03:00',
                                 '2017-03-30 14:05:00',
                                 '2017-03-30 14:08:00']*2),
         'exp_time': np.array([0., 3., 5., 8., 0., 3., 5., 8.]),
         'exp_ind': np.array([0, 3, 5, 8, 0, 3, 5, 8], dtype=int),
         'zeit': np.array([0., 3., 5., 8., 0., 3., 5., 8.]),
         'zeit_ind': np.array([0, 3, 5, 8, 0, 3, 5, 8], dtype=int),
         'acquisition': np.ones(8, dtype=int),
         'instrument': np.ones(8, dtype=int),
         'trial': np.ones(8, dtype=int)})
    assert_frame_equal(re_df, correct_df)

def test_tidy_data():
    # Test that it will not overwrite existing file
    with pytest.raises(RuntimeError) as excinfo:
        fishact.parse.tidy_data('test.csv', 'test_geno.txt', 'test.csv')
    excinfo.match("Cowardly refusing to overwrite input file.")

    with pytest.raises(RuntimeError) as excinfo:
        fishact.parse.tidy_data('test.csv', 'test_geno.txt', 'test_geno.txt')
    excinfo.match("Cowardly refusing to overwrite input file.")

    with pytest.raises(RuntimeError) as excinfo:
        fishact.parse.tidy_data('test.csv', 'test_geno.txt',
                                'tests/empty_file_for_tests.csv')
    excinfo.match("tests/empty_file_for_tests.csv already exists, cowardly refusing to overwrite.")

    ## TO DO: integration test: make sure output CSV is as expected.


"""
=================================================
Hyper-parameters of Approximate Nearest Neighbors
=================================================

This example demonstrates the behaviour of the
accuracy of the nearest neighbor queries of Locality Sensitive Hashing
Forest as the number of candidates and the number of estimators (trees)
vary.

In the first plot, accuracy is measured with the number of candidates. Here,
the term "number of candidates" refers to maximum bound for the number of
distinct points retrieved from each tree to calculate the distances. Nearest
neighbors are selected from this pool of candidates. Number of estimators is
maintained at three fixed levels (1, 5, 10).

In the second plot, the number of candidates is fixed at 50. Number of trees
is varied and the accuracy is plotted against those values. To measure the
accuracy, the true nearest neighbors are required, therefore
:class:`sklearn.neighbors.NearestNeighbors` is used to compute the exact
neighbors.
"""
from __future__ import division
print(__doc__)

# Author: Maheshakya Wijewardena <maheshakya.10@cse.mrt.ac.lk>
#
# License: BSD 3 clause

###############################################################################
import numpy as np
from sklearn.datasets.samples_generator import make_blobs
from sklearn.neighbors import LSHForest
from sklearn.neighbors import NearestNeighbors
import matplotlib.pyplot as plt

# Initialize size of the database, iterations and required neighbors.
n_samples = 10000
n_features = 100
n_queries = 30
rng = np.random.RandomState(42)

# Generate sample data
X, _ = make_blobs(n_samples=n_samples + n_queries,
                  n_features=n_features, centers=10,
                  random_state=0)
X_index = X[:n_samples]
X_query = X[n_samples:]
# Get exact neighbors
nbrs = NearestNeighbors(n_neighbors=1, algorithm='brute',
                        metric='cosine').fit(X_index)
neighbors_exact = nbrs.kneighbors(X_query, return_distance=False)

# Set `n_candidate` values
n_candidates_values = np.linspace(10, 500, 5).astype(np.int)
n_estimators_for_candidate_value = [1, 5, 10]
n_iter = 10
stds_accuracies = np.zeros((len(n_estimators_for_candidate_value),
                            n_candidates_values.shape[0]),
                           dtype=float)
accuracies_c = np.zeros((len(n_estimators_for_candidate_value),
                         n_candidates_values.shape[0]), dtype=float)

# LSH Forest is a stochastic index: perform several iteration to estimate
# expected accuracy and standard deviation displayed as error bars in
# the plots
for j, value in enumerate(n_estimators_for_candidate_value):
    for i, n_candidates in enumerate(n_candidates_values):
        accuracy_c = []
        for seed in range(n_iter):
            lshf = LSHForest(n_estimators=value,
                             n_candidates=n_candidates, n_neighbors=1,
                             random_state=seed)
            # Build the LSH Forest index
            lshf.fit(X_index)
            # Get neighbors
            neighbors_approx = lshf.kneighbors(X_query,
                                               return_distance=False)
            accuracy_c.append(np.sum(np.equal(neighbors_approx,
                                              neighbors_exact)) /
                              n_queries)

        stds_accuracies[j, i] = np.std(accuracy_c)
        accuracies_c[j, i] = np.mean(accuracy_c)

# Set `n_estimators` values
n_estimators_values = [1, 5, 10, 20, 30, 40, 50]
accuracies_trees = np.zeros(len(n_estimators_values), dtype=float)

# Calculate average accuracy for each value of `n_estimators`
for i, n_estimators in enumerate(n_estimators_values):
    lshf = LSHForest(n_estimators=n_estimators, n_neighbors=1)
    # Build the LSH Forest index
    lshf.fit(X_index)
    # Get neighbors
    neighbors_approx = lshf.kneighbors(X_query, return_distance=False)
    accuracies_trees[i] = np.sum(np.equal(neighbors_approx,
                                          neighbors_exact))/n_queries

###############################################################################
# Plot the accuracy variation with `n_candidates`
plt.figure()
colors = ['c', 'm', 'y']
for i, n_estimators in enumerate(n_estimators_for_candidate_value):
    label = 'n_estimators = %d ' % n_estimators
    plt.plot(n_candidates_values, accuracies_c[i, :],
             'o-', c=colors[i], label=label)
    plt.errorbar(n_candidates_values, accuracies_c[i, :],
                 stds_accuracies[i, :], c=colors[i])

plt.legend(loc='upper left', fontsize='small')
plt.ylim([0, 1.2])
plt.xlim(min(n_candidates_values), max(n_candidates_values))
plt.ylabel("Accuracy")
plt.xlabel("n_candidates")
plt.grid(which='both')
plt.title("Accuracy variation with n_candidates")

# Plot the accuracy variation with `n_estimators`
plt.figure()
plt.scatter(n_estimators_values, accuracies_trees, c='k')
plt.plot(n_estimators_values, accuracies_trees, c='g')
plt.ylim([0, 1.2])
plt.xlim(min(n_estimators_values), max(n_estimators_values))
plt.ylabel("Accuracy")
plt.xlabel("n_estimators")
plt.grid(which='both')
plt.title("Accuracy variation with n_estimators")

plt.show()


"""
=============================================
A demo of the mean-shift clustering algorithm
=============================================

Reference:

Dorin Comaniciu and Peter Meer, "Mean Shift: A robust approach toward
feature space analysis". IEEE Transactions on Pattern Analysis and
Machine Intelligence. 2002. pp. 603-619.

"""
print(__doc__)

import numpy as np
from sklearn.cluster import MeanShift, estimate_bandwidth
from sklearn.datasets import make_blobs

# #############################################################################
# Generate sample data
centers = [[1, 1], [-1, -1], [1, -1]]
X, _ = make_blobs(n_samples=10000, centers=centers, cluster_std=0.6)

# #############################################################################
# Compute clustering with MeanShift

# The following bandwidth can be automatically detected using
bandwidth = estimate_bandwidth(X, quantile=0.2, n_samples=500)

ms = MeanShift(bandwidth=bandwidth, bin_seeding=True)
ms.fit(X)
labels = ms.labels_
cluster_centers = ms.cluster_centers_

labels_unique = np.unique(labels)
n_clusters_ = len(labels_unique)

print("number of estimated clusters : %d" % n_clusters_)

# #############################################################################
# Plot result
import matplotlib.pyplot as plt
from itertools import cycle

plt.figure(1)
plt.clf()

colors = cycle('bgrcmykbgrcmykbgrcmykbgrcmyk')
for k, col in zip(range(n_clusters_), colors):
    my_members = labels == k
    cluster_center = cluster_centers[k]
    plt.plot(X[my_members, 0], X[my_members, 1], col + '.')
    plt.plot(cluster_center[0], cluster_center[1], 'o', markerfacecolor=col,
             markeredgecolor='k', markersize=14)
plt.title('Estimated number of clusters: %d' % n_clusters_)
plt.show()


"""
=================================================
Anomaly detection with Local Outlier Factor (LOF)
=================================================

This example presents the Local Outlier Factor (LOF) estimator. The LOF
algorithm is an unsupervised outlier detection method which computes the local
density deviation of a given data point with respect to its neighbors.
It considers as outlier samples that have a substantially lower density than
their neighbors.

The number of neighbors considered, (parameter n_neighbors) is typically
chosen 1) greater than the minimum number of objects a cluster has to contain,
so that other objects can be local outliers relative to this cluster, and 2)
smaller than the maximum number of close by objects that can potentially be
local outliers.
In practice, such informations are generally not available, and taking
n_neighbors=20 appears to work well in general.
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.neighbors import LocalOutlierFactor
print(__doc__)

np.random.seed(42)

# Generate train data
X = 0.3 * np.random.randn(100, 2)
# Generate some abnormal novel observations
X_outliers = np.random.uniform(low=-4, high=4, size=(20, 2))
X = np.r_[X + 2, X - 2, X_outliers]

# fit the model
clf = LocalOutlierFactor(n_neighbors=20)
y_pred = clf.fit_predict(X)
y_pred_outliers = y_pred[200:]

# plot the level sets of the decision function
xx, yy = np.meshgrid(np.linspace(-5, 5, 50), np.linspace(-5, 5, 50))
Z = clf._decision_function(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)

plt.title("Local Outlier Factor (LOF)")
plt.contourf(xx, yy, Z, cmap=plt.cm.Blues_r)

a = plt.scatter(X[:200, 0], X[:200, 1], c='white')
b = plt.scatter(X[200:, 0], X[200:, 1], c='red')
plt.axis('tight')
plt.xlim((-5, 5))
plt.ylim((-5, 5))
plt.legend([a, b],
           ["normal observations",
            "abnormal observations"],
           loc="upper left")
plt.show()


from __future__ import division, print_function, absolute_import

import pytest
from numpy.testing import assert_, assert_array_equal
from scipy._lib._numpy_compat import suppress_warnings

try:
    import matplotlib
    matplotlib.rcParams['backend'] = 'Agg'
    import matplotlib.pyplot as plt
    from matplotlib.collections import LineCollection
    from matplotlib import MatplotlibDeprecationWarning
    has_matplotlib = True
except:
    has_matplotlib = False

from scipy.spatial import \
     delaunay_plot_2d, voronoi_plot_2d, convex_hull_plot_2d, \
     Delaunay, Voronoi, ConvexHull

@pytest.mark.skipif(not has_matplotlib, reason="Matplotlib not available")
class TestPlotting:
    points = [(0,0), (0,1), (1,0), (1,1)]

    def test_delaunay(self):
        # Smoke test
        fig = plt.figure()
        obj = Delaunay(self.points)
        s_before = obj.simplices.copy()
        with suppress_warnings as sup:
            # filter can be removed when matplotlib 1.x is dropped
            sup.filter(message="The ishold function was deprecated in version")
            r = delaunay_plot_2d(obj, ax=fig.gca())
        assert_array_equal(obj.simplices, s_before)  # shouldn't modify
        assert_(r is fig)
        delaunay_plot_2d(obj, ax=fig.gca())

    def test_voronoi(self):
        # Smoke test
        fig = plt.figure()
        obj = Voronoi(self.points)
        with suppress_warnings as sup:
            # filter can be removed when matplotlib 1.x is dropped
            sup.filter(message="The ishold function was deprecated in version")
            r = voronoi_plot_2d(obj, ax=fig.gca())
        assert_(r is fig)
        voronoi_plot_2d(obj)
        voronoi_plot_2d(obj, show_vertices=False)

    def test_convex_hull(self):
        # Smoke test
        fig = plt.figure()
        tri = ConvexHull(self.points)
        with suppress_warnings as sup:
            # filter can be removed when matplotlib 1.x is dropped
            sup.filter(message="The ishold function was deprecated in version")
            r = convex_hull_plot_2d(tri, ax=fig.gca())
        assert_(r is fig)
        convex_hull_plot_2d(tri)


import warnings
import itertools
from copy import copy
from functools import partial
from collections.abc import Iterable, Sequence, Mapping
from numbers import Number
from datetime import datetime

import numpy as np
import pandas as pd
import matplotlib as mpl

from ._decorators import (
    share_init_params_with_map,
)
from .palettes import (
    QUAL_PALETTES,
    color_palette,
    cubehelix_palette,
    _parse_cubehelix_args,
)
from .utils import (
    get_color_cycle,
    remove_na,
)

class SemanticMapping:
    """Base class for mapping data values to plot attributes."""

    # -- Default attributes that all SemanticMapping subclasses must set

    # Whether the mapping is numeric, categorical, or datetime
    map_type = None

    # Ordered list of unique values in the input data
    levels = None

    # A mapping from the data values to corresponding plot attributes
    lookup_table = None

    def __init__(self, plotter):

        # TODO Putting this here so we can continue to use a lot of the
        # logic that's built into the library, but the idea of this class
        # is to move towards semantic mappings that are agnositic about the
        # kind of plot they're going to be used to draw.
        # Fully achieving that is going to take some thinking.
        self.plotter = plotter

    def map(cls, plotter, *args, **kwargs):
        # This method is assigned the __init__ docstring
        method_name = "_{}_map".format(cls.__name__[:-7].lower())
        setattr(plotter, method_name, cls(plotter, *args, **kwargs))
        return plotter

    def _lookup_single(self, key):
        """Apply the mapping to a single data value."""
        return self.lookup_table[key]

    def __call__(self, key, *args, **kwargs):
        """Get the attribute(s) values for the data key."""
        if isinstance(key, (list, np.ndarray, pd.Series)):
            return [self._lookup_single(k, *args, **kwargs) for k in key]
        else:
            return self._lookup_single(key, *args, **kwargs)

@share_init_params_with_map
class HueMapping(SemanticMapping):
    """Mapping that sets artist colors according to data values."""
    # A specification of the colors that should appear in the plot
    palette = None

    # An object that normalizes data values to [0, 1] range for color mapping
    norm = None

    # A continuous colormap object for interpolating in a numeric context
    cmap = None

    def __init__(
        self, plotter, palette=None, order=None, norm=None,
    ):
        """Map the levels of the `hue` variable to distinct colors.

        Parameters
        ----------
        # TODO add generic parameters

        """
        super().__init__(plotter)

        data = plotter.plot_data["hue"]

        if data.notna().any():

            map_type = self.infer_map_type(
                palette, norm, plotter.input_format, plotter.var_types["hue"]
            )

            # Our goal is to end up with a dictionary mapping every unique
            # value in `data` to a color. We will also keep track of the
            # metadata about this mapping we will need for, e.g., a legend

            # --- Option 1: numeric mapping with a matplotlib colormap

            if map_type == "numeric":

                data = pd.to_numeric(data)
                levels, lookup_table, norm, cmap = self.numeric_mapping(
                    data, palette, norm,
                )

            # --- Option 2: categorical mapping using seaborn palette

            elif map_type == "categorical":

                cmap = norm = None
                levels, lookup_table = self.categorical_mapping(
                    data, palette, order,
                )

            # --- Option 3: datetime mapping

            else:
                # TODO this needs actual implementation
                cmap = norm = None
                levels, lookup_table = self.categorical_mapping(
                    # Casting data to list to handle differences in the way
                    # pandas and numpy represent datetime64 data
                    list(data), palette, order,
                )

            self.map_type = map_type
            self.lookup_table = lookup_table
            self.palette = palette
            self.levels = levels
            self.norm = norm
            self.cmap = cmap

    def _lookup_single(self, key):
        """Get the color for a single value, using colormap to interpolate."""
        try:
            # Use a value that's in the original data vector
            value = self.lookup_table[key]
        except KeyError:
            # Use the colormap to interpolate between existing datapoints
            # (e.g. in the context of making a continuous legend)
            normed = self.norm(key)
            if np.ma.is_masked(normed):
                normed = np.nan
            value = self.cmap(normed)
        return value

    def infer_map_type(self, palette, norm, input_format, var_type):
        """Determine how to implement the mapping."""
        if palette in QUAL_PALETTES:
            map_type = "categorical"
        elif norm is not None:
            map_type = "numeric"
        elif isinstance(palette, (dict, list)):
            map_type = "categorical"
        elif input_format == "wide":
            map_type = "categorical"
        else:
            map_type = var_type

        return map_type

    def categorical_mapping(self, data, palette, order):
        """Determine colors when the hue mapping is categorical."""
        # -- Identify the order and name of the levels

        levels = categorical_order(data, order)
        n_colors = len(levels)

        # -- Identify the set of colors to use

        if isinstance(palette, dict):

            missing = set(levels) - set(palette)
            if any(missing):
                err = "The palette dictionary is missing keys: {}"
                raise ValueError(err.format(missing))

            lookup_table = palette

        else:

            if palette is None:
                if n_colors <= len(get_color_cycle()):
                    colors = color_palette(None, n_colors)
                else:
                    colors = color_palette("husl", n_colors)
            elif isinstance(palette, list):
                if len(palette) != n_colors:
                    err = "The palette list has the wrong number of colors."
                    raise ValueError(err)
                colors = palette
            else:
                colors = color_palette(palette, n_colors)

            lookup_table = dict(zip(levels, colors))

        return levels, lookup_table

    def numeric_mapping(self, data, palette, norm):
        """Determine colors when the hue variable is quantitative."""
        if isinstance(palette, dict):

            # The presence of a norm object overrides a dictionary of hues
            # in specifying a numeric mapping, so we need to process it here.
            levels = list(sorted(palette))
            colors = [palette[k] for k in sorted(palette)]
            cmap = mpl.colors.ListedColormap(colors)
            lookup_table = palette.copy()

        else:

            # The levels are the sorted unique values in the data
            levels = list(np.sort(remove_na(data.unique())))

            # --- Sort out the colormap to use from the palette argument

            # Default numeric palette is our default cubehelix palette
            # TODO do we want to do something complicated to ensure contrast?
            palette = "ch:" if palette is None else palette

            if isinstance(palette, mpl.colors.Colormap):
                cmap = palette
            elif str(palette).startswith("ch:"):
                args, kwargs = _parse_cubehelix_args(palette)
                cmap = cubehelix_palette(0, *args, as_cmap=True, **kwargs)
            else:
                try:
                    cmap = mpl.cm.get_cmap(palette)
                except (ValueError, TypeError):
                    err = "Palette {} not understood"
                    raise ValueError(err)

            # Now sort out the data normalization
            if norm is None:
                norm = mpl.colors.Normalize()
            elif isinstance(norm, tuple):
                norm = mpl.colors.Normalize(*norm)
            elif not isinstance(norm, mpl.colors.Normalize):
                err = "``hue_norm`` must be None, tuple, or Normalize object."
                raise ValueError(err)

            if not norm.scaled():
                norm(np.asarray(data.dropna()))

            lookup_table = dict(zip(levels, cmap(norm(levels))))

        return levels, lookup_table, norm, cmap

@share_init_params_with_map
class SizeMapping(SemanticMapping):
    """Mapping that sets artist sizes according to data values."""
    # An object that normalizes data values to [0, 1] range
    norm = None

    def __init__(
        self, plotter, sizes=None, order=None, norm=None,
    ):
        """Map the levels of the `size` variable to distinct values.

        Parameters
        ----------
        # TODO add generic parameters

        """
        super().__init__(plotter)

        data = plotter.plot_data["size"]

        if data.notna().any():

            map_type = self.infer_map_type(
                norm, sizes, plotter.var_types["size"]
            )

            # --- Option 1: numeric mapping

            if map_type == "numeric":

                levels, lookup_table, norm = self.numeric_mapping(
                    data, sizes, norm,
                )

            # --- Option 2: categorical mapping

            elif map_type == "categorical":

                levels, lookup_table = self.categorical_mapping(
                    data, sizes, order,
                )

            # --- Option 3: datetime mapping

            # TODO this needs an actual implementation
            else:

                levels, lookup_table = self.categorical_mapping(
                    # Casting data to list to handle differences in the way
                    # pandas and numpy represent datetime64 data
                    list(data), sizes, order,
                )

            self.map_type = map_type
            self.levels = levels
            self.norm = norm
            self.sizes = sizes
            self.lookup_table = lookup_table

    def infer_map_type(self, norm, sizes, var_type):

        if norm is not None:
            map_type = "numeric"
        elif isinstance(sizes, (dict, list)):
            map_type = "categorical"
        else:
            map_type = var_type

        return map_type

    def _lookup_single(self, key):

        try:
            value = self.lookup_table[key]
        except KeyError:
            normed = self.norm(key)
            if np.ma.is_masked(normed):
                normed = np.nan
            size_values = self.lookup_table.values()
            size_range = min(size_values), max(size_values)
            value = size_range[0] + normed * np.ptp(size_range)
        return value

    def categorical_mapping(self, data, sizes, order):

        levels = categorical_order(data, order)

        if isinstance(sizes, dict):

            # Dict inputs map existing data values to the size attribute
            missing = set(levels) - set(sizes)
            if any(missing):
                err = f"Missing sizes for the following levels: {missing}"
                raise ValueError(err)
            lookup_table = sizes.copy()

        elif isinstance(sizes, list):

            # List inputs give size values in the same order as the levels
            if len(sizes) != len(levels):
                err = "The `sizes` list has the wrong number of values."
                raise ValueError(err)

            lookup_table = dict(zip(levels, sizes))

        else:

            if isinstance(sizes, tuple):

                # Tuple input sets the min, max size values
                if len(sizes) != 2:
                    err = "A `sizes` tuple must have only 2 values"
                    raise ValueError(err)

            elif sizes is not None:

                err = f"Value for `sizes` not understood: {sizes}"
                raise ValueError(err)

            else:

                # Otherwise, we need to get the min, max size values from
                # the plotter object we are attached to.

                # TODO this is going to cause us trouble later, because we
                # want to restructure things so that the plotter is generic
                # across the visual representation of the data. But at this
                # point, we don't know the visual representation. Likely we
                # want to change the logic of this Mapping so that it gives
                # points on a nornalized range that then gets unnormalized
                # when we know what we're drawing. But given the way the
                # package works now, this way is cleanest.
                sizes = self.plotter._default_size_range

            # For categorical sizes, use regularly-spaced linear steps
            # between the minimum and maximum sizes
            sizes = np.linspace(*sizes, len(levels))
            lookup_table = dict(zip(levels, sizes))

        return levels, lookup_table

    def numeric_mapping(self, data, sizes, norm):

        if isinstance(sizes, dict):
            # The presence of a norm object overrides a dictionary of sizes
            # in specifying a numeric mapping, so we need to process it
            # dictionary here
            levels = list(np.sort(list(sizes)))
            size_values = sizes.values()
            size_range = min(size_values), max(size_values)

        else:

            # The levels here will be the unique values in the data
            levels = list(np.sort(remove_na(data.unique())))

            if isinstance(sizes, tuple):

                # For numeric inputs, the size can be parametrized by
                # the minimum and maximum artist values to map to. The
                # norm object that gets set up next specifies how to
                # do the mapping.

                if len(sizes) != 2:
                    err = "A `sizes` tuple must have only 2 values"
                    raise ValueError(err)

                size_range = sizes

            elif sizes is not None:

                err = f"Value for `sizes` not understood: {sizes}"
                raise ValueError(err)

            else:

                # When not provided, we get the size range from the plotter
                # object we are attached to. See the note in the categorical
                # method about how this is suboptimal for future development.:
                size_range = self.plotter._default_size_range

        # Now that we know the minimum and maximum sizes that will get drawn,
        # we need to map the data values that we have into that range. We will
        # use a matplotlib Normalize class, which is typically used for numeric
        # color mapping but works fine here too. It takes data values and maps
        # them into a [0, 1] interval, potentially nonlinear-ly.

        if norm is None:
            # Default is a linear function between the min and max data values
            norm = mpl.colors.Normalize()
        elif isinstance(norm, tuple):
            # It is also possible to give different limits in data space
            norm = mpl.colors.Normalize(*norm)
        elif not isinstance(norm, mpl.colors.Normalize):
            err = f"Value for size `norm` parameter not understood: {norm}"
            raise ValueError(err)
        else:
            # If provided with Normalize object, copy it so we can modify
            norm = copy(norm)

        # Set the mapping so all output values are in [0, 1]
        norm.clip = True

        # If the input range is not set, use the full range of the data
        if not norm.scaled():
            norm(levels)

        # Map from data values to [0, 1] range
        sizes_scaled = norm(levels)

        # Now map from the scaled range into the artist units
        if isinstance(sizes, dict):
            lookup_table = sizes
        else:
            lo, hi = size_range
            sizes = lo + sizes_scaled * (hi - lo)
            lookup_table = dict(zip(levels, sizes))

        return levels, lookup_table, norm

@share_init_params_with_map
class StyleMapping(SemanticMapping):
    """Mapping that sets artist style according to data values."""

    # Style mapping is always treated as categorical
    map_type = "categorical"

    def __init__(
        self, plotter, markers=None, dashes=None, order=None,
    ):
        """Map the levels of the `style` variable to distinct values.

        Parameters
        ----------
        # TODO add generic parameters

        """
        super().__init__(plotter)

        data = plotter.plot_data["style"]

        if data.notna().any():

            # Cast to list to handle numpy/pandas datetime quirks
            if variable_type(data) == "datetime":
                data = list(data)

            # Find ordered unique values
            levels = categorical_order(data, order)

            markers = self._map_attributes(
                markers, levels, unique_markers(len(levels)), "markers",
            )
            dashes = self._map_attributes(
                dashes, levels, unique_dashes(len(levels)), "dashes",
            )

            # Build the paths matplotlib will use to draw the markers
            paths = {}
            filled_markers = []
            for k, m in markers.items():
                if not isinstance(m, mpl.markers.MarkerStyle):
                    m = mpl.markers.MarkerStyle(m)
                paths[k] = m.get_path().transformed(m.get_transform())
                filled_markers.append(m.is_filled())

            # Mixture of filled and unfilled markers will show line art markers
            # in the edge color, which defaults to white. This can be handled,
            # but there would be additional complexity with specifying the
            # weight of the line art markers without overwhelming the filled
            # ones with the edges. So for now, we will disallow mixtures.
            if any(filled_markers) and not all(filled_markers):
                err = "Filled and line art markers cannot be mixed"
                raise ValueError(err)

            lookup_table = {}
            for key in levels:
                lookup_table[key] = {}
                if markers:
                    lookup_table[key]["marker"] = markers[key]
                    lookup_table[key]["path"] = paths[key]
                if dashes:
                    lookup_table[key]["dashes"] = dashes[key]

            self.levels = levels
            self.lookup_table = lookup_table

    def _lookup_single(self, key, attr=None):
        """Get attribute(s) for a given data point."""
        if attr is None:
            value = self.lookup_table[key]
        else:
            value = self.lookup_table[key][attr]
        return value

    def _map_attributes(self, arg, levels, defaults, attr):
        """Handle the specification for a given style attribute."""
        if arg is True:
            lookup_table = dict(zip(levels, defaults))
        elif isinstance(arg, dict):
            missing = set(levels) - set(arg)
            if missing:
                err = f"These `{attr}` levels are missing values: {missing}"
                raise ValueError(err)
            lookup_table = arg
        elif isinstance(arg, Sequence):
            if len(levels) != len(arg):
                err = f"The `{attr}` argument has the wrong number of values"
                raise ValueError(err)
            lookup_table = dict(zip(levels, arg))
        elif arg:
            err = f"This `{attr}` argument was not understood: {arg}"
            raise ValueError(err)
        else:
            lookup_table = {}

        return lookup_table

# =========================================================================== #

class VectorPlotter:
    """Base class for objects underlying *plot functions."""

    _semantic_mappings = {
        "hue": HueMapping,
        "size": SizeMapping,
        "style": StyleMapping,
    }

    # TODO units is another example of a non-mapping "semantic"
    # we need a general name for this and separate handling
    semantics = "x", "y", "hue", "size", "style", "units"
    wide_structure = {
        "x": "index", "y": "values", "hue": "columns", "style": "columns",
    }
    flat_structure = {"x": "index", "y": "values"}

    _default_size_range = 1, 2  # Unused but needed in tests, ugh

    def __init__(self, data=None, variables={}):

        self.assign_variables(data, variables)

        for var, cls in self._semantic_mappings.items():
            if var in self.semantics:

                # Create the mapping function
                map_func = partial(cls.map, plotter=self)
                setattr(self, f"map_{var}", map_func)

                # Call the mapping function to initialize with default values
                getattr(self, f"map_{var}")()

    @classmethod
    def get_semantics(cls, kwargs):
        """Subset a dictionary` arguments with known semantic variables."""
        return {k: kwargs[k] for k in cls.semantics}

    def assign_variables(self, data=None, variables={}):
        """Define plot variables, optionally using lookup from `data`."""
        x = variables.get("x", None)
        y = variables.get("y", None)

        if x is None and y is None:
            self.input_format = "wide"
            plot_data, variables = self._assign_variables_wideform(
                data, **variables,
            )
        else:
            self.input_format = "long"
            plot_data, variables = self._assign_variables_longform(
                data, **variables,
            )

        self.plot_data = plot_data
        self.variables = variables
        self.var_types = {
            v: variable_type(
                plot_data[v],
                boolean_type="numeric" if v in "xy" else "categorical"
            )
            for v in variables
        }

        return self

    def _assign_variables_wideform(self, data=None, **kwargs):
        """Define plot variables given wide-form data.

        Parameters
        ----------
        data : flat vector or collection of vectors
            Data can be a vector or mapping that is coerceable to a Series
            or a sequence- or mapping-based collection of such vectors, or a
            rectangular numpy array, or a Pandas DataFrame.
        kwargs : variable -> data mappings
            Behavior with keyword arguments is currently undefined.

        Returns
        -------
        plot_data : :class:`pandas.DataFrame`
            Long-form data object mapping seaborn variables (x, y, hue, ...)
            to data vectors.
        variables : dict
            Keys are defined seaborn variables; values are names inferred from
            the inputs (or None when no name can be determined).

        """
        # TODO raise here if any kwarg values are not None,
        # # if we decide for "structure-only" wide API

        # First, determine if the data object actually has any data in it
        empty = data is None or not len(data)

        # Then, determine if we have "flat" data (a single vector)
        if isinstance(data, dict):
            values = data.values()
        else:
            values = np.atleast_1d(data)
        flat = not any(
            isinstance(v, Iterable) and not isinstance(v, (str, bytes))
            for v in values
        )

        if empty:

            # Make an object with the structure of plot_data, but empty
            plot_data = pd.DataFrame(columns=self.semantics)
            variables = {}

        elif flat:

            # Handle flat data by converting to pandas Series and using the
            # index and/or values to define x and/or y
            # (Could be accomplished with a more general to_series() interface)
            flat_data = pd.Series(data).copy()
            names = {
                "values": flat_data.name,
                "index": flat_data.index.name
            }

            plot_data = {}
            variables = {}

            for var in ["x", "y"]:
                if var in self.flat_structure:
                    attr = self.flat_structure[var]
                    plot_data[var] = getattr(flat_data, attr)
                    variables[var] = names[self.flat_structure[var]]

            plot_data = pd.DataFrame(plot_data).reindex(columns=self.semantics)

        else:

            # Otherwise assume we have some collection of vectors.

            # Handle Python sequences such that entries end up in the columns,
            # not in the rows, of the intermediate wide DataFrame.
            # One way to accomplish this is to convert to a dict of Series.
            if isinstance(data, Sequence):
                data_dict = {}
                for i, var in enumerate(data):
                    key = getattr(var, "name", i)
                    # TODO is there a safer/more generic way to ensure Series?
                    # sort of like np.asarray, but for pandas?
                    data_dict[key] = pd.Series(var)

                data = data_dict

            # Pandas requires that dict values either be Series objects
            # or all have the same length, but we want to allow "ragged" inputs
            if isinstance(data, Mapping):
                data = {key: pd.Series(val) for key, val in data.items()}

            # Otherwise, delegate to the pandas DataFrame constructor
            # This is where we'd prefer to use a general interface that says
            # "give me this data as a pandas DataFrame", so we can accept
            # DataFrame objects from other libraries
            wide_data = pd.DataFrame(data, copy=True)

            # At this point we should reduce the dataframe to numeric cols
            numeric_cols = wide_data.apply(variable_type) == "numeric"
            wide_data = wide_data.loc[:, numeric_cols]

            # Now melt the data to long form
            melt_kws = {"var_name": "columns", "value_name": "values"}
            if "index" in self.wide_structure.values():
                melt_kws["id_vars"] = "index"
                wide_data["index"] = wide_data.index.to_series()
            plot_data = wide_data.melt(**melt_kws)

            # Assign names corresponding to plot semantics
            for var, attr in self.wide_structure.items():
                plot_data[var] = plot_data[attr]
            plot_data = plot_data.reindex(columns=self.semantics)

            # Define the variable names
            variables = {}
            for var, attr in self.wide_structure.items():
                obj = getattr(wide_data, attr)
                variables[var] = getattr(obj, "name", None)

        return plot_data, variables

    def _assign_variables_longform(self, data=None, **kwargs):
        """Define plot variables given long-form data and/or vector inputs.

        Parameters
        ----------
        data : dict-like collection of vectors
            Input data where variable names map to vector values.
        kwargs : variable -> data mappings
            Keys are seaborn variables (x, y, hue, ...) and values are vectors
            in any format that can construct a :class:`pandas.DataFrame` or
            names of columns or index levels in ``data``.

        Returns
        -------
        plot_data : :class:`pandas.DataFrame`
            Long-form data object mapping seaborn variables (x, y, hue, ...)
            to data vectors.
        variables : dict
            Keys are defined seaborn variables; values are names inferred from
            the inputs (or None when no name can be determined).

        Raises
        ------
        ValueError
            When variables are strings that don't appear in ``data``.

        """
        plot_data = {}
        variables = {}

        # Data is optional; all variables can be defined as vectors
        if data is None:
            data = {}

        # TODO should we try a data.to_dict() or similar here to more
        # generally accept objects with that interface?
        # Note that dict(df) also works for pandas, and gives us what we
        # want, whereas DataFrame.to_dict() gives a nested dict instead of
        # a dict of series.

        # Variables can also be extraced from the index attribute
        # TODO is this the most general way to enable it?
        # There is no index.to_dict on multiindex, unfortunately
        try:
            index = data.index.to_frame()
        except AttributeError:
            index = {}

        # The caller will determine the order of variables in plot_data
        for key, val in kwargs.items():

            if isinstance(val, (str, bytes)):
                # String inputs trigger __getitem__
                if val in data:
                    # First try to get an entry in the data object
                    plot_data[key] = data[val]
                    variables[key] = val
                elif val in index:
                    # Failing that, try to get an entry in the index object
                    plot_data[key] = index[val]
                    variables[key] = val
                else:
                    # We don't know what this name means
                    err = f"Could not interpret input '{val}'"
                    raise ValueError(err)

            else:

                # Otherwise, assume the value is itself a vector of data
                # TODO check for 1D here or let pd.DataFrame raise?
                plot_data[key] = val
                # Try to infer the name of the variable
                variables[key] = getattr(val, "name", None)

        # Construct a tidy plot DataFrame. This will convert a number of
        # types automatically, aligning on index in case of pandas objects
        plot_data = pd.DataFrame(plot_data, columns=self.semantics)

        # Reduce the variables dictionary to fields with valid data
        variables = {
            var: name
            for var, name in variables.items()
            if plot_data[var].notnull().any()
        }

        return plot_data, variables

    def _semantic_subsets(
        self, grouping_semantics, reverse=False, from_comp_data=False,
    ):
        """Generator for getting subsets of data defined by semantic variables.

        Parameters
        ----------
        grouping_semantics : list of strings
            Semantic variables that define the subsets of data.
        reverse : bool, optional
            If True, reverse the order of iteration.
        from_comp_data : bool, optional
            If True, use self.comp_data rather than self.plot_data

        Yields
        ------
        sub_vars : dict
            Keys are semantic names, values are the level of that semantic.
        sub_data : :class:`pandas.DataFrame`
            Subset of ``plot_data`` for this combination of semantic values.

        """
        if isinstance(grouping_semantics, str):
            grouping_semantics = [grouping_semantics]

        # Reduce to the semantics used in this plot
        grouping_semantics = [
            var for var in grouping_semantics if var in self.variables
        ]

        if from_comp_data:
            data = self.comp_data
        else:
            data = self.plot_data

        if grouping_semantics:

            grouped_data = data.groupby(
                grouping_semantics, sort=False, as_index=False
            )

            grouping_keys = []
            for var in grouping_semantics:
                # TODO this is messy, add "semantic levels" property?
                map_obj = getattr(self, f"_{var}_map")
                grouping_keys.append(map_obj.levels)

            iter_keys = itertools.product(*grouping_keys)
            if reverse:
                iter_keys = reversed(list(iter_keys))

            for key in iter_keys:

                # Pandas fails with singleton tuple inputs
                pd_key = key[0] if len(key) == 1 else key

                try:
                    data_subset = grouped_data.get_group(pd_key)
                except KeyError:
                    continue

                yield dict(zip(grouping_semantics, key)), data_subset

        else:

            yield {}, data

    @property
    def comp_data(self):
        """Dataframe with numeric x and y, after unit conversion and log scaling."""
        if not hasattr(self, "ax"):
            # Probably a good idea, but will need a bunch of tests updated
            # Most of these tests should just use the external interface
            # Then this can be reeneabled.
            # raise AttributeError("No Axes attached to plotter")
            return self.plot_data

        if not hasattr(self, "_comp_data"):

            comp_data = self.plot_data.copy(deep=False)
            for var in "xy":
                axis = getattr(self.ax, f"{var}axis")
                comp_var = axis.convert_units(self.plot_data[var])
                if axis.get_scale() == "log":
                    comp_var = np.log10(comp_var)
                comp_data[var] = comp_var
            self._comp_data = comp_data

        return self._comp_data

    def _attach(self, ax, allowed_types=None, log_scale=None):
        """Associate the plotter with a matplotlib Axes and initialize its units.

        Parameters
        ----------
        ax : :class:`matplotlib.axes.Axes`
            Axes object that we will eventually plot onto.
        allowed_types : str or list of str
            If provided, raise when either the x or y variable does not have
            one of the declared seaborn types.
        log_scale : bool, number, or pair of bools or numbers
            If not False, set the axes to use log scaling, with the given
            base or defaulting to 10. If a tuple, interpreted as separate
            arguments for the x and y axes.

        """
        if allowed_types is None:
            # TODO should we define this default somewhere?
            allowed_types = ["numeric", "datetime", "categorical"]
        elif isinstance(allowed_types, str):
            allowed_types = [allowed_types]

        for var in set("xy").intersection(self.variables):

            # Check types of x/y variables
            var_type = self.var_types[var]
            if var_type not in allowed_types:
                err = (
                    f"The {var} variable is {var_type}, but one of "
                    f"{allowed_types} is required"
                )
                raise TypeError(err)

            # Register with the matplotlib unit conversion machinery
            # TODO do we want to warn or raise if mixing units?
            axis = getattr(ax, f"{var}axis")
            seed_data = self.plot_data[var]
            if var_type == "categorical":
                seed_data = categorical_order(seed_data)
            axis.update_units(seed_data)

        # Possibly log-scale one or both axes
        if log_scale is not None:
            # Allow single value or x, y tuple
            try:
                scalex, scaley = log_scale
            except TypeError:
                scalex = log_scale if "x" in self.variables else False
                scaley = log_scale if "y" in self.variables else False

            for axis, scale in zip("xy", (scalex, scaley)):
                if scale:
                    set_scale = getattr(ax, f"set_{axis}scale")
                    if scale is True:
                        set_scale("log")
                    else:
                        set_scale("log", **{f"base{axis}": scale})

        self.ax = ax

    def _add_axis_labels(self, ax, default_x="", default_y=""):
        """Add axis labels from internal variable names if not already existing."""
        if not ax.get_xlabel():
            ax.set_xlabel(self.variables.get("x", default_x))
        if not ax.get_ylabel():
            ax.set_ylabel(self.variables.get("y", default_y))

def variable_type(vector, boolean_type="numeric"):
    """Determine whether a vector contains numeric, categorical, or dateime data.

    This function differs from the pandas typing API in two ways:

    - Python sequences or object-typed PyData objects are considered numeric if
      all of their entries are numeric.
    - String or mixed-type data are considered categorical even if not
      explicitly represented as a :class:pandas.api.types.CategoricalDtype`.

    Parameters
    ----------
    vector : :func:`pandas.Series`, :func:`numpy.ndarray`, or Python sequence
        Input data to test.
    binary_type : 'numeric' or 'categorical'
        Type to use for vectors containing only 0s and 1s (and NAs).

    Returns
    -------
    var_type : 'numeric', 'categorical', or 'datetime'
        Name identifying the type of data in the vector.

    """
    # Special-case all-na data, which is always "numeric"
    if pd.isna(vector).all():
        return "numeric"

    # Special-case binary/boolean data, allow caller to determine
    # This triggers a numpy warning when vector has strings/objects
    # https://github.com/numpy/numpy/issues/6784
    # Because we reduce with .all(), we are agnostic about whether the
    # comparison returns a scalar or vector, so we will ignore the warning.
    # It triggers a separate DeprecationWarning when the vector has datetimes:
    # https://github.com/numpy/numpy/issues/13548
    # This is considered a bug by numpy and will likely go away.
    with warnings.catch_warnings():
        warnings.simplefilter(
            action='ignore', category=(FutureWarning, DeprecationWarning)
        )
        if np.isin(vector, [0, 1, np.nan]).all():
            return boolean_type

    # Defer to positive pandas tests
    if pd.api.types.is_numeric_dtype(vector):
        return "numeric"

    if pd.api.types.is_categorical_dtype(vector):
        return "categorical"

    if pd.api.types.is_datetime64_dtype(vector):
        return "datetime"

    # --- If we get to here, we need to check the entries

    # Check for a collection where everything is a number

    def all_numeric(x):
        for x_i in x:
            if not isinstance(x_i, Number):
                return False
        return True

    if all_numeric(vector):
        return "numeric"

    # Check for a collection where everything is a datetime

    def all_datetime(x):
        for x_i in x:
            if not isinstance(x_i, (datetime, np.datetime64)):
                return False
        return True

    if all_datetime(vector):
        return "datetime"

    # Otherwise, our final fallback is to consider things categorical

    return "categorical"

def infer_orient(x=None, y=None, orient=None, require_numeric=True):
    """Determine how the plot should be oriented based on the data.

    For historical reasons, the convention is to call a plot "horizontally"
    or "vertically" oriented based on the axis representing its dependent
    variable. Practically, this is used when determining the axis for
    numerical aggregation.

    Paramters
    ---------
    x, y : Vector data or None
        Positional data vectors for the plot.
    orient : string or None
        Specified orientation, which must start with "v" or "h" if not None.
    require_numeric : bool
        If set, raise when the implied dependent variable is not numeric.

    Returns
    -------
    orient : "v" or "h"

    Raises
    ------
    ValueError: When `orient` is not None and does not start with "h" or "v"
    TypeError: When dependant variable is not numeric, with `require_numeric`

    """

    x_type = None if x is None else variable_type(x)
    y_type = None if y is None else variable_type(y)

    nonnumeric_dv_error = "{} orientation requires numeric `{}` variable."
    single_var_warning = "{} orientation ignored with only `{}` specified."

    if x is None:
        if str(orient).startswith("h"):
            warnings.warn(single_var_warning.format("Horizontal", "y"))
        if require_numeric and y_type != "numeric":
            raise TypeError(nonnumeric_dv_error.format("Vertical", "y"))
        return "v"

    elif y is None:
        if str(orient).startswith("v"):
            warnings.warn(single_var_warning.format("Vertical", "x"))
        if require_numeric and x_type != "numeric":
            raise TypeError(nonnumeric_dv_error.format("Horizontal", "x"))
        return "h"

    elif str(orient).startswith("v"):
        if require_numeric and y_type != "numeric":
            raise TypeError(nonnumeric_dv_error.format("Vertical", "y"))
        return "v"

    elif str(orient).startswith("h"):
        if require_numeric and x_type != "numeric":
            raise TypeError(nonnumeric_dv_error.format("Horizontal", "x"))
        return "h"

    elif orient is not None:
        raise ValueError(f"Value for `orient` not understood: {orient}")

    elif x_type != "numeric" and y_type == "numeric":
        return "v"

    elif x_type == "numeric" and y_type != "numeric":
        return "h"

    elif require_numeric and "numeric" not in (x_type, y_type):
        err = "Neither the `x` nor `y` variable appears to be numeric."
        raise TypeError(err)

    else:
        return "v"

def unique_dashes(n):
    """Build an arbitrarily long list of unique dash styles for lines.

    Parameters
    ----------
    n : int
        Number of unique dash specs to generate.

    Returns
    -------
    dashes : list of strings or tuples
        Valid arguments for the ``dashes`` parameter on
        :class:`matplotlib.lines.Line2D`. The first spec is a solid
        line (``""``), the remainder are sequences of long and short
        dashes.

    """
    # Start with dash specs that are well distinguishable
    dashes = [
        "",
        (4, 1.5),
        (1, 1),
        (3, 1.25, 1.5, 1.25),
        (5, 1, 1, 1),
    ]

    # Now programatically build as many as we need
    p = 3
    while len(dashes) < n:

        # Take combinations of long and short dashes
        a = itertools.combinations_with_replacement([3, 1.25], p)
        b = itertools.combinations_with_replacement([4, 1], p)

        # Interleave the combinations, reversing one of the streams
        segment_list = itertools.chain(*zip(
            list(a)[1:-1][::-1],
            list(b)[1:-1]
        ))

        # Now insert the gaps
        for segments in segment_list:
            gap = min(segments)
            spec = tuple(itertools.chain(*((seg, gap) for seg in segments)))
            dashes.append(spec)

        p += 1

    return dashes[:n]

def unique_markers(n):
    """Build an arbitrarily long list of unique marker styles for points.

    Parameters
    ----------
    n : int
        Number of unique marker specs to generate.

    Returns
    -------
    markers : list of string or tuples
        Values for defining :class:`matplotlib.markers.MarkerStyle` objects.
        All markers will be filled.

    """
    # Start with marker specs that are well distinguishable
    markers = [
        "o",
        "X",
        (4, 0, 45),
        "P",
        (4, 0, 0),
        (4, 1, 0),
        "^",
        (4, 1, 45),
        "v",
    ]

    # Now generate more from regular polygons of increasing order
    s = 5
    while len(markers) < n:
        a = 360 / (s + 1) / 2
        markers.extend([
            (s + 1, 1, a),
            (s + 1, 0, a),
            (s, 1, 0),
            (s, 0, 0),
        ])
        s += 1

    # Convert to MarkerStyle object, using only exactly what we need
    # markers = [mpl.markers.MarkerStyle(m) for m in markers[:n]]

    return markers[:n]

def categorical_order(vector, order=None):
    """Return a list of unique data values.

    Determine an ordered list of levels in ``values``.

    Parameters
    ----------
    vector : list, array, Categorical, or Series
        Vector of "categorical" values
    order : list-like, optional
        Desired order of category levels to override the order determined
        from the ``values`` object.

    Returns
    -------
    order : list
        Ordered list of category levels not including null values.

    """
    if order is None:
        if hasattr(vector, "categories"):
            order = vector.categories
        else:
            try:
                order = vector.cat.categories
            except (TypeError, AttributeError):

                try:
                    order = vector.unique()
                except AttributeError:
                    order = pd.unique(vector)

                if variable_type(vector) == "numeric":
                    order = np.sort(order)

        order = filter(pd.notnull, order)
    return list(order)


#!/usr/bin/python

# electronicdos.py v0.5 5-16-2012 Jeff Doak jeff.w.doak@gmail.com
import numpy as np
from scipy.interpolate import UnivariateSpline
from scipy.integrate import quad
from scipy.optimize import fsolve
import sys, subprocess

BOLTZCONST = 8.617e-5 #eV/K

class ElectronicDOS:
    """
    Class to calculate equilibrium carrier concentrations, as well as
    equilibrium thermodynamic properties of an electronic density of states.

    Class constants of ElectronicDOS:
    - BOLTZCONST - Boltzmann's Constant (eV/K)

    Instance attributes of ElectronicDOS:
    - n_atoms - number of atoms of the unit cell from which the DOS was
        calculated
    - energy - numpy array of energies at which DOS is calculated (eV)
    - dos_tot - numpy array of density of spin up and spin down allowed electron
        states at each energy in the array energy (# states/eV/atom)
    - dos_spin - numpy array of the difference in density between spin up and
        spin down states (# states/eV/atom)
    - e_min - minimum energy in numpy array energy (eV)
    - e_max - maximum energy in numpy array energy (eV)
    - e_fermi - zero Kelvin fermi energy for the electronic DOS (eV)
    - step_size - energy difference between two consecutive points in the DOSCAR
        file (eV)
    - vbm - valence band maximum, set to e_fermi for metals (eV)
    - cbm - conduction band minimum, set to e_fermi for metals (eV)
    - band_gap -  band gap around the fermi energy, zero for metals (eV)
    - temp - numpy array of temperatures at which to calculate equilibrium
        electron chemical potentials, electron concentrations, and hole
        concentrations (K)
    - mu_e - numpy array of electron chemical potentials calculated at each
        temperature in temp (eV)
    - num_e - numpy array of equilibrium electron concentrations calculated at
        each temperature in temp (# e's/atom)
    - num_h - numpy array of equilibrium hole concentrations calculated at each
        temperature in temp (# h's/atom)
    - E_el - numpy array of electronic energy calculated at each temperature
        in temp (eV/atom)
    - S_el - numpy array of electronic entropy calculated at each temperature
        in temp (kB/atom)
    - F_el - numpy array of electronic free energy calculated at each
        temperature in temp (eV/atom)
    """

    def __init__(self,input_,format=None):
        if isinstance(input_,str):
            try:
                input_ = open(input_,'r')
            except IOError:
                print "Error reading input file."
                print "Program will now exit!"
                sys.exit(1)
        if isinstance(input_,file):
            if format == "ezvasp":
                self.read_ezvasp_dos(input_)
            else:
                self.read_doscar(input_)
                nelec = subprocess.Popen("grep NELECT OUTCAR",
                    shell=True,stdin=None,stdout=subprocess.PIPE).communicate()[0]
                self.nelec = int(float(nelec.split()[2]))
            self.get_bandgap()
        # Calculate finite temperature properties
        self.temp = np.linspace(0,2000,21)
        self.mu_e = np.zeros_like(self.temp)
        self.num_e = np.zeros_like(self.temp)
        self.num_h = np.zeros_like(self.temp)
        self.E_el = np.zeros_like(self.temp)
        self.S_el = np.zeros_like(self.temp)
        self.F_el = np.zeros_like(self.temp)
        # Calculate E_el_0
        self.E_el_0 = None
        tol = 1e-5
        for i in range(len(self.temp)):
            if i < tol:
                self.mu_e[i] = self.e_fermi
                self.E_el[i] = 0.0
                self.S_el[i] = 0.0
                self.num_e[i] = 0.0
                self.num_h[i] = 0.0
            elif i > 0.0:
                self.mu_e[i] = self.calc_mu_e(self.temp[i])
                if self.E_el_0 == None:
                    self.E_el_0 = self.calc_E_el(self.mu_e[i],self.temp[i])
                self.num_e[i] = self.n(self.mu_e[i],self.temp[i])
                self.num_h[i] = self.p(self.mu_e[i],self.temp[i])
                self.E_el[i] = (self.calc_E_el(self.mu_e[i],self.temp[i]))
                self.S_el[i] = self.calc_S_el(self.mu_e[i],self.temp[i])

        self.E_el[1:] = self.E_el[1:] - self.E_el_0
        self.F_el = self.E_el - self.temp*BOLTZCONST*self.S_el

    def read_doscar(self,input_):
        """
        Reads in a doscar file to grab the density of states as a function of
        energy. The argument input_ is assumed to be a file object.
        """
        self.n_atoms = int(input_.readline().split()[0])
        # Discard header information
        for i in range(4):
            input_.readline()
        # Read in Fermi Energy
        line = input_.readline().split()
        self.e_max = float(line[0])
        self.e_min = float(line[1])
        self.e_fermi = float(line[3])
        energy = []; dos_tot = []; dos_spin = []
        for line in input_:
            line = line.split()
            energy.append(float(line[0]))
            if len(line) == 3:
                dos_tot.append(float(line[1]))  # DOS includes spin up and down
                dos_spin.append(0.0)
            elif len(line) == 5:
                dos_tot.append(float(line[1])+float(line[2]))
                dos_spin.append(float(line[1])-float(line[2]))
        self.energy = np.array(energy)
        #self.dos_tot = np.array(dos_tot)/float(self.n_atoms)
        self.dos_tot = np.array(dos_tot)
        #self.dos_spin = np.array(dos_spin)/float(self.n_atoms)
        self.dos_spin = np.array(dos_spin)
        self.dos_spline = UnivariateSpline(self.energy,self.dos_tot)

    def read_ezvasp_dos(self,input_):
        """
        Reads an ezvasp-formatted dos.out file to get the electronic density of
        states. The argument input_ is assumned to be a file object.
        """
        nions = subprocess.Popen("grep NIONS OUTCAR",
                shell=True,stdin=None,stdout=subprocess.PIPE).communicate()[0]
        self.n_atoms = int(float(nions.split()[-1]))
        self.e_min = 0.0
        line = input_.readline().split()
        self.nelec = int(float(line[0]))
        self.step_size = float(line[1])
        self.scale = float(line[2])
        energy = []; dos_tot = []
        i = 0
        for line in input_:
            line = line.split()
            dos_tot.append(float(line[0]))
            energy.append(float(i)*self.step_size)
            i += 1
        self.energy = np.array(energy)
        self.dos_tot = np.array(dos_tot)
        self.dos_spin = np.zeros_like(self.dos_tot) # Change this for spin-polar
        self.dos_spline = UnivariateSpline(self.energy,self.dos_tot)
        self.e_max = self.energy[-1]
        # Find the 0 Kelvin 'Fermi Energy' using ATAT's method
        ne = 0.0
        for i in range(len(self.dos_tot)):
            ne += self.dos_tot[i]*self.step_size
            e_fermi = self.energy[i]
            if ne >= self.nelec:
                break
        self.e_fermi = e_fermi

    def get_bandgap(self):
        """
        Finds the band gap of a DOS around the fermi energy.
        """
        self.step_size = self.energy[1] - self.energy[0]
        i = 0
        not_found = True
        while not_found:
            if self.energy[i] < self.e_fermi and self.dos_tot[i] > 1e-3:
                bot = self.energy[i]
            elif self.energy[i] > self.e_fermi and self.dos_tot[i] > 1e-3:
                top = self.energy[i]
                not_found = False
            i += 1
        if top - bot < 2*self.step_size:
            self.vbm = self.cbm = self.e_fermi
            self.band_gap = 0.0
        else:
            self.vbm = bot; self.cbm = top
            self.band_gap = top - bot

    def shift_energy(self,new_ref):
        """
        Change the reference energy for all of the energy attributes.
        """
        self.energy = self.energy - new_ref
        self.e_min = self.e_min - new_ref
        self.e_max = self.e_max - new_ref
        self.e_fermi = self.e_fermi - new_ref
        self.vbm = self.vbm - new_ref
        self.cbm = self.cbm - new_ref
        self.mu_e = self.mu_e - new_ref

    #def sum_dos(self,weight,start,end,args=None):
    def sum_dos(self,weight,start,end,args=None):
        """
        Sums the density of states, dos, in the energy range [start,end], weighted
        by the function weight, which takes as inputs energy and args.
        """
        flag = False
        sum = 0.
        for i in range(len(self.energy)):
            if flag:
                sum += self.step_size*self.dos_tot[i]*weight(
                        self.energy[i],args)
                if self.energy[i] > end:
                    break
            elif self.energy[i] >= start:
                flag = True
        return sum

    #def integrate_dos(self,weight,start,end,args=None,threshold=0.1):
    def ium_dos(self,weight,start,end,args=None,threshold=0.1):
        """
        Takes numpy arrays containing the energy and dos and integrates them over
        the range [start,end] with the weighting function weight. Weight should take
        as an argument the integrated energy and a list of other arguements args.
        """
        def integrand(x,weight,args):
            return self.dos_spline(x)*weight(x,args)
        result = quad(
                integrand,start,end,args=(weight,args),full_output=1,limit=350)
        integral = result[0]
        error = result[1]
        #if error > integral*threshold:
        #    print "Numerical integration error is greater than"
        #    print str(threshold)+" of the integrated value."
        #    sys.exit(1)
        return integral

    def n(self,mu_e,T):
        """
        Calculate the intrinsic number of conduction electrons per atom at an
        electron chemical potential mu_e and temperature T.
        """
        def fermi(x,args):
            mu = args[0]; T = args[1]
            return 1./(np.exp((x-mu)/(BOLTZCONST*T))+1.)
        #n = self.integrate_dos(fermi,self.cbm,self.e_max,args=(mu_e,T))
        #n = self.sum_dos(fermi,self.cbm,self.e_max,args=(mu_e,T))
        n = self.sum_dos(fermi,mu_e,self.e_max,args=(mu_e,T))
        return n

    def p(self,mu_e,T):
        """
        Calculate the intrinsic number of valence holes per atom at an electron
        chemical potential of mu_e and temperature T.
        """
        def fermi(x,args):
            mu = args[0]; T = args[1]
            return 1./(np.exp((mu-x)/(BOLTZCONST*T))+1.)
        #p = self.integrate_dos(fermi,self.e_min,self.vbm,args=(mu_e,T))
        #p = self.sum_dos(fermi,self.e_min,self.vbm,args=(mu_e,T))
        p = self.sum_dos(fermi,self.e_min,mu_e,args=(mu_e,T))
        return p

    def charge_neut2(self,mu_e,args):
        def fermi(x,args):
            mu = args[0]; T = args[1]
            return 1./(np.exp((x-mu)/(BOLTZCONST*T))+1.)
        T = args
        n_sum = self.sum_dos(fermi,self.e_min,self.e_max,args=(mu_e,T))
        return self.nelec - n_sum

    def charge_neutrality(self,mu_e,args):
        """
        Condition for charge neutrality for intrinsic doping in a perfect
        semiconductor. This function should be overwritten for a more
        complicated case.
        """
        T = args  # Args could also include atomic chemical potentials.
        return self.p(mu_e,T) - self.n(mu_e,T)

    def calc_mu_e(self,temp):
        """
        Calculate the electron chemical potential at temperature temp using the
        condition of charge neutrality.
        """
        #mu_e = fsolve(self.charge_neutrality,self.e_fermi,args=(temp))
        mu_e = fsolve(self.charge_neut2,self.e_fermi,args=(temp))
        return mu_e

    def calc_E_el(self,mu_e,T):
        """
        Calculate the electronic energy at a temperature T and electron chemical
        potential mu_e.
        """
        def energy(x,args):
            return x
        def fermi_energy(x,args):
            mu = args[0]; T = args[1]
            if x-mu < -30.0*BOLTZCONST*T:
                return x
            elif x-mu > 30.0*BOLTZCONST*T:
                return 0.0
            else:
                return x/(np.exp((x-mu)/(BOLTZCONST*T))+1.)
        #E = self.integrate_dos(fermi_energy,self.e_min,self.e_max,args=(mu_e,T))
        #E_0 = self.integrate_dos(
        #        fermi_energy,self.e_min,self.e_max,args=(mu_e,T))
        E = self.sum_dos(fermi_energy,self.e_min,self.e_max,args=(mu_e,T))
        #E_0 = self.sum_dos(energy,self.e_min,self.e_fermi,args=None)
        return E

    def calc_S_el(self,mu_e,T):
        """
        Calculate the electronic entropy at an electron chemical potential mu_e
        and temperature T.
        """
        def weight(x,args):
            mu = args[0]; T = args[1]
            x = (x - mu)/(BOLTZCONST*T)
            f = 1.0/(np.exp(x)+1)
            if f > 1e-5 and (1.0 - f) > 1e-5:
                return -f*np.log(f)-(1.-f)*np.log(1.-f)
            else:
                return 0.0
            #f = -np.log(np.exp(x)+1)/(np.exp(x)+1)
            #f += -np.log(np.exp(-x)+1)/(np.exp(-x)+1)
            #return f
        #S = self.integrate_dos(weight,self.e_min,self.e_max,args=(mu_e,T))
        S = self.sum_dos(weight,self.e_min,self.e_max,args=(mu_e,T))
        return S

def fermi_dirac_dist(x,args):
    """
    Calculates the Fermi-Dirac distribution for an energy x, temperature
    args[0], and electron chemical potential args[1].
    """
    T = args[0]; mu = args[1]
    return 1./(np.exp((x-mu)/(BOLTZCONST*T))+1.)

def test2(argv):
    doscar = ElectronicDOS(open(str(argv[0]),'r'))
    T = 500
    #n = doscar.integrate_dos(
    #        fermi_dirac_dist,doscar.cbm,doscar.e_max,args=(T,doscar.e_fermi))
    p = doscar.p(doscar.e_fermi,T)
    print p

def test3(argv):
    format = None
    if len(argv) > 1:
        format = str(argv[1])
    doscar = ElectronicDOS(open(str(argv[0]),'r'),format)
    print doscar.temp
    print doscar.num_e
    print doscar.num_h
    print doscar.E_el
    print doscar.S_el
    print doscar.F_el

def atat_test(argv):
    format = None
    if len(argv) > 1:
        format = str(argv[1])
    doscar = ElectronicDOS(open(str(argv[0]),'r'),format)
    print doscar.E_el_0
    for i in range(len(doscar.temp)):
        print doscar.temp[i],doscar.mu_e[i],doscar.E_el[i],doscar.S_el[i],doscar.F_el[i]

def test1(argv):
    import matplotlib.pyplot as plt
    doscar = ElectronicDOS(open(str(argv[0]),'r'))
    plt.plot(doscar.energy,doscar.dos_tot)
    plt.show()

def main(argv):
    import matplotlib.pyplot as plt
    doscar = open(str(argv[0]))
    e_fermi,energy,n_tot,n_spin = read_doscar(doscar)
    plt.plot(energy,n_tot)
    if len(argv) > 1:
        doscar2 = open(str(argv[1]))
        e_fermi2,energy2,n_tot2,n_spin2 = read_doscar(doscar2)
        plt.plot(energy2,n_tot2)
    plt.show()

if __name__ == "__main__":
    import sys
    #test3(sys.argv[1:])
    atat_test(sys.argv[1:])