import datasets
import pandas as pd
from tqdm.notebook import tqdm
import json
import os

def load_dataset(dataset_name, **kwargs):
    if dataset_name == "Stocknet":
        root_path = r"../../../stocknet-dataset/tweet/raw"
        stock_lists = os.listdir(root_path)
        all = pd.DataFrame()
        for stock in tqdm(stock_lists, desc="Loading Stocknet dataset..."):
            stock_path = os.path.join(root_path, stock)
            date_files = os.listdir(stock_path)
            for date in date_files:
                with open(os.path.join(stock_path, date_files[0])) as f:
                    json_list = f.readlines()
                tmp_json = []
                for json_str in json_list:
                    tmp_json.append(json.loads(json_str))
                tmp_json = pd.DataFrame(tmp_json)
                all = pd.concat([all, tmp_json], axis=0)
        all = all.reset_index(drop=True)
        all = datasets.Dataset.from_pandas(all)
        return all

    else:
        raise NotImplementedError("Only support Stocknet dataset for now")