#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   reader_utils.py    
@Contact :   wangcc@csc.com.cn
@License :   (C)Copyright 2017-2018, CSC

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2021/10/15 10:44   wangcc     1.0         None
'''


import json
from .config import data_path, cal_config
from joblib import Parallel, delayed
import copy
import os
import re
import pandas as pd
import numpy as np
import time
import h5py
import numba as nb
from pathlib import Path
from .tools import bisect_left, bisect_right, to_intdate
from pyarrow.dataset import field as ffield
from pyarrow import dataset as ds
markets = ["zz500", "hs300", "zz800", "zz1000", "all"]


class BaseProvider:
    """ data prov"""

    def calendar(self, start_time=None, end_time=None, freq="Tday"):
        """
        :param start_time
        :param end_time
        :param freq: includdes Aday for natural day; Tday for trading day; report_period for financial statement
        :return:
        """
        return Cal.calendar(start_time, end_time, freq)

    def list_instruments(self, instruments, start_time=None, end_time=None, freq="Tday", as_list=False):
        """
        :param instruments: instruments list or market (ei. "zz500", "hs300", "zz800", "zz1000", "all"
        :param start_time:
        :param end_time:
        :param freq: includdes Aday for natural day; Tday for trading day; report_period for financial statement
        :param as_list: if as_list, return list of instruments, else return dict with values of inday-outday span
        :return:
        """
        return Inst.list_instruments(instruments, start_time, end_time, freq, as_list)

    def list_features(self, freq="*"):
        """
        :param freq: includdes Aday for natural day; Tday for trading day; report_period for financial statement
        :return: list of features
        """
        return list_features(freq)

    def features(self, instruments, fields, start_time=None, end_time=None, freq="Tday", return_mat=False):
        """
        :param instruments: instruments list or market (ei. "zz500", "hs300", "zz800", "zz1000", "all"
        :param fields: features str or list
        :param start_time:
        :param end_time:
        :param freq: includdes Aday for natural day; Tday for trading day; report_period for financial statement
        :param return_mat: if return_mat, return nd-ns numpy array
        :return:
        """
        if isinstance(fields, str):
            fields = [fields, ]
        if freq == "report_period":
            fields += ["ann_date", ]
        fields = list(fields)  # In case of tuple.
        return FeatureD.features(instruments, fields, start_time, end_time, freq, return_mat)

    def features_tday(self, instruments, fields, start_time=None, end_time=None, return_mat=False):
        """
        :param instruments: instruments list or market (ei. "zz500", "hs300", "zz800", "zz1000", "all"
        :param fields:  features str or list
        :param start_time:
        :param end_time:
        :param return_mat:
        :return:
        """
        return self.features(instruments, fields, start_time, end_time, "Tday", return_mat)

    def features_aday(self, instruments, fields, start_time=None, end_time=None, return_mat=False):
        """
        :param instruments: instruments list or market (ei. "zz500", "hs300", "zz800", "zz1000", "all"
        :param fields:  features str or list
        :param start_time:
        :param end_time:
        :param return_mat:
        :return:
        """
        return self.features(instruments, fields, start_time, end_time, "Aday", return_mat)

    def features_rpday(self, instruments, fields, start_time=None, end_time=None, return_mat=False):
        """
        :param instruments: instruments list or market (ei. "zz500", "hs300", "zz800", "zz1000", "all"
        :param fields:  features str or list
        :param start_time:
        :param end_time:
        :param return_mat:
        :return:
        """
        return self.features(instruments, fields, start_time, end_time, "report_period", return_mat)

    def list_datasets(self):
        """
        :return:  list of dataset name
        """
        return datasets_ls

    def dataset_info(self, dataset):
        """
        :param dataset: dataset name
        :return: the columns with dtype info of the dataset
        """
        return DatasetD.dataset_info(dataset)

    def datasets(self, instruments, fields, start_time=None, end_time=None, dataset=None):
        """
        :param instruments: instruments list or market (ei. "zz500", "hs300", "zz800", "zz1000", "all"
        :param fields:  features str or list
        :param start_time:
        :param end_time:
        :param dataset: dataset name
        :return:
        """
        if isinstance(fields, str):
            fields = [fields, ]
        if fields is None:
            _, _, others, _ = DatasetD.dataset_info(dataset)
            fields = list(others.keys())
        fields = list(fields)
        return DatasetD.datasets(instruments, fields, start_time, end_time, dataset)

    def datasets_fl(self, filt, fields, dataset=None, **kwargs):
        """
        :param filt: conditions to filt the dataset
        :param fields:
        :param dataset: dataset name
        :param kwargs: the args to evaluate the filter statement
        :return:
        """
        if isinstance(fields, str):
            fields = [fields, ]
        if fields is None:
            _, _, others, _ = DatasetD.dataset_info(dataset)
            fields = list(others.keys())
        fields = list(fields)
        return DatasetD.datasets_fl(filt, fields, dataset, **kwargs)


class CalendarProvider:
    """Calendar provider base class: Provide calendar data."""

    def __init__(self):
        self.cal_cache = {}
        self._uri_cal = os.path.join(data_path, "calendars", "{}.txt")

    def cache_cal(self, freq):
        _calendar = np.array(self.load_calendar(freq))
        _calendar_index = {x: i for i, x in enumerate(_calendar)}  # for fast search
        self.cal_cache[freq] = (_calendar, _calendar_index)

    def load_calendar(self, freq):
        """Load original int calendar from file."""

        fname = self._uri_cal.format(freq)
        if not os.path.exists(fname):
            raise ValueError("calendar not exists for freq " + freq)
        with open(fname) as f:
            return [int(x.strip()) for x in f]

    def calendar(self, start_time=None, end_time=None, freq="Tday"):
        _calendar, _ = self._get_calendar(freq)
        if start_time == "None":
            start_time = None
        if end_time == "None":
            end_time = None
        # strip
        if start_time:
            start_time = to_intdate(start_time)
            if start_time > _calendar[-1]:
                return np.array([])
        else:
            start_time = _calendar[0]
        if end_time:
            end_time = to_intdate(end_time)
            if end_time < _calendar[0]:
                return np.array([])
        else:
            end_time = _calendar[-1]
        si, ei = self.locate_index(start_time, end_time, freq)
        return _calendar[si: ei + 1]

    def locate_index(self, start_time, end_time, freq='Tday'):
        """Locate the start time index and end time index in a calendar under certain frequency."""
        try:
            start_time = to_intdate(start_time)
        except:
            pass
        end_time = to_intdate(end_time)
        calendar, calendar_index = self._get_calendar(freq=freq)
        if start_time not in calendar_index:
            try:
                start_time = calendar[bisect_left(calendar, start_time)]
            except IndexError:
                raise IndexError(
                    f"`start_time` doesnt exist in {freq} calendar`"
                )
        start_index = calendar_index[start_time]
        if end_time not in calendar_index:
            end_time = calendar[bisect_right(calendar, end_time) - 1]
        end_index = calendar_index[end_time]
        return start_index, end_index

    def _get_calendar(self, freq):
        """Load calendar using memcache."""
        if freq in self.cal_cache:
            return self.cal_cache.get(freq)
        else:
            self.cache_cal(freq)
            return self.cal_cache.get(freq)

    def predays(self, day, shifts, freq):
        calendar, calendar_index = self._get_calendar(freq=freq)
        if shifts > 0:
            if day not in calendar_index:
                try:
                    day = calendar[bisect_left(calendar, day)]
                except IndexError:
                    raise IndexError("`start_time` not exist")
            index = calendar_index[day]
        else:
            if day not in calendar_index:
                day = calendar[bisect_right(calendar, day) - 1]
            index = calendar_index[day]
        return calendar[index + shifts]


class InstrumentProvider:
    """Instrument provider base class:Provide instrument data."""

    @property
    def _uri_inst(self):
        """Instrument file uri."""
        return os.path.join(data_path, "instruments", "{}.txt")

    def _load_instruments(self, market):
        fname = self._uri_inst.format(market)
        if not os.path.exists(fname):
            raise ValueError("instruments not exists for market " + market)

        _instruments = dict()
        df = pd.read_csv(
            fname,
            sep="\t",
            usecols=[0, 1, 2],
            names=["inst", "start_datetime", "end_datetime"],
            dtype={"inst": str},
            # parse_dates=["start_datetime", "end_datetime"],
        )
        for row in df.itertuples(index=False):
            _instruments.setdefault(row[0], []).append((row[1], row[2]))
        return _instruments

    def list_instruments(self, instruments, start_time=None, end_time=None, freq="Tday", as_list=False):
        if isinstance(instruments, str):
            market = instruments
        else:
            market = instruments["market"]

        _instruments = self._load_instruments(market)
        # strip
        # use calendar boundary
        cal = Cal.calendar(freq=freq)
        start_time = to_intdate(start_time or cal[0])
        end_time = to_intdate(end_time or cal[-1])
        _instruments_filtered = {
            inst: list(
                filter(
                    lambda x: x[0] <= x[1],
                    [(max(start_time, x[0]), min(end_time, x[1])) for x in spans],
                )
            )
            for inst, spans in _instruments.items()
        }
        _instruments_filtered = {key: value for key, value in _instruments_filtered.items() if value}
        # as list
        if as_list:
            return list(_instruments_filtered)
        return _instruments_filtered


class BaseDataProvider:

    @staticmethod
    def get_instruments_d(instruments, start_time, end_time, freq):
        """
        Parse different types of input instruments to output instruments_d
        Wrong format of input instruments will lead to exception.

        """
        if isinstance(instruments, dict):
            if "market" in instruments:
                # dict of stockpool config
                instruments_d = Inst.list_instruments(instruments=instruments, start_time=start_time,
                                                      end_time=end_time, freq=freq, as_list=False)
            else:
                # dict of instruments and timestamp
                instruments_d = instruments
        elif isinstance(instruments, str):
            if instruments in markets:
                instruments_d = Inst.list_instruments(instruments=instruments, start_time=start_time,
                                                      end_time=end_time, freq=freq, as_list=False)
            else:
                raise ValueError("market not found")

        elif isinstance(instruments, (list, tuple, pd.Index, np.ndarray)):
            # list or tuple of a group of instruments
            instruments_d = list(instruments)
        else:
            raise ValueError("Unsupported input type for param `instrument`")
        return instruments_d

    @staticmethod
    def get_column_names(fields):
        """
        Get column names from input fields

        """
        if len(fields) == 0:
            raise ValueError("fields cannot be empty")
        fields = fields.copy()
        column_names = [str(f) for f in fields]
        return column_names

    @staticmethod
    def remove_repeat_field(fields):
        """remove repeat field

        :param fields: list; features fields
        :return: list
        """
        fields = copy.deepcopy(fields)
        _fields = set(fields)
        return sorted(_fields, key=fields.index)

    @staticmethod
    def normalize_cache_fields(fields: [list, tuple]):
        """normalize cache fields

        :param fields: features fields
        :return: list
        """
        fields = [i for i in fields if i is not None]
        return BaseDataProvider.remove_repeat_field(fields)


class FeatureProvider(BaseDataProvider):
    """Dataset provider class: Provide Dataset data."""

    @property
    def _uri_data(self):
        return os.path.join(data_path, "features", "{}", "{}.h5")

    def features(self, instruments, fields, start_time=None, end_time=None, freq="Tday", return_mat=False):
        instruments_d = self.get_instruments_d(instruments, start_time, end_time, freq)
        column_names = self.get_column_names(fields)
        cal = Cal.calendar(start_time, end_time, freq)
        if len(cal) == 0:
            return pd.DataFrame(columns=column_names)
        start_time = cal[0]
        end_time = cal[-1]
        data = self.feature_processor(self._uri_data, instruments_d, column_names, start_time, end_time, freq, return_mat)
        return data

    @staticmethod
    def feature_processor(fpath, instruments_d, column_names, start_time, end_time, freq, return_mat=False):
        """
        Load and process the data, return the data set.
        - default using multi-kernel method.

        """
        # One process for one task, so that the memory will be freed quicker.
        if isinstance(instruments_d, dict):
            instruments_d = list(instruments_d.keys())
        start_index, end_index = Cal.locate_index(start_time, end_time, freq)
        dtname = cal_config[freq]["col_name"]
        timeindex, mask, codes = read_mask(fpath.format(freq, dtname), instruments_d, start_index, end_index, return_mat)
        if mask is None:
            return pd.DataFrame(columns=column_names)

        workers = min(kernels, len(instruments_d))
        normalize_column_names = FeatureD.normalize_cache_fields(column_names)
        start = time.perf_counter()

        field_l = []
        tasks_l = []
        for field in normalize_column_names:
            field_l.append(field)
            tasks_l.append(
                delayed(read_mmap)(
                    fpath.format(freq, field), 'data', mask
                )
            )

        data = dict(
            zip(
                field_l,
                Parallel(n_jobs=workers, backend="threading")(tasks_l),
            )
        )
        if return_mat:
            data = {
                k: v.reshape((len(instruments_d), -1)).T for k, v in data.items()
            }
            data["stockcode"] = codes
            data[dtname] = Cal.calendar(start_time, end_time, freq)
        else:
            data["stockcode"] = codes
            data[dtname] = timeindex
            data = pd.DataFrame(data).set_index(["stockcode", dtname]).dropna(how="all")
        print(time.perf_counter() - start)
        return data


class DatasetProvider(BaseDataProvider):
    """Dataset provider class: Provide Dataset data."""

    @property
    def _uri_data(self):
        return os.path.join(data_path, "datasets", "{}")

    def datasets(self, instruments, fields, start_time=None, end_time=None, dataset=None):
        instruments_d = self.get_instruments_d(instruments, start_time, end_time, freq="Aday")
        column_names = self.get_column_names(fields)
        cal = Cal.calendar(start_time, end_time, "Aday")
        if len(cal) == 0:
            return pd.DataFrame(columns=column_names)
        start_time = cal[0]
        end_time = cal[-1]
        fpath = self._uri_data.format(dataset)
        data = self.dataset_processor(fpath, instruments_d, column_names, start_time, end_time)
        return data

    def datasets_fl(self, filt, fields, dataset, **kwargs):
        column_names = self.get_column_names(fields)
        fpath = self._uri_data.format(dataset)
        data = self.dataset_processor_fl(fpath, filt, column_names, **kwargs)
        return data

    def dataset_info(self, dataset):
        config_path = self._uri_data.format(dataset) + "_config.json"
        return DatasetD.read_config(config_path)

    @staticmethod
    def dataset_processor(fpath, instruments_d, column_names, start_time=None, end_time=None):
        """
        :param fpath:
        :param instruments_d:
        :param column_names:
        :param start_time:
        :param end_time:
        :return:
        """
        # One process for one task, so that the memory will be freed quicker.
        normalize_column_names = DatasetD.normalize_cache_fields(column_names)
        config_path = fpath + "_config.json"
        code_flt, time_flt, _, is_annrepo = DatasetD.read_config(config_path)
        primary_keys = [code_flt, time_flt]            # 插入primary key
        if is_annrepo:
            primary_keys += ["report_period"]
        normalize_column_names = primary_keys + normalize_column_names
        code_flt = ffield(code_flt) if code_flt is not None else None
        time_flt = ffield(time_flt) if time_flt is not None else None
        dset = ds.dataset(fpath, format="feather")                        # 还差不规则特征

        if code_flt is None:
            if time_flt is None:
                return dset.to_table(columns=normalize_column_names).to_pandas()
            else:
                condition = (time_flt >= start_time) & (time_flt <= end_time)
        else:
            if time_flt is None:
                condition = code_flt.isin(instruments_d)
            else:
                condition = (time_flt >= start_time) & (time_flt <= end_time) & (code_flt.isin(instruments_d))
        start = time.perf_counter()
        data = dset.to_table(columns=normalize_column_names, filter=condition).to_pandas()
        print(time.perf_counter() - start)
        return data

    @staticmethod
    def dataset_processor_fl(fpath, filt, column_names, **kwargs):
        """
        :param fpath:
        :param filt:
        :param column_names:
        :param kwargs:
        :return:
        """

        # One process for one task, so that the memory will be freed quicker.
        config_path = fpath + "_config.json"
        code_flt, time_flt, _, is_annrepo = DatasetD.read_config(config_path)
        primary_keys = [code_flt, time_flt]            # 插入primary key
        if is_annrepo:
            primary_keys += ["report_period"]
        column_names = primary_keys + column_names

        normalize_column_names = DatasetD.normalize_cache_fields(column_names)
        dset = ds.dataset(fpath, format="feather")
        # 还差不规则特征
        start = time.perf_counter()
        if filt is None:
            data = dset.to_table(columns=normalize_column_names).to_pandas()
        else:
            for k in kwargs.keys():
                exec(f"{k}=kwargs['{k}']")
            filt=eval(parse_field(filt))
            data = dset.to_table(columns=normalize_column_names, filter=filt).to_pandas()
        print(time.perf_counter() - start)
        return data

    @staticmethod
    def read_config(config_path):
        with open(config_path, "r") as f:
            json_cfg = json.load(f)
            code_flt = json_cfg["code_flt"]
            time_flt = json_cfg["time_flt"]
            others = json_cfg["others"]
            is_annrepo = json_cfg["is_annrepo"]
        return code_flt, time_flt, others, is_annrepo


def list_features(freq="*"):
    features = (Path(data_path) / "features").glob(f"**/{freq}/*.h5")
    features = [i.name.replace(".h5", "") for i in features]
    return features


def list_datasets():
    datasets = (Path(data_path) / "datasets").glob("*")
    datasets = [i.name for i in datasets if i.is_dir()]
    return datasets


def read_mmap(pth, key, mask=None):
    """numpy 数据结构提取, mmap加速"""
    with h5py.File(pth, 'r') as f:
        ds = f[key]
        offset = ds.id.get_offset()
        assert ds.chunks is None
        assert ds.compression is None
        assert offset > 0
        dtype = ds.dtype
        shape = ds.shape
    arr = np.memmap(pth, mode='r', shape=shape, offset=offset, dtype=dtype)
    if not mask is None:
        return arr[mask]
    return arr


def read_index(pth):
    with h5py.File(pth, 'r') as f:
        instlist = f['instlist'][:].astype(str)  # inst code list
        instloc = f['instloc'][:]
        rsilist = f['rsilist'][:]
        reilist = f['reilist'][:]
        cidx = f['data'].shape[0]
    return instlist, instloc, rsilist, reilist, cidx


def read_mask(path, ori_instruments, start_index, end_index, mat_mode=False):
    instlist, instloc, rsilist, reilist, cidx = read_index(path)
    timeindex = read_mmap(path, "data")
    mask, codes = None, []
    if ori_instruments is None:
        codes = ori_instruments
        infoloc = np.arange(len(instloc))
    else:
        codes, infoloc = get_availables(ori_instruments, instlist)
    if len(infoloc) > 0:
        mask, amounts = gen_idxs(cidx, infoloc, instloc, rsilist, reilist, start_index, end_index)
        if not mat_mode:                             # 返回矩阵形式
            codes = np.repeat(codes, amounts)
            timeindex = timeindex[mask]
    return timeindex, mask, codes


def get_availables(instruments, availables):
    series = pd.Series(np.arange(len(instruments)), index=instruments).reindex(availables).dropna()
    return series.index.values, series.values.astype(int)


@nb.jit(nopython=True)
def gen_idxs(cidx, infoloc, instloc, rsilist, reilist, start_index, end_index):
    mask = np.zeros(cidx, dtype=np.bool_)
    amounts = []
    for i in infoloc:
        if (end_index < rsilist[i]) | (start_index > reilist[i]) | (start_index > end_index):
            amounts.append(0)
            continue
        instsiloc = instloc[i]
        tsi = max(start_index - rsilist[i], 0)
        tei = min(end_index - rsilist[i], reilist[i] - rsilist[i])
        h5si, h5ei = tsi + instsiloc, tei + instsiloc
        amounts.append(h5ei - h5si + 1)
        mask[h5si:h5ei + 1] = True
        assert amounts[-1]>0
    return mask, amounts


def parse_field(field):
    if not isinstance(field, str):
        field = str(field)
    return re.sub(r"\$(\w+)", r'ffield("\1")', field)


features_ls = {**{a: "Aday" for a in list_features("Aday")}, **{t: "Tday" for t in list_features("Tday")},
               **{rp: "report_period" for rp in list_features("report_period")}}
datasets_ls = list_datasets()
DatasetD = DatasetProvider()
FeatureD = FeatureProvider()
Cal = CalendarProvider()
Inst = InstrumentProvider()
D = BaseProvider()

kernels = 20
