# -*- coding: utf-8 -*-
# pylint: disable=W0102

from datetime import datetime, date

import nose
import numpy as np

import re
import itertools
from pandas import Index, MultiIndex, DataFrame, DatetimeIndex, Series, Categorical
from pandas.compat import OrderedDict, lrange
from pandas.sparse.array import SparseArray
from pandas.core.internals import (BlockPlacement, SingleBlockManager, make_block,
                                   BlockManager)
import pandas.core.common as com
import pandas.core.internals as internals
import pandas.util.testing as tm
import pandas as pd
from pandas.util.testing import (
    assert_almost_equal, assert_frame_equal, randn, assert_series_equal)
from pandas.compat import zip, u

def assert_block_equal(left, right):
    assert_almost_equal(left.values, right.values)
    assert(left.dtype == right.dtype)
    assert_almost_equal(left.mgr_locs, right.mgr_locs)

def get_numeric_mat(shape):
    arr = np.arange(shape[0])
    return np.lib.stride_tricks.as_strided(
        x=arr, shape=shape,
        strides=(arr.itemsize,) + (0,) * (len(shape) - 1)).copy()

N = 10

def create_block(typestr, placement, item_shape=None, num_offset=0):
    """
    Supported typestr:

        * float, f8, f4, f2
        * int, i8, i4, i2, i1
        * uint, u8, u4, u2, u1
        * complex, c16, c8
        * bool
        * object, string, O
        * datetime, dt, M8[ns], M8[ns, tz]
        * timedelta, td, m8[ns]
        * sparse (SparseArray with fill_value=0.0)
        * sparse_na (SparseArray with fill_value=np.nan)
        * category, category2

    """
    placement = BlockPlacement(placement)
    num_items = len(placement)

    if item_shape is None:
        item_shape = (N,)

    shape = (num_items,) + item_shape

    mat = get_numeric_mat(shape)

    if typestr in ('float', 'f8', 'f4', 'f2',
                   'int', 'i8', 'i4', 'i2', 'i1',
                   'uint', 'u8', 'u4', 'u2', 'u1'):
        values = mat.astype(typestr) + num_offset
    elif typestr in ('complex', 'c16', 'c8'):
        values = 1.j * (mat.astype(typestr) + num_offset)
    elif typestr in ('object', 'string', 'O'):
        values = np.reshape(['A%d' % i for i in mat.ravel() + num_offset],
                            shape)
    elif typestr in ('b','bool',):
        values = np.ones(shape, dtype=np.bool_)
    elif typestr in ('datetime', 'dt', 'M8[ns]'):
        values = (mat * 1e9).astype('M8[ns]')
    elif typestr.startswith('M8[ns'):
        # datetime with tz
        m = re.search('M8\[ns,\s*(\w+\/?\w*)\]', typestr)
        assert m is not None, "incompatible typestr -> {0}".format(typestr)
        tz = m.groups()[0]
        assert num_items == 1, "must have only 1 num items for a tz-aware"
        values = DatetimeIndex(np.arange(N) * 1e9, tz=tz)
    elif typestr in ('timedelta', 'td', 'm8[ns]'):
        values = (mat * 1).astype('m8[ns]')
    elif typestr in ('category',):
        values = Categorical([1,1,2,2,3,3,3,3,4,4])
    elif typestr in ('category2',):
        values = Categorical(['a','a','a','a','b','b','c','c','c','d'])
    elif typestr in ('sparse', 'sparse_na'):
        # FIXME: doesn't support num_rows != 10
        assert shape[-1] == 10
        assert all(s == 1 for s in shape[:-1])
        if typestr.endswith('_na'):
            fill_value = np.nan
        else:
            fill_value = 0.0
        values = SparseArray([fill_value, fill_value, 1, 2, 3, fill_value,
                              4, 5, fill_value, 6], fill_value=fill_value)
        arr = values.sp_values.view()
        arr += (num_offset - 1)
    else:
        raise ValueError('Unsupported typestr: "%s"' % typestr)

    return make_block(values, placement=placement, ndim=len(shape))

def create_single_mgr(typestr, num_rows=None):
    if num_rows is None:
        num_rows = N

    return SingleBlockManager(
        create_block(typestr, placement=slice(0, num_rows), item_shape=()),
        np.arange(num_rows))

def create_mgr(descr, item_shape=None):
    """
    Construct BlockManager from string description.

    String description syntax looks similar to np.matrix initializer.  It looks
    like this::

        a,b,c: f8; d,e,f: i8

    Rules are rather simple:

    * see list of supported datatypes in `create_block` method
    * components are semicolon-separated
    * each component is `NAME,NAME,NAME: DTYPE_ID`
    * whitespace around colons & semicolons are removed
    * components with same DTYPE_ID are combined into single block
    * to force multiple blocks with same dtype, use '-SUFFIX'::

        'a:f8-1; b:f8-2; c:f8-foobar'

    """
    if item_shape is None:
        item_shape = (N,)

    offset = 0
    mgr_items = []
    block_placements = OrderedDict()
    for d in descr.split(';'):
        d = d.strip()
        names, blockstr = d.partition(':')[::2]
        blockstr = blockstr.strip()
        names = names.strip().split(',')

        mgr_items.extend(names)
        placement = list(np.arange(len(names)) + offset)
        try:
            block_placements[blockstr].extend(placement)
        except KeyError:
            block_placements[blockstr] = placement
        offset += len(names)

    mgr_items = Index(mgr_items)

    blocks = []
    num_offset = 0
    for blockstr, placement in block_placements.items():
        typestr = blockstr.split('-')[0]
        blocks.append(create_block(typestr, placement, item_shape=item_shape,
                                   num_offset=num_offset,))
        num_offset += len(placement)

    return BlockManager(sorted(blocks, key=lambda b: b.mgr_locs[0]),
                        [mgr_items] + [np.arange(n) for n in item_shape])

class TestBlock(tm.TestCase):

    _multiprocess_can_split_ = True

    def setUp(self):
        # self.fblock = get_float_ex()  # a,c,e
        # self.cblock = get_complex_ex() #
        # self.oblock = get_obj_ex()
        # self.bool_block = get_bool_ex()
        # self.int_block = get_int_ex()

        self.fblock = create_block('float', [0, 2, 4])
        self.cblock = create_block('complex', [7])
        self.oblock = create_block('object', [1, 3])
        self.bool_block = create_block('bool', [5])
        self.int_block = create_block('int', [6])

    def test_constructor(self):
        int32block = create_block('i4', [0])
        self.assertEqual(int32block.dtype, np.int32)

    def test_pickle(self):

        def _check(blk):
            assert_block_equal(self.round_trip_pickle(blk), blk)

        _check(self.fblock)
        _check(self.cblock)
        _check(self.oblock)
        _check(self.bool_block)

    def test_mgr_locs(self):
        assert_almost_equal(self.fblock.mgr_locs, [0, 2, 4])

    def test_attrs(self):
        self.assertEqual(self.fblock.shape, self.fblock.values.shape)
        self.assertEqual(self.fblock.dtype, self.fblock.values.dtype)
        self.assertEqual(len(self.fblock), len(self.fblock.values))

    def test_merge(self):
        avals = randn(2, 10)
        bvals = randn(2, 10)

        ref_cols = Index(['e', 'a', 'b', 'd', 'f'])

        ablock = make_block(avals,
                            ref_cols.get_indexer(['e', 'b']))
        bblock = make_block(bvals,
                            ref_cols.get_indexer(['a', 'd']))
        merged = ablock.merge(bblock)
        assert_almost_equal(merged.mgr_locs, [0, 1, 2, 3])
        assert_almost_equal(merged.values[[0, 2]], avals)
        assert_almost_equal(merged.values[[1, 3]], bvals)

        # TODO: merge with mixed type?

    def test_copy(self):
        cop = self.fblock.copy()
        self.assertIsNot(cop, self.fblock)
        assert_block_equal(self.fblock, cop)

    def test_reindex_index(self):
        pass

    def test_reindex_cast(self):
        pass

    def test_insert(self):
        pass

    def test_delete(self):
        newb = self.fblock.copy()
        newb.delete(0)
        assert_almost_equal(newb.mgr_locs, [2, 4])
        self.assertTrue((newb.values[0] == 1).all())

        newb = self.fblock.copy()
        newb.delete(1)
        assert_almost_equal(newb.mgr_locs, [0, 4])
        self.assertTrue((newb.values[1] == 2).all())

        newb = self.fblock.copy()
        newb.delete(2)
        assert_almost_equal(newb.mgr_locs, [0, 2])
        self.assertTrue((newb.values[1] == 1).all())

        newb = self.fblock.copy()
        self.assertRaises(Exception, newb.delete, 3)

    def test_split_block_at(self):

        # with dup column support this method was taken out
        # GH3679
        raise nose.SkipTest("skipping for now")

        bs = list(self.fblock.split_block_at('a'))
        self.assertEqual(len(bs), 1)
        self.assertTrue(np.array_equal(bs[0].items, ['c', 'e']))

        bs = list(self.fblock.split_block_at('c'))
        self.assertEqual(len(bs), 2)
        self.assertTrue(np.array_equal(bs[0].items, ['a']))
        self.assertTrue(np.array_equal(bs[1].items, ['e']))

        bs = list(self.fblock.split_block_at('e'))
        self.assertEqual(len(bs), 1)
        self.assertTrue(np.array_equal(bs[0].items, ['a', 'c']))

        bblock = get_bool_ex(['f'])
        bs = list(bblock.split_block_at('f'))
        self.assertEqual(len(bs), 0)

    def test_get(self):
        pass

    def test_set(self):
        pass

    def test_fillna(self):
        pass

    def test_repr(self):
        pass

class TestDatetimeBlock(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_try_coerce_arg(self):
        block = create_block('datetime', [0])

        # coerce None
        none_coerced = block._try_coerce_args(block.values, None)[2]
        self.assertTrue(pd.Timestamp(none_coerced) is pd.NaT)

        # coerce different types of date bojects
        vals = (np.datetime64('2010-10-10'),
                datetime(2010, 10, 10),
                date(2010, 10, 10))
        for val in vals:
            coerced = block._try_coerce_args(block.values, val)[2]
            self.assertEqual(np.int64, type(coerced))
            self.assertEqual(pd.Timestamp('2010-10-10'), pd.Timestamp(coerced))

class TestBlockManager(tm.TestCase):
    _multiprocess_can_split_ = True

    def setUp(self):
        self.mgr = create_mgr('a: f8; b: object; c: f8; d: object; e: f8;'
                              'f: bool; g: i8; h: complex')

    def test_constructor_corner(self):
        pass

    def test_attrs(self):
        mgr = create_mgr('a,b,c: f8-1; d,e,f: f8-2')
        self.assertEqual(mgr.nblocks, 2)
        self.assertEqual(len(mgr), 6)

    def test_is_mixed_dtype(self):
        self.assertFalse(create_mgr('a,b:f8').is_mixed_type)
        self.assertFalse(create_mgr('a:f8-1; b:f8-2').is_mixed_type)

        self.assertTrue(create_mgr('a,b:f8; c,d: f4').is_mixed_type)
        self.assertTrue(create_mgr('a,b:f8; c,d: object').is_mixed_type)

    def test_is_indexed_like(self):
        mgr1 = create_mgr('a,b: f8')
        mgr2 = create_mgr('a:i8; b:bool')
        mgr3 = create_mgr('a,b,c: f8')
        self.assertTrue(mgr1._is_indexed_like(mgr1))
        self.assertTrue(mgr1._is_indexed_like(mgr2))
        self.assertTrue(mgr1._is_indexed_like(mgr3))

        self.assertFalse(mgr1._is_indexed_like(
            mgr1.get_slice(slice(-1), axis=1)))

    def test_duplicate_ref_loc_failure(self):
        tmp_mgr = create_mgr('a:bool; a: f8')

        axes, blocks = tmp_mgr.axes, tmp_mgr.blocks

        blocks[0].mgr_locs = np.array([0])
        blocks[1].mgr_locs = np.array([0])
        # test trying to create block manager with overlapping ref locs
        self.assertRaises(AssertionError, BlockManager, blocks, axes)

        blocks[0].mgr_locs = np.array([0])
        blocks[1].mgr_locs = np.array([1])
        mgr = BlockManager(blocks, axes)
        mgr.iget(1)

    def test_contains(self):
        self.assertIn('a', self.mgr)
        self.assertNotIn('baz', self.mgr)

    def test_pickle(self):

        mgr2 = self.round_trip_pickle(self.mgr)
        assert_frame_equal(DataFrame(self.mgr), DataFrame(mgr2))

        # share ref_items
        # self.assertIs(mgr2.blocks[0].ref_items, mgr2.blocks[1].ref_items)

        # GH2431
        self.assertTrue(hasattr(mgr2, "_is_consolidated"))
        self.assertTrue(hasattr(mgr2, "_known_consolidated"))

        # reset to False on load
        self.assertFalse(mgr2._is_consolidated)
        self.assertFalse(mgr2._known_consolidated)

    def test_non_unique_pickle(self):

        mgr = create_mgr('a,a,a:f8')
        mgr2 = self.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

        mgr = create_mgr('a: f8; a: i8')
        mgr2 = self.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

    def test_categorical_block_pickle(self):
        mgr = create_mgr('a: category')
        mgr2 = self.round_trip_pickle(mgr)
        assert_frame_equal(DataFrame(mgr), DataFrame(mgr2))

        smgr = create_single_mgr('category')
        smgr2 = self.round_trip_pickle(smgr)
        assert_series_equal(Series(smgr), Series(smgr2))

    def test_get_scalar(self):
        for item in self.mgr.items:
            for i, index in enumerate(self.mgr.axes[1]):
                res = self.mgr.get_scalar((item, index))
                exp = self.mgr.get(item, fastpath=False)[i]
                assert_almost_equal(res, exp)
                exp = self.mgr.get(item).internal_values()[i]
                assert_almost_equal(res, exp)

    def test_get(self):
        cols = Index(list('abc'))
        values = np.random.rand(3, 3)
        block = make_block(values=values.copy(),
                           placement=np.arange(3))
        mgr = BlockManager(blocks=[block], axes=[cols, np.arange(3)])

        assert_almost_equal(mgr.get('a', fastpath=False), values[0])
        assert_almost_equal(mgr.get('b', fastpath=False), values[1])
        assert_almost_equal(mgr.get('c', fastpath=False), values[2])
        assert_almost_equal(mgr.get('a').internal_values(), values[0])
        assert_almost_equal(mgr.get('b').internal_values(), values[1])
        assert_almost_equal(mgr.get('c').internal_values(), values[2])

    def test_set(self):
        mgr = create_mgr('a,b,c: int', item_shape=(3,))

        mgr.set('d', np.array(['foo'] * 3))
        mgr.set('b', np.array(['bar'] * 3))
        assert_almost_equal(mgr.get('a').internal_values(), [0] * 3)
        assert_almost_equal(mgr.get('b').internal_values(), ['bar'] * 3)
        assert_almost_equal(mgr.get('c').internal_values(), [2] * 3)
        assert_almost_equal(mgr.get('d').internal_values(), ['foo'] * 3)

    def test_insert(self):
        self.mgr.insert(0, 'inserted', np.arange(N))

        self.assertEqual(self.mgr.items[0], 'inserted')
        assert_almost_equal(self.mgr.get('inserted'), np.arange(N))

        for blk in self.mgr.blocks:
            yield self.assertIs, self.mgr.items, blk.ref_items

    def test_set_change_dtype(self):
        self.mgr.set('baz', np.zeros(N, dtype=bool))

        self.mgr.set('baz', np.repeat('foo', N))
        self.assertEqual(self.mgr.get('baz').dtype, np.object_)

        mgr2 = self.mgr.consolidate()
        mgr2.set('baz', np.repeat('foo', N))
        self.assertEqual(mgr2.get('baz').dtype, np.object_)

        mgr2.set('quux', randn(N).astype(int))
        self.assertEqual(mgr2.get('quux').dtype, np.int_)

        mgr2.set('quux', randn(N))
        self.assertEqual(mgr2.get('quux').dtype, np.float_)

    def test_set_change_dtype_slice(self): # GH8850
        cols = MultiIndex.from_tuples([('1st','a'), ('2nd','b'), ('3rd','c')])
        df = DataFrame([[1.0, 2, 3], [4.0, 5, 6]], columns=cols)
        df['2nd'] = df['2nd'] * 2.0

        self.assertEqual(sorted(df.blocks.keys()), ['float64', 'int64'])
        assert_frame_equal(df.blocks['float64'],
                DataFrame([[1.0, 4.0], [4.0, 10.0]], columns=cols[:2]))
        assert_frame_equal(df.blocks['int64'],
                DataFrame([[3], [6]], columns=cols[2:]))

    def test_copy(self):
        shallow = self.mgr.copy(deep=False)

        # we don't guaranteee block ordering
        for blk in self.mgr.blocks:
            found = False
            for cp_blk in shallow.blocks:
                if cp_blk.values is blk.values:
                    found = True
                    break
            self.assertTrue(found)

    def test_sparse(self):
        mgr = create_mgr('a: sparse-1; b: sparse-2')
        # what to test here?
        self.assertEqual(mgr.as_matrix().dtype, np.float64)

    def test_sparse_mixed(self):
        mgr = create_mgr('a: sparse-1; b: sparse-2; c: f8')
        self.assertEqual(len(mgr.blocks), 3)
        self.assertIsInstance(mgr, BlockManager)

        # what to test here?

    def test_as_matrix_float(self):
        mgr = create_mgr('c: f4; d: f2; e: f8')
        self.assertEqual(mgr.as_matrix().dtype, np.float64)

        mgr = create_mgr('c: f4; d: f2')
        self.assertEqual(mgr.as_matrix().dtype, np.float32)

    def test_as_matrix_int_bool(self):
        mgr = create_mgr('a: bool-1; b: bool-2')
        self.assertEqual(mgr.as_matrix().dtype, np.bool_)

        mgr = create_mgr('a: i8-1; b: i8-2; c: i4; d: i2; e: u1')
        self.assertEqual(mgr.as_matrix().dtype, np.int64)

        mgr = create_mgr('c: i4; d: i2; e: u1')
        self.assertEqual(mgr.as_matrix().dtype, np.int32)

    def test_as_matrix_datetime(self):
        mgr = create_mgr('h: datetime-1; g: datetime-2')
        self.assertEqual(mgr.as_matrix().dtype, 'M8[ns]')

    def test_as_matrix_datetime_tz(self):
        mgr = create_mgr('h: M8[ns, US/Eastern]; g: M8[ns, CET]')
        self.assertEqual(mgr.get('h').dtype, 'datetime64[ns, US/Eastern]')
        self.assertEqual(mgr.get('g').dtype, 'datetime64[ns, CET]')
        self.assertEqual(mgr.as_matrix().dtype, 'object')

    def test_astype(self):
        # coerce all
        mgr = create_mgr('c: f4; d: f2; e: f8')
        for t in ['float16', 'float32', 'float64', 'int32', 'int64']:
            t = np.dtype(t)
            tmgr = mgr.astype(t)
            self.assertEqual(tmgr.get('c').dtype.type, t)
            self.assertEqual(tmgr.get('d').dtype.type, t)
            self.assertEqual(tmgr.get('e').dtype.type, t)

        # mixed
        mgr = create_mgr('a,b: object; c: bool; d: datetime;'
                         'e: f4; f: f2; g: f8')
        for t in ['float16', 'float32', 'float64', 'int32', 'int64']:
            t = np.dtype(t)
            tmgr = mgr.astype(t, raise_on_error=False)
            self.assertEqual(tmgr.get('c').dtype.type, t)
            self.assertEqual(tmgr.get('e').dtype.type, t)
            self.assertEqual(tmgr.get('f').dtype.type, t)
            self.assertEqual(tmgr.get('g').dtype.type, t)

            self.assertEqual(tmgr.get('a').dtype.type, np.object_)
            self.assertEqual(tmgr.get('b').dtype.type, np.object_)
            if t != np.int64:
                self.assertEqual(tmgr.get('d').dtype.type, np.datetime64)
            else:
                self.assertEqual(tmgr.get('d').dtype.type, t)

    def test_convert(self):
        def _compare(old_mgr, new_mgr):
            """ compare the blocks, numeric compare ==, object don't """
            old_blocks = set(old_mgr.blocks)
            new_blocks = set(new_mgr.blocks)
            self.assertEqual(len(old_blocks), len(new_blocks))

            # compare non-numeric
            for b in old_blocks:
                found = False
                for nb in new_blocks:
                    if (b.values == nb.values).all():
                        found = True
                        break
                self.assertTrue(found)

            for b in new_blocks:
                found = False
                for ob in old_blocks:
                    if (b.values == ob.values).all():
                        found = True
                        break
                self.assertTrue(found)

        # noops
        mgr = create_mgr('f: i8; g: f8')
        new_mgr = mgr.convert()
        _compare(mgr,new_mgr)

        mgr = create_mgr('a, b: object; f: i8; g: f8')
        new_mgr = mgr.convert()
        _compare(mgr,new_mgr)

        # convert
        mgr = create_mgr('a,b,foo: object; f: i8; g: f8')
        mgr.set('a', np.array(['1'] * N, dtype=np.object_))
        mgr.set('b', np.array(['2.'] * N, dtype=np.object_))
        mgr.set('foo', np.array(['foo.'] * N, dtype=np.object_))
        new_mgr = mgr.convert(numeric=True)
        self.assertEqual(new_mgr.get('a').dtype, np.int64)
        self.assertEqual(new_mgr.get('b').dtype, np.float64)
        self.assertEqual(new_mgr.get('foo').dtype, np.object_)
        self.assertEqual(new_mgr.get('f').dtype, np.int64)
        self.assertEqual(new_mgr.get('g').dtype, np.float64)

        mgr = create_mgr('a,b,foo: object; f: i4; bool: bool; dt: datetime;'
                         'i: i8; g: f8; h: f2')
        mgr.set('a', np.array(['1'] * N, dtype=np.object_))
        mgr.set('b', np.array(['2.'] * N, dtype=np.object_))
        mgr.set('foo', np.array(['foo.'] * N, dtype=np.object_))
        new_mgr = mgr.convert(numeric=True)
        self.assertEqual(new_mgr.get('a').dtype, np.int64)
        self.assertEqual(new_mgr.get('b').dtype, np.float64)
        self.assertEqual(new_mgr.get('foo').dtype, np.object_)
        self.assertEqual(new_mgr.get('f').dtype, np.int32)
        self.assertEqual(new_mgr.get('bool').dtype, np.bool_)
        self.assertEqual(new_mgr.get('dt').dtype.type, np.datetime64)
        self.assertEqual(new_mgr.get('i').dtype, np.int64)
        self.assertEqual(new_mgr.get('g').dtype, np.float64)
        self.assertEqual(new_mgr.get('h').dtype, np.float16)

    def test_interleave(self):

        # self
        for dtype in ['f8','i8','object','bool','complex','M8[ns]','m8[ns]']:
            mgr = create_mgr('a: {0}'.format(dtype))
            self.assertEqual(mgr.as_matrix().dtype,dtype)
            mgr = create_mgr('a: {0}; b: {0}'.format(dtype))
            self.assertEqual(mgr.as_matrix().dtype,dtype)

        # will be converted according the actual dtype of the underlying
        mgr = create_mgr('a: category')
        self.assertEqual(mgr.as_matrix().dtype,'i8')
        mgr = create_mgr('a: category; b: category')
        self.assertEqual(mgr.as_matrix().dtype,'i8'),
        mgr = create_mgr('a: category; b: category2')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: category2')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: category2; b: category2')
        self.assertEqual(mgr.as_matrix().dtype,'object')

        # combinations
        mgr = create_mgr('a: f8')
        self.assertEqual(mgr.as_matrix().dtype,'f8')
        mgr = create_mgr('a: f8; b: i8')
        self.assertEqual(mgr.as_matrix().dtype,'f8')
        mgr = create_mgr('a: f4; b: i8')
        self.assertEqual(mgr.as_matrix().dtype,'f4')
        mgr = create_mgr('a: f4; b: i8; d: object')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: bool; b: i8')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: complex')
        self.assertEqual(mgr.as_matrix().dtype,'complex')
        mgr = create_mgr('a: f8; b: category')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: M8[ns]; b: category')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: M8[ns]; b: bool')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: M8[ns]; b: i8')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: m8[ns]; b: bool')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: m8[ns]; b: i8')
        self.assertEqual(mgr.as_matrix().dtype,'object')
        mgr = create_mgr('a: M8[ns]; b: m8[ns]')
        self.assertEqual(mgr.as_matrix().dtype,'object')

    def test_interleave_non_unique_cols(self):
        df = DataFrame([
            [pd.Timestamp('20130101'), 3.5],
            [pd.Timestamp('20130102'), 4.5]],
            columns=['x', 'x'],
            index=[1, 2])

        df_unique = df.copy()
        df_unique.columns = ['x', 'y']
        self.assertEqual(df_unique.values.shape, df.values.shape)
        tm.assert_numpy_array_equal(df_unique.values[0], df.values[0])
        tm.assert_numpy_array_equal(df_unique.values[1], df.values[1])

    def test_consolidate(self):
        pass

    def test_consolidate_ordering_issues(self):
        self.mgr.set('f', randn(N))
        self.mgr.set('d', randn(N))
        self.mgr.set('b', randn(N))
        self.mgr.set('g', randn(N))
        self.mgr.set('h', randn(N))

        cons = self.mgr.consolidate()
        self.assertEqual(cons.nblocks, 1)
        assert_almost_equal(cons.blocks[0].mgr_locs,
                            np.arange(len(cons.items)))

    def test_reindex_index(self):
        pass

    def test_reindex_items(self):
        # mgr is not consolidated, f8 & f8-2 blocks
        mgr = create_mgr('a: f8; b: i8; c: f8; d: i8; e: f8;'
                         'f: bool; g: f8-2')

        reindexed = mgr.reindex_axis(['g', 'c', 'a', 'd'], axis=0)
        self.assertEqual(reindexed.nblocks, 2)
        assert_almost_equal(reindexed.items, ['g', 'c', 'a', 'd'])
        assert_almost_equal(mgr.get('g',fastpath=False), reindexed.get('g',fastpath=False))
        assert_almost_equal(mgr.get('c',fastpath=False), reindexed.get('c',fastpath=False))
        assert_almost_equal(mgr.get('a',fastpath=False), reindexed.get('a',fastpath=False))
        assert_almost_equal(mgr.get('d',fastpath=False), reindexed.get('d',fastpath=False))
        assert_almost_equal(mgr.get('g').internal_values(), reindexed.get('g').internal_values())
        assert_almost_equal(mgr.get('c').internal_values(), reindexed.get('c').internal_values())
        assert_almost_equal(mgr.get('a').internal_values(), reindexed.get('a').internal_values())
        assert_almost_equal(mgr.get('d').internal_values(), reindexed.get('d').internal_values())

    def test_multiindex_xs(self):
        mgr = create_mgr('a,b,c: f8; d,e,f: i8')

        index = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux'],
                                   ['one', 'two', 'three']],
                           labels=[[0, 0, 0, 1, 1, 2, 2, 3, 3, 3],
                                   [0, 1, 2, 0, 1, 1, 2, 0, 1, 2]],
                           names=['first', 'second'])

        mgr.set_axis(1, index)
        result = mgr.xs('bar', axis=1)
        self.assertEqual(result.shape, (6, 2))
        self.assertEqual(result.axes[1][0], ('bar', 'one'))
        self.assertEqual(result.axes[1][1], ('bar', 'two'))

    def test_get_numeric_data(self):
        mgr = create_mgr('int: int; float: float; complex: complex;'
                         'str: object; bool: bool; obj: object; dt: datetime',
                         item_shape=(3,))
        mgr.set('obj', np.array([1, 2, 3], dtype=np.object_))

        numeric = mgr.get_numeric_data()
        assert_almost_equal(numeric.items, ['int', 'float', 'complex', 'bool'])
        assert_almost_equal(mgr.get('float',fastpath=False), numeric.get('float',fastpath=False))
        assert_almost_equal(mgr.get('float').internal_values(), numeric.get('float').internal_values())

        # Check sharing
        numeric.set('float', np.array([100., 200., 300.]))
        assert_almost_equal(mgr.get('float',fastpath=False), np.array([100., 200., 300.]))
        assert_almost_equal(mgr.get('float').internal_values(), np.array([100., 200., 300.]))

        numeric2 = mgr.get_numeric_data(copy=True)
        assert_almost_equal(numeric.items, ['int', 'float', 'complex', 'bool'])
        numeric2.set('float', np.array([1000., 2000., 3000.]))
        assert_almost_equal(mgr.get('float',fastpath=False), np.array([100., 200., 300.]))
        assert_almost_equal(mgr.get('float').internal_values(), np.array([100., 200., 300.]))

    def test_get_bool_data(self):
        mgr = create_mgr('int: int; float: float; complex: complex;'
                         'str: object; bool: bool; obj: object; dt: datetime',
                         item_shape=(3,))
        mgr.set('obj', np.array([True, False, True], dtype=np.object_))

        bools = mgr.get_bool_data()
        assert_almost_equal(bools.items, ['bool'])
        assert_almost_equal(mgr.get('bool',fastpath=False), bools.get('bool',fastpath=False))
        assert_almost_equal(mgr.get('bool').internal_values(), bools.get('bool').internal_values())

        bools.set('bool', np.array([True, False, True]))
        assert_almost_equal(mgr.get('bool',fastpath=False), [True, False, True])
        assert_almost_equal(mgr.get('bool').internal_values(), [True, False, True])

        # Check sharing
        bools2 = mgr.get_bool_data(copy=True)
        bools2.set('bool', np.array([False, True, False]))
        assert_almost_equal(mgr.get('bool',fastpath=False), [True, False, True])
        assert_almost_equal(mgr.get('bool').internal_values(), [True, False, True])

    def test_unicode_repr_doesnt_raise(self):
        str_repr = repr(create_mgr(u('b,\u05d0: object')))

    def test_missing_unicode_key(self):
        df = DataFrame({"a": [1]})
        try:
            df.ix[:, u("\u05d0")]  # should not raise UnicodeEncodeError
        except KeyError:
            pass  # this is the expected exception

    def test_equals(self):
        # unique items
        bm1 = create_mgr('a,b,c: i8-1; d,e,f: i8-2')
        bm2 = BlockManager(bm1.blocks[::-1], bm1.axes)
        self.assertTrue(bm1.equals(bm2))

        bm1 = create_mgr('a,a,a: i8-1; b,b,b: i8-2')
        bm2 = BlockManager(bm1.blocks[::-1], bm1.axes)
        self.assertTrue(bm1.equals(bm2))

    def test_equals_block_order_different_dtypes(self):
        # GH 9330

        mgr_strings = [
            "a:i8;b:f8", # basic case
            "a:i8;b:f8;c:c8;d:b", # many types
            "a:i8;e:dt;f:td;g:string", # more types
            "a:i8;b:category;c:category2;d:category2", # categories
            "c:sparse;d:sparse_na;b:f8", # sparse
            ]

        for mgr_string in mgr_strings:
            bm = create_mgr(mgr_string)
            block_perms = itertools.permutations(bm.blocks)
            for bm_perm in block_perms:
                bm_this = BlockManager(bm_perm, bm.axes)
                self.assertTrue(bm.equals(bm_this))
                self.assertTrue(bm_this.equals(bm))

    def test_single_mgr_ctor(self):
        mgr = create_single_mgr('f8', num_rows=5)
        self.assertEqual(mgr.as_matrix().tolist(), [0., 1., 2., 3., 4.])

class TestIndexing(object):
    # Nosetests-style data-driven tests.
    #
    # This test applies different indexing routines to block managers and
    # compares the outcome to the result of same operations on np.ndarray.
    #
    # NOTE: sparse (SparseBlock with fill_value != np.nan) fail a lot of tests
    #       and are disabled.

    MANAGERS = [
        create_single_mgr('f8', N),
        create_single_mgr('i8', N),
        #create_single_mgr('sparse', N),
        create_single_mgr('sparse_na', N),

        # 2-dim
        create_mgr('a,b,c,d,e,f: f8', item_shape=(N,)),
        create_mgr('a,b,c,d,e,f: i8', item_shape=(N,)),
        create_mgr('a,b: f8; c,d: i8; e,f: string', item_shape=(N,)),
        create_mgr('a,b: f8; c,d: i8; e,f: f8', item_shape=(N,)),
        #create_mgr('a: sparse', item_shape=(N,)),
        create_mgr('a: sparse_na', item_shape=(N,)),

        # 3-dim
        create_mgr('a,b,c,d,e,f: f8', item_shape=(N, N)),
        create_mgr('a,b,c,d,e,f: i8', item_shape=(N, N)),
        create_mgr('a,b: f8; c,d: i8; e,f: string', item_shape=(N, N)),
        create_mgr('a,b: f8; c,d: i8; e,f: f8', item_shape=(N, N)),
        # create_mgr('a: sparse', item_shape=(1, N)),
    ]

    # MANAGERS = [MANAGERS[6]]

    def test_get_slice(self):
        def assert_slice_ok(mgr, axis, slobj):
            # import pudb; pudb.set_trace()
            mat = mgr.as_matrix()

            # we maybe using an ndarray to test slicing and
            # might not be the full length of the axis
            if isinstance(slobj, np.ndarray):
                ax = mgr.axes[axis]
                if len(ax) and len(slobj) and len(slobj) != len(ax):
                    slobj = np.concatenate([slobj, np.zeros(len(ax)-len(slobj),dtype=bool)])
            sliced = mgr.get_slice(slobj, axis=axis)
            mat_slobj = (slice(None),) * axis + (slobj,)
            assert_almost_equal(mat[mat_slobj], sliced.as_matrix())
            assert_almost_equal(mgr.axes[axis][slobj], sliced.axes[axis])

        for mgr in self.MANAGERS:
            for ax in range(mgr.ndim):
                # slice
                yield assert_slice_ok, mgr, ax, slice(None)
                yield assert_slice_ok, mgr, ax, slice(3)
                yield assert_slice_ok, mgr, ax, slice(100)
                yield assert_slice_ok, mgr, ax, slice(1, 4)
                yield assert_slice_ok, mgr, ax, slice(3, 0, -2)

                # boolean mask
                yield assert_slice_ok, mgr, ax, np.array([], dtype=np.bool_)
                yield (assert_slice_ok, mgr, ax,
                       np.ones(mgr.shape[ax], dtype=np.bool_))
                yield (assert_slice_ok, mgr, ax,
                       np.zeros(mgr.shape[ax], dtype=np.bool_))

                if mgr.shape[ax] >= 3:
                    yield (assert_slice_ok, mgr, ax,
                           np.arange(mgr.shape[ax]) % 3 == 0)
                    yield (assert_slice_ok, mgr, ax,
                           np.array([True, True, False], dtype=np.bool_))

                # fancy indexer
                yield assert_slice_ok, mgr, ax, []
                yield assert_slice_ok, mgr, ax, lrange(mgr.shape[ax])

                if mgr.shape[ax] >= 3:
                    yield assert_slice_ok, mgr, ax, [0, 1, 2]
                    yield assert_slice_ok, mgr, ax, [-1, -2, -3]

    def test_take(self):
        def assert_take_ok(mgr, axis, indexer):
            mat = mgr.as_matrix()
            taken = mgr.take(indexer, axis)
            assert_almost_equal(np.take(mat, indexer, axis),
                                taken.as_matrix())
            assert_almost_equal(mgr.axes[axis].take(indexer),
                                taken.axes[axis])

        for mgr in self.MANAGERS:
            for ax in range(mgr.ndim):
                # take/fancy indexer
                yield assert_take_ok, mgr, ax, []
                yield assert_take_ok, mgr, ax, [0, 0, 0]
                yield assert_take_ok, mgr, ax, lrange(mgr.shape[ax])

                if mgr.shape[ax] >= 3:
                    yield assert_take_ok, mgr, ax, [0, 1, 2]
                    yield assert_take_ok, mgr, ax, [-1, -2, -3]

    def test_reindex_axis(self):
        def assert_reindex_axis_is_ok(mgr, axis, new_labels,
                                      fill_value):
            mat = mgr.as_matrix()
            indexer = mgr.axes[axis].get_indexer_for(new_labels)

            reindexed = mgr.reindex_axis(new_labels, axis,
                                         fill_value=fill_value)
            assert_almost_equal(com.take_nd(mat, indexer, axis,
                                            fill_value=fill_value),
                                reindexed.as_matrix())
            assert_almost_equal(reindexed.axes[axis], new_labels)

        for mgr in self.MANAGERS:
            for ax in range(mgr.ndim):
                for fill_value in (None, np.nan, 100.):
                    yield assert_reindex_axis_is_ok, mgr, ax, [], fill_value
                    yield (assert_reindex_axis_is_ok, mgr, ax,
                           mgr.axes[ax], fill_value)
                    yield (assert_reindex_axis_is_ok, mgr, ax,
                           mgr.axes[ax][[0, 0, 0]], fill_value)
                    yield (assert_reindex_axis_is_ok, mgr, ax,
                           ['foo', 'bar', 'baz'], fill_value)
                    yield (assert_reindex_axis_is_ok, mgr, ax,
                           ['foo', mgr.axes[ax][0], 'baz'], fill_value)

                    if mgr.shape[ax] >= 3:
                        yield (assert_reindex_axis_is_ok, mgr, ax,
                               mgr.axes[ax][:-3], fill_value)
                        yield (assert_reindex_axis_is_ok, mgr, ax,
                               mgr.axes[ax][-3::-1], fill_value)
                        yield (assert_reindex_axis_is_ok, mgr, ax,
                               mgr.axes[ax][[0, 1, 2, 0, 1, 2]], fill_value)

    def test_reindex_indexer(self):
        def assert_reindex_indexer_is_ok(mgr, axis, new_labels, indexer,
                                         fill_value):
            mat = mgr.as_matrix()
            reindexed_mat = com.take_nd(mat, indexer, axis,
                                        fill_value=fill_value)
            reindexed = mgr.reindex_indexer(new_labels, indexer, axis,
                                            fill_value=fill_value)
            assert_almost_equal(reindexed_mat, reindexed.as_matrix())
            assert_almost_equal(reindexed.axes[axis], new_labels)

        for mgr in self.MANAGERS:
            for ax in range(mgr.ndim):
                for fill_value in (None, np.nan, 100.):
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           [], [], fill_value)
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           mgr.axes[ax], np.arange(mgr.shape[ax]), fill_value)
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           ['foo'] * mgr.shape[ax], np.arange(mgr.shape[ax]),
                           fill_value)

                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           mgr.axes[ax][::-1], np.arange(mgr.shape[ax]),
                           fill_value)
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           mgr.axes[ax], np.arange(mgr.shape[ax])[::-1],
                           fill_value)
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           ['foo', 'bar', 'baz'], [0, 0, 0], fill_value)
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           ['foo', 'bar', 'baz'], [-1, 0, -1], fill_value)
                    yield (assert_reindex_indexer_is_ok, mgr, ax,
                           ['foo', mgr.axes[ax][0], 'baz'], [-1, -1, -1],
                           fill_value)

                    if mgr.shape[ax] >= 3:
                        yield (assert_reindex_indexer_is_ok, mgr, ax,
                               ['foo', 'bar', 'baz'], [0, 1, 2], fill_value)

    # test_get_slice(slice_like, axis)
    # take(indexer, axis)
    # reindex_axis(new_labels, axis)
    # reindex_indexer(new_labels, indexer, axis)

class TestBlockPlacement(tm.TestCase):
    _multiprocess_can_split_ = True

    def test_slice_len(self):
        self.assertEqual(len(BlockPlacement(slice(0, 4))), 4)
        self.assertEqual(len(BlockPlacement(slice(0, 4, 2))), 2)
        self.assertEqual(len(BlockPlacement(slice(0, 3, 2))), 2)

        self.assertEqual(len(BlockPlacement(slice(0, 1, 2))), 1)
        self.assertEqual(len(BlockPlacement(slice(1, 0, -1))), 1)

    def test_zero_step_raises(self):
        self.assertRaises(ValueError, BlockPlacement, slice(1, 1, 0))
        self.assertRaises(ValueError, BlockPlacement, slice(1, 2, 0))

    def test_unbounded_slice_raises(self):
        def assert_unbounded_slice_error(slc):
            # assertRaisesRegexp is not available in py2.6
            # self.assertRaisesRegexp(ValueError, "unbounded slice",
            #                         lambda: BlockPlacement(slc))
            self.assertRaises(ValueError, BlockPlacement, slc)

        assert_unbounded_slice_error(slice(None, None))
        assert_unbounded_slice_error(slice(10, None))
        assert_unbounded_slice_error(slice(None, None, -1))
        assert_unbounded_slice_error(slice(None, 10, -1))

        # These are "unbounded" because negative index will change depending on
        # container shape.
        assert_unbounded_slice_error(slice(-1, None))
        assert_unbounded_slice_error(slice(None, -1))
        assert_unbounded_slice_error(slice(-1, -1))
        assert_unbounded_slice_error(slice(-1, None, -1))
        assert_unbounded_slice_error(slice(None, -1, -1))
        assert_unbounded_slice_error(slice(-1, -1, -1))

    def test_not_slice_like_slices(self):
        def assert_not_slice_like(slc):
            self.assertTrue(not BlockPlacement(slc).is_slice_like)

        assert_not_slice_like(slice(0, 0))
        assert_not_slice_like(slice(100, 0))

        assert_not_slice_like(slice(100, 100, -1))
        assert_not_slice_like(slice(0, 100, -1))

        self.assertTrue(not BlockPlacement(slice(0, 0)).is_slice_like)
        self.assertTrue(not BlockPlacement(slice(100, 100)).is_slice_like)

    def test_array_to_slice_conversion(self):
        def assert_as_slice_equals(arr, slc):
            self.assertEqual(BlockPlacement(arr).as_slice, slc)

        assert_as_slice_equals([0], slice(0, 1, 1))
        assert_as_slice_equals([100], slice(100, 101, 1))

        assert_as_slice_equals([0, 1, 2], slice(0, 3, 1))
        assert_as_slice_equals([0, 5, 10], slice(0, 15, 5))
        assert_as_slice_equals([0, 100], slice(0, 200, 100))

        assert_as_slice_equals([2, 1], slice(2, 0, -1))
        assert_as_slice_equals([2, 1, 0], slice(2, None, -1))
        assert_as_slice_equals([100, 0], slice(100, None, -100))

    def test_not_slice_like_arrays(self):
        def assert_not_slice_like(arr):
            self.assertTrue(not BlockPlacement(arr).is_slice_like)

        assert_not_slice_like([])
        assert_not_slice_like([-1])
        assert_not_slice_like([-1, -2, -3])
        assert_not_slice_like([-10])
        assert_not_slice_like([-1])
        assert_not_slice_like([-1, 0, 1, 2])
        assert_not_slice_like([-2, 0, 2, 4])
        assert_not_slice_like([1, 0, -1])
        assert_not_slice_like([1, 1, 1])

    def test_slice_iter(self):
        self.assertEqual(list(BlockPlacement(slice(0, 3))), [0, 1, 2])
        self.assertEqual(list(BlockPlacement(slice(0, 0))), [])
        self.assertEqual(list(BlockPlacement(slice(3, 0))), [])

        self.assertEqual(list(BlockPlacement(slice(3, 0, -1))), [3, 2, 1])
        self.assertEqual(list(BlockPlacement(slice(3, None, -1))),
                          [3, 2, 1, 0])

    def test_slice_to_array_conversion(self):
        def assert_as_array_equals(slc, asarray):
            tm.assert_numpy_array_equal(
                BlockPlacement(slc).as_array,
                np.asarray(asarray))

        assert_as_array_equals(slice(0, 3), [0, 1, 2])
        assert_as_array_equals(slice(0, 0), [])
        assert_as_array_equals(slice(3, 0), [])

        assert_as_array_equals(slice(3, 0, -1), [3, 2, 1])
        assert_as_array_equals(slice(3, None, -1), [3, 2, 1, 0])
        assert_as_array_equals(slice(31, None, -10), [31, 21, 11, 1])

    def test_blockplacement_add(self):
        bpl = BlockPlacement(slice(0, 5))
        self.assertEqual(bpl.add(1).as_slice, slice(1, 6, 1))
        self.assertEqual(bpl.add(np.arange(5)).as_slice,
                          slice(0, 10, 2))
        self.assertEqual(list(bpl.add(np.arange(5, 0, -1))),
                          [5, 5, 5, 5, 5])

    def test_blockplacement_add_int(self):
        def assert_add_equals(val, inc, result):
            self.assertEqual(list(BlockPlacement(val).add(inc)),
                              result)

        assert_add_equals(slice(0, 0), 0, [])
        assert_add_equals(slice(1, 4), 0, [1, 2, 3])
        assert_add_equals(slice(3, 0, -1), 0, [3, 2, 1])
        assert_add_equals(slice(2, None, -1), 0, [2, 1, 0])
        assert_add_equals([1, 2, 4], 0, [1, 2, 4])

        assert_add_equals(slice(0, 0), 10, [])
        assert_add_equals(slice(1, 4), 10, [11, 12, 13])
        assert_add_equals(slice(3, 0, -1), 10, [13, 12, 11])
        assert_add_equals(slice(2, None, -1), 10, [12, 11, 10])
        assert_add_equals([1, 2, 4], 10, [11, 12, 14])

        assert_add_equals(slice(0, 0), -1, [])
        assert_add_equals(slice(1, 4), -1, [0, 1, 2])
        assert_add_equals(slice(3, 0, -1), -1, [2, 1, 0])
        assert_add_equals([1, 2, 4], -1, [0, 1, 3])

        self.assertRaises(ValueError,
                          lambda: BlockPlacement(slice(1, 4)).add(-10))
        self.assertRaises(ValueError,
                          lambda: BlockPlacement([1, 2, 4]).add(-10))
        self.assertRaises(ValueError,
                          lambda: BlockPlacement(slice(2, None, -1)).add(-1))

    # def test_blockplacement_array_add(self):

    #     assert_add_equals(slice(0, 2), [0, 1, 1], [0, 2, 3])
    #     assert_add_equals(slice(2, None, -1), [1, 1, 0], [3, 2, 0])

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)


"""
Optical flow
============

This tutorial offers a short overview of the optical flow routines available in
pysteps and it will cover how to compute and plot the motion field from a
sequence of radar images.
"""

from datetime import datetime
from pprint import pprint
import matplotlib.pyplot as plt
import numpy as np

from pysteps import io, motion, rcparams
from pysteps.utils import conversion, transformation
from pysteps.visualization import plot_precip_field, quiver

################################################################################
# Read the radar input images
# ---------------------------
#
# First, we will import the sequence of radar composites.
# You need the pysteps-data archive downloaded and the pystepsrc file
# configured with the data_source paths pointing to data folders.

# Selected case
date = datetime.strptime("201505151630", "%Y%m%d%H%M")
data_source = rcparams.data_sources["mch"]

###############################################################################
# Load the data from the archive
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

root_path = data_source["root_path"]
path_fmt = data_source["path_fmt"]
fn_pattern = data_source["fn_pattern"]
fn_ext = data_source["fn_ext"]
importer_name = data_source["importer"]
importer_kwargs = data_source["importer_kwargs"]
timestep = data_source["timestep"]

# Find the input files from the archive
fns = io.archive.find_by_date(
    date, root_path, path_fmt, fn_pattern, fn_ext, timestep=5, num_prev_files=9
)

# Read the radar composites
importer = io.get_method(importer_name, "importer")
R, quality, metadata = io.read_timeseries(fns, importer, **importer_kwargs)

del quality  # Not used

###############################################################################
# Preprocess the data
# ~~~~~~~~~~~~~~~~~~~

# Convert to mm/h
R, metadata = conversion.to_rainrate(R, metadata)

# Store the reference frame
R_ = R[-1, :, :].copy()

# Log-transform the data [dBR]
R, metadata = transformation.dB_transform(R, metadata, threshold=0.1, zerovalue=-15.0)

# Nicely print the metadata
pprint(metadata)

################################################################################
# Lucas-Kanade (LK)
# -----------------
#
# The Lucas-Kanade optical flow method implemented in pysteps is a local
# tracking approach that relies on the OpenCV package.
# Local features are tracked in a sequence of two or more radar images. The
# scheme includes a final interpolation step in order to produce a smooth
# field of motion vectors.

oflow_method = motion.get_method("LK")
V1 = oflow_method(R[-3:, :, :])

# Plot the motion field on top of the reference frame
plot_precip_field(R_, geodata=metadata, title="LK")
quiver(V1, geodata=metadata, step=25)
plt.show()

################################################################################
# Variational echo tracking (VET)
# -------------------------------
#
# This module implements the VET algorithm presented
# by Laroche and Zawadzki (1995) and used in the McGill Algorithm for
# Prediction by Lagrangian Extrapolation (MAPLE) described in
# Germann and Zawadzki (2002).
# The approach essentially consists of a global optimization routine that seeks
# at minimizing a cost function between the displaced and the reference image.

oflow_method = motion.get_method("VET")
V2 = oflow_method(R[-3:, :, :])

# Plot the motion field
plot_precip_field(R_, geodata=metadata, title="VET")
quiver(V2, geodata=metadata, step=25)
plt.show()

################################################################################
# Dynamic and adaptive radar tracking of storms (DARTS)
# -----------------------------------------------------
#
# DARTS uses a spectral approach to optical flow that is based on the discrete
# Fourier transform (DFT) of a temporal sequence of radar fields.
# The level of truncation of the DFT coefficients controls the degree of
# smoothness of the estimated motion field, allowing for an efficient
# motion estimation. DARTS requires a longer sequence of radar fields for
# estimating the motion, here we are going to use all the available 10 fields.

oflow_method = motion.get_method("DARTS")
R[~np.isfinite(R)] = metadata["zerovalue"]
V3 = oflow_method(R)  # needs longer training sequence

# Plot the motion field
plot_precip_field(R_, geodata=metadata, title="DARTS")
quiver(V3, geodata=metadata, step=25)
plt.show()

################################################################################
# Anisotropic diffusion method (Proesmans et al 1994)
# ---------------------------------------------------
#
# This module implements the anisotropic diffusion method presented in Proesmans
# et al. (1994), a robust optical flow technique which employs the notion of
# inconsitency during the solution of the optical flow equations.

oflow_method = motion.get_method("proesmans")
R[~np.isfinite(R)] = metadata["zerovalue"]
V4 = oflow_method(R[-2:, :, :])

# Plot the motion field
plot_precip_field(R_, geodata=metadata, title="Proesmans")
quiver(V4, geodata=metadata, step=25)
plt.show()

# sphinx_gallery_thumbnail_number = 1


"""
A mockup "Foo" units class which supports
conversion and different tick formatting depending on the "unit".
Here the "unit" is just a scalar conversion factor, but this example shows mpl is
entirely agnostic to what kind of units client packages use
"""

import matplotlib

from matplotlib.cbook import iterable
import matplotlib.units as units
import matplotlib.ticker as ticker
from pylab import figure, show

class Foo:
    def __init__( self, val, unit=1.0 ):
        self.unit = unit
        self._val = val * unit

    def value( self, unit ):
        if unit is None: unit = self.unit
        return self._val / unit

class FooConverter:

    @staticmethod
    def axisinfo(unit, axis):
        'return the Foo AxisInfo'
        if unit==1.0 or unit==2.0:
            return units.AxisInfo(
                majloc = ticker.IndexLocator( 8, 0 ),
                majfmt = ticker.FormatStrFormatter("VAL: %s"),
                label='foo',
                )

        else:
            return None

    @staticmethod
    def convert(obj, unit, axis):
        """
        convert obj using unit.  If obj is a sequence, return the
        converted sequence
        """
        if units.ConversionInterface.is_numlike(obj):
            return obj

        if iterable(obj):
            return [o.value(unit) for o in obj]
        else:
            return obj.value(unit)

    @staticmethod
    def default_units(x, axis):
        'return the default unit for x or None'
        if iterable(x):
            for thisx in x:
                return thisx.unit
        else:
            return x.unit

units.registry[Foo] = FooConverter()

# create some Foos
x = []
for val in range( 0, 50, 2 ):
    x.append( Foo( val, 1.0 ) )

# and some arbitrary y data
y = [i for i in range( len(x) ) ]

# plot specifying units
fig = figure()
fig.suptitle("Custom units")
fig.subplots_adjust(bottom=0.2)
ax = fig.add_subplot(1,2,2)
ax.plot( x, y, 'o', xunits=2.0 )
for label in ax.get_xticklabels():
    label.set_rotation(30)
    label.set_ha('right')
ax.set_title("xunits = 2.0")

# plot without specifying units; will use the None branch for axisinfo
ax = fig.add_subplot(1,2,1)
ax.plot( x, y ) # uses default units
ax.set_title('default units')
for label in ax.get_xticklabels():
    label.set_rotation(30)
    label.set_ha('right')

show()


# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for dnn_linear_combined.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import shutil
import tempfile

import numpy as np
import six

from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.estimator.canned import dnn_linear_combined
from tensorflow.python.estimator.canned import dnn_testing_utils
from tensorflow.python.estimator.canned import linear_testing_utils
from tensorflow.python.estimator.canned import prediction_keys
from tensorflow.python.estimator.export import export
from tensorflow.python.estimator.inputs import numpy_io
from tensorflow.python.estimator.inputs import pandas_io
from tensorflow.python.feature_column import feature_column
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
from tensorflow.python.summary.writer import writer_cache
from tensorflow.python.training import checkpoint_utils
from tensorflow.python.training import gradient_descent
from tensorflow.python.training import input as input_lib
from tensorflow.python.training import optimizer as optimizer_lib

try:
  # pylint: disable=g-import-not-at-top
  import pandas as pd
  HAS_PANDAS = True
except IOError:
  # Pandas writes a temporary file during import. If it fails, don't use pandas.
  HAS_PANDAS = False
except ImportError:
  HAS_PANDAS = False

class DNNOnlyModelFnTest(dnn_testing_utils.BaseDNNModelFnTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNModelFnTest.__init__(self, self._dnn_only_model_fn)

  def _dnn_only_model_fn(self,
                         features,
                         labels,
                         mode,
                         head,
                         hidden_units,
                         feature_columns,
                         optimizer='Adagrad',
                         activation_fn=nn.relu,
                         dropout=None,
                         input_layer_partitioner=None,
                         config=None):
    return dnn_linear_combined._dnn_linear_combined_model_fn(
        features=features,
        labels=labels,
        mode=mode,
        head=head,
        linear_feature_columns=[],
        dnn_hidden_units=hidden_units,
        dnn_feature_columns=feature_columns,
        dnn_optimizer=optimizer,
        dnn_activation_fn=activation_fn,
        dnn_dropout=dropout,
        input_layer_partitioner=input_layer_partitioner,
        config=config)

# A function to mimic linear-regressor init reuse same tests.
def _linear_regressor_fn(feature_columns,
                         model_dir=None,
                         label_dimension=1,
                         weight_column=None,
                         optimizer='Ftrl',
                         config=None,
                         partitioner=None):
  return dnn_linear_combined.DNNLinearCombinedRegressor(
      model_dir=model_dir,
      linear_feature_columns=feature_columns,
      linear_optimizer=optimizer,
      label_dimension=label_dimension,
      weight_column=weight_column,
      input_layer_partitioner=partitioner,
      config=config)

class LinearOnlyRegressorPartitionerTest(
    linear_testing_utils.BaseLinearRegressorPartitionerTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearRegressorPartitionerTest.__init__(
        self, _linear_regressor_fn)

class LinearOnlyRegressorEvaluationTest(
    linear_testing_utils.BaseLinearRegressorEvaluationTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearRegressorEvaluationTest.__init__(
        self, _linear_regressor_fn)

class LinearOnlyRegressorPredictTest(
    linear_testing_utils.BaseLinearRegressorPredictTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearRegressorPredictTest.__init__(
        self, _linear_regressor_fn)

class LinearOnlyRegressorIntegrationTest(
    linear_testing_utils.BaseLinearRegressorIntegrationTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearRegressorIntegrationTest.__init__(
        self, _linear_regressor_fn)

class LinearOnlyRegressorTrainingTest(
    linear_testing_utils.BaseLinearRegressorTrainingTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearRegressorTrainingTest.__init__(
        self, _linear_regressor_fn)

def _linear_classifier_fn(feature_columns,
                          model_dir=None,
                          n_classes=2,
                          weight_column=None,
                          label_vocabulary=None,
                          optimizer='Ftrl',
                          config=None,
                          partitioner=None):
  return dnn_linear_combined.DNNLinearCombinedClassifier(
      model_dir=model_dir,
      linear_feature_columns=feature_columns,
      linear_optimizer=optimizer,
      n_classes=n_classes,
      weight_column=weight_column,
      label_vocabulary=label_vocabulary,
      input_layer_partitioner=partitioner,
      config=config)

class LinearOnlyClassifierTrainingTest(
    linear_testing_utils.BaseLinearClassifierTrainingTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearClassifierTrainingTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn)

class LinearOnlyClassifierClassesEvaluationTest(
    linear_testing_utils.BaseLinearClassifierEvaluationTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearClassifierEvaluationTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn)

class LinearOnlyClassifierPredictTest(
    linear_testing_utils.BaseLinearClassifierPredictTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearClassifierPredictTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn)

class LinearOnlyClassifierIntegrationTest(
    linear_testing_utils.BaseLinearClassifierIntegrationTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    linear_testing_utils.BaseLinearClassifierIntegrationTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn)

class DNNLinearCombinedRegressorIntegrationTest(test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      writer_cache.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _test_complete_flow(
      self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
      label_dimension, batch_size):
    linear_feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))]
    dnn_feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))]
    feature_columns = linear_feature_columns + dnn_feature_columns
    est = dnn_linear_combined.DNNLinearCombinedRegressor(
        linear_feature_columns=linear_feature_columns,
        dnn_hidden_units=(2, 2),
        dnn_feature_columns=dnn_feature_columns,
        label_dimension=label_dimension,
        model_dir=self._model_dir)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predictions = np.array([
        x[prediction_keys.PredictionKeys.PREDICTIONS]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, label_dimension), predictions.shape)

    # EXPORT
    feature_spec = feature_column.make_parse_example_spec(feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_savedmodel(tempfile.mkdtemp(),
                                       serving_input_receiver_fn)
    self.assertTrue(gfile.Exists(export_dir))

  def test_numpy_input_fn(self):
    """Tests complete flow with numpy_input_fn."""
    label_dimension = 2
    batch_size = 10
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)
    # learn y = x
    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size)

  def test_pandas_input_fn(self):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return
    label_dimension = 1
    batch_size = 10
    data = np.linspace(0., 2., batch_size, dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(data)
    train_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x,
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size)

  def test_input_fn_from_parse_example(self):
    """Tests complete flow with input_fn constructed from parse_example."""
    label_dimension = 2
    batch_size = 10
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)

    serialized_examples = []
    for datum in data:
      example = example_pb2.Example(features=feature_pb2.Features(
          feature={
              'x': feature_pb2.Feature(
                  float_list=feature_pb2.FloatList(value=datum)),
              'y': feature_pb2.Feature(
                  float_list=feature_pb2.FloatList(value=datum)),
          }))
      serialized_examples.append(example.SerializeToString())

    feature_spec = {
        'x': parsing_ops.FixedLenFeature([label_dimension], dtypes.float32),
        'y': parsing_ops.FixedLenFeature([label_dimension], dtypes.float32),
    }
    def _train_input_fn():
      feature_map = parsing_ops.parse_example(serialized_examples, feature_spec)
      features = linear_testing_utils.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _eval_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _predict_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils.queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

    self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size)

# A function to mimic dnn-classifier init reuse same tests.
def _dnn_classifier_fn(hidden_units,
                       feature_columns,
                       model_dir=None,
                       n_classes=2,
                       weight_column=None,
                       label_vocabulary=None,
                       optimizer='Adagrad',
                       config=None,
                       input_layer_partitioner=None):
  return dnn_linear_combined.DNNLinearCombinedClassifier(
      model_dir=model_dir,
      dnn_hidden_units=hidden_units,
      dnn_feature_columns=feature_columns,
      dnn_optimizer=optimizer,
      n_classes=n_classes,
      weight_column=weight_column,
      label_vocabulary=label_vocabulary,
      input_layer_partitioner=input_layer_partitioner,
      config=config)

class DNNOnlyClassifierEvaluateTest(
    dnn_testing_utils.BaseDNNClassifierEvaluateTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
        self, _dnn_classifier_fn)

class DNNOnlyClassifierPredictTest(
    dnn_testing_utils.BaseDNNClassifierPredictTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
        self, _dnn_classifier_fn)

class DNNOnlyClassifierTrainTest(
    dnn_testing_utils.BaseDNNClassifierTrainTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
        self, _dnn_classifier_fn)

# A function to mimic dnn-regressor init reuse same tests.
def _dnn_regressor_fn(hidden_units,
                      feature_columns,
                      model_dir=None,
                      label_dimension=1,
                      weight_column=None,
                      optimizer='Adagrad',
                      config=None,
                      input_layer_partitioner=None):
  return dnn_linear_combined.DNNLinearCombinedRegressor(
      model_dir=model_dir,
      dnn_hidden_units=hidden_units,
      dnn_feature_columns=feature_columns,
      dnn_optimizer=optimizer,
      label_dimension=label_dimension,
      weight_column=weight_column,
      input_layer_partitioner=input_layer_partitioner,
      config=config)

class DNNOnlyRegressorEvaluateTest(
    dnn_testing_utils.BaseDNNRegressorEvaluateTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
        self, _dnn_regressor_fn)

class DNNOnlyRegressorPredictTest(
    dnn_testing_utils.BaseDNNRegressorPredictTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
        self, _dnn_regressor_fn)

class DNNOnlyRegressorTrainTest(
    dnn_testing_utils.BaseDNNRegressorTrainTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
        self, _dnn_regressor_fn)

class DNNLinearCombinedClassifierIntegrationTest(test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      writer_cache.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _as_label(self, data_in_float):
    return np.rint(data_in_float).astype(np.int64)

  def _test_complete_flow(
      self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
      n_classes, batch_size):
    linear_feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))]
    dnn_feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))]
    feature_columns = linear_feature_columns + dnn_feature_columns
    est = dnn_linear_combined.DNNLinearCombinedClassifier(
        linear_feature_columns=linear_feature_columns,
        dnn_hidden_units=(2, 2),
        dnn_feature_columns=dnn_feature_columns,
        n_classes=n_classes,
        model_dir=self._model_dir)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predicted_proba = np.array([
        x[prediction_keys.PredictionKeys.PROBABILITIES]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)

    # EXPORT
    feature_spec = feature_column.make_parse_example_spec(feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_savedmodel(tempfile.mkdtemp(),
                                       serving_input_receiver_fn)
    self.assertTrue(gfile.Exists(export_dir))

  def test_numpy_input_fn(self):
    """Tests complete flow with numpy_input_fn."""
    n_classes = 3
    input_dimension = 2
    batch_size = 10
    data = np.linspace(
        0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
    x_data = data.reshape(batch_size, input_dimension)
    y_data = self._as_label(np.reshape(data[:batch_size], (batch_size, 1)))
    # learn y = x
    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        y=y_data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        y=y_data,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size)

  def test_pandas_input_fn(self):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return
    input_dimension = 1
    n_classes = 2
    batch_size = 10
    data = np.linspace(0., n_classes - 1., batch_size, dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(self._as_label(data))
    train_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x,
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size)

  def test_input_fn_from_parse_example(self):
    """Tests complete flow with input_fn constructed from parse_example."""
    input_dimension = 2
    n_classes = 3
    batch_size = 10
    data = np.linspace(0., n_classes-1., batch_size * input_dimension,
                       dtype=np.float32)
    data = data.reshape(batch_size, input_dimension)

    serialized_examples = []
    for datum in data:
      example = example_pb2.Example(features=feature_pb2.Features(
          feature={
              'x':
                  feature_pb2.Feature(float_list=feature_pb2.FloatList(
                      value=datum)),
              'y':
                  feature_pb2.Feature(int64_list=feature_pb2.Int64List(
                      value=self._as_label(datum[:1]))),
          }))
      serialized_examples.append(example.SerializeToString())

    feature_spec = {
        'x': parsing_ops.FixedLenFeature([input_dimension], dtypes.float32),
        'y': parsing_ops.FixedLenFeature([1], dtypes.int64),
    }
    def _train_input_fn():
      feature_map = parsing_ops.parse_example(serialized_examples, feature_spec)
      features = linear_testing_utils.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _eval_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _predict_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils.queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

    self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size)

class DNNLinearCombinedTests(test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      shutil.rmtree(self._model_dir)

  def _mock_optimizer(self, real_optimizer, var_name_prefix):
    """Verifies global_step is None and var_names start with given prefix."""

    def _minimize(loss, global_step=None, var_list=None):
      self.assertIsNone(global_step)
      trainable_vars = var_list or ops.get_collection(
          ops.GraphKeys.TRAINABLE_VARIABLES)
      var_names = [var.name for var in trainable_vars]
      self.assertTrue(
          all([name.startswith(var_name_prefix) for name in var_names]))
      # var is used to check this op called by training.
      var = variables_lib.Variable(0., name=(var_name_prefix + '_called'))
      with ops.control_dependencies([var.assign(100.)]):
        return real_optimizer.minimize(loss, global_step, var_list)

    optimizer_mock = test.mock.NonCallableMagicMock(
        spec=optimizer_lib.Optimizer, wraps=real_optimizer)
    optimizer_mock.minimize = test.mock.MagicMock(wraps=_minimize)

    return optimizer_mock

  def test_train_op_calls_both_dnn_and_linear(self):
    opt = gradient_descent.GradientDescentOptimizer(1.)
    x_column = feature_column.numeric_column('x')
    input_fn = numpy_io.numpy_input_fn(
        x={'x': np.array([[0.], [1.]])},
        y=np.array([[0.], [1.]]),
        batch_size=1,
        shuffle=False)
    est = dnn_linear_combined.DNNLinearCombinedClassifier(
        linear_feature_columns=[x_column],
        # verifies linear_optimizer is used only for linear part.
        linear_optimizer=self._mock_optimizer(opt, 'linear'),
        dnn_hidden_units=(2, 2),
        dnn_feature_columns=[x_column],
        # verifies dnn_optimizer is used only for linear part.
        dnn_optimizer=self._mock_optimizer(opt, 'dnn'),
        model_dir=self._model_dir)
    est.train(input_fn, steps=1)
    # verifies train_op fires linear minimize op
    self.assertEqual(100.,
                     checkpoint_utils.load_variable(
                         self._model_dir, 'binary_logistic_head/linear_called'))
    # verifies train_op fires dnn minimize op
    self.assertEqual(100.,
                     checkpoint_utils.load_variable(
                         self._model_dir, 'binary_logistic_head/dnn_called'))

  def test_dnn_and_linear_logits_are_added(self):
    with ops.Graph().as_default():
      variables_lib.Variable([[1.0]], name='linear/linear_model/x/weights')
      variables_lib.Variable([2.0], name='linear/linear_model/bias_weights')
      variables_lib.Variable([[3.0]], name='dnn/hiddenlayer_0/kernel')
      variables_lib.Variable([4.0], name='dnn/hiddenlayer_0/bias')
      variables_lib.Variable([[5.0]], name='dnn/logits/kernel')
      variables_lib.Variable([6.0], name='dnn/logits/bias')
      variables_lib.Variable(1, name='global_step', dtype=dtypes.int64)
      linear_testing_utils.save_variables_to_ckpt(self._model_dir)

    x_column = feature_column.numeric_column('x')
    est = dnn_linear_combined.DNNLinearCombinedRegressor(
        linear_feature_columns=[x_column],
        dnn_hidden_units=[1],
        dnn_feature_columns=[x_column],
        model_dir=self._model_dir)
    input_fn = numpy_io.numpy_input_fn(
        x={'x': np.array([[10.]])}, batch_size=1, shuffle=False)
    # linear logits = 10*1 + 2 = 12
    # dnn logits = (10*3 + 4)*5 + 6 = 176
    # logits = dnn + linear = 176 + 12 = 188
    self.assertAllClose(
        {
            prediction_keys.PredictionKeys.PREDICTIONS: [188.],
        },
        next(est.predict(input_fn=input_fn)))

if __name__ == '__main__':
  test.main()


#!/usr/bin/env python
#adapted from Nathan Salomonis: http://code.activestate.com/recipes/578175-hierarchical-clustering-heatmap-python/

import matplotlib as mpl
#pick non-x display
mpl.use('Agg')
import matplotlib.pyplot as pylab
import scipy
import scipy.cluster.hierarchy as sch
import scipy.spatial.distance as dist
import numpy
import string
import time
import sys, os
import getopt
import numpy as np
import pandas as ps
from .PhenotypeCollection import PhenotypeCollection
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
#ignore these warnings
#/usr/lib/pymodules/python2.7/matplotlib/collections.py:548: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison
#  if self._edgecolors == 'face':
#/usr/lib/pymodules/python2.7/matplotlib/backends/backend_pdf.py:2184: FutureWarning: comparison to `None` will result in an elementwise object comparison in the future.
#  different = bool(ours != theirs)

################# Perform the hierarchical clustering #################

def heatmap(x, row_header, column_header, primary_pt_models, color_f, row_method,
            column_method, row_metric, column_metric,
            filename, sample_f, secondary_pt_models):

    print "\nrunning hiearchical clustering using %s for columns and %s for rows" % (column_metric,row_metric)

    """
    This below code is based in large part on the protype methods:
    http://old.nabble.com/How-to-plot-heatmap-with-matplotlib--td32534593.html
    http://stackoverflow.com/questions/7664826/how-to-get-flat-clustering-corresponding-to-color-clusters-in-the-dendrogram-cre

    x is an m by n ndarray, m observations, n genes
    """

    ### Define the color gradient to use based on the provided name
    #if color_gradient == 'red_white_blue':
    #    cmap=pylab.cm.bwr
    #if color_gradient == 'red_black_sky':
    #    cmap=RedBlackSkyBlue()
    #if color_gradient == 'red_black_blue':
    #    cmap=RedBlackBlue()
    #if color_gradient == 'red_black_green':
    #    cmap=RedBlackGreen()
    #if color_gradient == 'yellow_black_blue':
    #    cmap=YellowBlackBlue()
    #if color_gradient == 'seismic':
    #    cmap=pylab.cm.seismic
    #if color_gradient == 'green_white_purple':
    #    cmap=pylab.cm.PiYG_r
    #if color_gradient == 'coolwarm':
    #    cmap=pylab.cm.coolwarm

    ### Scale the max and min colors so that 0 is white/black
    #vmin=x.min()
    #vmax=x.max()
    #vmax = max([vmax,abs(vmin)])
    #vmin = vmax*-1
    #norm = mpl.colors.Normalize(vmin/2, vmax/2) ### adjust the max and min to scale these colors

    ### Scale the Matplotlib window size
    default_window_hight = 10.5
    default_window_width = 10
    fig = pylab.figure(figsize=(default_window_width,default_window_hight)) ### could use m,n to scale here
    color_bar_w = 0.015 ### Sufficient size to show
    color_bar_w = 0.015 ### Sufficient size to show

    ## calculate positions for all elements
    # ax1, placement of dendrogram 1, on the left of the heatmap
    #if row_method != None: w1 =
    [ax1_x, ax1_y, ax1_w, ax1_h] = [0.05,0.42,0.2,0.4]   ### The second value controls the position of the matrix relative to the bottom of the view
    width_between_ax1_axr = 0.004
    height_between_ax1_axc = 0.004 ### distance between the top color bar axis and the matrix

    # axr, placement of row side colorbar
    [axr_x, axr_y, axr_w, axr_h] = [0.31,0.1,color_bar_w,0.6] ### second to last controls the width of the side color bar - 0.015 when showing
    axr_x = ax1_x + ax1_w + width_between_ax1_axr
    axr_y = ax1_y; axr_h = ax1_h
    width_between_axr_axm = 0.004

    # axc, placement of column side colorbar
    [axc_x, axc_y, axc_w, axc_h] = [0.4,0.63,0.5,color_bar_w] ### last one controls the hight of the top color bar - 0.015 when showing
    axc_x = axr_x + axr_w + width_between_axr_axm
    axc_y = ax1_y + ax1_h + height_between_ax1_axc
    height_between_axc_ax2 = 0.004

    # axm, placement of heatmap for the data matrix
    [axm_x, axm_y, axm_w, axm_h] = [0.4,0.9,2.5,0.5]
    axm_x = axr_x + axr_w + width_between_axr_axm
    axm_y = ax1_y; axm_h = ax1_h
    axm_w = axc_w

    # ax2, placement of dendrogram 2, on the top of the heatmap
    [ax2_x, ax2_y, ax2_w, ax2_h] = [0.3,0.72,0.6,0.15] ### last one controls hight of the dendrogram
    ax2_x = axr_x + axr_w + width_between_axr_axm
    ax2_y = ax1_y + ax1_h + height_between_ax1_axc + axc_h + height_between_axc_ax2
    ax2_w = axc_w

    # placement of the phenotype legend
    [axpl_x, axpl_y, axpl_w, axpl_h] = [0.78,0.84,0.05,0.13]
    # placement of the sample legend

    # axcb - placement of the sample legend
    [axsl_x, axsl_y, axsl_w, axsl_h] = [0.05,0.29,0.05,0.09]

    # axcb - placement of the color legend
    [axcb_x, axcb_y, axcb_w, axcb_h] = [0.05, 0.88,0.05,0.09]

    # Compute and plot top dendrogram
    if not column_method is None and x.shape[1] > 1:
        start_time = time.time()
        d2 = dist.pdist(x.T)
        D2 = dist.squareform(d2)
        ax2 = fig.add_axes([ax2_x, ax2_y, ax2_w, ax2_h], frame_on=True)
        Y2 = sch.linkage(D2, method=column_method, metric=column_metric) ### array-clustering metric - 'average', 'single', 'centroid', 'complete'
        Z2 = sch.dendrogram(Y2)
        ind2 = sch.fcluster(Y2,0.7*max(Y2[:,2]),'distance') ### This is the default behavior of dendrogram
        time_diff = str(round(time.time()-start_time,1))
        ax2.set_xticks([]) ### Hides ticks
        ax2.set_yticks([])
        #print 'Column clustering completed in %s seconds' % time_diff
    else:
        ind2 = ['NA']*len(column_header) ### Used for exporting the flat cluster data

    # Compute and plot left dendrogram.
    if not row_method is None and x.shape[0] > 1:
        start_time = time.time()
        x_bin = x.copy()
        x_bin[x_bin > 0] = 1
        d1 = dist.pdist(x_bin)
        D1 = dist.squareform(d1)  # full matrix
        ax1 = fig.add_axes([ax1_x, ax1_y, ax1_w, ax1_h], frame_on=True) # frame_on may be False
        ax1.set_xticks([]) ### Hides ticks
        ax1.set_yticks([])
        Y1 = sch.linkage(D1, method=row_method, metric=row_metric) ### gene-clustering metric - 'average', 'single', 'centroid', 'complete'
        Z1 = sch.dendrogram(Y1, orientation='right')
        ind1 = sch.fcluster(Y1,0.7*max(Y1[:,2]),'distance') ### This is the default behavior of dendrogram
        time_diff = str(round(time.time()-start_time,1))
        #print 'Row clustering completed in %s seconds' % time_diff
    else:
        ind1 = ['NA']*len(row_header) ### Used for exporting the flat cluster data

    # Plot heatmap color legend
    n = len(x[0]); m = len(x)
    if secondary_pt_models is not None:
        cmaplist = np.array([[247,247,247],[166,206,227],[178,223,138],[31,120,180]])/256.0
    else:
        cmaplist = np.array([[247,247,247],[31,120,180]])/256.0
    cmap = mpl.colors.ListedColormap(cmaplist)
    axcb = fig.add_axes([axcb_x, axcb_y, axcb_w, axcb_h], frame_on=False)  # axes for colorbar
    #cb = mpl.colorbar.ColorbarBase(axcb, cmap=cmap, orientation='horizontal')
    bounds = numpy.linspace(0, len(cmaplist), len(cmaplist) + 1)
    norm = mpl.colors.BoundaryNorm(bounds, len(cmaplist))
    cb = mpl.colorbar.ColorbarBase(axcb, cmap=cmap, norm=norm, spacing='proportional', ticks=bounds, boundaries=bounds)
    if secondary_pt_models is not None:
        axcb.set_yticklabels(["negative", "%s positive" % primary_pt_models.get_name(), "%s positive" % secondary_pt_models.get_name(), "both predictors positive"], fontsize = 8)
        axcb.yaxis.set_ticks([0.125, 0.375, 0.625, 0.875])
    else:
        axcb.set_yticklabels(["%s negative" % primary_pt_models.get_name(), "%s positive" % primary_pt_models.get_name()], fontsize = 8)
        axcb.yaxis.set_ticks([0.25, 0.75])
    axcb.set_title("Heatmap colorkey", fontsize = 10, loc = "left")

    # Plot distance matrix.
    axm = fig.add_axes([axm_x, axm_y, axm_w, axm_h])  # axes for the data matrix
    xt = x

    if not column_method is None and x.shape[1] > 1:
        idx2 = Z2['leaves'] ### apply the clustering for the array-dendrograms to the actual matrix data
        xt = xt[:,idx2]
        ind2 = ind2[idx2] ### reorder the flat cluster to match the order of the leaves the dendrogram
        pass
    if not row_method is None and x.shape[0] > 1 :
        idx1 = Z1['leaves'] ### apply the clustering for the gene-dendrograms to the actual matrix data
        xt = xt[idx1,:]   # xt is transformed x
        ind1 = ind1[idx1] ### reorder the flat cluster to match the order of the leaves the dendrogram
    ### taken from http://stackoverflow.com/questions/2982929/plotting-results-of-hierarchical-clustering-ontop-of-a-matrix-of-data-in-python/3011894#3011894
    im = axm.matshow(xt, aspect='auto', origin='lower', cmap=cmap, norm=norm) ### norm=norm added to scale coloring of expression with zero = white or black

    axm.set_xticks([]) ### Hides x-ticks
    axm.set_yticks([])

    # Add text
    new_row_header=[]
    new_column_header=[]
    for i in range(x.shape[0]):
        margin = 0
        if len(row_header) > 0 :
            fontdict = {'fontsize': 7}
        if len(row_header) > 30 :
            fontdict = {'fontsize': 7}
            margin = 0.5
        if len(row_header) > 50 :
            fontdict = {'fontsize': 4}
        if len(row_header) > 100 :
            fontdict = {'fontsize': 2}
        if len(row_header) > 200:
            fontdict = {'fontsize': 1}
        #if len(row_header)<100: ### Don't visualize gene associations when more than 100 rows
        axm.plot([-0.5, len(column_header)], [i - 0.5, i - 0.5], color = 'black', ls = '-')
        if x.shape[0] > 1 and row_method is not None:
            label = row_header[idx1[i]]
        else:
            label = row_header[i]
        fontdict.items
        axm.text(x.shape[1] + 0.2, i - margin , '  ' + label, fontdict = fontdict)
        new_row_header.append(label)

    for i in range(x.shape[1]):
        if not column_method is None and x.shape[1] > 1:
            axm.plot([i-0.5, i-0.5], [-0.5, len(row_header) - 0.5], color = 'black', ls = '-')
            axm.text(i-0.5, -0.5, ' '+ column_header[idx2[i]], fontdict = {'fontsize': 7}, rotation=270, verticalalignment="top") # rotation could also be degrees
            new_column_header.append(column_header[idx2[i]])
        else: ### When not clustering columns
            axm.plot([i-0.5, i-0.5], [-0.5, len(row_header) - 0.5], color = 'black', ls = '-')
            axm.text(i-0.5, -0.8, ' '+column_header[i], fontdict = {'fontsize': 7}, rotation=270, verticalalignment="top")
            new_column_header.append(column_header[i])

    pt2acc = primary_pt_models.get_pt2acc()
    pt2acc.index = pt2acc.loc[:, "accession"]
    #colors
    colors = ps.read_csv(color_f, index_col = None, sep = "\t")
    if "category" in pt2acc.columns:
        #assign categories to colors
        import sets
        #get unique categories in the order they appear in the pt mapping table
        cats = sorted(set(pt2acc.loc[:, "category"].tolist()), key=lambda x: pt2acc.loc[:, "category"].tolist().index(x))
        if not colors.shape[0] < len(cats):
            # Plot phenotype legend
            axpl = fig.add_axes([axpl_x, axpl_y, axpl_w, axpl_h], frame_on=False)  # axes for colorbar
            #for i in pt2cat2col.index:
            #    if pt2cat2col.loc[i,"Category"] not in cat2col:
            #        cat2col[pt2cat2col.loc[i,"Category"]] = pt2cat2col.loc[i, ["r", "g", "b"]]
            #        col2id[pt2cat2col.loc[i,"Category"]] = j
            #        j += 1
            pt2cat = dict([(pt2acc.loc[i, "accession"], pt2acc.loc[i, "category"]) for i in pt2acc.index])
            cat2id = dict([(cats[i - 1], i) for i in range(1, len(cats) + 1)])
            cmaplist = ps.DataFrame(colors.iloc[:len(cats),])
            cmaplist.index = cats
            cmaplist = cmaplist / 256.0
            cmap_p = mpl.colors.ListedColormap(cmaplist.values)
            bounds = numpy.linspace(0, cmaplist.shape[0], cmaplist.shape[0] + 1)
            norm = mpl.colors.BoundaryNorm(bounds, cmaplist.shape[0])
            cb = mpl.colorbar.ColorbarBase(axpl, cmap=cmap_p, norm=norm, spacing='proportional', ticks=bounds, boundaries=bounds)
            axpl.set_yticklabels([i for i in cats], fontsize = 6)
            axpl.yaxis.set_ticks(np.arange(1.0 / len(cats) / 2, 1,  1.0 / len(cats)))
            axpl.set_title("Phenotype colorkey", fontsize =  10, loc = "left")
            # Plot colside colors
            # axc --> axes for column side colorbar
            axc = fig.add_axes([axc_x, axc_y, axc_w, axc_h])  # axes for column side colorbar
            dc = numpy.array([cat2id[pt2cat[i]]  for i in column_header]).T
            if x.shape[1] > 1 and column_method is not None:
                dc = dc[idx2]
            dc.shape = (1, x.shape[1])
            im_c = axc.matshow(dc, aspect='auto', origin='lower', cmap=cmap_p)
            axc.set_xticks([]) ### Hides ticks
            axc.set_yticks([])

    # Plot rowside colors
    if sample_f is not None and x.shape[0] > 1:
        samples = ps.read_csv(sample_f, sep = "\t", index_col = "sample_name")
        if "category" in samples.columns:
            #get unique sample categories and sort according to the order they appear in the sampling file
            sample_cats = sorted(set(samples.loc[:, "category"].tolist()), key = lambda x: samples.loc[:, "category"].tolist().index(x))
            cat2col = dict([(sample_cats[i - 1], i) for i in range(1, len(sample_cats) + 1)])
            cmaplist = ps.DataFrame(colors.iloc[:len(sample_cats),]) / 256.0
            cmap_p = mpl.colors.ListedColormap(cmaplist.values)
            axr = fig.add_axes([axr_x, axr_y, axr_w, axr_h])  # axes for row side colorbar
            dr = numpy.array([cat2col[samples.loc[i, "category"]]  for i in row_header]).T
            if row_method is not None:
                dr = dr[idx1]
            dr.shape = (samples.shape[0], 1)
            #cmap_r = mpl.colors.ListedColormap(['r', 'g', 'b', 'y', 'w', 'k', 'm'])
            im_r = axr.matshow(dr, aspect='auto', origin='lower', cmap=cmap_p)
            axr.set_xticks([]) ### Hides ticks
            axr.set_yticks([])
            # Plot sample legend
            axsl = fig.add_axes([axsl_x, axsl_y, axsl_w, axsl_h], frame_on=False)  # axes for colorbar
            bounds = numpy.linspace(0, len(sample_cats), len(sample_cats) + 1)
            norm = mpl.colors.BoundaryNorm(bounds, len(sample_cats))
            cb = mpl.colorbar.ColorbarBase(axsl, cmap=cmap_p, norm=norm, spacing='proportional', ticks=bounds, boundaries=bounds)
            axsl.yaxis.set_ticks(np.arange(1.0 / len(sample_cats) / 2, 1,  1.0 / len(sample_cats)))
            axsl.set_yticklabels([i for i in sample_cats], fontsize = 6)
            axsl.set_title("Sample colorkey", loc = "left", fontsize = 10)

    #exportFlatClusterData(filename, new_row_header,new_column_header,xt,ind1,ind2)

    ### Render the graphic
    if len(row_header)>50 or len(column_header)>50:
        pylab.rcParams['font.size'] = 6
    else:
        pylab.rcParams['font.size'] = 8

    pylab.savefig(filename)
    pylab.savefig(filename, dpi=300) #,dpi=200
    #pylab.show()

def getColorRange(x):
    """ Determines the range of colors, centered at zero, for normalizing cmap """
    vmax=x.max()
    vmin=x.min()
    if vmax<0 and vmin<0: direction = 'negative'
    elif vmax>0 and vmin>0: direction = 'positive'
    else: direction = 'both'
    if direction == 'both':
        vmax = max([vmax,abs(vmin)])
        vmin = -1*vmax
        return vmax,vmin
    else:
        return vmax,vmin

################# Export the flat cluster data #################

def exportFlatClusterData(filename, new_row_header,new_column_header,xt,ind1,ind2):
    """ Export the clustered results as a text file, only indicating the flat-clusters rather than the tree """

    filename = string.replace(filename,'.pdf','.txt')
    export_text = open(filename,'w')
    column_header = string.join(['UID','row_clusters-flat']+new_column_header,'\t')+'\n' ### format column-names for export
    export_text.write(column_header)
    column_clusters = string.join(['column_clusters-flat','']+ map(str, ind2),'\t')+'\n' ### format column-flat-clusters for export
    export_text.write(column_clusters)

    ### The clusters, dendrogram and flat clusters are drawn bottom-up, so we need to reverse the order to match
    new_row_header = new_row_header[::-1]
    xt = xt[::-1]

    ### Export each row in the clustered data matrix xt
    i=0
    for row in xt:
        export_text.write(string.join([new_row_header[i],str(ind1[i])]+map(str, row),'\t')+'\n')
        i+=1
    export_text.close()

    ### Export as CDT file
    filename = string.replace(filename,'.txt','.cdt')
    export_cdt = open(filename,'w')
    column_header = string.join(['UNIQID','NAME','GWEIGHT']+new_column_header,'\t')+'\n' ### format column-names for export
    export_cdt.write(column_header)
    eweight = string.join(['EWEIGHT','','']+ ['1']*len(new_column_header),'\t')+'\n' ### format column-flat-clusters for export
    export_cdt.write(eweight)

    ### Export each row in the clustered data matrix xt
    i=0
    for row in xt:
        export_cdt.write(string.join([new_row_header[i]]*2+['1']+map(str, row),'\t')+'\n')
        i+=1
    export_cdt.close()

################# Create Custom Color Gradients #################
#http://matplotlib.sourceforge.net/examples/pylab_examples/custom_cmap.html

def RedBlackSkyBlue():
    cdict = {'red':   ((0.0, 0.0, 0.0),
                       (0.5, 0.0, 0.1),
                       (1.0, 1.0, 1.0)),

             'green': ((0.0, 0.0, 0.9),
                       (0.5, 0.1, 0.0),
                       (1.0, 0.0, 0.0)),

             'blue':  ((0.0, 0.0, 1.0),
                       (0.5, 0.1, 0.0),
                       (1.0, 0.0, 0.0))
            }

    my_cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    return my_cmap

def RedBlackBlue():
    cdict = {'red':   ((0.0, 0.0, 0.0),
                       (0.5, 0.0, 0.1),
                       (1.0, 1.0, 1.0)),

             'green': ((0.0, 0.0, 0.0),
                       (1.0, 0.0, 0.0)),

             'blue':  ((0.0, 0.0, 1.0),
                       (0.5, 0.1, 0.0),
                       (1.0, 0.0, 0.0))
            }

    my_cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    return my_cmap

def RedBlackGreen():
    cdict = {'red':   ((0.0, 0.0, 0.0),
                       (0.5, 0.0, 0.1),
                       (1.0, 1.0, 1.0)),

             'blue': ((0.0, 0.0, 0.0),
                       (1.0, 0.0, 0.0)),

             'green':  ((0.0, 0.0, 1.0),
                       (0.5, 0.1, 0.0),
                       (1.0, 0.0, 0.0))
            }

    my_cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    return my_cmap

def YellowBlackBlue():
    cdict = {'red':   ((0.0, 0.0, 0.0),
                       (0.5, 0.0, 0.1),
                       (1.0, 1.0, 1.0)),

             'green': ((0.0, 0.0, 0.8),
                       (0.5, 0.1, 0.0),
                       (1.0, 1.0, 1.0)),

             'blue':  ((0.0, 0.0, 1.0),
                       (0.5, 0.1, 0.0),
                       (1.0, 0.0, 0.0))
            }
    ### yellow is created by adding y = 1 to RedBlackSkyBlue green last tuple
    ### modulate between blue and cyan using the last y var in the first green tuple
    my_cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    return my_cmap