
import datetime
import traceback
import json
import numpy as np
import pandas as pd
from typing import List, Dict, Optional, Tuple, Union, Set

from ....constant import SectorType, IndexPriceSource, INDEX_PRICE_EXTRA_TRADE_DAYS, FUND_NAV_EXTRA_TRADE_DAYS
from ...api.raw import RawDataApi
from ...api.basic import BasicDataApi
from ...view.basic_models import *
from ...manager.fund_tag_update import fund_tag
# from ..raw.raw_data_helper import RawDataHelper
from .basic_data_helper import BasicDataHelper


class BasicDataPart1:

    _REPORT_DATE_LIST = ['0331', '0630', '0930', '1231']

    def __init__(self, data_helper: BasicDataHelper):
        self._data_helper = data_helper
        self._raw_data_api = RawDataApi()
        self._basic_data_api = BasicDataApi()
        self._index_info_df = self._basic_data_api.get_index_info()
        self._fund_info_df = self._basic_data_api.get_fund_info()

    def _status_mapping(self, x):
        if x == 'Open':
            return 0
        elif x == 'Suspended':
            return 1
        elif x == 'Limited':
            return 2
        elif x == 'Close':
            return 3
        else:
            return None

    def _find_tag(self, symbol, wind_class_II):
        if '沪深300' in symbol and wind_class_II in ['普通股票型基金', '增强指数型基金', '被动指数型基金']:
            return 'A股大盘'
        elif '中证500' in symbol and wind_class_II in ['普通股票型基金', '增强指数型基金', '被动指数型基金']:
            return 'A股中盘'
        elif '标普500' in symbol:
            return '美股大盘'
        elif '创业板' in symbol and wind_class_II in ['普通股票型基金', '增强指数型基金', '被动指数型基金']:
            return '创业板'
        elif '德国' in symbol:
            return '德国大盘'
        elif '日本' in symbol or '日经' in symbol:
            return '日本大盘'
        elif (('国债' in symbol or '利率' in symbol or '金融债' in symbol)
              and wind_class_II in ['短期纯债型基金', '中长期纯债型基金', '被动指数型债券基金']):
            return '利率债'
        elif (('信用' in symbol or '企债' in symbol or '企业债' in symbol)
              and wind_class_II in ['短期纯债型基金', '中长期纯债型基金', '被动指数型债券基金']):
            return '信用债'
        elif '黄金' in symbol:
            return '黄金'
        elif '原油' in symbol or '石油' in symbol or '油气' in symbol:
            return '原油'
        elif ('地产' in symbol or '金融' in symbol) and ('美国'not in symbol):
            return '房地产'
        else:
            return 'null'

    def fund_info(self):
        # Update manually
        # Not verified
        track_index_df = pd.read_csv('./data/fund_track_index.csv', index_col=0)
        wind_fund_info = self._raw_data_api.get_wind_fund_info()
        fund_fee = self._raw_data_api.get_fund_fee()
        wind_fund_info['order_book_id'] = [_.split('!')[0].split('.')[0] for _ in wind_fund_info['wind_id'] ]
        res = []
        for i in wind_fund_info['wind_id']:
            if not '!' in i:
                res.append(0)
            else:
                res_i = int(i.split('!')[1].split('.')[0])
                res.append(res_i)
        wind_fund_info['transition'] = res
        wind_fund_info['fund_id'] = [o + '!' + str(t) for o, t in zip(wind_fund_info['order_book_id'], wind_fund_info['transition'])]
        wind_fund_info = wind_fund_info.set_index('fund_id')
        fund_fee = fund_fee.drop(['id','desc_name'], axis = 1).set_index('fund_id')
        wind_fund_info = wind_fund_info.join(fund_fee)
        wind_fund_info['fund_id'] = wind_fund_info.index
        wind_fund_info['asset_type'] = [self._find_tag(symbol, wind_class_II) for symbol, wind_class_II in zip(wind_fund_info['desc_name'],wind_fund_info['wind_class_2'])]
        wind_fund_info = wind_fund_info[[i.split('.')[1] == 'OF' for i in wind_fund_info['wind_id']]]
        wind_fund_info = wind_fund_info.set_index('wind_id')
        dic = {k:v for k,v in zip(self._index_info_df['desc_name'], self._index_info_df['index_id'])}
        wind_fund_info = wind_fund_info.join(track_index_df[['track_index']])
        wind_fund_info['wind_id'] = wind_fund_info.index
        wind_fund_info['track_index'] = wind_fund_info['track_index'].map(lambda x: dic.get(x,'none'))
        self._data_helper._upload_basic(wind_fund_info, FundInfo.__table__.name)

    def fund_info_from_raw(self, fund_id_list: Set[str], benchmark_profile: pd.Series):
        def _invest_type_class_to_wind_class_1(x: pd.Series) -> str:
            if x.invest_type_class == '其他基金':
                if 'QDII' in x.desc_name:
                    return '国际(QDII)基金'
            return x.invest_type_class

        def _invest_type_class_to_wind_class_2(invest_type_class: str) -> str:
            '''Choice的基金二级类型转换为wind的'''

            if invest_type_class == '混合债券型基金（一级）':
                return '混合债券型一级基金'
            elif invest_type_class == '长期纯债型基金':
                return '中长期纯债型基金'
            elif invest_type_class == '中短期纯债型基金':
                return '短期纯债型基金'
            elif invest_type_class == '混合债券型基金（二级）':
                return '混合债券型二级基金'
            elif invest_type_class == '被动指数型基金':
                return '被动指数型债券基金'
            elif invest_type_class in ('QDII基金', '封闭式基金', '保本型基金'):
                # error case
                return
            else:
                return invest_type_class

        def _lambda_structure_type(x: pd.Series) -> int:
            if x.base_fund_id is None:
                return 0
            if x.base_fund_id == x.wind_id:
                return 1
            # 到这里应该必然是分级基金，所以直接判断A,B即可
            elif 'A' in x.desc_name:
                return 2
            elif 'B' in x.desc_name:
                return 3
            else:
                # 分级基金A,B份额可能名字里不带A,B，之后需要手工修改下
                return 4

        def _lambda_index_id_new(x: pd.Series, benchmark: pd.DataFrame) -> Optional[str]:
            try:
                return fund_tag(x.base_fund_id, x.desc_name, x.is_open, x.currency, x.wind_class_2, benchmark.at[x.wind_id, 'benchmark_s_raw'], x.fund_id)
            except KeyError:
                print(f'{x.fund_id} has no benchmark info')
                return

        def _lambda_fund_manager(x: str) -> pd.Series:
            manager_list: List[str] = x.split(',')
            return pd.Series(['\r\n'.join(manager_list), manager_list[0]])

        def _lambda_benchmark(x: str) -> pd.Series:
            def _parse_index(part: str):
                if '×' in part:
                    es = part.split('×')
                elif '*' in part:
                    es = part.split('*')
                elif '╳' in part:
                    es = part.split('╳')
                else:
                    return part

                if '%' in es[0]:
                    return es[1]
                else:
                    return es[0]

            if x is None:
                return
            benchmarks: List[str] = x.split('+')
            if len(benchmarks) == 1:
                second = None
            else:
                second = _parse_index(benchmarks[1])
            return pd.Series([_parse_index(benchmarks[0]), second])

        def _lambda_fee(x: str, em_fund_fee: pd.DataFrame) -> pd.Series:
            def _parse_fee(fee: str):
                if fee is None:
                    return
                cheapest_fee = fee.split(';')[0]
                es = cheapest_fee.split(':')
                if len(es) > 1:
                    value = es[1]
                else:
                    value = es[0]
                if value[-1] == '%':
                    value = float(value[:-1]) / 100
                else:
                    try:
                        value = float(value) / 100
                    except ValueError:
                        return
                return value

            fee = em_fund_fee.loc[x, ['manage_fee', 'trustee_fee', 'purchase_fee', 'redeem_fee']]
            return fee.transform(_parse_fee)

        # 获取raw中的基金基础信息
        df = self._raw_data_api.get_em_fund_info(fund_id_list).drop(columns='_update_time')
        # 其中一部分列可以直接使用，改一下名字即可
        df = df.rename(columns={'em_id': 'wind_id', 'name': 'desc_name', 'manager_id': 'pre_fund_manager'})
        # is open
        df['is_open'] = 1
        key_words = ['封闭','定期','定开']
        key_words = '|'.join(key_words)
        _wind_ids = df[df['desc_name'].str.contains(key_words)].wind_id
        df = df.set_index('wind_id')
        df.loc[_wind_ids,'is_open'] = 0
        df = df.reset_index()
        # 没有end_date的话给一个默认值
        df['end_date'] = df['end_date'].where(df.end_date.notna(), '2040-12-31')
        # 这里可能还需要加上港币
        df['currency'] = df.currency.map({'人民币元': 'CNY', '美元': 'USD'})
        # 基金经理
        df[['manager_id', 'fund_manager']] = df.pre_fund_manager.apply(_lambda_fund_manager)
        # 新基金都是!0，这里需要检查是否有相同order_book_id的旧基金
        df['order_book_id'] = df.wind_id.transform(lambda x: x.split('.')[0])
        df['transition'] = 0
        df['fund_id'] = df.order_book_id.transform(lambda x: x + '!0')
        # Choice的基金一二级分类转换为wind的
        df['wind_class_1'] = df[['desc_name', 'invest_type_class_2']].transform(_invest_type_class_to_wind_class_1, axis=1)
        df['wind_class_2'] = df.invest_type_class_2.transform(_invest_type_class_to_wind_class_2)
        df['asset_type'] = df.apply(lambda x: self._find_tag(x.desc_name, x.wind_class_2), axis=1)
        df['structure_type'] = df.apply(_lambda_structure_type, axis=1)
        df['is_etf'] = df.desc_name.transform(lambda x: 1 if ('ETF' in x and '联接' not in x) else 0)
        fund_benchmark = self._basic_data_api.get_fund_benchmark(fund_id_list).set_index('em_id')
        df['index_id_new'] = df.apply(_lambda_index_id_new, axis=1, benchmark=fund_benchmark)
        em_fund_fee = self._raw_data_api.get_em_fund_fee(fund_id_list).drop(columns='_update_time').set_index('em_id')
        df[['manage_fee', 'trustee_fee', 'purchase_fee', 'redeem_fee']] = df.wind_id.transform(_lambda_fee, em_fund_fee=em_fund_fee)
        df = df.set_index('wind_id').join(em_fund_fee[['subscr_fee', 'purchase_fee', 'redeem_fee', 'service_fee']].rename(columns={'subscr_fee': 'subscr_fee_detail', 'purchase_fee': 'purchase_fee_detail', 'redeem_fee': 'redeem_fee_detail'}))
        df['benchmark'] = benchmark_profile
        # 拆解上边的benchmark
        df[['benchmark_1', 'benchmark_2']] = df.benchmark.apply(_lambda_benchmark)
        df = df.drop(columns='invest_type_class_2').reset_index()
        # 将场内变为场外的
        df['wind_id'] = df.wind_id.transform(lambda x: x.split('.')[0] + '.OF' if '.OF' not in x else x)
        self._data_helper._upload_basic(df, FundInfo.__table__.name)

    def index_info(self):
        # Update manually
        # Not verified
        df = pd.read_csv('./data/index_info.csv')
        self._data_helper._upload_basic(df, IndexInfo.__table__.name)

    def index_info_from_raw(self, index_list: Union[Tuple[str], List[str]], index_ids: List[str]):
        assert len(index_list) == len(index_ids), 'index list length should match the ids'

        df = self._raw_data_api.get_em_index_info(index_list).drop(columns=['name', 'publish_date', '_update_time'])
        assert df.shape[0] == len(index_ids), 'index info length should match the ids'
        df = df.rename(columns={'short_name': 'desc_name'})
        # 先全体置为false，后续可以手工修改
        df['is_select'] = False
        df['tag_method'] = 'PE百分位'
        # index_id是人工起的，所以只能作为参数传进来
        df['index_id'] = index_ids
        self._data_helper._upload_basic(df, IndexInfo.__table__.name)

    def index_component(self):
        def _lambda_func_index_component_2(em_id_to_index_id: Dict[str, str], x: str):
            info = json.loads(x)
            r = []
            for one in info:
                em_id = one[0]
                if em_id in ['000915.SH', '000916.SH']:
                    em_id = em_id.split('.')[0] + '.CSI'
                r.append((em_id_to_index_id[em_id], one[1]))
            return json.dumps(r)

        def _lambda_func_index_component_3(wind_id_to_id: Dict[str, str], x: str):
            if x is None:
                return x
            r = []
            info = json.loads(x)
            for one in info:
                try:
                    r.append(wind_id_to_id[one + '.OF'])
                except KeyError:
                    r.append(one + '!0')
            return json.dumps(r)

        index_info_df = self._index_info_df[self._index_info_df.em_id.notna()]
        index_info_df = index_info_df.set_index('em_id')
        em_id_to_index_id: Dict[str, str] = index_info_df['index_id'].to_dict()

        df = self._raw_data_api.get_cs_index_component().drop(columns=['_update_time'])
        df['index_id'] = df['index_id'].map(em_id_to_index_id)
        df['sector'] = df['sector'].transform(lambda x: _lambda_func_index_component_2(em_id_to_index_id, x))
        fund_info_df = self._fund_info_df.set_index('wind_id')
        wind_id_to_id: Dict[str, str] = fund_info_df['fund_id'].to_dict()
        df['related_funds'] = df['related_funds'].transform(lambda x: _lambda_func_index_component_3(wind_id_to_id, x))
        self._data_helper._upload_basic(df, IndexComponent.__table__.name)

    def fund_nav_from_rq(self, start_date, end_date):
        df = self._raw_data_api.get_rq_fund_nav(start_date, end_date)
        df['fund_id'] = df.apply(
            lambda x: self._data_helper._get_fund_id_from_order_book_id(x['order_book_id'], x['datetime']), axis=1)
        df['subscribe_status'] = df['subscribe_status'].map(self._status_mapping)
        df['redeem_status'] = df['redeem_status'].map(self._status_mapping)
        df = df[df['fund_id'].notna()]
        self._data_helper._upload_basic(df, FundNav.__table__.name)

    def fund_nav(self, start_date, end_date):
        try:
            raw_start_date_dt: datetime.date = pd.to_datetime(start_date, infer_datetime_format=True).date()
            end_date = pd.to_datetime(end_date, infer_datetime_format=True).date()

            # 多取N个交易日的数据
            start_date = raw_start_date_dt - datetime.timedelta(days=40)
            trading_day_df = BasicDataApi().get_trading_day_list(start_date, end_date)
            trading_day_df = trading_day_df[trading_day_df.datetime <= raw_start_date_dt]
            assert trading_day_df.datetime.size >= FUND_NAV_EXTRA_TRADE_DAYS, f'should at least have {FUND_NAV_EXTRA_TRADE_DAYS} trading days'
            # del_date在start_date后边一天，确保可以正确计算ret等daily difference
            start_date = trading_day_df.datetime.array[-FUND_NAV_EXTRA_TRADE_DAYS]
            del_date = trading_day_df.datetime.array[-(FUND_NAV_EXTRA_TRADE_DAYS-1)]
            print(f'[fund_nav] (start_date){start_date} (del_date){del_date} (end_date){end_date}')

            # start_date = start_date.strftime('%Y%m%d')
            df = self._raw_data_api.get_em_fund_nav(start_date, end_date).drop(columns='_update_time')
            df['fund_id'] = df.apply(
                lambda x: self._data_helper._get_fund_id_from_order_book_id(x['CODES'].split('.')[0], x['DATES']), axis=1)
            df = df.drop(['CODES'], axis=1)
            df = df[df['fund_id'].notna()]
            df = df.rename(columns={
                    'DATES': 'datetime',
                    'ORIGINALUNIT': 'unit_net_value',
                    'ORIGINALNAVACCUM': 'acc_net_value',
                    'ADJUSTEDNAV': 'adjusted_net_value',
                    'UNITYIELD10K': 'daily_profit',
                    'YIELDOF7DAYS': 'weekly_yield'
                })
            df['weekly_yield'] = df['weekly_yield'].map(lambda x: x / 100.0)

            anv_df = df.pivot_table(index='datetime', columns='fund_id', values='adjusted_net_value').fillna(method='ffill')
            df = df.set_index(['datetime', 'fund_id'])
            df['change_rate'] = (anv_df / anv_df.shift(1) - 1).stack()

            df = df[df.index.get_level_values(level='datetime').date >= del_date]
            # TODO: 最好应该原子提交下边两步
            BasicDataApi().delete_fund_nav(del_date, end_date)
            self._data_helper._upload_basic(df.reset_index(), FundNav.__table__.name)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def fund_benchmark(self):
        df = pd.read_csv('./data/benchmark_res2.csv')
        df = df[['fund_id', 'fund_em_id', 'benchmark', 'index_text', 'assets', 'industry']].rename(columns={'fund_em_id': 'em_id'})
        self._data_helper._upload_basic(df, FundBenchmark.__table__.name)

    def index_price(self, start_date, end_date):
        try:
            raw_start_date_dt: datetime.date = pd.to_datetime(start_date, infer_datetime_format=True).date()
            # 多取N个交易日的数据
            start_date = raw_start_date_dt - datetime.timedelta(days=90)
            trading_day_df = BasicDataApi().get_trading_day_list(start_date.strftime('%Y%m%d'), end_date)
            trading_day_df = trading_day_df[trading_day_df.datetime <= raw_start_date_dt]
            assert trading_day_df.datetime.size >= INDEX_PRICE_EXTRA_TRADE_DAYS, f'should at least have {INDEX_PRICE_EXTRA_TRADE_DAYS} trading days'
            # del_date在start_date后边一天，确保可以正确计算ret等daily difference
            start_date = trading_day_df.datetime.array[-INDEX_PRICE_EXTRA_TRADE_DAYS]
            del_date = trading_day_df.datetime.array[-(INDEX_PRICE_EXTRA_TRADE_DAYS-1)]
            print(f'[index_price] (start_date){start_date} (del_date){del_date} (end_date){end_date}')
            # 准备多个数据源
            cm_index_price = self._raw_data_api.get_raw_cm_index_price_df(start_date, end_date).drop(columns=['_update_time']).set_index('datetime').fillna(method='ffill')
            cxindex_index_price = self._raw_data_api.get_cxindex_index_price_df(start_date, end_date).drop(columns=['_update_time'])
            yahoo_index_price = self._raw_data_api.get_yahoo_index_price_df(start_date, end_date).drop(columns=['_update_time'])
            index_info_df = self._index_info_df[self._index_info_df.price_source == IndexPriceSource.default]
            em_id_list: Dict[str, str] = index_info_df[['em_id', 'index_id']].dropna().set_index('em_id').to_dict()['index_id']
            em_index_price = self._raw_data_api.get_em_index_price(start_date, end_date, em_id_list.keys()).drop(columns=['_update_time'])
            df_list = []
            # 处理Choice数据
            index_id_list = []
            for em_id, index_id in em_id_list.items():
                df = em_index_price[em_index_price['em_id'] == em_id]
                df = df.drop(columns=['em_id'])
                df['index_id'] = index_id
                df['ret'] = df.close / df.close.shift(1) - 1
                df = df[df.datetime >= pd.to_datetime(del_date).date()]
                df_list.append(df)
                if not df.empty:
                    index_id_list.append(index_id)
                # if em_id == 'SPX.GI':
                #     df['index_id'] = 'sp500'
                #     df_list.append(df.copy())
                #     if not df.empty:
                #         index_id_list.append('sp500')

                #     df = df.set_index('datetime').join(cm_index_price[['usd_central_parity_rate']])
                #     df['index_id'] = 'sp500rmb'
                #     df['open'] *= df['usd_central_parity_rate']
                #     df['close'] *= df['usd_central_parity_rate']
                #     df['low'] *= df['usd_central_parity_rate']
                #     df['high'] *= df['usd_central_parity_rate']
                #     df['ret'] = df.close / df.close.shift(1) - 1
                #     df = df.copy().reset_index().drop('usd_central_parity_rate', axis=1)
                #     df_list.append(df.copy())
                #     if not df.empty:
                #         index_id_list.append('sp500rmb')
                # else:
                #     df_list.append(df)
                #     if not df.empty:
                #         index_id_list.append(index_id)

            end_date_dt: datetime.date = pd.to_datetime(end_date, infer_datetime_format=True).date()
            index_list = ['sp500']  #, 'dax30', 'n225']
            cm_index_list = ['sp500rmb']  #, 'dax30rmb', 'n225rmb']
            for i, c in zip(index_list, cm_index_list):
                df = yahoo_index_price[yahoo_index_price['index_id'] == i]
                df = df.set_index('datetime').reindex(cm_index_price.index).fillna(method='ffill')
                # 存进去人民币计价的数据
                df = df[['close', 'open', 'low', 'high', 'volume']]
                df['close'] *= cm_index_price['usd_central_parity_rate']
                df['open'] *= cm_index_price['usd_central_parity_rate']
                df['low'] *= cm_index_price['usd_central_parity_rate']
                df['high'] *= cm_index_price['usd_central_parity_rate']
                df['index_id'] = c
                df['ret'] = df.close / df.close.shift(1) - 1
                df = df.loc[raw_start_date_dt:end_date_dt, :]
                df_list.append(df.reset_index())

            df_list.append(cxindex_index_price[cxindex_index_price.datetime == end_date_dt])
            # 存进去美元计价的数据
            df_list.append(yahoo_index_price[yahoo_index_price.datetime.between(raw_start_date_dt, end_date_dt)])

            df = pd.concat(df_list, ignore_index=True)
            # TODO: 改成原子提交下边两个操作
            BasicDataApi().delete_index_price(index_id_list, del_date, end_date)
            self._data_helper._upload_basic(df.replace([np.Inf, -np.Inf], np.nan), IndexPrice.__table__.name)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def fund_rating_latest(self):
        try:
            fund_rating_df = self._raw_data_api.get_fund_rating()
            fund_info = self._fund_info_df[['fund_id', 'order_book_id', 'start_date', 'end_date']]
            score_name = ['zs', 'sh3', 'sh5', 'jajx']
            score_dic = {}
            for s in score_name:
                score_df = fund_rating_df[['order_book_id', 'datetime', s]]
                datelist = sorted(list(set(score_df.dropna().datetime.tolist())))
                score_dic[s] = datelist[-1]
            res = []
            for s, d in score_dic.items():
                try:
                    df = fund_rating_df[['order_book_id', 'datetime', s]]
                    df = df[df['datetime'] == d]
                    con1 = fund_info['start_date'] <= d
                    con2 = fund_info['end_date'] >= d
                    fund_info_i = fund_info[con1 & con2]
                    dic = {row['order_book_id']: row['fund_id']
                        for index, row in fund_info_i.iterrows()}
                    df['fund_id'] = df['order_book_id'].map(lambda x: dic[x])
                    df = df[['fund_id', s]].copy().set_index('fund_id')
                    res.append(df)
                except Exception as e:
                    print(e)
            df = pd.concat(res, axis=1, sort=False)
            df['fund_id'] = df.index
            df['update_time'] = datetime.date.today()
            self._data_helper._upload_basic(df, FundRatingLatest.__table__.name, to_truncate=True)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def rq_stock_price(self, start_date, end_date):
        try:
            stock_price_df = self._raw_data_api.get_rq_stock_price(start_date, end_date)
            stock_post_price_df = self._raw_data_api.get_rq_stock_post_price(start_date, end_date)

            stock_post_price_df['adj_close'] = stock_post_price_df['close']
            stock_post_price_df = stock_post_price_df.filter(items=['adj_close', 'datetime', 'order_book_id'],
                axis='columns')

            stock_price_merge_df = stock_price_df.merge(stock_post_price_df, how='left', on=['datetime', 'order_book_id'])
            stock_price_merge_df['post_adj_factor'] = stock_price_merge_df['adj_close'] / stock_price_merge_df['close']
            stock_price_merge_df = stock_price_merge_df.rename(columns={'order_book_id': 'stock_id'})
            self._data_helper._upload_basic(stock_price_merge_df, StockPrice.__table__.name)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def rq_fund_ret(self, start_date, end_date):
        try:
            column_rename_dict = {
                'sharp_ratio': 'sharpe_ratio',
                'last_week_return': 'w1_ret',
                'last_month_return': 'm1_ret',
                'last_three_month_return': 'm3_ret',
                'last_six_month_return': 'm6_ret',
                'last_twelve_month_return': 'y1_ret',
                'max_drop_down': 'mdd',
                'annualized_returns': 'annual_ret',
                'average_size': 'avg_size',
                'information_ratio': 'info_ratio',
                'to_date_return':'to_date_ret'
            }
            column_drop_list = ['order_book_id', 'year_to_date_return', 'annualized_risk']

            df = self._raw_data_api.get_rq_fund_indicator(start_date, end_date)
            df['fund_id'] = df.apply(
                lambda x: self._data_helper._get_fund_id_from_order_book_id(x['order_book_id'], x['datetime']), axis=1)
            df = df.rename(columns=column_rename_dict).drop(columns=column_drop_list)
            df = df[df['fund_id'].notnull()]

            self._data_helper._upload_basic(df, FundRet.__table__.name)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def stock_info(self):
        try:
            df = self._raw_data_api.get_stock_info()
            self._data_helper._upload_basic(df, StockInfo.__table__.name, to_truncate=True)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def sector_info(self):
        try:
            df = pd.read_csv('./data/sector_info.csv', index_col=0)
            df['sector_type'] = df['sector_type'].map({'industry': SectorType.industry.name, 'topic': SectorType.topic.name})
            self._data_helper._upload_basic(df, SectorInfo.__table__.name, to_truncate=True)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def sector_funds(self):
        def _lambda_sector_funds_1(em_id_map: Dict[str, str], x: str):
            try:
                return em_id_map[x]
            except KeyError:
                return x

        try:
            index_info_df = self._index_info_df[self._index_info_df.em_id.notna()]
            index_info_df = index_info_df.set_index('em_id')
            em_id_list: List[str] = index_info_df.index.to_list()
            em_id_to_index_id: Dict[str, str] = index_info_df['index_id'].to_dict()

            df = pd.read_csv('./data/sector_funds.csv', index_col=0)
            df['csi_index_id'] = df['csi_index_id'].transform(lambda x: _lambda_sector_funds_1({em_id.split('.')[0]: em_id for em_id in em_id_list if em_id is not None}, x))
            df['index_id'] = df['csi_index_id'].map(em_id_to_index_id)
            self._data_helper._upload_basic(df, SectorFunds.__table__.name, to_truncate=True)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def fund_size(self, start_date, end_date):
        try:
            # Since fund_size in basic db is latest snapshot, we only use end_date as param
            df = self._raw_data_api.get_em_fund_scale(end_date=end_date)
            if df is None:
                print('failed to get em fund scale for fund size')
                return False
            # 将场内变为场外的
            df['CODES'] = df.CODES.transform(lambda x: x.split('.')[0] + '.OF' if '.OF' not in x else x)
            df = df.pivot_table(index='DATES', columns='CODES', values='FUNDSCALE').fillna(method='ffill')
            df = df.iloc[[-1], :].apply(lambda x: pd.Series({'fund_id': self._data_helper._get_fund_id_from_order_book_id(x.name.split('.')[0], x.index.array[0]), 'latest_size': float(x.array[0])})).T
            df = df.reset_index(drop=True)
            df = df[df.notnull().all(axis=1)]
            self._data_helper._upload_basic(df, FundSize.__table__.name, to_truncate=True)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def fund_status_latest(self, end_date: str):
        try:
            df: Optional[pd.DataFrame] = self._raw_data_api.get_em_fund_status(end_date=end_date)
            if df is None:
                print('failed to get em fund scale for fund status latest')
                return False
            # 将场内变为场外的
            df['CODES'] = df.CODES.transform(lambda x: x.split('.')[0] + '.OF' if '.OF' not in x else x)
            df = df.pivot(index='DATES', columns='CODES', values=['PURCHSTATUS', 'REDEMSTATUS', 'LPLIMIT', 'OTCLPLIMITJG', 'OTCLPLIMITGR', 'TRADESTATUS']).fillna(method='ffill')
            df = df.iloc[[-1], :].stack().reset_index()
            df['fund_id'] = df[['DATES', 'CODES']].apply(lambda x: self._data_helper._get_fund_id_from_order_book_id(x.CODES.split('.')[0], x.DATES), axis=1)
            df = df[df.fund_id.notna()]
            df = df.drop(columns=['DATES', 'CODES']).rename(columns=lambda x: FundStatusLatest.__getattribute__(FundStatusLatest, x).name)
            self._data_helper._upload_basic(df, FundStatusLatest.__table__.name, to_truncate=True)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def fund_rate(self, end_date):
        def _count_star_history(x):
            color_change = {
                '红':'蓝',
                '蓝':'白',
                '白':'红',
            }
            if x is None:
                return None
            if '星' in x:
                return str(x.count('★') + x.count('☆') + 1) + color_change[x[-2]]
            else:
                return str(x.count('★') + x.count('☆'))

        def _count_A_rating(x: Optional[str]) -> Optional[str]:
            if x is None:
                return
            return str(x.count('A'))

        def _count_star_recent(x):
            if x is None:
                return None
            if '星' in x:
                return str(x.count('★') + x.count('☆') + 1) + x[-2]
            else:
                return str(x.count('★') + x.count('☆'))
        try:
            df = self._raw_data_api.get_em_fund_rate(end_date)
            df['fund_id'] = df.apply( lambda x : self._data_helper._get_fund_id_from_order_book_id(x.CODES.split('.')[0], x.DATES), axis=1)
            no_star_columns = ['CODES','DATES','_update_time','fund_id']
            A_rate_columns = ['TXCOMRATRAT']
            for col in df:
                if col in no_star_columns:
                    continue
                if col in A_rate_columns:
                    df[col] = df[col].map(_count_A_rating)
                else:
                    df[col] = df[col].map(lambda x: _count_star_recent(x))
            df = df.drop(['CODES','_update_time'], axis=1)
            name_dic = {
                'DATES':'datetime',
                'MKTCOMPRE3YRAT':'mkt_3y_comp',
                'MORNSTAR3YRAT':'mornstar_3y',
                'MORNSTAR5YRAT':'mornstar_5y',
                'MERCHANTS3YRAT':'merchant_3y',
                'TXCOMRATRAT':'tx_comp',
                'SHSTOCKSTAR3YCOMRAT':'sh_3y_comp',
                'SHSTOCKSTAR3YSTOCKRAT':'sh_3y_stock',
                'SHSTOCKSTAR3YTIMERAT':'sh_3y_timeret',
                'SHSTOCKSTAR3YSHARPERAT':'sh_3y_sharpe',
                'SHSTOCKSTAR5YCOMRATRAT':'sh_5y_comp',
                'SHSTOCKSTAR5YSTOCKRAT':'sh_5y_stock',
                'SHSTOCKSTAR5YTIMERAT':'sh_5y_timeret',
                'SHSTOCKSTAR5YSHARPRAT':'sh_5y_sharpe',
                'JAJXCOMRAT':'jian_comp',
                'JAJXEARNINGPOWERRAT':'jian_earn',
                'JAJXACHIEVESTABILITYRAT':'jian_stable',
                'JAJXANTIRISKRAT':'jian_risk',
                'JAJXSTOCKSELECTIONRAT':'jian_stock',
                'JAJXTIMESELECTIONRAT':'jian_timeret',
                'JAJXBENCHMARKTRACKINGRAT':'jian_track',
                'JAJXEXCESSEARNINGSRAT':'jian_alpha',
                'JAJXTOTALFEERAT':'jian_fee',
                'fund_id':'fund_id',
            }
            df = df.rename(columns = name_dic)
            df = df.dropna(subset=['fund_id','datetime'], axis=0)
            # 寻找离目前最近且没超过当前时间的财报日期，将数据写在这个日期上（先删掉这个日期上的旧数据）
            end_date_dt = pd.to_datetime(end_date, infer_datetime_format=True)
            md = end_date_dt.strftime('%m%d')
            sentry_date = pd.Series(self._REPORT_DATE_LIST)
            sentry = sentry_date[sentry_date <= md]
            if sentry.empty:
                # 这里取前一年的12月31
                real_date = datetime.date(end_date_dt.year - 1, 12, 31)
            else:
                real_date = pd.to_datetime(str(end_date_dt.year)+sentry.array[-1], infer_datetime_format=True).date()
            df['datetime'] = real_date
            # TODO: 最好原子提交下边两步
            BasicDataApi().delete_fund_rate(real_date, df.fund_id.to_list())
            self._data_helper._upload_basic(df, FundRate.__table__.name)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def trading_day_list(self, start_date, end_date, is_automatic_update=True):
        try:
            # 如果是每日自动更新，我们需要去存T+1日的数据
            if is_automatic_update:
                start_date = end_date
                end_date = ''
            df = self._raw_data_api.get_em_tradedates(start_date, end_date)
            df = df.rename(columns={'TRADEDATES': 'datetime'})
            if is_automatic_update:
                df = df.iloc[[1], :]
            self._data_helper._upload_basic(df, TradingDayList.__table__.name)
            return True
        except Exception as e:
            print(e)
            traceback.print_exc()
            return False

    def fund_size_hold_rate_history(self):
        start_date = '20200630'
        end_date = '20200630'

        fund_holding_rate = self._raw_data_api.get_em_fund_holding_rate(start_date=start_date, end_date=end_date)
        fund_holding_rate['fund_id'] = fund_holding_rate.apply(
            lambda x: self._data_helper._get_fund_id_from_order_book_id(x['CODES'].split('.')[0], x['DATES']), axis=1)
        fund_holding_rate = fund_holding_rate[fund_holding_rate.fund_id.notna()]
        fund_holding_rate = fund_holding_rate.drop(['CODES', '_update_time'], axis=1).rename(columns={
            'DATES': 'datetime',
            'HOLDPERSONALHOLDINGPCT': 'personal_holds',
            'HOLDINSTIHOLDINGPCT': 'institution_holds',
            'HOLDNUM': 'hold_num',
        })
        fund_holding_rate = fund_holding_rate[fund_holding_rate.drop(columns='fund_id').notna().any(axis=1)]

        fund_scale = self._raw_data_api.get_em_fund_scale(start_date=start_date, end_date=end_date)
        fund_scale = fund_scale.pivot(index='DATES', columns='CODES', values='FUNDSCALE').fillna(method='ffill').stack()
        fund_scale = pd.DataFrame(fund_scale, columns=['FUNDSCALE']).reset_index()
        fund_scale['fund_id'] = fund_scale.apply(
            lambda x: self._data_helper._get_fund_id_from_order_book_id(x['CODES'].split('.')[0], x['DATES']), axis=1)
        fund_scale = fund_scale[fund_scale.fund_id.notna()]
        fund_scale = fund_scale.drop(['CODES'], axis=1).rename(columns={'DATES': 'datetime', 'FUNDSCALE': 'size'})
        fund_scale = fund_scale.dropna().drop_duplicates(subset=['size', 'fund_id'], keep='first')

        df = fund_scale.set_index(['fund_id', 'datetime']).join(fund_holding_rate.set_index(['fund_id', 'datetime']), how='outer').reset_index()
        self._data_helper._upload_basic(df, Fund_size_and_hold_rate.__table__.name)

    def fund_hold_asset(self, start_date: str, end_date: str):
        # history, update only on report date
        em_hold_asset = self._raw_data_api.get_em_fund_hold_asset(start_date, end_date)
        em_hold_asset['CODES'] = em_hold_asset.apply(lambda x:
                            self._data_helper._get_fund_id_from_order_book_id(x.CODES.split('.')[0], x.DATES), axis=1)
        em_hold_asset = em_hold_asset.dropna(subset=['CODES']).drop(columns='_update_time')
        name_dict = {
            'CODES': 'fund_id',
            'DATES': 'datetime',
            'PRTSTOCKTONAV': 'stock_nav_ratio',
            'PRTBONDTONAV': 'bond_nav_ratio',
            'PRTCASHTONAV': 'cash_nav_ratio',
            'PRTOTHERTONAV': 'other_nav_ratio',
            'MMFFIRSTREPOTONAV': 'first_repo_to_nav',
            'MMFAVGPTM': 'avg_ptm',
        }
        em_hold_asset = em_hold_asset.rename(columns=name_dict)
        self._data_helper._upload_basic(em_hold_asset, FundHoldAsset.__table__.name)

    def fund_hold_industry(self, end_date: str):
        # history, update only on report date
        em_hold_industry = self._raw_data_api.get_em_fund_hold_industry(end_date)
        em_hold_industry['CODES'] = em_hold_industry.apply(lambda x:
                            self._data_helper._get_fund_id_from_order_book_id(x.CODES.split('.')[0], x.DATES), axis=1)
        em_hold_industry = em_hold_industry.dropna(subset=['CODES']).drop(columns='_update_time')
        name_dict = {
            'CODES': 'fund_id',
            'DATES': 'datetime',
        }
        em_hold_industry = em_hold_industry.rename(columns=name_dict)
        self._data_helper._upload_basic(em_hold_industry, FundHoldIndustry.__table__.name)

    def fund_hold_stock(self, end_date: str):
        #history, update only on report date
        em_hold_stock = self._raw_data_api.get_em_fund_hold_stock(end_date)
        em_hold_stock['CODES'] = em_hold_stock.apply(lambda x :
                            self._data_helper._get_fund_id_from_order_book_id(x.CODES.split('.')[0], x.DATES), axis=1)
        em_hold_stock = em_hold_stock.dropna(subset=['CODES']).drop('_update_time', axis=1)
        name_dict = {
            'CODES':'fund_id',
            'DATES':'datetime',
        }
        em_hold_stock = em_hold_stock.rename(columns=name_dict)
        self._data_helper._upload_basic(em_hold_stock, FundHoldStock.__table__.name)

    def fund_hold_bond(self, end_date: str):
        #history, update only on report date
        em_hold_bond = self._raw_data_api.get_em_fund_hold_bond(end_date)
        em_hold_bond['CODES'] = em_hold_bond.apply(lambda x :
                            self._data_helper._get_fund_id_from_order_book_id(x.CODES.split('.')[0], x.DATES), axis=1)
        em_hold_bond = em_hold_bond.dropna(subset=['CODES']).drop('_update_time', axis=1)
        name_dict = {
            'CODES':'fund_id',
            'DATES':'datetime',
        }
        em_hold_bond = em_hold_bond.rename(columns=name_dict)
        self._data_helper._upload_basic(em_hold_bond, FundHoldBond.__table__.name)

    def mng_info(self):
        df = self._raw_data_api.get_em_mng_info()
        df['fund_id'] = df.apply(lambda x : 
                    self._data_helper._get_fund_id_from_order_book_id(x.code.split('.')[0], datetime.date(2020,8,10)), axis=1)
        df = df.drop('code',axis=1).dropna(subset=['fund_id'])
        self._data_helper._upload_basic(df, MngInfo.__table__.name)

    def mng_change(self):
        df = self._raw_data_api.get_em_fund_mng_change()
        df['fund_id'] = df.apply(lambda x : 
                    self._data_helper._get_fund_id_from_order_book_id(x.code.split('.')[0], datetime.date(2020,8,10)), axis=1)
        df = df.drop('code',axis=1).dropna(subset=['fund_id'])
        self._data_helper._upload_basic(df, FundMngChange.__table__.name)

    def fund_comp_mng_change(self):
        df = self._raw_data_api.get_fund_com_mng_change()
        self._data_helper._upload_basic(df, FundCompMngChange.__table__.name)

    def fund_comp_core_mng(self):
        df = self._raw_data_api.get_em_fund_com_core_mng()
        df['fund_id'] = df.apply(lambda x : 
                    self._data_helper._get_fund_id_from_order_book_id(x.code.split('.')[0], datetime.date(2020,8,10)), axis=1)
        df = df.drop('code',axis=1).dropna(subset=['fund_id'])
        self._data_helper._upload_basic(df, FundCompCoreMng.__table__.name)

    def fund_open_info(self):
        df = self._raw_data_api.get_em_fund_open()
        df['fund_id'] = df.apply(
            lambda x : self._data_helper._get_fund_id_from_order_book_id(x.em_id.split('.')[0], datetime.date(2020,9,15)), axis = 1)       
        df = df.dropna(subset=['fund_id'])
        df = df.drop(['em_id','_update_time'],axis=1)
        self._data_helper._upload_basic(df, FundOpenInfo.__table__.name)

    def process_all(self, start_date, end_date):
        failed_tasks = []
        # Depend on RQ
        # if not self.rq_fund_ret(start_date, end_date):
        #     failed_tasks.append('rq_fund_ret')

        # # Depend on RQ
        # if not self.stock_info():
        #     failed_tasks.append('stock_info')

        # Depend on RQ
        # if not self.rq_stock_price(start_date, end_date):
        #     failed_tasks.append('rq_stock_price')

        if not self.index_price(start_date, end_date):
            failed_tasks.append('index_price')

        # new_fund_list, _ = RawDataHelper.get_new_and_delisted_fund_list(end_date)
        # if new_fund_list is None:
        #     failed_tasks.append('get new_fund_list')
        # else:
        #     if not self.fund_info_from_raw(new_fund_list):
        #         failed_tasks.append('fund_info')

        # TODO: Depend on RQ, change to EM
        if not self.fund_rating_latest():
            failed_tasks.append('fund_rating_latest')

        if not self.fund_size(start_date, end_date):
            failed_tasks.append('fund_size')

        if not self.fund_nav(start_date, end_date):
            failed_tasks.append('fund_nav')

        if not self.trading_day_list(start_date, end_date):
            failed_tasks.append('trading_day_list')

        if not self.fund_rate(end_date):
            failed_tasks.append('fund_rate')

        if not self.fund_status_latest(end_date):
            failed_tasks.append('fund_status_latest')

        return failed_tasks

if __name__ == '__main__':
    bdp = BasicDataPart1(BasicDataHelper())
    # start_date = '20200428'
    # end_date = '20200428'
    # bdp.fund_nav(start_date, end_date)
    # bdp.index_price(start_date, end_date)
    # bdp.rq_stock_price(start_date, end_date)
    # bdp.rq_fund_ret(start_date, end_date)
    # bdp.fund_rating_latest()
    # bdp.fund_size('20200513', '20200513')
    # bdp.trading_day_list('20200124', '20200301', False)
    # bdp.fund_benchmark()
    # bdp.sector_info()
    # bdp.sector_funds()
    # bdp.index_component()
    # bdp.fund_rate('20200812')
    # bdp.fund_size_hold_rate_history()
