import requests
import pandas as pd

CN_CODE_TYPE = {
    "600":'sh',
    "601":'sh',
    "603":'sh',
    "605":'sh',
    "000":'sz',
    "300":'sz',
    "688":'sh',
    "002":'sz',
    "50":"sh",
    "51":"sh",
    "52":"sh",
    "16":"sz",
    "15":"sz",
    "18":"sz"
    }



def get_history_k(code, start_date='19700101', end_date=None):

    if '.' in code:
        code = code.split('.')[0]

    stock_code = ''

    if code[0:2] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:2]]
    elif code[0:3] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:3]]
    print(stock_code)
    if stock_code != '':

        url = "https://api.zhituapi.com/hs/history/%s/d/f?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)

        if end_date is not None:
            url += "&end_date=%s"%(end_date)

        response = requests.get(url)

        data = response.json()
        print(data)
        df = pd.DataFrame(data)

        df['t'] = df['t'].apply(pd.to_datetime)
        df.rename({'t': 'date', 'o': 'open', 'h': 'high', 'l': 'low', 'c': 'close', 'v': 'volume', 'a': 'amount'},
                  axis='columns', inplace=True)
        df = df[['date', 'open', 'high', 'low', 'close', 'volume', 'amount']].set_index('date')

    else:
        df = pd.DataFrame()


    return df


def get_5min_k(code, start_date='19700101', end_date=None):
    stock_code = ''

    if code[0:2] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:2]]
    elif code[0:3] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:3]]

    if stock_code != '':

        url = "https://api.zhituapi.com/hs/history/%s/5/n?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)

        if end_date is not None:
            url += "&end_date=%s"%(end_date)

        response = requests.get(url)

        data = response.json()

        df = pd.DataFrame(data)
        df['t'] = df['t'].apply(pd.to_datetime)
        df.rename({'t': 'date', 'o': 'open', 'h': 'high', 'l': 'low', 'c': 'close', 'v': 'volume', 'a': 'amount'},
                  axis='columns', inplace=True)
        df = df[['date','close', 'volume']].set_index('date')

    else:
        df = pd.DataFrame()


    return df