"""
生成预测数据，  3月14日 和 3月17日的小时级数据
             3月18日至 3月21日的daily 数据
"""

import os
import json
from tqdm import tqdm
from functions import generate_hourly_strings


def generate_hour_data(s, e, flag):
    """
    从 temp_data读取小时级别数据作为训练数据和预测数据
    :param s: 开始日期
    :param e: 结束日期
    :param flag: train / pred
    :return:
    """
    target_string_list = generate_hourly_strings(start_date=s, end_date=e)
    path = "data/temp_data"
    L = []
    for file in tqdm(target_string_list):
        json_path = os.path.join(path, "hour_" + file + ".json")
        with open(json_path, encoding="utf-8") as f:
            data = json.loads(f.read())
        for obj in data:
            L.append(obj)

    with open("data/{}_data/{}_{}_{}.json".format(flag, flag, s, e), "w", encoding="utf-8") as f:
        f.write(json.dumps(L, ensure_ascii=False))


if __name__ == "__main__":
    iii = int(input("请输入标识符，输入 1 生成训练数据， 输入 2 生成预测数据: \n"))
    if iii == 1:
        f = "train"
    elif iii == 2:
        f = "pred"
    else:
        print("输入错误")
    start = str(input("请输入开始字符串， 格式为 yyyymmddhh:\n"))
    end = str(input("请输入结束字符串, 格式为 yyyymmddhh: \n"))
    generate_hour_data(s=start, e=end, flag=f)