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


class DUC2003(Dataset):

    def __init__(self, summary_no=0):

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

        if summary_no > 3:
            raise Exception("summary_no has to be 0~3 (4 different writer).")
        self.summary_no = summary_no

    @property
    def extract_targets(self):
        return ["DUC2003/input.txt",
                "DUC2003/task1_ref{}.txt".format(self.summary_no)]

    def prepare(self, dataset_root, extracted_path):
        news_path = self.__get_extracted_path(extracted_path, 0)
        summary_path = self.__get_extracted_path(extracted_path, 1)
        write_file_name = "duc2003_ref{}.txt".format(self.summary_no)
        write_path = os.path.join(dataset_root, write_file_name)

        write_file = open(write_path, "w", encoding="utf-8")
        with open(news_path, encoding="utf-8") as n:
            with open(summary_path, encoding="utf-8") as s:
                for nline, sline in zip(n, s):
                    line = "\t".join([
                        nline.strip(), sline.strip()
                        ]) + "\n"
                    write_file.write(line)
        write_file.close()
        return write_path

    def __get_extracted_path(self, extracted_path, index):
        return os.path.join(
                extracted_path, 
                os.path.basename(self.extract_targets[index]))

    def make_resource(self, data_root):
        return Resource(
                data_root,
                columns=["news", "summary"],
                target="summary")
