import json
import os

import luigi
import pandas as pd


class DatasetParameter(luigi.Parameter):
    pass


class DatasetTask(luigi.Task):
    dataset = DatasetParameter()


class DataFrame(luigi.LocalTarget):
    def __init__(self, path=None, format=None, is_tmp=False):
        if os.path.splitext(path)[-1] == "":
            path += ".pd"   # pd is short for Pandas
        super().__init__(path=path, format=format, is_tmp=is_tmp)

    def save(self, data_frame: pd.DataFrame):
        self.makedirs()
        data_frame.to_pickle(self.path)

    def load(self):
        data_frame = pd.read_pickle(self.path)
        return data_frame


class JSON(luigi.LocalTarget):
    def save(self, input_dict: dict):
        with open(self.path, "w") as f:
            json.dump(input_dict, f, indent=4)

    def load(self):
        with open(self.path, "r") as f:
            ret = json.load(f)
        return ret