# -*- coding: utf-8 -*-
# pylint: disable-msg=W0612,E1101

import pytest

from pandas import DataFrame
import pandas as pd

from numpy import nan
import numpy as np

from pandas import melt, lreshape, wide_to_long
import pandas.util.testing as tm
from pandas.compat import range


class TestMelt(object):

    def setup_method(self, method):
        self.df = tm.makeTimeDataFrame()[:10]
        self.df['id1'] = (self.df['A'] > 0).astype(np.int64)
        self.df['id2'] = (self.df['B'] > 0).astype(np.int64)

        self.var_name = 'var'
        self.value_name = 'val'

        self.df1 = pd.DataFrame([[1.067683, -1.110463, 0.20867
                                  ], [-1.321405, 0.368915, -1.055342],
                                 [-0.807333, 0.08298, -0.873361]])
        self.df1.columns = [list('ABC'), list('abc')]
        self.df1.columns.names = ['CAP', 'low']

    def test_top_level_method(self):
        result = melt(self.df)
        assert result.columns.tolist() == ['variable', 'value']

    def test_method_signatures(self):
        tm.assert_frame_equal(self.df.melt(),
                              melt(self.df))

        tm.assert_frame_equal(self.df.melt(id_vars=['id1', 'id2'],
                                           value_vars=['A', 'B']),
                              melt(self.df,
                                   id_vars=['id1', 'id2'],
                                   value_vars=['A', 'B']))

        tm.assert_frame_equal(self.df.melt(var_name=self.var_name,
                                           value_name=self.value_name),
                              melt(self.df,
                                   var_name=self.var_name,
                                   value_name=self.value_name))

        tm.assert_frame_equal(self.df1.melt(col_level=0),
                              melt(self.df1, col_level=0))

    def test_default_col_names(self):
        result = self.df.melt()
        assert result.columns.tolist() == ['variable', 'value']

        result1 = self.df.melt(id_vars=['id1'])
        assert result1.columns.tolist() == ['id1', 'variable', 'value']

        result2 = self.df.melt(id_vars=['id1', 'id2'])
        assert result2.columns.tolist() == ['id1', 'id2', 'variable', 'value']

    def test_value_vars(self):
        result3 = self.df.melt(id_vars=['id1', 'id2'], value_vars='A')
        assert len(result3) == 10

        result4 = self.df.melt(id_vars=['id1', 'id2'], value_vars=['A', 'B'])
        expected4 = DataFrame({'id1': self.df['id1'].tolist() * 2,
                               'id2': self.df['id2'].tolist() * 2,
                               'variable': ['A'] * 10 + ['B'] * 10,
                               'value': (self.df['A'].tolist() +
                                         self.df['B'].tolist())},
                              columns=['id1', 'id2', 'variable', 'value'])
        tm.assert_frame_equal(result4, expected4)

    def test_value_vars_types(self):
        # GH 15348
        expected = DataFrame({'id1': self.df['id1'].tolist() * 2,
                              'id2': self.df['id2'].tolist() * 2,
                              'variable': ['A'] * 10 + ['B'] * 10,
                              'value': (self.df['A'].tolist() +
                                        self.df['B'].tolist())},
                             columns=['id1', 'id2', 'variable', 'value'])

        for type_ in (tuple, list, np.array):
            result = self.df.melt(id_vars=['id1', 'id2'],
                                  value_vars=type_(('A', 'B')))
            tm.assert_frame_equal(result, expected)

    def test_vars_work_with_multiindex(self):
        expected = DataFrame({
            ('A', 'a'): self.df1[('A', 'a')],
            'CAP': ['B'] * len(self.df1),
            'low': ['b'] * len(self.df1),
            'value': self.df1[('B', 'b')],
        }, columns=[('A', 'a'), 'CAP', 'low', 'value'])

        result = self.df1.melt(id_vars=[('A', 'a')], value_vars=[('B', 'b')])
        tm.assert_frame_equal(result, expected)

    def test_tuple_vars_fail_with_multiindex(self):
        # melt should fail with an informative error message if
        # the columns have a MultiIndex and a tuple is passed
        # for id_vars or value_vars.
        tuple_a = ('A', 'a')
        list_a = [tuple_a]
        tuple_b = ('B', 'b')
        list_b = [tuple_b]

        for id_vars, value_vars in ((tuple_a, list_b), (list_a, tuple_b),
                                    (tuple_a, tuple_b)):
            with tm.assert_raises_regex(ValueError, r'MultiIndex'):
                self.df1.melt(id_vars=id_vars, value_vars=value_vars)

    def test_custom_var_name(self):
        result5 = self.df.melt(var_name=self.var_name)
        assert result5.columns.tolist() == ['var', 'value']

        result6 = self.df.melt(id_vars=['id1'], var_name=self.var_name)
        assert result6.columns.tolist() == ['id1', 'var', 'value']

        result7 = self.df.melt(id_vars=['id1', 'id2'], var_name=self.var_name)
        assert result7.columns.tolist() == ['id1', 'id2', 'var', 'value']

        result8 = self.df.melt(id_vars=['id1', 'id2'], value_vars='A',
                               var_name=self.var_name)
        assert result8.columns.tolist() == ['id1', 'id2', 'var', 'value']

        result9 = self.df.melt(id_vars=['id1', 'id2'], value_vars=['A', 'B'],
                               var_name=self.var_name)
        expected9 = DataFrame({'id1': self.df['id1'].tolist() * 2,
                               'id2': self.df['id2'].tolist() * 2,
                               self.var_name: ['A'] * 10 + ['B'] * 10,
                               'value': (self.df['A'].tolist() +
                                         self.df['B'].tolist())},
                              columns=['id1', 'id2', self.var_name, 'value'])
        tm.assert_frame_equal(result9, expected9)

    def test_custom_value_name(self):
        result10 = self.df.melt(value_name=self.value_name)
        assert result10.columns.tolist() == ['variable', 'val']

        result11 = self.df.melt(id_vars=['id1'], value_name=self.value_name)
        assert result11.columns.tolist() == ['id1', 'variable', 'val']

        result12 = self.df.melt(id_vars=['id1', 'id2'],
                                value_name=self.value_name)
        assert result12.columns.tolist() == ['id1', 'id2', 'variable', 'val']

        result13 = self.df.melt(id_vars=['id1', 'id2'], value_vars='A',
                                value_name=self.value_name)
        assert result13.columns.tolist() == ['id1', 'id2', 'variable', 'val']

        result14 = self.df.melt(id_vars=['id1', 'id2'], value_vars=['A', 'B'],
                                value_name=self.value_name)
        expected14 = DataFrame({'id1': self.df['id1'].tolist() * 2,
                                'id2': self.df['id2'].tolist() * 2,
                                'variable': ['A'] * 10 + ['B'] * 10,
                                self.value_name: (self.df['A'].tolist() +
                                                  self.df['B'].tolist())},
                               columns=['id1', 'id2', 'variable',
                                        self.value_name])
        tm.assert_frame_equal(result14, expected14)

    def test_custom_var_and_value_name(self):

        result15 = self.df.melt(var_name=self.var_name,
                                value_name=self.value_name)
        assert result15.columns.tolist() == ['var', 'val']

        result16 = self.df.melt(id_vars=['id1'], var_name=self.var_name,
                                value_name=self.value_name)
        assert result16.columns.tolist() == ['id1', 'var', 'val']

        result17 = self.df.melt(id_vars=['id1', 'id2'],
                                var_name=self.var_name,
                                value_name=self.value_name)
        assert result17.columns.tolist() == ['id1', 'id2', 'var', 'val']

        result18 = self.df.melt(id_vars=['id1', 'id2'], value_vars='A',
                                var_name=self.var_name,
                                value_name=self.value_name)
        assert result18.columns.tolist() == ['id1', 'id2', 'var', 'val']

        result19 = self.df.melt(id_vars=['id1', 'id2'], value_vars=['A', 'B'],
                                var_name=self.var_name,
                                value_name=self.value_name)
        expected19 = DataFrame({'id1': self.df['id1'].tolist() * 2,
                                'id2': self.df['id2'].tolist() * 2,
                                self.var_name: ['A'] * 10 + ['B'] * 10,
                                self.value_name: (self.df['A'].tolist() +
                                                  self.df['B'].tolist())},
                               columns=['id1', 'id2', self.var_name,
                                        self.value_name])
        tm.assert_frame_equal(result19, expected19)

        df20 = self.df.copy()
        df20.columns.name = 'foo'
        result20 = df20.melt()
        assert result20.columns.tolist() == ['foo', 'value']

    def test_col_level(self):
        res1 = self.df1.melt(col_level=0)
        res2 = self.df1.melt(col_level='CAP')
        assert res1.columns.tolist() == ['CAP', 'value']
        assert res2.columns.tolist() == ['CAP', 'value']

    def test_multiindex(self):
        res = self.df1.melt()
        assert res.columns.tolist() == ['CAP', 'low', 'value']


class TestLreshape(object):

    def test_pairs(self):
        data = {'birthdt': ['08jan2009', '20dec2008', '30dec2008', '21dec2008',
                            '11jan2009'],
                'birthwt': [1766, 3301, 1454, 3139, 4133],
                'id': [101, 102, 103, 104, 105],
                'sex': ['Male', 'Female', 'Female', 'Female', 'Female'],
                'visitdt1': ['11jan2009', '22dec2008', '04jan2009',
                             '29dec2008', '20jan2009'],
                'visitdt2':
                ['21jan2009', nan, '22jan2009', '31dec2008', '03feb2009'],
                'visitdt3': ['05feb2009', nan, nan, '02jan2009', '15feb2009'],
                'wt1': [1823, 3338, 1549, 3298, 4306],
                'wt2': [2011.0, nan, 1892.0, 3338.0, 4575.0],
                'wt3': [2293.0, nan, nan, 3377.0, 4805.0]}

        df = DataFrame(data)

        spec = {'visitdt': ['visitdt%d' % i for i in range(1, 4)],
                'wt': ['wt%d' % i for i in range(1, 4)]}
        result = lreshape(df, spec)

        exp_data = {'birthdt':
                    ['08jan2009', '20dec2008', '30dec2008', '21dec2008',
                     '11jan2009', '08jan2009', '30dec2008', '21dec2008',
                     '11jan2009', '08jan2009', '21dec2008', '11jan2009'],
                    'birthwt': [1766, 3301, 1454, 3139, 4133, 1766, 1454, 3139,
                                4133, 1766, 3139, 4133],
                    'id': [101, 102, 103, 104, 105, 101, 103, 104, 105, 101,
                           104, 105],
                    'sex': ['Male', 'Female', 'Female', 'Female', 'Female',
                            'Male', 'Female', 'Female', 'Female', 'Male',
                            'Female', 'Female'],
                    'visitdt': ['11jan2009', '22dec2008', '04jan2009',
                                '29dec2008', '20jan2009', '21jan2009',
                                '22jan2009', '31dec2008', '03feb2009',
                                '05feb2009', '02jan2009', '15feb2009'],
                    'wt': [1823.0, 3338.0, 1549.0, 3298.0, 4306.0, 2011.0,
                           1892.0, 3338.0, 4575.0, 2293.0, 3377.0, 4805.0]}
        exp = DataFrame(exp_data, columns=result.columns)
        tm.assert_frame_equal(result, exp)

        result = lreshape(df, spec, dropna=False)
        exp_data = {'birthdt':
                    ['08jan2009', '20dec2008', '30dec2008', '21dec2008',
                     '11jan2009', '08jan2009', '20dec2008', '30dec2008',
                     '21dec2008', '11jan2009', '08jan2009', '20dec2008',
                     '30dec2008', '21dec2008', '11jan2009'],
                    'birthwt': [1766, 3301, 1454, 3139, 4133, 1766, 3301, 1454,
                                3139, 4133, 1766, 3301, 1454, 3139, 4133],
                    'id': [101, 102, 103, 104, 105, 101, 102, 103, 104, 105,
                           101, 102, 103, 104, 105],
                    'sex': ['Male', 'Female', 'Female', 'Female', 'Female',
                            'Male', 'Female', 'Female', 'Female', 'Female',
                            'Male', 'Female', 'Female', 'Female', 'Female'],
                    'visitdt': ['11jan2009', '22dec2008', '04jan2009',
                                '29dec2008', '20jan2009', '21jan2009', nan,
                                '22jan2009', '31dec2008', '03feb2009',
                                '05feb2009', nan, nan, '02jan2009',
                                '15feb2009'],
                    'wt': [1823.0, 3338.0, 1549.0, 3298.0, 4306.0, 2011.0, nan,
                           1892.0, 3338.0, 4575.0, 2293.0, nan, nan, 3377.0,
                           4805.0]}
        exp = DataFrame(exp_data, columns=result.columns)
        tm.assert_frame_equal(result, exp)

        spec = {'visitdt': ['visitdt%d' % i for i in range(1, 3)],
                'wt': ['wt%d' % i for i in range(1, 4)]}
        pytest.raises(ValueError, lreshape, df, spec)


class TestWideToLong(object):

    def test_simple(self):
        np.random.seed(123)
        x = np.random.randn(3)
        df = pd.DataFrame({"A1970": {0: "a",
                                     1: "b",
                                     2: "c"},
                           "A1980": {0: "d",
                                     1: "e",
                                     2: "f"},
                           "B1970": {0: 2.5,
                                     1: 1.2,
                                     2: .7},
                           "B1980": {0: 3.2,
                                     1: 1.3,
                                     2: .1},
                           "X": dict(zip(
                               range(3), x))})
        df["id"] = df.index
        exp_data = {"X": x.tolist() + x.tolist(),
                    "A": ['a', 'b', 'c', 'd', 'e', 'f'],
                    "B": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
                    "year": ['1970', '1970', '1970', '1980', '1980', '1980'],
                    "id": [0, 1, 2, 0, 1, 2]}
        exp_frame = DataFrame(exp_data)
        exp_frame = exp_frame.set_index(['id', 'year'])[["X", "A", "B"]]
        long_frame = wide_to_long(df, ["A", "B"], i="id", j="year")
        tm.assert_frame_equal(long_frame, exp_frame)

    def test_stubs(self):
        # GH9204
        df = pd.DataFrame([[0, 1, 2, 3, 8], [4, 5, 6, 7, 9]])
        df.columns = ['id', 'inc1', 'inc2', 'edu1', 'edu2']
        stubs = ['inc', 'edu']

        # TODO: unused?
        df_long = pd.wide_to_long(df, stubs, i='id', j='age')  # noqa

        assert stubs == ['inc', 'edu']

    def test_separating_character(self):
        # GH14779
        np.random.seed(123)
        x = np.random.randn(3)
        df = pd.DataFrame({"A.1970": {0: "a",
                                      1: "b",
                                      2: "c"},
                           "A.1980": {0: "d",
                                      1: "e",
                                      2: "f"},
                           "B.1970": {0: 2.5,
                                      1: 1.2,
                                      2: .7},
                           "B.1980": {0: 3.2,
                                      1: 1.3,
                                      2: .1},
                           "X": dict(zip(
                               range(3), x))})
        df["id"] = df.index
        exp_data = {"X": x.tolist() + x.tolist(),
                    "A": ['a', 'b', 'c', 'd', 'e', 'f'],
                    "B": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
                    "year": ['1970', '1970', '1970', '1980', '1980', '1980'],
                    "id": [0, 1, 2, 0, 1, 2]}
        exp_frame = DataFrame(exp_data)
        exp_frame = exp_frame.set_index(['id', 'year'])[["X", "A", "B"]]
        long_frame = wide_to_long(df, ["A", "B"], i="id", j="year", sep=".")
        tm.assert_frame_equal(long_frame, exp_frame)

    def test_escapable_characters(self):
        np.random.seed(123)
        x = np.random.randn(3)
        df = pd.DataFrame({"A(quarterly)1970": {0: "a",
                                                1: "b",
                                                2: "c"},
                           "A(quarterly)1980": {0: "d",
                                                1: "e",
                                                2: "f"},
                           "B(quarterly)1970": {0: 2.5,
                                                1: 1.2,
                                                2: .7},
                           "B(quarterly)1980": {0: 3.2,
                                                1: 1.3,
                                                2: .1},
                           "X": dict(zip(
                               range(3), x))})
        df["id"] = df.index
        exp_data = {"X": x.tolist() + x.tolist(),
                    "A(quarterly)": ['a', 'b', 'c', 'd', 'e', 'f'],
                    "B(quarterly)": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
                    "year": ['1970', '1970', '1970', '1980', '1980', '1980'],
                    "id": [0, 1, 2, 0, 1, 2]}
        exp_frame = DataFrame(exp_data)
        exp_frame = exp_frame.set_index(
            ['id', 'year'])[["X", "A(quarterly)", "B(quarterly)"]]
        long_frame = wide_to_long(df, ["A(quarterly)", "B(quarterly)"],
                                  i="id", j="year")
        tm.assert_frame_equal(long_frame, exp_frame)

    def test_unbalanced(self):
        # test that we can have a varying amount of time variables
        df = pd.DataFrame({'A2010': [1.0, 2.0],
                           'A2011': [3.0, 4.0],
                           'B2010': [5.0, 6.0],
                           'X': ['X1', 'X2']})
        df['id'] = df.index
        exp_data = {'X': ['X1', 'X1', 'X2', 'X2'],
                    'A': [1.0, 3.0, 2.0, 4.0],
                    'B': [5.0, np.nan, 6.0, np.nan],
                    'id': [0, 0, 1, 1],
                    'year': ['2010', '2011', '2010', '2011']}
        exp_frame = pd.DataFrame(exp_data)
        exp_frame = exp_frame.set_index(['id', 'year'])[["X", "A", "B"]]
        long_frame = wide_to_long(df, ['A', 'B'], i='id', j='year')
        tm.assert_frame_equal(long_frame, exp_frame)

    def test_character_overlap(self):
        # Test we handle overlapping characters in both id_vars and value_vars
        df = pd.DataFrame({
            'A11': ['a11', 'a22', 'a33'],
            'A12': ['a21', 'a22', 'a23'],
            'B11': ['b11', 'b12', 'b13'],
            'B12': ['b21', 'b22', 'b23'],
            'BB11': [1, 2, 3],
            'BB12': [4, 5, 6],
            'BBBX': [91, 92, 93],
            'BBBZ': [91, 92, 93]
        })
        df['id'] = df.index
        exp_frame = pd.DataFrame({
            'BBBX': [91, 92, 93, 91, 92, 93],
            'BBBZ': [91, 92, 93, 91, 92, 93],
            'A': ['a11', 'a22', 'a33', 'a21', 'a22', 'a23'],
            'B': ['b11', 'b12', 'b13', 'b21', 'b22', 'b23'],
            'BB': [1, 2, 3, 4, 5, 6],
            'id': [0, 1, 2, 0, 1, 2],
            'year': ['11', '11', '11', '12', '12', '12']})
        exp_frame = exp_frame.set_index(['id', 'year'])[
            ['BBBX', 'BBBZ', 'A', 'B', 'BB']]
        long_frame = wide_to_long(df, ['A', 'B', 'BB'], i='id', j='year')
        tm.assert_frame_equal(long_frame.sort_index(axis=1),
                              exp_frame.sort_index(axis=1))

    def test_invalid_separator(self):
        # if an invalid separator is supplied a empty data frame is returned
        sep = 'nope!'
        df = pd.DataFrame({'A2010': [1.0, 2.0],
                           'A2011': [3.0, 4.0],
                           'B2010': [5.0, 6.0],
                           'X': ['X1', 'X2']})
        df['id'] = df.index
        exp_data = {'X': '',
                    'A2010': [],
                    'A2011': [],
                    'B2010': [],
                    'id': [],
                    'year': [],
                    'A': [],
                    'B': []}
        exp_frame = pd.DataFrame(exp_data)
        exp_frame = exp_frame.set_index(['id', 'year'])[[
            'X', 'A2010', 'A2011', 'B2010', 'A', 'B']]
        exp_frame.index.set_levels([[0, 1], []], inplace=True)
        long_frame = wide_to_long(df, ['A', 'B'], i='id', j='year', sep=sep)
        tm.assert_frame_equal(long_frame.sort_index(axis=1),
                              exp_frame.sort_index(axis=1))

    def test_num_string_disambiguation(self):
        # Test that we can disambiguate number value_vars from
        # string value_vars
        df = pd.DataFrame({
            'A11': ['a11', 'a22', 'a33'],
            'A12': ['a21', 'a22', 'a23'],
            'B11': ['b11', 'b12', 'b13'],
            'B12': ['b21', 'b22', 'b23'],
            'BB11': [1, 2, 3],
            'BB12': [4, 5, 6],
            'Arating': [91, 92, 93],
            'Arating_old': [91, 92, 93]
        })
        df['id'] = df.index
        exp_frame = pd.DataFrame({
            'Arating': [91, 92, 93, 91, 92, 93],
            'Arating_old': [91, 92, 93, 91, 92, 93],
            'A': ['a11', 'a22', 'a33', 'a21', 'a22', 'a23'],
            'B': ['b11', 'b12', 'b13', 'b21', 'b22', 'b23'],
            'BB': [1, 2, 3, 4, 5, 6],
            'id': [0, 1, 2, 0, 1, 2],
            'year': ['11', '11', '11', '12', '12', '12']})
        exp_frame = exp_frame.set_index(['id', 'year'])[
            ['Arating', 'Arating_old', 'A', 'B', 'BB']]
        long_frame = wide_to_long(df, ['A', 'B', 'BB'], i='id', j='year')
        tm.assert_frame_equal(long_frame.sort_index(axis=1),
                              exp_frame.sort_index(axis=1))

    def test_invalid_suffixtype(self):
        # If all stubs names end with a string, but a numeric suffix is
        # assumed,  an empty data frame is returned
        df = pd.DataFrame({'Aone': [1.0, 2.0],
                           'Atwo': [3.0, 4.0],
                           'Bone': [5.0, 6.0],
                           'X': ['X1', 'X2']})
        df['id'] = df.index
        exp_data = {'X': '',
                    'Aone': [],
                    'Atwo': [],
                    'Bone': [],
                    'id': [],
                    'year': [],
                    'A': [],
                    'B': []}
        exp_frame = pd.DataFrame(exp_data)
        exp_frame = exp_frame.set_index(['id', 'year'])[[
            'X', 'Aone', 'Atwo', 'Bone', 'A', 'B']]
        exp_frame.index.set_levels([[0, 1], []], inplace=True)
        long_frame = wide_to_long(df, ['A', 'B'], i='id', j='year')
        tm.assert_frame_equal(long_frame.sort_index(axis=1),
                              exp_frame.sort_index(axis=1))

    def test_multiple_id_columns(self):
        # Taken from http://www.ats.ucla.edu/stat/stata/modules/reshapel.htm
        df = pd.DataFrame({
            'famid': [1, 1, 1, 2, 2, 2, 3, 3, 3],
            'birth': [1, 2, 3, 1, 2, 3, 1, 2, 3],
            'ht1': [2.8, 2.9, 2.2, 2, 1.8, 1.9, 2.2, 2.3, 2.1],
            'ht2': [3.4, 3.8, 2.9, 3.2, 2.8, 2.4, 3.3, 3.4, 2.9]
        })
        exp_frame = pd.DataFrame({
            'ht': [2.8, 3.4, 2.9, 3.8, 2.2, 2.9, 2.0, 3.2, 1.8,
                   2.8, 1.9, 2.4, 2.2, 3.3, 2.3, 3.4, 2.1, 2.9],
            'famid': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3],
            'birth': [1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3],
            'age': ['1', '2', '1', '2', '1', '2', '1', '2', '1',
                    '2', '1', '2', '1', '2', '1', '2', '1', '2']
        })
        exp_frame = exp_frame.set_index(['famid', 'birth', 'age'])[['ht']]
        long_frame = wide_to_long(df, 'ht', i=['famid', 'birth'], j='age')
        tm.assert_frame_equal(long_frame, exp_frame)

    def test_non_unique_idvars(self):
        # GH16382
        # Raise an error message if non unique id vars (i) are passed
        df = pd.DataFrame({
            'A_A1': [1, 2, 3, 4, 5],
            'B_B1': [1, 2, 3, 4, 5],
            'x': [1, 1, 1, 1, 1]
        })
        with pytest.raises(ValueError):
            wide_to_long(df, ['A_A', 'B_B'], i='x', j='colname')
