# Path: alphaware/enums.py
# class OutputDataFormat(IntEnum):
#     MULTI_INDEX_DF = 0
#     PITVOT_TABLE_DF = 1
# 
# class FreqType(StrEnum):
#     EOD = 'D'
#     EOW = 'W'
#     EOM = 'M'
#     EOQ = 'Q'
#     EOSY = 'S'
#     EOY = 'Y'
# 
# Path: alphaware/const.py
# INDEX_FACTOR = namedtuple('INDEX_FACTOR', ['date_index', 'sec_index', 'full_index', 'col_score'])
# 
# Path: alphaware/utils/pandas_utils.py
# @expect_types(data=(pd.Series, pd.DataFrame))
# def convert_df_format(data, target_format=OutputDataFormat.MULTI_INDEX_DF, col_name='factor',
#                       multi_index=INDEX_FACTOR):
#     if target_format == OutputDataFormat.MULTI_INDEX_DF:
#         tmp = data.stack()
#         data_ = pd.DataFrame(tmp)
#         data_.index.names = multi_index.full_index
#         data_.columns = [col_name]
#     else:
#         tmp = data.unstack()
#         index = tmp.index
#         columns = tmp.columns.get_level_values(multi_index.sec_index).tolist()
#         data_ = pd.DataFrame(tmp.values, index=index, columns=columns)
# 
#     return data_
# 
# @expect_types(df=(pd.Series, pd.DataFrame))
# def top(df, column=None, n=5):
#     if isinstance(df, pd.Series):
#         ret = df.sort_values(ascending=False)[:n]
#     else:
#         py_assert(column is not None, "Specify the col name or use pandas Series type of data")
#         ret = df.sort_values(by=column, ascending=False)[:n]
# 
#     return ret
# 
# @expect_types(data=(pd.DataFrame, pd.Series), freq=optional(FreqType, str))
# def group_by_freq(data, freq=FreqType.EOM):
#     data_ = pd.DataFrame(data) if isinstance(data, pd.Series) else data
#     if freq == FreqType.EOD:
#         return data_.groupby([lambda x: x.year, lambda x: x.month, lambda x: x.day])
#     elif freq == FreqType.EOM:
#         return data_.groupby(pd.TimeGrouper(freq='M'))
#     elif freq == FreqType.EOQ:
#         return data_.groupby(pd.TimeGrouper(freq='Q'))
#     elif freq == FreqType.EOY:
#         return data_.groupby(pd.TimeGrouper(freq='A'))
# 
# @expect_types(x=(pd.DataFrame, pd.Series))
# def fwd_return(x, date_index=INDEX_FACTOR.date_index, sec_index=INDEX_FACTOR.sec_index, period=1):
#     """
#     每个日期和股票代码对应的未来收益
#     """
#     dates = sorted(set(x.index.get_level_values(date_index)))
#     ret = pd.DataFrame()
#     for i in range(len(dates) - period):
#         shift_date = dates[i + period]
#         data_concat = x.loc[shift_date].reset_index()
#         data_concat[date_index] = [dates[i]] * len(data_concat)
#         ret = pd.concat([ret, data_concat], axis=0)
#     ret.set_index([date_index, sec_index], inplace=True)
#     ret.columns = ['fwd_return']
#     return ret
# 
# @preprocess(data=ensure_pd_df)
# def weighted_rank(data, ascend_order=None, weight=None, out_df=False):
#     nb_col = len(data.columns)
#     ascend_order = [1] * nb_col if ascend_order is None else ascend_order
#     weight = [1.0 / nb_col] * nb_col if weight is None else weight
#     rank = np.zeros([len(data), nb_col])
#     for i in range(len(data.columns)):
#         rank[:, i] = np.argsort(data[data.columns[i]]) if ascend_order[i] else np.argsort(data[data.columns[i]])[::-1]
#     weight_rank = np.dot(rank, np.array(weight))
#     return pd.DataFrame(weight_rank, index=data.index) if out_df else weight_rank


from unittest import TestCase
from parameterized import parameterized
from numpy.testing.utils import assert_array_equal
from pandas import (MultiIndex,
                    Index)
from pandas.util.testing import assert_frame_equal, assert_series_equal
from alphaware.enums import OutputDataFormat, FreqType
from alphaware.const import INDEX_FACTOR
from alphaware.utils import (convert_df_format,
                             top,
                             group_by_freq,
                             fwd_return,
                             weighted_rank)
from datetime import datetime as dt
import pandas as pd
import numpy as np
# -*- coding: utf-8 -*-



class TestPandasUtils(TestCase):
    @parameterized.expand([(pd.DataFrame({'001': [1, 2, 3], '002': [2, 3, 4]}, index=['2014', '2015', '2016']),
                            OutputDataFormat.MULTI_INDEX_DF,
                            'test_factor',
                            INDEX_FACTOR,
                            pd.DataFrame(index=MultiIndex(levels=[['2014', '2015', '2016'], ['001', '002']],
                                                          labels=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]],
                                                          names=['trade_date', 'ticker']),
                                         data=[1, 2, 2, 3, 3, 4],
                                         columns=['test_factor']))])
    def test_convert_df_format_1(self, data, target_format, col_name, multi_index, expected):
        calculated = convert_df_format(data, target_format, col_name, multi_index)
        assert_frame_equal(calculated, expected)

    @parameterized.expand(
        [(pd.DataFrame(
            index=MultiIndex.from_product([['2014', '2015', '2016'], ['001', '002']], names=['trade_date', 'ticker']),
            data=[1, 2, 3, 4, 5, 6],
            columns=['factor']),
          OutputDataFormat.PITVOT_TABLE_DF,
          'factor',

# Path: alphaware/enums.py
# class OutputDataFormat(IntEnum):
#     MULTI_INDEX_DF = 0
#     PITVOT_TABLE_DF = 1
# 
# class FreqType(StrEnum):
#     EOD = 'D'
#     EOW = 'W'
#     EOM = 'M'
#     EOQ = 'Q'
#     EOSY = 'S'
#     EOY = 'Y'
# 
# Path: alphaware/const.py
# INDEX_FACTOR = namedtuple('INDEX_FACTOR', ['date_index', 'sec_index', 'full_index', 'col_score'])
# 
# Path: alphaware/utils/pandas_utils.py
# @expect_types(data=(pd.Series, pd.DataFrame))
# def convert_df_format(data, target_format=OutputDataFormat.MULTI_INDEX_DF, col_name='factor',
#                       multi_index=INDEX_FACTOR):
#     if target_format == OutputDataFormat.MULTI_INDEX_DF:
#         tmp = data.stack()
#         data_ = pd.DataFrame(tmp)
#         data_.index.names = multi_index.full_index
#         data_.columns = [col_name]
#     else:
#         tmp = data.unstack()
#         index = tmp.index
#         columns = tmp.columns.get_level_values(multi_index.sec_index).tolist()
#         data_ = pd.DataFrame(tmp.values, index=index, columns=columns)
# 
#     return data_
# 
# @expect_types(df=(pd.Series, pd.DataFrame))
# def top(df, column=None, n=5):
#     if isinstance(df, pd.Series):
#         ret = df.sort_values(ascending=False)[:n]
#     else:
#         py_assert(column is not None, "Specify the col name or use pandas Series type of data")
#         ret = df.sort_values(by=column, ascending=False)[:n]
# 
#     return ret
# 
# @expect_types(data=(pd.DataFrame, pd.Series), freq=optional(FreqType, str))
# def group_by_freq(data, freq=FreqType.EOM):
#     data_ = pd.DataFrame(data) if isinstance(data, pd.Series) else data
#     if freq == FreqType.EOD:
#         return data_.groupby([lambda x: x.year, lambda x: x.month, lambda x: x.day])
#     elif freq == FreqType.EOM:
#         return data_.groupby(pd.TimeGrouper(freq='M'))
#     elif freq == FreqType.EOQ:
#         return data_.groupby(pd.TimeGrouper(freq='Q'))
#     elif freq == FreqType.EOY:
#         return data_.groupby(pd.TimeGrouper(freq='A'))
# 
# @expect_types(x=(pd.DataFrame, pd.Series))
# def fwd_return(x, date_index=INDEX_FACTOR.date_index, sec_index=INDEX_FACTOR.sec_index, period=1):
#     """
#     每个日期和股票代码对应的未来收益
#     """
#     dates = sorted(set(x.index.get_level_values(date_index)))
#     ret = pd.DataFrame()
#     for i in range(len(dates) - period):
#         shift_date = dates[i + period]
#         data_concat = x.loc[shift_date].reset_index()
#         data_concat[date_index] = [dates[i]] * len(data_concat)
#         ret = pd.concat([ret, data_concat], axis=0)
#     ret.set_index([date_index, sec_index], inplace=True)
#     ret.columns = ['fwd_return']
#     return ret
# 
# @preprocess(data=ensure_pd_df)
# def weighted_rank(data, ascend_order=None, weight=None, out_df=False):
#     nb_col = len(data.columns)
#     ascend_order = [1] * nb_col if ascend_order is None else ascend_order
#     weight = [1.0 / nb_col] * nb_col if weight is None else weight
#     rank = np.zeros([len(data), nb_col])
#     for i in range(len(data.columns)):
#         rank[:, i] = np.argsort(data[data.columns[i]]) if ascend_order[i] else np.argsort(data[data.columns[i]])[::-1]
#     weight_rank = np.dot(rank, np.array(weight))
#     return pd.DataFrame(weight_rank, index=data.index) if out_df else weight_rank


# Path: alphaware/tests/utils/test_pandas_utils.py
from unittest import TestCase
from parameterized import parameterized
from numpy.testing.utils import assert_array_equal
from pandas import (MultiIndex,
                    Index)
from pandas.util.testing import assert_frame_equal, assert_series_equal
from alphaware.enums import OutputDataFormat, FreqType
from alphaware.const import INDEX_FACTOR
from alphaware.utils import (convert_df_format,
                             top,
                             group_by_freq,
                             fwd_return,
                             weighted_rank)
from datetime import datetime as dt
import pandas as pd
import numpy as np

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



class TestPandasUtils(TestCase):
    @parameterized.expand([(pd.DataFrame({'001': [1, 2, 3], '002': [2, 3, 4]}, index=['2014', '2015', '2016']),
                            OutputDataFormat.MULTI_INDEX_DF,
                            'test_factor',
                            INDEX_FACTOR,
                            pd.DataFrame(index=MultiIndex(levels=[['2014', '2015', '2016'], ['001', '002']],
                                                          labels=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]],
                                                          names=['trade_date', 'ticker']),
                                         data=[1, 2, 2, 3, 3, 4],
                                         columns=['test_factor']))])
    def test_convert_df_format_1(self, data, target_format, col_name, multi_index, expected):
        calculated = convert_df_format(data, target_format, col_name, multi_index)
        assert_frame_equal(calculated, expected)

    @parameterized.expand(
        [(pd.DataFrame(
            index=MultiIndex.from_product([['2014', '2015', '2016'], ['001', '002']], names=['trade_date', 'ticker']),
            data=[1, 2, 3, 4, 5, 6],
            columns=['factor']),
          OutputDataFormat.PITVOT_TABLE_DF,
          'factor',

          INDEX_FACTOR,