import os
from chazutsu.datasets.framework.dataset import Dataset
from chazutsu.datasets.framework.resource import Resource


class DUC2004(Dataset):

    def __init__(self):

        super().__init__(
            name="DUC2004",
            site_url="http://duc.nist.gov/duc2004/tasks.html",
            download_url="https://s3-ap-northeast-1.amazonaws.com/dev.tech-sketch.jp/chakki/chazutsu/DUC2004.zip",  # noqa
            description="DUC2004 Summarization dataset. " \
                        "This dataset has 500 news, and its summary " \
                        "generated by 4 different humans."
            )

    @property
    def extract_targets(self):
        refs = ["DUC2004/task1_ref{}.txt".format(n) for n in range(4)]
        return ["DUC2004/input.txt"] + refs

    def prepare(self, dataset_root, extracted_path):
        paths = self.__get_extracted_pathes(extracted_path)
        write_path = os.path.join(dataset_root, "duc2004.txt")

        data_files = [open(p, encoding="utf-8") for p in paths]
        files_iter = zip(*data_files)
        with open(write_path, "w", encoding="utf-8") as f:
            for lines in files_iter:
                lines = [ln.strip() for ln in lines]
                line = "\t".join(lines) + "\n"
                f.write(line)
            for _f in data_files:
                _f.close()
        return write_path

    def __get_extracted_pathes(self, extracted_path):
        paths = []
        for e in self.extract_targets:
            p = os.path.join(extracted_path, os.path.basename(e))
            paths.append(p)
        return paths

    def make_resource(self, data_root):
        return Resource(
                data_root,
                columns=["news", "summary1", "summary2",
                                 "summary3", "summary4"],
                target="summary1")
