import os
import warnings
from pkg_resources import resource_filename
import tempfile

import numpy
from numpy import nan
import pandas
from shapely.geometry import Polygon
import geopandas

import pytest
import numpy.testing as nptest
import pandas.util.testing as pdtest

from pygridtools import core
from pygridgen.tests import raises
from . import utils

BASELINE_IMAGES = 'baseline_files/test_core'
try:
    import pygridgen
    HASPGG = True
except ImportError:
    HASPGG = False

@pytest.fixture
def A():
    return numpy.arange(12).reshape(4, 3).astype(float)

@pytest.fixture
def B():
    return numpy.arange(8).reshape(2, 4).astype(float)

@pytest.fixture
def C():
    return numpy.arange(25).reshape(5, 5).astype(float)

@pytest.mark.parametrize('fxn', [numpy.fliplr, numpy.flipud, numpy.fliplr])
def test_transform(A, fxn):
    result = core.transform(A, fxn)
    expected = fxn(A)
    nptest.assert_array_equal(result, expected)

@pytest.mark.parametrize(('index', 'axis', 'first', 'second'), [
    (3, 0, 'top', 'bottom'),
    (2, 1, 'left', 'right'),
    (5, 0, None, None),
    (5, 1, None, None),
])
def test_split_rows(C, index, axis, first, second):
    expected = {
        'top': numpy.array([
            [ 0.0,  1.0,  2.0,  3.0,  4.0],
            [ 5.0,  6.0,  7.0,  8.0,  9.0],
            [10.0, 11.0, 12.0, 13.0, 14.0],
        ]),
        'bottom': numpy.array([
            [15., 16., 17., 18., 19.],
            [20., 21., 22., 23., 24.],
        ]),
        'left': numpy.array([
            [ 0.,  1.],
            [ 5.,  6.],
            [10., 11.],
            [15., 16.],
            [20., 21.],
        ]),
        'right': numpy.array([
            [ 2.,  3.,  4.],
            [ 7.,  8.,  9.],
            [12., 13., 14.],
            [17., 18., 19.],
            [22., 23., 24.],
        ]),
    }
    if first and second:
        a, b = core.split(C, index, axis)
        nptest.assert_array_equal(a, expected[first])
        nptest.assert_array_equal(b, expected[second])
    else:
        with raises(ValueError):
            left, right = core.split(C, index, axis=axis)

@pytest.mark.parametrize('N', [1, 3, None])
def test__interp_between_vectors(N):
    index = numpy.arange(0, 4)
    vector1 = -1 * index**2 - 1
    vector2 = 2 * index**2 + 2

    expected = {
        1: numpy.array([
            [ -1.0,  -2.0,  -5.0, -10.0],
            [  0.5,   1.0,   2.5,   5.0],
            [  2.0,   4.0,  10.0,  20.0],
        ]),
        3: numpy.array([
            [ -1.00,  -2.00,  -5.00, -10.00],
            [ -0.25,  -0.50,  -1.25,  -2.50],
            [  0.50,   1.00,   2.50,   5.00],
            [  1.25,   2.50,   6.25,  12.50],
            [  2.00,   4.00,  10.00,  20.00],
        ])
    }

    if N:
        result = core._interp_between_vectors(vector1, vector2, n_nodes=N)
        nptest.assert_array_equal(result, expected[N])
    else:
        with raises(ValueError):
            core._interp_between_vectors(vector1, vector2, n_nodes=0)

@pytest.mark.parametrize(('n', 'axis'), [
    (1, 0), (4, 0), (1, 1), (3, 1)
])
def test_insert(C, n, axis):
    expected = {
        (1, 0): numpy.array([
            [ 0.0,  1.0,  2.0,  3.0,  4.0],
            [ 5.0,  6.0,  7.0,  8.0,  9.0],
            [ 7.5,  8.5,  9.5, 10.5, 11.5],
            [10.0, 11.0, 12.0, 13.0, 14.0],
            [15.0, 16.0, 17.0, 18.0, 19.0],
            [20.0, 21.0, 22.0, 23.0, 24.0],
        ]),
        (4, 0): numpy.array([
            [ 0.0,  1.0,  2.0,  3.0,  4.0],
            [ 5.0,  6.0,  7.0,  8.0,  9.0],
            [ 6.0,  7.0,  8.0,  9.0, 10.0],
            [ 7.0,  8.0,  9.0, 10.0, 11.0],
            [ 8.0,  9.0, 10.0, 11.0, 12.0],
            [ 9.0, 10.0, 11.0, 12.0, 13.0],
            [10.0, 11.0, 12.0, 13.0, 14.0],
            [15.0, 16.0, 17.0, 18.0, 19.0],
            [20.0, 21.0, 22.0, 23.0, 24.0],
        ]),
        (1, 1): numpy.array([
            [ 0.0,  1.0,  1.5,  2.0,  3.0,  4.0],
            [ 5.0,  6.0,  6.5,  7.0,  8.0,  9.0],
            [10.0, 11.0, 11.5, 12.0, 13.0, 14.0],
            [15.0, 16.0, 16.5, 17.0, 18.0, 19.0],
            [20.0, 21.0, 21.5, 22.0, 23.0, 24.0],
        ]),
        (3, 1): numpy.array([
            [ 0.00,  1.00,  1.25,  1.50,  1.75,  2.00,  3.00,  4.00],
            [ 5.00,  6.00,  6.25,  6.50,  6.75,  7.00,  8.00,  9.00],
            [10.00, 11.00, 11.25, 11.50, 11.75, 12.00, 13.00, 14.00],
            [15.00, 16.00, 16.25, 16.50, 16.75, 17.00, 18.00, 19.00],
            [20.00, 21.00, 21.25, 21.50, 21.75, 22.00, 23.00, 24.00],
        ])
    }
    result = core.insert(C, 2, axis=axis, n_nodes=n)
    nptest.assert_array_equal(result, expected[(n, axis)])

@pytest.mark.parametrize('how', ['h', 'v'])
@pytest.mark.parametrize('where', ['+', '-'])
@pytest.mark.parametrize('shift', [0, 2, -1])
def test_merge(A, B, how, where, shift):
    expected = {
        ('v', '+', 0): numpy.array([
            [0.,  1.,  2., nan],
            [3.,  4.,  5., nan],
            [6.,  7.,  8., nan],
            [9., 10., 11., nan],
            [0.,  1.,  2.,  3.],
            [4.,  5.,  6.,  7.]
        ]),
        ('v', '-', 0): numpy.array([
            [0.,  1.,  2.,  3.],
            [4.,  5.,  6.,  7.],
            [0.,  1.,  2., nan],
            [3.,  4.,  5., nan],
            [6.,  7.,  8., nan],
            [9., 10., 11., nan]
        ]),
        ('v', '+', 2): numpy.array([
            [ 0.,   1.,   2., nan,   nan, nan],
            [ 3.,   4.,   5., nan,   nan, nan],
            [ 6.,   7.,   8., nan,   nan, nan],
            [ 9.,  10.,  11., nan,   nan, nan],
            [nan,  nan,   0.,   1.,   2.,  3.],
            [nan,  nan,   4.,   5.,   6.,  7.]
        ]),
        ('v', '-', 2): numpy.array([
            [nan, nan,  0.,  1.,  2.,  3.],
            [nan, nan,  4.,  5.,  6.,  7.],
            [ 0.,  1.,  2., nan, nan, nan],
            [ 3.,  4.,  5., nan, nan, nan],
            [ 6.,  7.,  8., nan, nan, nan],
            [ 9., 10., 11., nan, nan, nan]
        ]),
        ('v', '+', -1): numpy.array([
            [nan, 0., 1.,   2.],
            [nan, 3., 4.,   5.],
            [nan, 6., 7.,   8.],
            [nan, 9., 10., 11.],
            [ 0., 1., 2.,   3.],
            [ 4., 5., 6.,   7.]
        ]),
        ('v', '-', -1): numpy.array([
            [ 0., 1., 2.,   3.],
            [ 4., 5., 6.,   7.],
            [nan, 0., 1.,   2.],
            [nan, 3., 4.,   5.],
            [nan, 6., 7.,   8.],
            [nan, 9., 10., 11.]
        ]),
        ('h', '+', 0): numpy.array([
            [0.,  1.,  2.,  0.,  1.,  2.,  3.],
            [3.,  4.,  5.,  4.,  5.,  6.,  7.],
            [6.,  7.,  8., nan, nan, nan, nan],
            [9., 10., 11., nan, nan, nan, nan]
        ]),
        ('h', '-', 0): numpy.array([
            [ 0.,  1.,  2.,  3., 0.,  1.,  2.],
            [ 4.,  5.,  6.,  7., 3.,  4.,  5.],
            [nan, nan, nan, nan, 6.,  7.,  8.],
            [nan, nan, nan, nan, 9., 10., 11.]
        ]),
        ('h', '+', 2): numpy.array([
            [0.,  1.,  2., nan, nan, nan, nan],
            [3.,  4.,  5., nan, nan, nan, nan],
            [6.,  7.,  8.,  0.,  1.,  2.,  3.],
            [9., 10., 11.,  4.,  5.,  6.,  7.]
        ]),
        ('h', '-', 2): numpy.array([
            [nan, nan, nan, nan, 0.,  1.,  2.],
            [nan, nan, nan, nan, 3.,  4.,  5.],
            [ 0.,  1.,  2.,  3., 6.,  7.,  8.],
            [ 4.,  5.,  6.,  7., 9., 10., 11.]
        ]),
        ('h', '+', -1): numpy.array([
            [nan, nan, nan,  0.,  1.,  2.,  3.],
            [ 0.,  1.,  2.,  4.,  5.,  6.,  7.],
            [ 3.,  4.,  5., nan, nan, nan, nan],
            [ 6.,  7.,  8., nan, nan, nan, nan],
            [ 9., 10., 11., nan, nan, nan, nan]
        ]),
        ('h', '-', -1): numpy.array([
            [ 0.,  1.,  2.,  3., nan, nan, nan],
            [ 4.,  5.,  6.,  7.,  0.,  1.,  2.],
            [nan, nan, nan, nan,  3.,  4.,  5.],
            [nan, nan, nan, nan,  6.,  7.,  8.],
            [nan, nan, nan, nan,  9., 10., 11.]
        ]),
    }
    result = core.merge(A, B, how=how, where=where, shift=shift)
    nptest.assert_array_equal(result, expected[(how, where, shift)])

@pytest.fixture
def g1(simple_nodes):
    xn, yn = simple_nodes
    g = core.ModelGrid(xn[:, :3], yn[:, :3])

    mask = g.cell_mask
    mask[:2, :2] = True
    g.cell_mask = mask
    return g

@pytest.fixture
def g2(simple_nodes):
    xn, yn = simple_nodes
    g = core.ModelGrid(xn[2:5, 3:], yn[2:5, 3:])
    return g

@pytest.fixture
def polyverts():
    return geopandas.GeoSeries(Polygon([(2.4, 0.9), (3.6, 0.9), (3.6, 2.4), (2.4, 2.4)]))

def test_ModelGrid_bad_shapes(simple_cells):
    xc, yc = simple_cells
    with raises(ValueError):
        mg = core.ModelGrid(xc, yc[2:, 2:])

def test_ModelGrid_nodes_and_cells(g1, simple_cells):
    xc, yc = simple_cells
    assert (isinstance(g1.nodes_x, numpy.ndarray))
    assert (isinstance(g1.nodes_y, numpy.ndarray))
    assert (isinstance(g1.cells_x, numpy.ndarray))
    nptest.assert_array_equal(g1.cells_x, xc[:, :2])
    assert (isinstance(g1.cells_y, numpy.ndarray))
    nptest.assert_array_equal(g1.cells_y, yc[:, :2])

def test_ModelGrid_counts_and_shapes(g1):
    expected_rows = 9
    expected_cols = 3

    assert (g1.icells == expected_cols - 1)
    assert (g1.jcells == expected_rows - 1)

    assert (g1.inodes == expected_cols)
    assert (g1.jnodes == expected_rows)

    assert (g1.shape == (expected_rows, expected_cols))
    assert (g1.cell_shape == (expected_rows - 1, expected_cols - 1))

def test_ModelGrid_cell_mask(g1):
    expected_mask = numpy.array([
        [1, 1], [1, 1], [0, 0], [0, 0],
        [0, 0], [0, 0], [0, 0], [0, 0],
    ])
    nptest.assert_array_equal(g1.cell_mask, expected_mask)

@pytest.mark.parametrize(('usemask', 'which', 'error'), [
    (True, 'nodes', ValueError),
    (False, 'nodes', None),
    (True, 'cells', None),
])
def test_ModelGrid_to_dataframe(g1, usemask, which, error):
    def name_cols(df):
        df.columns.names = ['coord', 'ii']
        df.index.names = ['jj']
        return df

    if error:
        with raises(ValueError):
            g1.to_dataframe(usemask=usemask, which=which)
    else:

        expected = {
            (False, 'nodes'): pandas.DataFrame({
                ('easting', 0): {
                    0: 1.0, 1: 1.0, 2: 1.0, 3: 1.0, 4: 1.0,
                    5: 1.0, 6: 1.0, 7: 1.0, 8: 1.0
                }, ('easting', 1): {
                    0: 1.5, 1: 1.5, 2: 1.5, 3: 1.5, 4: 1.5,
                    5: 1.5, 6: 1.5, 7: 1.5, 8: 1.5
                }, ('easting', 2): {
                    0: 2.0, 1: 2.0, 2: 2.0, 3: 2.0, 4: 2.0,
                    5: 2.0, 6: 2.0, 7: 2.0, 8: 2.0
                }, ('northing', 0): {
                    0: 0.0, 1: 0.5, 2: 1.0, 3: 1.5, 4: 2.0,
                    5: 2.5, 6: 3.0, 7: 3.5, 8: 4.0
                }, ('northing', 1): {
                    0: 0.0, 1: 0.5, 2: 1.0, 3: 1.5, 4: 2.0,
                    5: 2.5, 6: 3.0, 7: 3.5, 8: 4.0
                }, ('northing', 2): {
                    0: 0.0, 1: 0.5, 2: 1.0, 3: 1.5, 4: 2.0,
                    5: 2.5, 6: 3.0, 7: 3.5, 8: 4.0}
            }).pipe(name_cols),
            (True, 'cells'): pandas.DataFrame({
                ('easting', 0): {
                    0: nan, 1: nan, 2: 1.25, 3: 1.25, 4: 1.25,
                    5: 1.25, 6: 1.25, 7: 1.25
                }, ('easting', 1): {
                    0: nan, 1: nan, 2: 1.75, 3: 1.75, 4: 1.75,
                    5: 1.75, 6: 1.75, 7: 1.75
                }, ('northing', 0): {
                    0: nan, 1: nan, 2: 1.25, 3: 1.75, 4: 2.25,
                    5: 2.75, 6: 3.25, 7: 3.75
                }, ('northing', 1): {
                    0: nan, 1: nan, 2: 1.25, 3: 1.75, 4: 2.25,
                    5: 2.75, 6: 3.25, 7: 3.75
                }
            }).pipe(name_cols),
        }

        result = g1.to_dataframe(usemask=usemask, which=which)
        pdtest.assert_frame_equal(result, expected[(usemask, which)], check_names=False)
        pdtest.assert_index_equal(result.columns, expected[(usemask, which)].columns)

@pytest.mark.parametrize(('usemask', 'which', 'error'), [
    (True, 'nodes', ValueError),
    (False, 'nodes', None),
    (True, 'cells', None),
    (False, 'cells', None),
])
def test_ModelGrid_to_coord_pairs(g1, usemask, which, error):
    if error:
        with raises(error):
            g1.to_coord_pairs(usemask=usemask, which=which)
    else:

        expected = {
            ('nodes', False): numpy.array([
                [1.0, 0.0], [1.5, 0.0], [2.0, 0.0], [1.0, 0.5],
                [1.5, 0.5], [2.0, 0.5], [1.0, 1.0], [1.5, 1.0],
                [2.0, 1.0], [1.0, 1.5], [1.5, 1.5], [2.0, 1.5],
                [1.0, 2.0], [1.5, 2.0], [2.0, 2.0], [1.0, 2.5],
                [1.5, 2.5], [2.0, 2.5], [1.0, 3.0], [1.5, 3.0],
                [2.0, 3.0], [1.0, 3.5], [1.5, 3.5], [2.0, 3.5],
                [1.0, 4.0], [1.5, 4.0], [2.0, 4.0]
            ]),
            ('cells', False): numpy.array([
                [1.25, 0.25], [1.75, 0.25], [1.25, 0.75], [1.75, 0.75],
                [1.25, 1.25], [1.75, 1.25], [1.25, 1.75], [1.75, 1.75],
                [1.25, 2.25], [1.75, 2.25], [1.25, 2.75], [1.75, 2.75],
                [1.25, 3.25], [1.75, 3.25], [1.25, 3.75], [1.75, 3.75]
            ]),
            ('cells', True): numpy.array([
                [nan, nan], [nan, nan], [nan, nan], [nan, nan],
                [1.25, 1.25], [1.75, 1.25], [1.25, 1.75], [1.75, 1.75],
                [1.25, 2.25], [1.75, 2.25], [1.25, 2.75], [1.75, 2.75],
                [1.25, 3.25], [1.75, 3.25], [1.25, 3.75], [1.75, 3.75]
            ])
        }

        result = g1.to_coord_pairs(usemask=usemask, which=which)
        nptest.assert_array_equal(result, expected[which, usemask])

def test_ModelGrid_transform(mg, simple_nodes):
    xn, yn = simple_nodes
    g = mg.transform(lambda x: x * 10)
    nptest.assert_array_equal(g.xn, xn * 10)
    nptest.assert_array_equal(g.yn, yn * 10)

def test_ModelGrid_transform_x(mg, simple_nodes):
    xn, yn = simple_nodes
    g = mg.transform_x(lambda x: x * 10)
    nptest.assert_array_equal(g.xn, xn * 10)
    nptest.assert_array_equal(g.yn, yn)

def test_ModelGrid_transform_y(mg, simple_nodes):
    xn, yn = simple_nodes
    g = mg.transform_y(lambda y: y * 10)
    nptest.assert_array_equal(g.xn, xn)
    nptest.assert_array_equal(g.yn, yn * 10)

def test_ModelGrid_transpose(mg, simple_nodes):
    xn, yn = simple_nodes
    g = mg.transpose()
    nptest.assert_array_equal(g.xn, xn.T)
    nptest.assert_array_equal(g.yn, yn.T)

def test_ModelGrid_fliplr(mg, simple_nodes):
    xn, yn = simple_nodes
    g = mg.fliplr()
    nptest.assert_array_equal(g.xn, numpy.fliplr(xn))
    nptest.assert_array_equal(g.yn, numpy.fliplr(yn))

def test_ModelGrid_flipud(mg, simple_nodes):
    xn, yn = simple_nodes
    g = mg.flipud()
    nptest.assert_array_equal(g.xn, numpy.flipud(xn))
    nptest.assert_array_equal(g.yn, numpy.flipud(yn))

def test_ModelGrid_split_ax0(mg, simple_nodes):
    xn, yn = simple_nodes
    mgtop, mgbottom = mg.split(3, axis=0)
    nptest.assert_array_equal(mgtop.nodes_x, xn[:3, :])
    nptest.assert_array_equal(mgtop.nodes_y, yn[:3, :])
    nptest.assert_array_equal(mgbottom.nodes_x, xn[3:, :])
    nptest.assert_array_equal(mgbottom.nodes_y, yn[3:, :])

def test_ModelGrid_node_mask(simple_nodes):
    g = core.ModelGrid(*simple_nodes).update_cell_mask()
    expected = numpy.array([
        [0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 1, 1, 1, 1]
    ]).astype(bool)
    nptest.assert_array_equal(expected, g.node_mask)

def test_ModelGrid_merge(g1, g2, simple_nodes):
    g3 = g1.merge(g2, how='horiz', where='+', shift=2)
    g4 = core.ModelGrid(*simple_nodes).update_cell_mask()

    nptest.assert_array_equal(g3.xn, g4.xn)
    nptest.assert_array_equal(g3.xc, g4.xc)

@pytest.mark.mpl_image_compare(baseline_dir=BASELINE_IMAGES, tolerance=15)
def test_ModelGrid_merge_with_mask(simple_nodes):
    mg1 = core.ModelGrid(*simple_nodes).update_cell_mask()
    mg2 = (
        mg1.transform_x(lambda x: x + 1)
           .transform_y(lambda y: y + 5)
           .update_cell_mask(mask=mg1.cell_mask)
    )

    merged = mg1.merge(mg2, where='+', shift=1, min_nodes=1)
    expected = numpy.array([
        [0, 0, 1, 1, 1, 1, 1],
        [0, 0, 1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 1],
        [0, 0, 0, 0, 0, 0, 1],
        [0, 0, 1, 1, 1, 1, 1],
        [0, 0, 1, 1, 1, 1, 1],
        [0, 0, 1, 1, 1, 1, 1],
        [0, 0, 1, 1, 1, 1, 1],
        [1, 0, 1, 1, 1, 1, 1],
        [1, 0, 0, 1, 1, 1, 1],
        [1, 0, 0, 1, 1, 1, 1],
        [1, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 1, 1, 1, 1],
        [1, 0, 0, 1, 1, 1, 1],
        [1, 0, 0, 1, 1, 1, 1],
        [1, 0, 0, 1, 1, 1, 1]
    ]).astype(bool)
    nptest.assert_array_equal(merged.cell_mask, expected)
    fig, artists = merged.plot_cells()
    return fig

def test_ModelGrid_insert_3_ax0(mg):
    known_xnodes = numpy.ma.masked_invalid(numpy.array([
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0],
        [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0],
        [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
        [1.0, 1.5, 2.0, nan, nan, nan, nan],
    ]))

    known_ynodes = numpy.ma.masked_invalid(numpy.array([
        [0.000, 0.000, 0.000,   nan,   nan,   nan,   nan],
        [0.500, 0.500, 0.500,   nan,   nan,   nan,   nan],
        [0.625, 0.625, 0.625,   nan,   nan,   nan,   nan],
        [0.750, 0.750, 0.750,   nan,   nan,   nan,   nan],
        [0.875, 0.875, 0.875,   nan,   nan,   nan,   nan],
        [1.000, 1.000, 1.000, 1.000, 1.000, 1.000, 1.000],
        [1.500, 1.500, 1.500, 1.500, 1.500, 1.500, 1.500],
        [2.000, 2.000, 2.000, 2.000, 2.000, 2.000, 2.000],
        [2.500, 2.500, 2.500,   nan,   nan,   nan,   nan],
        [3.000, 3.000, 3.000,   nan,   nan,   nan,   nan],
        [3.500, 3.500, 3.500,   nan,   nan,   nan,   nan],
        [4.000, 4.000, 4.000,   nan,   nan,   nan,   nan],
    ]))

    result = mg.insert(2, axis=0, n_nodes=3)
    nptest.assert_array_equal(result.nodes_x, known_xnodes)
    nptest.assert_array_equal(result.nodes_y, known_ynodes)

def test_ModelGrid_insert_3_ax1(mg):
    known_xnodes = numpy.ma.masked_invalid(numpy.array([
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000,   nan,   nan,   nan,   nan],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000,   nan,   nan,   nan,   nan],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000, 2.500, 3.000, 3.500, 4.000],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000, 2.500, 3.000, 3.500, 4.000],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000, 2.500, 3.000, 3.500, 4.000],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000,   nan,   nan,   nan,   nan],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000,   nan,   nan,   nan,   nan],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000,   nan,   nan,   nan,   nan],
        [1.000, 1.500, 1.625, 1.750, 1.875, 2.000,   nan,   nan,   nan,   nan]
    ]))

    known_ynodes = numpy.ma.masked_invalid(numpy.array([
        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, nan, nan, nan, nan],
        [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, nan, nan, nan, nan],
        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
        [1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5],
        [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0],
        [2.5, 2.5, 2.5, 2.5, 2.5, 2.5, nan, nan, nan, nan],
        [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, nan, nan, nan, nan],
        [3.5, 3.5, 3.5, 3.5, 3.5, 3.5, nan, nan, nan, nan],
        [4.0, 4.0, 4.0, 4.0, 4.0, 4.0, nan, nan, nan, nan],
    ]))

    result = mg.insert(2, axis=1, n_nodes=3)
    nptest.assert_array_equal(result.nodes_x, known_xnodes)
    nptest.assert_array_equal(result.nodes_y, known_ynodes)

def test_extract(mg, simple_nodes):
    xn, yn = simple_nodes
    result = mg.extract(jstart=2, jend=5, istart=3, iend=6)
    nptest.assert_array_equal(result.nodes_x, xn[2:5, 3:6])
    nptest.assert_array_equal(result.nodes_y, yn[2:5, 3:6])

@pytest.mark.parametrize(('where', 'use_existing'), [
    ('inside', False),
    ('inside', True),
    ('outside', False)
])
def test_ModelGrid_mask_centroids(mg, polyverts, where, use_existing):
    expected = {
        ('inside', False): numpy.array([
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0]
        ]),
        ('inside', True): numpy.array([
            [0, 0, 1, 1, 1, 1],
            [0, 0, 1, 1, 1, 1],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 1, 1, 1, 1],
            [0, 0, 1, 1, 1, 1],
            [0, 0, 1, 1, 1, 1],
            [0, 0, 1, 1, 1, 1]
        ]),
        ('outside', False):  numpy.array([
            [1, 1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1, 1],
            [1, 1, 1, 0, 0, 1],
            [1, 1, 1, 0, 0, 1],
            [1, 1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1, 1]
        ])
    }

    result = mg.mask_centroids(**{where: polyverts}, use_existing=use_existing)

    nptest.assert_array_equal(
        result.cell_mask.astype(int),
        expected[(where, use_existing)].astype(int)
    )

@pytest.mark.parametrize(('kwargs', 'error'), [
    [dict(min_nodes=0), ValueError],
    [dict(min_nodes=5), ValueError],
    [dict(triangles=True), NotImplementedError],
])
def test_ModelGrid_mask_nodes_errors(mg, polyverts, kwargs, error):
    with raises(error):
        mg.mask_nodes(inside=polyverts, **kwargs)

def test_masks_no_polys(mg):
    with raises(ValueError):
        mg.mask_nodes()

    with raises(ValueError):
        mg.mask_centroids()

def test_ModelGrid_to_point_geodataframe(g1):
    expectedfile = resource_filename('pygridtools.tests.baseline_files',  'mgshp_nomask_nodes_points.shp')
    expected = geopandas.read_file(expectedfile)
    result = g1.to_point_geodataframe(which='nodes', usemask=False)
    utils.assert_gdfs_equal(expected.drop(columns=['river', 'reach']), result)

@pytest.mark.xfail
@pytest.mark.parametrize('usemask', [True, False])
def test_ModelGrid_to_gis_cells(g1, usemask):
    expectedfile = {
        True: 'mgshp_mask_cells_polys.shp',
        False: 'mgshp_nomask_cells_polys.shp',
    }
    expectedfile = resource_filename('pygridtools.tests.baseline_files',
                                     expectedfile[usemask])
    expected = geopandas.read_file(expectedfile)
    result = g1.to_polygon_geodataframe(usemask=usemask)
    utils.assert_gdfs_equal(expected.drop(columns=['river', 'reach']), result)

@pytest.mark.parametrize(('which', 'usemask', 'error'), [
    ('nodes', True, ValueError),
    ('junk', False, ValueError),
    ('nodes', False, None),
    ('cells', False, None),
])
def test_ModelGrid__get_x_y_nodes_and_mask(g1, which, usemask, error):
    if error:
        with raises(error):
            g1._get_x_y(which, usemask=usemask)
    else:
        x, y = g1._get_x_y(which, usemask=usemask)
        nptest.assert_array_equal(x, getattr(g1, 'x' + which[0]))
        nptest.assert_array_equal(y, getattr(g1, 'y' + which[0]))

@pytest.mark.mpl_image_compare(baseline_dir=BASELINE_IMAGES, tolerance=15)
def test_ModelGrid_plots_basic(simple_nodes):
    mg = core.ModelGrid(*simple_nodes)
    mg.cell_mask = numpy.ma.masked_invalid(mg.xc).mask
    fig, artists = mg.plot_cells()
    return fig

@pytest.mark.mpl_image_compare(baseline_dir=BASELINE_IMAGES, tolerance=15)
def test_ModelGrid_plots_masked(river_grid, river_bathy):
    fig, artists = river_grid.plot_cells(cell_kws=dict(colors=river_bathy, cmap='Reds_r'))
    return fig

@pytest.mark.parametrize(('otherargs', 'gridtype'), [
    (dict(), None),
    (dict(verbose=True), None),
    (dict(rawgrid=False), core.ModelGrid)
])
@pytest.mark.skipif(not HASPGG, reason='pygridgen unavailabile')
def test_make_grid(simple_boundary_gdf, otherargs, gridtype):
    if not gridtype:
        gridtype = pygridgen.Gridgen

    gridparams = {'nnodes': 12, 'verbose': False, 'ul_idx': 0}
    gridparams.update(otherargs)
    grid = core.make_grid(9, 7, domain=simple_boundary_gdf, **gridparams)
    assert (isinstance(grid, gridtype))


"""Create a mosaic plot from a contingency table.

It allows to visualize multivariate categorical data in a rigorous
and informative way.

see the docstring of the mosaic function for more informations.
"""
# Author: Enrico Giampieri - 21 Jan 2013

from __future__ import division
from statsmodels.compat.python import (iteritems, iterkeys, lrange, string_types, lzip,
                                itervalues, zip, range)
import numpy as np
from statsmodels.compat.collections import OrderedDict
from itertools import product

from numpy import iterable, r_, cumsum, array
from statsmodels.graphics import utils
from pandas import DataFrame

__all__ = ["mosaic"]

def _normalize_split(proportion):
    """
    return a list of proportions of the available space given the division
    if only a number is given, it will assume a split in two pieces
    """
    if not iterable(proportion):
        if proportion == 0:
            proportion = array([0.0, 1.0])
        elif proportion >= 1:
            proportion = array([1.0, 0.0])
        elif proportion < 0:
            raise ValueError("proportions should be positive,"
                              "given value: {}".format(proportion))
        else:
            proportion = array([proportion, 1.0 - proportion])
    proportion = np.asarray(proportion, dtype=float)
    if np.any(proportion < 0):
        raise ValueError("proportions should be positive,"
                          "given value: {}".format(proportion))
    if np.allclose(proportion, 0):
        raise ValueError("at least one proportion should be"
                          "greater than zero".format(proportion))
    # ok, data are meaningful, so go on
    if len(proportion) < 2:
        return array([0.0, 1.0])
    left = r_[0, cumsum(proportion)]
    left /= left[-1] * 1.0
    return left

def _split_rect(x, y, width, height, proportion, horizontal=True, gap=0.05):
    """
    Split the given rectangle in n segments whose proportion is specified
    along the given axis if a gap is inserted, they will be separated by a
    certain amount of space, retaining the relative proportion between them
    a gap of 1 correspond to a plot that is half void and the remaining half
    space is proportionally divided among the pieces.
    """
    x, y, w, h = float(x), float(y), float(width), float(height)
    if (w < 0) or (h < 0):
        raise ValueError("dimension of the square less than"
                          "zero w={} h=()".format(w, h))
    proportions = _normalize_split(proportion)

    # extract the starting point and the dimension of each subdivision
    # in respect to the unit square
    starting = proportions[:-1]
    amplitude = proportions[1:] - starting

    # how much each extrema is going to be displaced due to gaps
    starting += gap * np.arange(len(proportions) - 1)

    # how much the squares plus the gaps are extended
    extension = starting[-1] + amplitude[-1] - starting[0]

    # normalize everything for fit again in the original dimension
    starting /= extension
    amplitude /= extension

    # bring everything to the original square
    starting = (x if horizontal else y) + starting * (w if horizontal else h)
    amplitude = amplitude * (w if horizontal else h)

    # create each 4-tuple for each new block
    results = [(s, y, a, h) if horizontal else (x, s, w, a)
                for s, a in zip(starting, amplitude)]
    return results

def _reduce_dict(count_dict, partial_key):
    """
    Make partial sum on a counter dict.
    Given a match for the beginning of the category, it will sum each value.
    """
    L = len(partial_key)
    count = sum(v for k, v in iteritems(count_dict) if k[:L] == partial_key)
    return count

def _key_splitting(rect_dict, keys, values, key_subset, horizontal, gap):
    """
    Given a dictionary where each entry  is a rectangle, a list of key and
    value (count of elements in each category) it split each rect accordingly,
    as long as the key start with the tuple key_subset.  The other keys are
    returned without modification.
    """
    result = OrderedDict()
    L = len(key_subset)
    for name, (x, y, w, h) in iteritems(rect_dict):
        if key_subset == name[:L]:
            # split base on the values given
            divisions = _split_rect(x, y, w, h, values, horizontal, gap)
            for key, rect in zip(keys, divisions):
                result[name + (key,)] = rect
        else:
            result[name] = (x, y, w, h)
    return result

def _tuplify(obj):
    """convert an object in a tuple of strings (even if it is not iterable,
    like a single integer number, but keep the string healthy)
    """
    if np.iterable(obj) and not isinstance(obj, string_types):
        res = tuple(str(o) for o in obj)
    else:
        res = (str(obj),)
    return res

def _categories_level(keys):
    """use the Ordered dict to implement a simple ordered set
    return each level of each category
    [[key_1_level_1,key_2_level_1],[key_1_level_2,key_2_level_2]]
    """
    res = []
    for i in zip(*(keys)):
        tuplefied = _tuplify(i)
        res.append(list(OrderedDict([(j, None) for j in tuplefied])))
    return res

def _hierarchical_split(count_dict, horizontal=True, gap=0.05):
    """
    Split a square in a hierarchical way given a contingency table.

    Hierarchically split the unit square in alternate directions
    in proportion to the subdivision contained in the contingency table
    count_dict.  This is the function that actually perform the tiling
    for the creation of the mosaic plot.  If the gap array has been specified
    it will insert a corresponding amount of space (proportional to the
    unit lenght), while retaining the proportionality of the tiles.

    Parameters
    ----------
    count_dict : dict
        Dictionary containing the contingency table.
        Each category should contain a non-negative number
        with a tuple as index.  It expects that all the combination
        of keys to be representes; if that is not true, will
        automatically consider the missing values as 0
    horizontal : bool
        The starting direction of the split (by default along
        the horizontal axis)
    gap : float or array of floats
        The list of gaps to be applied on each subdivision.
        If the lenght of the given array is less of the number
        of subcategories (or if it's a single number) it will extend
        it with exponentially decreasing gaps

    Returns
    ----------
    base_rect : dict
        A dictionary containing the result of the split.
        To each key is associated a 4-tuple of coordinates
        that are required to create the corresponding rectangle:

            0 - x position of the lower left corner
            1 - y position of the lower left corner
            2 - width of the rectangle
            3 - height of the rectangle
    """
    # this is the unit square that we are going to divide
    base_rect = OrderedDict([(tuple(), (0, 0, 1, 1))])
    # get the list of each possible value for each level
    categories_levels = _categories_level(list(iterkeys(count_dict)))
    L = len(categories_levels)

    # recreate the gaps vector starting from an int
    if not np.iterable(gap):
        gap = [gap / 1.5 ** idx for idx in range(L)]
    # extend if it's too short
    if len(gap) < L:
        last = gap[-1]
        gap = list(*gap) + [last / 1.5 ** idx for idx in range(L)]
    # trim if it's too long
    gap = gap[:L]
    # put the count dictionay in order for the keys
    # this will allow some code simplification
    count_ordered = OrderedDict([(k, count_dict[k])
                        for k in list(product(*categories_levels))])
    for cat_idx, cat_enum in enumerate(categories_levels):
        # get the partial key up to the actual level
        base_keys = list(product(*categories_levels[:cat_idx]))
        for key in base_keys:
            # for each partial and each value calculate how many
            # observation we have in the counting dictionary
            part_count = [_reduce_dict(count_ordered, key + (partial,))
                            for partial in cat_enum]
            # reduce the gap for subsequents levels
            new_gap = gap[cat_idx]
            # split the given subkeys in the rectangle dictionary
            base_rect = _key_splitting(base_rect, cat_enum, part_count, key,
                                       horizontal, new_gap)
        horizontal = not horizontal
    return base_rect

def _single_hsv_to_rgb(hsv):
    """Transform a color from the hsv space to the rgb."""
    from matplotlib.colors import hsv_to_rgb
    return hsv_to_rgb(array(hsv).reshape(1, 1, 3)).reshape(3)

def _create_default_properties(data):
    """"Create the default properties of the mosaic given the data
    first it will varies the color hue (first category) then the color
    saturation (second category) and then the color value
    (third category).  If a fourth category is found, it will put
    decoration on the rectangle.  Doesn't manage more than four
    level of categories
    """
    categories_levels = _categories_level(list(iterkeys(data)))
    Nlevels = len(categories_levels)
    # first level, the hue
    L = len(categories_levels[0])
    # hue = np.linspace(1.0, 0.0, L+1)[:-1]
    hue = np.linspace(0.0, 1.0, L + 2)[:-2]
    # second level, the saturation
    L = len(categories_levels[1]) if Nlevels > 1 else 1
    saturation = np.linspace(0.5, 1.0, L + 1)[:-1]
    # third level, the value
    L = len(categories_levels[2]) if Nlevels > 2 else 1
    value = np.linspace(0.5, 1.0, L + 1)[:-1]
    # fourth level, the hatch
    L = len(categories_levels[3]) if Nlevels > 3 else 1
    hatch = ['', '/', '-', '|', '+'][:L + 1]
    # convert in list and merge with the levels
    hue = lzip(list(hue), categories_levels[0])
    saturation = lzip(list(saturation),
                     categories_levels[1] if Nlevels > 1 else [''])
    value = lzip(list(value),
                     categories_levels[2] if Nlevels > 2 else [''])
    hatch = lzip(list(hatch),
                     categories_levels[3] if Nlevels > 3 else [''])
    # create the properties dictionary
    properties = {}
    for h, s, v, t in product(hue, saturation, value, hatch):
        hv, hn = h
        sv, sn = s
        vv, vn = v
        tv, tn = t
        level = (hn,) + ((sn,) if sn else tuple())
        level = level + ((vn,) if vn else tuple())
        level = level + ((tn,) if tn else tuple())
        hsv = array([hv, sv, vv])
        prop = {'color': _single_hsv_to_rgb(hsv), 'hatch': tv, 'lw': 0}
        properties[level] = prop
    return properties

def _normalize_data(data, index):
    """normalize the data to a dict with tuples of strings as keys
    right now it works with:

        0 - dictionary (or equivalent mappable)
        1 - pandas.Series with simple or hierarchical indexes
        2 - numpy.ndarrays
        3 - everything that can be converted to a numpy array
        4 - pandas.DataFrame (via the _normalize_dataframe function)
    """
    # if data is a dataframe we need to take a completely new road
    # before coming back here. Use the hasattr to avoid importing
    # pandas explicitly
    if hasattr(data, 'pivot') and hasattr(data, 'groupby'):
        data = _normalize_dataframe(data, index)
        index = None
    # can it be used as a dictionary?
    try:
        items = list(iteritems(data))
    except AttributeError:
        # ok, I cannot use the data as a dictionary
        # Try to convert it to a numpy array, or die trying
        data = np.asarray(data)
        temp = OrderedDict()
        for idx in np.ndindex(data.shape):
            name = tuple(i for i in idx)
            temp[name] = data[idx]
        data = temp
        items = list(iteritems(data))
    # make all the keys a tuple, even if simple numbers
    data = OrderedDict([_tuplify(k), v] for k, v in items)
    categories_levels = _categories_level(list(iterkeys(data)))
    # fill the void in the counting dictionary
    indexes = product(*categories_levels)
    contingency = OrderedDict([(k, data.get(k, 0)) for k in indexes])
    data = contingency
    # reorder the keys order according to the one specified by the user
    # or if the index is None convert it into a simple list
    # right now it doesn't do any check, but can be modified in the future
    index = lrange(len(categories_levels)) if index is None else index
    contingency = OrderedDict()
    for key, value in iteritems(data):
        new_key = tuple(key[i] for i in index)
        contingency[new_key] = value
    data = contingency
    return data

def _normalize_dataframe(dataframe, index):
    """Take a pandas DataFrame and count the element present in the
    given columns, return a hierarchical index on those columns
    """
    #groupby the given keys, extract the same columns and count the element
    # then collapse them with a mean
    data = dataframe[index].dropna()
    grouped = data.groupby(index, sort=False)
    counted = grouped[index].count()
    averaged = counted.mean(axis=1)
    return averaged

def _statistical_coloring(data):
    """evaluate colors from the indipendence properties of the matrix
    It will encounter problem if one category has all zeros
    """
    data = _normalize_data(data, None)
    categories_levels = _categories_level(list(iterkeys(data)))
    Nlevels = len(categories_levels)
    total = 1.0 * sum(v for v in itervalues(data))
    # count the proportion of observation
    # for each level that has the given name
    # at each level
    levels_count = []
    for level_idx in range(Nlevels):
        proportion = {}
        for level in categories_levels[level_idx]:
            proportion[level] = 0.0
            for key, value in iteritems(data):
                if level == key[level_idx]:
                    proportion[level] += value
            proportion[level] /= total
        levels_count.append(proportion)
    # for each key I obtain the expected value
    # and it's standard deviation from a binomial distribution
    # under the hipothesys of independence
    expected = {}
    for key, value in iteritems(data):
        base = 1.0
        for i, k in enumerate(key):
            base *= levels_count[i][k]
        expected[key] = base * total, np.sqrt(total * base * (1.0 - base))
    # now we have the standard deviation of distance from the
    # expected value for each tile. We create the colors from this
    sigmas = dict((k, (data[k] - m) / s) for k, (m, s) in iteritems(expected))
    props = {}
    for key, dev in iteritems(sigmas):
        red = 0.0 if dev < 0 else (dev / (1 + dev))
        blue = 0.0 if dev > 0 else (dev / (-1 + dev))
        green = (1.0 - red - blue) / 2.0
        hatch = 'x' if dev > 2 else 'o' if dev < -2 else ''
        props[key] = {'color': [red, green, blue], 'hatch': hatch}
    return props

def _create_labels(rects, horizontal, ax, rotation):
    """find the position of the label for each value of each category

    right now it supports only up to the four categories

    ax: the axis on which the label should be applied
    rotation: the rotation list for each side
    """
    categories = _categories_level(list(iterkeys(rects)))
    if len(categories) > 4:
        msg = ("maximum of 4 level supported for axes labeling..and 4"
               "is alreay a lot of level, are you sure you need them all?")
        raise NotImplementedError(msg)
    labels = {}
    #keep it fixed as will be used a lot of times
    items = list(iteritems(rects))
    vertical = not horizontal

    #get the axis ticks and labels locator to put the correct values!
    ax2 = ax.twinx()
    ax3 = ax.twiny()
    #this is the order of execution for horizontal disposition
    ticks_pos = [ax.set_xticks, ax.set_yticks, ax3.set_xticks, ax2.set_yticks]
    ticks_lab = [ax.set_xticklabels, ax.set_yticklabels,
                 ax3.set_xticklabels, ax2.set_yticklabels]
    #for the vertical one, rotate it by one
    if vertical:
        ticks_pos = ticks_pos[1:] + ticks_pos[:1]
        ticks_lab = ticks_lab[1:] + ticks_lab[:1]
    #clean them
    for pos, lab in zip(ticks_pos, ticks_lab):
        pos([])
        lab([])
    #for each level, for each value in the level, take the mean of all
    #the sublevel that correspond to that partial key
    for level_idx, level in enumerate(categories):
        #this dictionary keep the labels only for this level
        level_ticks = dict()
        for value in level:
            #to which level it should refer to get the preceding
            #values of labels? it's rather a tricky question...
            #this is dependent on the side. It's a very crude management
            #but I couldn't think a more general way...
            if horizontal:
                if level_idx == 3:
                    index_select = [-1, -1, -1]
                else:
                    index_select = [+0, -1, -1]
            else:
                if level_idx == 3:
                    index_select = [+0, -1, +0]
                else:
                    index_select = [-1, -1, -1]
            #now I create the base key name and append the current value
            #It will search on all the rects to find the corresponding one
            #and use them to evaluate the mean position
            basekey = tuple(categories[i][index_select[i]]
                            for i in range(level_idx))
            basekey = basekey + (value,)
            subset = dict((k, v) for k, v in items
                          if basekey == k[:level_idx + 1])
            #now I extract the center of all the tiles and make a weighted
            #mean of all these center on the area of the tile
            #this should give me the (more or less) correct position
            #of the center of the category

            vals = list(itervalues(subset))
            W = sum(w * h for (x, y, w, h) in vals)
            x_lab = sum((x + w / 2.0) * w * h / W for (x, y, w, h) in vals)
            y_lab = sum((y + h / 2.0) * w * h / W for (x, y, w, h) in vals)
            #now base on the ordering, select which position to keep
            #needs to be written in a more general form of 4 level are enough?
            #should give also the horizontal and vertical alignment
            side = (level_idx + vertical) % 4
            level_ticks[value] = y_lab if side % 2 else x_lab
        #now we add the labels of this level to the correct axis

        ticks_pos[level_idx](list(itervalues(level_ticks)))
        ticks_lab[level_idx](list(iterkeys(level_ticks)),
                             rotation=rotation[level_idx])
    return labels

def mosaic(data, index=None, ax=None, horizontal=True, gap=0.005,
           properties=lambda key: None, labelizer=None,
           title='', statistic=False, axes_label=True,
           label_rotation=0.0):
    """Create a mosaic plot from a contingency table.

    It allows to visualize multivariate categorical data in a rigorous
    and informative way.

    Parameters
    ----------
    data : dict, pandas.Series, np.ndarray, pandas.DataFrame
        The contingency table that contains the data.
        Each category should contain a non-negative number
        with a tuple as index.  It expects that all the combination
        of keys to be representes; if that is not true, will
        automatically consider the missing values as 0.  The order
        of the keys will be the same as the one of insertion.
        If a dict of a Series (or any other dict like object)
        is used, it will take the keys as labels.  If a
        np.ndarray is provided, it will generate a simple
        numerical labels.
    index: list, optional
        Gives the preferred order for the category ordering. If not specified
        will default to the given order.  It doesn't support named indexes
        for hierarchical Series.  If a DataFrame is provided, it expects
        a list with the name of the columns.
    ax : matplotlib.Axes, optional
        The graph where display the mosaic. If not given, will
        create a new figure
    horizontal : bool, optional (default True)
        The starting direction of the split (by default along
        the horizontal axis)
    gap : float or array of floats
        The list of gaps to be applied on each subdivision.
        If the lenght of the given array is less of the number
        of subcategories (or if it's a single number) it will extend
        it with exponentially decreasing gaps
    labelizer : function (key) -> string, optional
        A function that generate the text to display at the center of
        each tile base on the key of that tile
    properties : function (key) -> dict, optional
        A function that for each tile in the mosaic take the key
        of the tile and returns the dictionary of properties
        of the generated Rectangle, like color, hatch or similar.
        A default properties set will be provided fot the keys whose
        color has not been defined, and will use color variation to help
        visually separates the various categories. It should return None
        to indicate that it should use the default property for the tile.
        A dictionary of the properties for each key can be passed,
        and it will be internally converted to the correct function
    statistic: bool, optional (default False)
        if true will use a crude statistical model to give colors to the plot.
        If the tile has a containt that is more than 2 standard deviation
        from the expected value under independence hipotesys, it will
        go from green to red (for positive deviations, blue otherwise) and
        will acquire an hatching when crosses the 3 sigma.
    title: string, optional
        The title of the axis
    axes_label: boolean, optional
        Show the name of each value of each category
        on the axis (default) or hide them.
    label_rotation: float or list of float
        the rotation of the axis label (if present). If a list is given
        each axis can have a different rotation

    Returns
    ----------
    fig : matplotlib.Figure
        The generate figure
    rects : dict
        A dictionary that has the same keys of the original
        dataset, that holds a reference to the coordinates of the
        tile and the Rectangle that represent it

    See Also
    ----------
    A Brief History of the Mosaic Display
        Michael Friendly, York University, Psychology Department
        Journal of Computational and Graphical Statistics, 2001

    Mosaic Displays for Loglinear Models.
        Michael Friendly, York University, Psychology Department
        Proceedings of the Statistical Graphics Section, 1992, 61-68.

    Mosaic displays for multi-way contingecy tables.
        Michael Friendly, York University, Psychology Department
        Journal of the american statistical association
        March 1994, Vol. 89, No. 425, Theory and Methods

    Examples
    ----------
    The most simple use case is to take a dictionary and plot the result

    >>> data = {'a': 10, 'b': 15, 'c': 16}
    >>> mosaic(data, title='basic dictionary')
    >>> pylab.show()

    A more useful example is given by a dictionary with multiple indices.
    In this case we use a wider gap to a better visual separation of the
    resulting plot

    >>> data = {('a', 'b'): 1, ('a', 'c'): 2, ('d', 'b'): 3, ('d', 'c'): 4}
    >>> mosaic(data, gap=0.05, title='complete dictionary')
    >>> pylab.show()

    The same data can be given as a simple or hierarchical indexed Series

    >>> rand = np.random.random
    >>> from itertools import product
    >>>
    >>> tuples = list(product(['bar', 'baz', 'foo', 'qux'], ['one', 'two']))
    >>> index = pd.MultiIndex.from_tuples(tuples, names=['first', 'second'])
    >>> data = pd.Series(rand(8), index=index)
    >>> mosaic(data, title='hierarchical index series')
    >>> pylab.show()

    The third accepted data structureis the np array, for which a
    very simple index will be created.

    >>> rand = np.random.random
    >>> data = 1+rand((2,2))
    >>> mosaic(data, title='random non-labeled array')
    >>> pylab.show()

    If you need to modify the labeling and the coloring you can give
    a function tocreate the labels and one with the graphical properties
    starting from the key tuple

    >>> data = {'a': 10, 'b': 15, 'c': 16}
    >>> props = lambda key: {'color': 'r' if 'a' in key else 'gray'}
    >>> labelizer = lambda k: {('a',): 'first', ('b',): 'second',
                               ('c',): 'third'}[k]
    >>> mosaic(data, title='colored dictionary',
                properties=props, labelizer=labelizer)
    >>> pylab.show()

    Using a DataFrame as source, specifying the name of the columns of interest
    >>> gender = ['male', 'male', 'male', 'female', 'female', 'female']
    >>> pet = ['cat', 'dog', 'dog', 'cat', 'dog', 'cat']
    >>> data = pandas.DataFrame({'gender': gender, 'pet': pet})
    >>> mosaic(data, ['pet', 'gender'])
    >>> pylab.show()
    """
    if isinstance(data, DataFrame) and index is None:
        raise ValueError("You must pass an index if data is a DataFrame."
                         " See examples.")

    from pylab import Rectangle
    fig, ax = utils.create_mpl_ax(ax)
    # normalize the data to a dict with tuple of strings as keys
    data = _normalize_data(data, index)
    # split the graph into different areas
    rects = _hierarchical_split(data, horizontal=horizontal, gap=gap)
    # if there is no specified way to create the labels
    # create a default one
    if labelizer is None:
        labelizer = lambda k: "\n".join(k)
    if statistic:
        default_props = _statistical_coloring(data)
    else:
        default_props = _create_default_properties(data)
    if isinstance(properties, dict):
        color_dict = properties
        properties = lambda key: color_dict.get(key, None)
    for k, v in iteritems(rects):
        # create each rectangle and put a label on it
        x, y, w, h = v
        conf = properties(k)
        props = conf if conf else default_props[k]
        text = labelizer(k)
        Rect = Rectangle((x, y), w, h, label=text, **props)
        ax.add_patch(Rect)
        ax.text(x + w / 2, y + h / 2, text, ha='center',
                 va='center', size='smaller')
    #creating the labels on the axis
    #o clearing it
    if axes_label:
        if np.iterable(label_rotation):
            rotation = label_rotation
        else:
            rotation = [label_rotation] * 4
        labels = _create_labels(rects, horizontal, ax, rotation)
    else:
        ax.set_xticks([])
        ax.set_xticklabels([])
        ax.set_yticks([])
        ax.set_yticklabels([])
    ax.set_title(title)
    return fig, rects


from __future__ import print_function
import fcl
import os
import time
from os.path import abspath, join, dirname, isfile
from fcl import kmeans
from fcl.datasets import load_sector_dataset, load_usps_dataset
from fcl.matrix.csr_matrix import get_csr_matrix_from_object, csr_matrix_to_libsvm_string
from sklearn.decomposition import TruncatedSVD, PCA
from scipy.sparse import csr_matrix
from sklearn.datasets import dump_svmlight_file
import numpy as np
import argparse

def get_pca_projection_csrmatrix(fcl_csr_input_matrix, component_ratio):
  n_components = int(fcl_csr_input_matrix.annz * component_ratio)
  p = TruncatedSVD(n_components = n_components)
  start = time.time()
  p.fit(fcl_csr_input_matrix.to_numpy())
  # convert to millis
  fin = (time.time() - start) * 1000
  (n_samples, n_dim) = fcl_csr_input_matrix.shape
  print("Truncated SVD took %.3fs to retrieve %s components for input_matrix with n_samples %d, n_dim %d" % (fin/1000.0, str(n_components), n_samples, n_dim))
  return get_csr_matrix_from_object(p.components_)

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='Create a pca matrix from an input matrix with given component ratio.')
  parser.add_argument('path_input_dataset', type=str, help='Path to the input libsvm dataset')
  parser.add_argument('path_output_dataset', type=str, help='Path to the input libsvm dataset')
  parser.add_argument('--component_ratio', default=0.1, type=float, help='Percentage of the average non zero values of the input dataset to use as components.')

  args = parser.parse_args()

  if not isfile(args.path_input_dataset):
    raise Exception("Unable to find path_input_dataset: %s" % args.path_input_dataset)

  print("Loading data from %s" % args.path_input_dataset)
  fcl_mtrx_input_dataset = get_csr_matrix_from_object(args.path_input_dataset)

  print("Retrieving the pca projection matrix")
  pca_mtrx = get_pca_projection_csrmatrix(fcl_mtrx_input_dataset, args.component_ratio)

  print("Convert pca projection matrix to libsvm string")
  pca_mtrx_lsvm_str = csr_matrix_to_libsvm_string(pca_mtrx)

  print("Writing pca projection matrix libsvm string to file %s" % args.path_output_dataset)
  with open(args.path_output_dataset, 'w') as f:
    f.write(pca_mtrx_lsvm_str)


"""
The :mod:`sklearn.pipeline` module implements utilities to build a composite
estimator, as a chain of transforms and estimators.
"""
# Author: Edouard Duchesnay
#         Gael Varoquaux
#         Virgile Fritsch
#         Alexandre Gramfort
#         Lars Buitinck
# Licence: BSD

from collections import defaultdict
from warnings import warn

import numpy as np
from scipy import sparse

from .base import BaseEstimator, TransformerMixin
from .externals.joblib import Parallel, delayed
from .externals import six
from .utils import tosequence
from .utils.metaestimators import if_delegate_has_method
from .externals.six import iteritems

__all__ = ['Pipeline', 'FeatureUnion']

class Pipeline(BaseEstimator):
    """Pipeline of transforms with a final estimator.

    Sequentially apply a list of transforms and a final estimator.
    Intermediate steps of the pipeline must be 'transforms', that is, they
    must implement fit and transform methods.
    The final estimator only needs to implement fit.

    The purpose of the pipeline is to assemble several steps that can be
    cross-validated together while setting different parameters.
    For this, it enables setting parameters of the various steps using their
    names and the parameter name separated by a '__', as in the example below.

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

    Parameters
    ----------
    steps : list
        List of (name, transform) tuples (implementing fit/transform) that are
        chained, in the order in which they are chained, with the last object
        an estimator.

    Attributes
    ----------
    named_steps : dict
        Read-only attribute to access any step parameter by user given name.
        Keys are step names and values are steps parameters.

    Examples
    --------
    >>> from sklearn import svm
    >>> from sklearn.datasets import samples_generator
    >>> from sklearn.feature_selection import SelectKBest
    >>> from sklearn.feature_selection import f_regression
    >>> from sklearn.pipeline import Pipeline
    >>> # generate some data to play with
    >>> X, y = samples_generator.make_classification(
    ...     n_informative=5, n_redundant=0, random_state=42)
    >>> # ANOVA SVM-C
    >>> anova_filter = SelectKBest(f_regression, k=5)
    >>> clf = svm.SVC(kernel='linear')
    >>> anova_svm = Pipeline([('anova', anova_filter), ('svc', clf)])
    >>> # You can set the parameters using the names issued
    >>> # For instance, fit using a k of 10 in the SelectKBest
    >>> # and a parameter 'C' of the svm
    >>> anova_svm.set_params(anova__k=10, svc__C=.1).fit(X, y)
    ...                                              # doctest: +ELLIPSIS
    Pipeline(steps=[...])
    >>> prediction = anova_svm.predict(X)
    >>> anova_svm.score(X, y)                        # doctest: +ELLIPSIS
    0.77...
    >>> # getting the selected features chosen by anova_filter
    >>> anova_svm.named_steps['anova'].get_support()
    ... # doctest: +NORMALIZE_WHITESPACE
    array([ True,  True,  True, False, False,  True, False,  True,  True, True,
           False, False,  True, False,  True, False, False, False, False,
           True], dtype=bool)
    """

    # BaseEstimator interface

    def __init__(self, steps):
        names, estimators = zip(*steps)
        if len(dict(steps)) != len(steps):
            raise ValueError("Provided step names are not unique: %s" % (names,))

        # shallow copy of steps
        self.steps = tosequence(steps)
        transforms = estimators[:-1]
        estimator = estimators[-1]

        for t in transforms:
            if (not (hasattr(t, "fit") or hasattr(t, "fit_transform")) or not
                    hasattr(t, "transform")):
                raise TypeError("All intermediate steps of the chain should "
                                "be transforms and implement fit and transform"
                                " '%s' (type %s) doesn't)" % (t, type(t)))

        if not hasattr(estimator, "fit"):
            raise TypeError("Last step of chain should implement fit "
                            "'%s' (type %s) doesn't)"
                            % (estimator, type(estimator)))

    @property
    def _estimator_type(self):
        return self.steps[-1][1]._estimator_type

    def get_params(self, deep=True):
        if not deep:
            return super(Pipeline, self).get_params(deep=False)
        else:
            out = self.named_steps
            for name, step in six.iteritems(self.named_steps):
                for key, value in six.iteritems(step.get_params(deep=True)):
                    out['%s__%s' % (name, key)] = value

            out.update(super(Pipeline, self).get_params(deep=False))
            return out

    @property
    def named_steps(self):
        return dict(self.steps)

    @property
    def _final_estimator(self):
        return self.steps[-1][1]

    # Estimator interface

    def _pre_transform(self, X, y=None, **fit_params):
        fit_params_steps = dict((step, {}) for step, _ in self.steps)
        for pname, pval in six.iteritems(fit_params):
            step, param = pname.split('__', 1)
            fit_params_steps[step][param] = pval
        Xt = X
        for name, transform in self.steps[:-1]:
            if hasattr(transform, "fit_transform"):
                Xt = transform.fit_transform(Xt, y, **fit_params_steps[name])
            else:
                Xt = transform.fit(Xt, y, **fit_params_steps[name]) \
                              .transform(Xt)
        return Xt, fit_params_steps[self.steps[-1][0]]

    def fit(self, X, y=None, **fit_params):
        """Fit all the transforms one after the other and transform the
        data, then fit the transformed data using the final estimator.

        Parameters
        ----------
        X : iterable
            Training data. Must fulfill input requirements of first step of the
            pipeline.
        y : iterable, default=None
            Training targets. Must fulfill label requirements for all steps of
            the pipeline.
        """
        Xt, fit_params = self._pre_transform(X, y, **fit_params)
        self.steps[-1][-1].fit(Xt, y, **fit_params)
        return self

    def fit_transform(self, X, y=None, **fit_params):
        """Fit all the transforms one after the other and transform the
        data, then use fit_transform on transformed data using the final
        estimator.

        Parameters
        ----------
        X : iterable
            Training data. Must fulfill input requirements of first step of the
            pipeline.

        y : iterable, default=None
            Training targets. Must fulfill label requirements for all steps of
            the pipeline.
        """
        Xt, fit_params = self._pre_transform(X, y, **fit_params)
        if hasattr(self.steps[-1][-1], 'fit_transform'):
            return self.steps[-1][-1].fit_transform(Xt, y, **fit_params)
        else:
            return self.steps[-1][-1].fit(Xt, y, **fit_params).transform(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def predict(self, X):
        """Applies transforms to the data, and the predict method of the
        final estimator. Valid only if the final estimator implements
        predict.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].predict(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def fit_predict(self, X, y=None, **fit_params):
        """Applies fit_predict of last step in pipeline after transforms.

        Applies fit_transforms of a pipeline to the data, followed by the
        fit_predict method of the final estimator in the pipeline. Valid
        only if the final estimator implements fit_predict.

        Parameters
        ----------
        X : iterable
            Training data. Must fulfill input requirements of first step of
            the pipeline.
        y : iterable, default=None
            Training targets. Must fulfill label requirements for all steps
            of the pipeline.
        """
        Xt, fit_params = self._pre_transform(X, y, **fit_params)
        return self.steps[-1][-1].fit_predict(Xt, y, **fit_params)

    @if_delegate_has_method(delegate='_final_estimator')
    def predict_proba(self, X):
        """Applies transforms to the data, and the predict_proba method of the
        final estimator. Valid only if the final estimator implements
        predict_proba.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].predict_proba(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def decision_function(self, X):
        """Applies transforms to the data, and the decision_function method of
        the final estimator. Valid only if the final estimator implements
        decision_function.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].decision_function(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def predict_log_proba(self, X):
        """Applies transforms to the data, and the predict_log_proba method of
        the final estimator. Valid only if the final estimator implements
        predict_log_proba.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].predict_log_proba(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def transform(self, X):
        """Applies transforms to the data, and the transform method of the
        final estimator. Valid only if the final estimator implements
        transform.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps:
            Xt = transform.transform(Xt)
        return Xt

    @if_delegate_has_method(delegate='_final_estimator')
    def inverse_transform(self, X):
        """Applies inverse transform to the data.
        Starts with the last step of the pipeline and applies ``inverse_transform`` in
        inverse order of the pipeline steps.
        Valid only if all steps of the pipeline implement inverse_transform.

        Parameters
        ----------
        X : iterable
            Data to inverse transform. Must fulfill output requirements of the
            last step of the pipeline.
        """
        if X.ndim == 1:
            warn("From version 0.19, a 1d X will not be reshaped in"
                 " pipeline.inverse_transform any more.", FutureWarning)
            X = X[None, :]
        Xt = X
        for name, step in self.steps[::-1]:
            Xt = step.inverse_transform(Xt)
        return Xt

    @if_delegate_has_method(delegate='_final_estimator')
    def score(self, X, y=None):
        """Applies transforms to the data, and the score method of the
        final estimator. Valid only if the final estimator implements
        score.

        Parameters
        ----------
        X : iterable
            Data to score. Must fulfill input requirements of first step of the
            pipeline.

        y : iterable, default=None
            Targets used for scoring. Must fulfill label requirements for all steps of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].score(Xt, y)

    @property
    def classes_(self):
        return self.steps[-1][-1].classes_

    @property
    def _pairwise(self):
        # check if first estimator expects pairwise input
        return getattr(self.steps[0][1], '_pairwise', False)

def _name_estimators(estimators):
    """Generate names for estimators."""

    names = [type(estimator).__name__.lower() for estimator in estimators]
    namecount = defaultdict(int)
    for est, name in zip(estimators, names):
        namecount[name] += 1

    for k, v in list(six.iteritems(namecount)):
        if v == 1:
            del namecount[k]

    for i in reversed(range(len(estimators))):
        name = names[i]
        if name in namecount:
            names[i] += "-%d" % namecount[name]
            namecount[name] -= 1

    return list(zip(names, estimators))

def make_pipeline(*steps):
    """Construct a Pipeline from the given estimators.

    This is a shorthand for the Pipeline constructor; it does not require, and
    does not permit, naming the estimators. Instead, their names will be set
    to the lowercase of their types automatically.

    Examples
    --------
    >>> from sklearn.naive_bayes import GaussianNB
    >>> from sklearn.preprocessing import StandardScaler
    >>> make_pipeline(StandardScaler(), GaussianNB())    # doctest: +NORMALIZE_WHITESPACE
    Pipeline(steps=[('standardscaler',
                     StandardScaler(copy=True, with_mean=True, with_std=True)),
                    ('gaussiannb', GaussianNB())])

    Returns
    -------
    p : Pipeline
    """
    return Pipeline(_name_estimators(steps))

def _fit_one_transformer(transformer, X, y):
    return transformer.fit(X, y)

def _transform_one(transformer, name, X, transformer_weights):
    if transformer_weights is not None and name in transformer_weights:
        # if we have a weight for this transformer, multiply output
        return transformer.transform(X) * transformer_weights[name]
    return transformer.transform(X)

def _fit_transform_one(transformer, name, X, y, transformer_weights,
                       **fit_params):
    if transformer_weights is not None and name in transformer_weights:
        # if we have a weight for this transformer, multiply output
        if hasattr(transformer, 'fit_transform'):
            X_transformed = transformer.fit_transform(X, y, **fit_params)
            return X_transformed * transformer_weights[name], transformer
        else:
            X_transformed = transformer.fit(X, y, **fit_params).transform(X)
            return X_transformed * transformer_weights[name], transformer
    if hasattr(transformer, 'fit_transform'):
        X_transformed = transformer.fit_transform(X, y, **fit_params)
        return X_transformed, transformer
    else:
        X_transformed = transformer.fit(X, y, **fit_params).transform(X)
        return X_transformed, transformer

class FeatureUnion(BaseEstimator, TransformerMixin):
    """Concatenates results of multiple transformer objects.

    This estimator applies a list of transformer objects in parallel to the
    input data, then concatenates the results. This is useful to combine
    several feature extraction mechanisms into a single transformer.

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

    Parameters
    ----------
    transformer_list: list of (string, transformer) tuples
        List of transformer objects to be applied to the data. The first
        half of each tuple is the name of the transformer.

    n_jobs: int, optional
        Number of jobs to run in parallel (default 1).

    transformer_weights: dict, optional
        Multiplicative weights for features per transformer.
        Keys are transformer names, values the weights.

    """
    def __init__(self, transformer_list, n_jobs=1, transformer_weights=None):
        self.transformer_list = transformer_list
        self.n_jobs = n_jobs
        self.transformer_weights = transformer_weights

    def get_feature_names(self):
        """Get feature names from all transformers.

        Returns
        -------
        feature_names : list of strings
            Names of the features produced by transform.
        """
        feature_names = []
        for name, trans in self.transformer_list:
            if not hasattr(trans, 'get_feature_names'):
                raise AttributeError("Transformer %s does not provide"
                                     " get_feature_names." % str(name))
            feature_names.extend([name + "__" + f for f in
                                  trans.get_feature_names()])
        return feature_names

    def fit(self, X, y=None):
        """Fit all transformers using X.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            Input data, used to fit transformers.
        """
        transformers = Parallel(n_jobs=self.n_jobs)(
            delayed(_fit_one_transformer)(trans, X, y)
            for name, trans in self.transformer_list)
        self._update_transformer_list(transformers)
        return self

    def fit_transform(self, X, y=None, **fit_params):
        """Fit all transformers using X, transform the data and concatenate
        results.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            Input data to be transformed.

        Returns
        -------
        X_t : array-like or sparse matrix, shape (n_samples, sum_n_components)
            hstack of results of transformers. sum_n_components is the
            sum of n_components (output dimension) over transformers.
        """
        result = Parallel(n_jobs=self.n_jobs)(
            delayed(_fit_transform_one)(trans, name, X, y,
                                        self.transformer_weights, **fit_params)
            for name, trans in self.transformer_list)

        Xs, transformers = zip(*result)
        self._update_transformer_list(transformers)
        if any(sparse.issparse(f) for f in Xs):
            Xs = sparse.hstack(Xs).tocsr()
        else:
            Xs = np.hstack(Xs)
        return Xs

    def transform(self, X):
        """Transform X separately by each transformer, concatenate results.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            Input data to be transformed.

        Returns
        -------
        X_t : array-like or sparse matrix, shape (n_samples, sum_n_components)
            hstack of results of transformers. sum_n_components is the
            sum of n_components (output dimension) over transformers.
        """
        Xs = Parallel(n_jobs=self.n_jobs)(
            delayed(_transform_one)(trans, name, X, self.transformer_weights)
            for name, trans in self.transformer_list)
        if any(sparse.issparse(f) for f in Xs):
            Xs = sparse.hstack(Xs).tocsr()
        else:
            Xs = np.hstack(Xs)
        return Xs

    def get_params(self, deep=True):
        if not deep:
            return super(FeatureUnion, self).get_params(deep=False)
        else:
            out = dict(self.transformer_list)
            for name, trans in self.transformer_list:
                for key, value in iteritems(trans.get_params(deep=True)):
                    out['%s__%s' % (name, key)] = value
            out.update(super(FeatureUnion, self).get_params(deep=False))
            return out

    def _update_transformer_list(self, transformers):
        self.transformer_list[:] = [
            (name, new)
            for ((name, old), new) in zip(self.transformer_list, transformers)
        ]

# XXX it would be nice to have a keyword-only n_jobs argument to this function,
# but that's not allowed in Python 2.x.
def make_union(*transformers):
    """Construct a FeatureUnion from the given transformers.

    This is a shorthand for the FeatureUnion constructor; it does not require,
    and does not permit, naming the transformers. Instead, they will be given
    names automatically based on their types. It also does not allow weighting.

    Examples
    --------
    >>> from sklearn.decomposition import PCA, TruncatedSVD
    >>> make_union(PCA(), TruncatedSVD())    # doctest: +NORMALIZE_WHITESPACE
    FeatureUnion(n_jobs=1,
                 transformer_list=[('pca', PCA(copy=True, n_components=None,
                                               whiten=False)),
                                   ('truncatedsvd',
                                    TruncatedSVD(algorithm='randomized',
                                                 n_components=2, n_iter=5,
                                                 random_state=None, tol=0.0))],
                 transformer_weights=None)

    Returns
    -------
    f : FeatureUnion
    """
    return FeatureUnion(_name_estimators(transformers))


#!/usr/bin/env python

#
# Arun Venkatraman (arunvenk@cs.cmu.edu)
# December 2016
#

# If we are not running from the build directory, then add lib to path from
# build assuming we are running from the python folder
import os
full_path = os.path.realpath(__file__)
if full_path.count("src/python") > 0:
    import sys
    to_add = os.path.abspath(os.path.join(os.path.split(full_path)[0], "../../build/"))
    sys.path.append(to_add)

from IPython import embed

import lib.ilqr_diffdrive as ilqr
import visualize_circle_world as vis

import numpy as np
import matplotlib.pyplot as plt

if __name__ == "__main__":
    obs_prior = [0.5, 0.5]
    world_dims = [-30, 30, -30, 30]

    w1 = ilqr.CircleWorld(world_dims)
    w2 = ilqr.CircleWorld(world_dims)

    obs_pos_1 = [-2, 0.0]
    obs_pos_2 = [2, 0.0]

    obs_radius = 10.0
    obstacle_1 = ilqr.Circle(obs_radius, obs_pos_1);
    obstacle_2 = ilqr.Circle(obs_radius, obs_pos_2);

    # add obstacle to world 1
    w1.add_obstacle(obstacle_1);
    # add obstacle to world 2
    w2.add_obstacle(obstacle_2);

    cost, states_true_1, obs_fname_1 = ilqr.control_diffdrive(ilqr.TRUE_ILQR,
            w1, w2, obs_prior, "true1", "true1")
    cost, states_true_2, obs_fname_2 = ilqr.control_diffdrive(ilqr.TRUE_ILQR,
            w2, w1, obs_prior, "true2", "true2")

    cost, states_weighted_1, obs_fname_3 =\
            ilqr.control_diffdrive(ilqr.PROB_WEIGHTED_CONTROL,
                w1, w2, obs_prior, "weight3", "weight3")
    cost, states_weighted_2, obs_fname_4 =\
            ilqr.control_diffdrive(ilqr.PROB_WEIGHTED_CONTROL,
                w2, w1, obs_prior, "weight4", "weight4")

    cost, states_hind_1, obs_fname_5 =\
            ilqr.control_diffdrive(ilqr.HINDSIGHT,
                w1, w2, obs_prior, "hind3", "hind3")
    cost, states_hind_2, obs_fname_6 =\
            ilqr.control_diffdrive(ilqr.HINDSIGHT,
                w2, w1, obs_prior, "hind4", "hind4")

    print("Drawing world 1")
    ax1 = vis.parse_draw_files([states_true_1, states_weighted_1, states_hind_1], obs_fname_1,
            show=False)
    plt.title('World 1')

    print("Drawing world 2")
    ax2 = vis.parse_draw_files([states_true_2, states_weighted_2, states_hind_2],
            obs_fname_2, show=False)
    plt.title('World 2')
    plt.show()

    embed()


'''
    Created on 08.06.2015

    @author: marinavidovic
    '''
import os
import pdb
import utils_svm
import pickle
import numpy as np
import copy
import genQ
import makePOIM
import view
import matplotlib
matplotlib.use('Agg')

if __name__ == '__main__':

    read_data = 1
    datapath = "/home/mvidovic/POIMall/data/real/human_acceptor_splice_data.txt"
    savepath = "/home/mvidovic/POIMall/data/real/human_acceptor_splice_data0.pkl"

    lines=1000
    if read_data:
        x,y=utils_svm.extractRealData(datapath,savepath,lines)
    else:
        fobj=open(savepath,'rb')
        x,y=pickle.load(fobj)
        fobj.close()
    num_pos = 100
    num_neg = 4*num_pos
    print "reduce samples"
    x_red,y_red = utils_svm.reduce_samples(copy.deepcopy(x),copy.deepcopy(y),num_pos,num_neg)
    nploci_letters,nploci_positions = utils_svm.non_polymorphic_loci(x_red)
    #read data
    experiment_name = "real1"
    if not os.path.exists(experiment_name):
        os.makedirs(experiment_name)
    poimpath=experiment_name+"/poim.pkl"
    tally=30
    positives=25
    sequenceno=100
    mutation_prob=0.0
    motif="ATTTT"
    mu=13
    x,y = makePOIM.gensequences(tally,positives,sequenceno,mutation_prob,motif,mu)

    #compute POIM
    poim_degree = 6
    kernel_degree = 8
    print "start poim computation"
    poims = makePOIM.computePOIM(x,y,poim_degree,kernel_degree,poimpath)
    Q2 = poims[0][1]
    #view.test()
    view.figurepoimsimple(Q2, "poim_pic", 0)


import os
import pandas as pd

def read_synchronisation_file(experiment_root):
    filepath = os.path.join(experiment_root, "labels", "synchronisation.csv")
    return pd.read_csv(filepath)

def convert_timestamps(experiment_root, timestamps, from_reference, to_reference):
    """
    Convert numeric timestamps (seconds for start of the video or posix timestamp) of a reference time (e.g. P3_eyetracker) to a different reference time (e.g. video time)

    Parameters
    ----------
    experiment_root: str
        Root of the current experiment (to find the right synchronisation matrix)
    timestamps: float or array like
        timestamps to be converted
    from_reference: str
        name of the reference of the original timestamps
    to_reference: str
        name of the reference time the timestamp has to be converted to

    Returns
    -------
        converted_timestamps: float or array like
            Timestamps given in to_reference time values
    """
    synchronisation_file = read_synchronisation_file(experiment_root)

    offset = synchronisation_file.loc[synchronisation_file["from"] == from_reference, to_reference].values[0]

    converted_timestamps = timestamps + offset

    return converted_timestamps

if __name__ == '__main__':
    exp_root = "/Volumes/DataDrive/igroups_recordings/igroups_experiment_8"

    print(convert_timestamps(exp_root, [1482326641, 1482326642], "P3_eyetracker", "video"))


# tmax_peakfind_example.py
"""
Demonstration of some of the primary functions in pygaero, including Tmax finding and elemental analysis.
"""

# Module import
from pygaero import pio
from pygaero import therm
from pygaero import gen_chem
import os
import matplotlib.pyplot as plt

def example1():
    # ------------------------------- File I/O and Data Cleaning Example -------------------------------- #
    indir = ""  # input directory (same folder as script by default)
    infiles = ['desorb1.csv', 'desorb2.csv']  # input files as a list of strings
    # Read in list of csvs with figaero desorptions
    df_desorbs_ls = pio.read_files(fdir=indir, flist=infiles)
    print('# of files imported: ', len(df_desorbs_ls))

    # Clean ion names from default A_CxHyOzI_Avg format (strip underscores '_' and remove iodide
    for df in df_desorbs_ls:
        print("Example of ion names before clean: ", df.columns.values[0:3])
        df.columns = gen_chem.cln_molec_names(idx_names=df.columns.values, delim="_")  # remove underscores
        df.columns = gen_chem.replace_group(molec_ls=df.columns.values, old_groups=["I"], new_group="")  # remove I
        print('Example of ion names after clean: ', df.columns.values[0:3])

    # Alternatively, one can just assign a single thermogram by df_example = pd.DataFrame.from_csv(indir+infile)
    # Adjust thermogram signals for 4.0 LPM figaero flow rate relative to nominal 2.0 LPM sample rate
    # print('Before flow rate adjust:', df_desorbs_ls[0].values[0:3, 5])
    therm.flow_correction(thermograms=df_desorbs_ls, aero_samp_rates=[4.0, 4.0])
    # print('After flow rate adjust:', df_desorbs_ls[0].values[0:3, 5])

    # ---------------------------------- Elemental Stats Example --------------------------------------- #
    # A. Calculate elemental statistics for species in each desorb CSV that was read in. Then append the DataFrames
    # containing these statistics into a list. Note, Iodide has been stripped from the names at this point, so
    # the parameter cluster_group=None
    ele_stats_ls = []
    for df in df_desorbs_ls:
        df_ele_temp = gen_chem.ele_stats(molec_ls=df.columns.values, ion_state=-1, cluster_group=None,
                                         clst_group_mw=0.0, xtra_elements=["Cl", "F"])
        ele_stats_ls.append(df_ele_temp)

    # -------------------------------- Peak Finding (TMax) Example --------------------------------------#
    # A. Smooth time series as step prior to Tmax (helps prevent mis-identification of TMax in noisy signals)
    for df in df_desorbs_ls:
        for series in df.columns.values:
            # print('series: ', series)
            df.ix[:, series] = therm.smooth(x=df.ix[:, series].values, window='hamming', window_len=15)
        plt.show()

    # B. Find TMax for all loaded thermograms. Returns a pandas DataFrame with ion names as index values and columns:
    # TMax1, MaxSig1, TMax2, MaxSig2, DubFlag (double peak flag - binary; -1 for no peaks found). Depending on the
    # specific data set, the [pk_thresh] and [pk_win] parameters may need to be optimized. See documentation for
    # function peakfind_df_ls in module therm.py for more details. Results are drastically improved by first
    # smoothing the time series, so that small fluctuations in signal are not mistaken for a peak.
    df_tmax_ls = therm.peakfind_df_ls(df_ls=df_desorbs_ls, pk_thresh=0.05, pk_win=20,
                                      min_temp=40.0, max_temp=190.0)

    # C. Quick plot to visualize Tmax values for 15 example ions
    # therm.plot_tmax(df=df_desorbs_ls[0], ions=df_tmax_ls[0].index.values[15:29],
    #                 tmax_temps=df_tmax_ls[0].ix[15:29, 'TMax1'], tmax_vals=df_tmax_ls[0].ix[15:29, 'MaxSig1'])
    therm.plot_tmax_double(df=df_desorbs_ls[0], ions=df_tmax_ls[0].index.values[15:29],
                           tmax_temps=df_tmax_ls[0].ix[15:29, 'TMax1'],
                           tmax_temps2=df_tmax_ls[0].ix[15:29, 'TMax2'],
                           tmax_vals=df_tmax_ls[0].ix[15:29, 'MaxSig1'],
                           tmax_vals2=df_tmax_ls[0].ix[15:29, 'MaxSig2'])

    # ----------------------------------- Saving Results Example -------------------------------------- #
    # Uncomment the following lines to save the example output
    # outdir = 'testout'
    # if outdir[-1] != '/':
    #     outdir += '/'
    # if not os.path.exists(outdir):
    #     os.makedirs(outdir)
    # # A. Save TMax data
    # for df, fname in zip(df_tmax_ls, ["desorb1_tmax", "desorb2_tmax"]):
    #     df.to_csv(outdir+fname+".csv")
    # # B. Save smoothed desorption thermogram time series
    # for df, fname in zip(df_desorbs_ls, ["desorb1_smth", "desorb2_smth"]):
    #     df.to_csv(outdir+fname+".csv")
    # # C. Save elemental stats for each desorption
    # for df, fname in zip(ele_stats_ls, ["desorb1_ele", "desorb2_ele"]):
    #     df.to_csv(outdir+fname+".csv")

    return 0

if __name__ == '__main__':
    example1()