import baostock as bs
import pandas as pd
import numpy as np
import datetime
import math
import pprint
import csv
import heapq
from collections import defaultdict
import json
import pickle as p
import multiprocessing as mp
import threading as td
from queue import Queue


# pd.set_option('display.max_columns', 1000)
# pd.set_option('display.width', 1000)
# pd.set_option('display.max_colwidth', 1000)
# pd.set_option('display.unicode.ambiguous_as_wide', True)
# pd.set_option('display.unicode.east_asian_width', True)


def choose_stack(code):
    rs = bs.query_history_k_data_plus(code, 'date,close,volume,amount,turn', start_date='2002-07-01', frequency='w',
                                      adjustflag='1')
    tmp_start = None
    start_time = None
    time_list = []
    tmp_amount = 0.
    tmp_turn = 0.
    start_amount = tmp_amount
    start_turn = tmp_turn
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        now_itm = rs.get_row_data()
        if not now_itm[3]:
            continue
        tmp_amount = float(now_itm[3]) if float(now_itm[3]) > 0. else tmp_amount
        if start_amount == 0.:
            start_amount = tmp_amount
        if not now_itm[4]:
            continue
        tmp_turn = float(now_itm[4]) if float(now_itm[4]) > 0. else tmp_turn
        if start_turn == 0.:
            start_turn = tmp_turn
        if tmp_start is None:
            tmp_start = now_itm
            start_time = datetime.datetime.strptime(tmp_start[0], '%Y-%m-%d').date()
            start_amount = tmp_amount
            start_turn = tmp_turn
            continue
        now_time = datetime.datetime.strptime(now_itm[0], '%Y-%m-%d').date()
        tmp_grow = float(now_itm[1]) / float(tmp_start[1])
        if tmp_grow < 1:
            tmp_start = now_itm
            start_time = now_time
            start_amount = tmp_amount
            start_turn = tmp_turn
            continue
        time_diff = (now_time - start_time).days
        tmp_avg_rate = math.log(tmp_grow) / time_diff
        if tmp_avg_rate < 0.0017:
            tmp_start = now_itm
            start_time = now_time
            start_amount = tmp_amount
            start_turn = tmp_turn
            continue
        if time_diff > 365:
            if time_list and time_list[-1][2] == start_time:
                if tmp_grow > time_list[-1][4]:
                    time_list[-1][3] = now_time
                    time_list[-1][4] = tmp_grow
                    time_list[-1][0] = tmp_avg_rate
                else:
                    if time_list[-1][4] / tmp_grow - 1 > 0.3:
                        tmp_start = now_itm
                        start_time = now_time
                        start_amount = tmp_amount
                        start_turn = tmp_turn
                        continue
            else:
                if start_turn == 0.:
                    tmp_start = now_itm
                    start_time = now_time
                    start_amount = tmp_amount
                    start_turn = tmp_turn
                    continue
                ltsz = start_amount / start_turn * 100
                # if ltsz < 10000000 * 1.0017 ** (
                #         start_time - datetime.datetime.strptime('2002-6-1', '%Y-%m-%d').date()).days * 2:
                #     tmp_start = now_itm
                #     start_time = now_time
                #     start_amount = tmp_amount
                #     start_turn = tmp_turn
                #     continue
                time_list.append([tmp_avg_rate, code, start_time, now_time, tmp_grow, ltsz])
    if time_list:
        return time_list
    return None


# df = choose_stack('sh.000011')
# df = choose_stack('sh.600519')
# pprint.pprint(df)
# df = choose_stack('sh.600150')
def main():
    start_dict = defaultdict(dict)
    end_dict = defaultdict(dict)
    max_len = 100
    with open("股票代码.txt", 'r') as f:
        for idx, tmp_code in enumerate(f):
            code = tmp_code[0:9]
            ret = choose_stack(code)
            print(idx + 1, code)
            if not isinstance(ret, type(None)):
                for itm in ret:
                    start_time = itm[2]
                    if start_time.month not in start_dict[start_time.year]:
                        start_dict[start_time.year][start_time.month] = [itm]
                    else:
                        if len(start_dict[start_time.year][start_time.month]) < max_len:
                            heapq.heappush(start_dict[start_time.year][start_time.month], itm)
                        else:
                            heapq.heappushpop(start_dict[start_time.year][start_time.month], itm)

                    end_time = itm[3]
                    if end_time.month not in end_dict[end_time.year]:
                        end_dict[end_time.year][end_time.month] = [itm]
                    else:
                        if len(end_dict[end_time.year][end_time.month]) < max_len:
                            heapq.heappush(end_dict[end_time.year][end_time.month], itm)
                        else:
                            heapq.heappushpop(end_dict[end_time.year][end_time.month], itm)
            if (idx + 1) % 100 == 0:
                all_dict = [start_dict, end_dict]
                store_dict = {'idx': idx, 'all_dict': all_dict}
                f = open(FILE_NAME, 'wb')
                p.dump(store_dict, f, 0)
                f.close()
                print('保存数据,now_code = {}'.format(code))
        all_dict = [start_dict, end_dict]
        store_dict = {'idx': idx, 'all_dict': all_dict}
        f = open(FILE_NAME, 'wb')
        p.dump(store_dict, f, 0)
        f.close()
        print('运行结束')


if __name__ == '__main__':
    FILE_NAME = 'p-1000-16.txt.txt'
    bs.login()
    main()
    bs.logout()