# -*- coding: utf-8 -*-

from collections import OrderedDict
import json
import os
import sys

attribute_names = [
    "Relation",
    "Section",
    "FileNumber",
    "Connective_SpanList",
    "Connective_GornList",
    "Connective_RawText",
    # "Connective_Trees",
    "Connective_StringPosition",
    "SentenceNumber",
    "ConnHead",
    "Conn1",
    "Conn2",
    "ConnHeadSemClass1",
    "ConnHeadSemClass2",
    "Conn2SemClass1",
    "Conn2SemClass2",
    "Attribution_Source",
    "Attribution_Type",
    "Attribution_Polarity",
    "Attribution_Determinacy",
    "Attribution_SpanList",
    "Attribution_GornList",
    "Attribution_RawText",
    "Arg1_SpanList",
    "Arg1_GornList",
    "Arg1_RawText",
    # "Arg1_Trees",
    "Arg1_Attribution_Source",
    "Arg1_Attribution_Type",
    "Arg1_Attribution_Polarity",
    "Arg1_Attribution_Determinacy",
    "Arg1_Attribution_SpanList",
    "Arg1_Attribution_GornList",
    "Arg1_Attribution_RawText",
    # "Arg1_Attribution_Trees",
    "Arg2_SpanList",
    "Arg2_GornList",
    "Arg2_RawText",
    # "Arg2_Trees",
    "Arg2_Attribution_Source",
    "Arg2_Attribution_Type",
    "Arg2_Attribution_Polarity",
    "Arg2_Attribution_Determinacy",
    "Arg2_Attribution_SpanList",
    "Arg2_Attribution_GornList",
    "Arg2_Attribution_RawText",
    "Sup1_SpanList",
    "Sup1_GornList",
    "Sup1_RawText",
    # "Sup1_Trees",
    "Sup2_SpanList",
    "Sup2_GornList",
    "Sup2_RawText",
    # "Sup2_Trees",
]
assert len(attribute_names) == 48

# Following the paper of Braud & Denis (2015).

sections_dict = {
    "Lin_etal_2009": {
        "train": ["%02d" % i for i in xrange(2, 22)],
        "val": ["22"],
        "test": ["23"],
    },
    "Ji_and_Eisenstein_2015": {
        "train": ["%02d" % i for i in xrange(2, 21)],
        "val": ["00", "01"],
        "test": ["21", "22"],
    },
    "Pitler_2009": {
        "train": ["%02d" % i for i in xrange(2, 21)],
        "val": ["21", "22"],
        "test": ["00", "01"],
    },
}

class PDTB2(object):
    """
    Penn Discourse Treebank 2.0
    """

    def __init__(self, path):
        """
        the ``path'' is assumed to be the root directory for *.pipe files
        *.pipe files can be generated by the script ``convert.pl'' provided by the dataset authors.
        """
        def load_section(path, relation, section):
            data = []
            dirname = section
            filenames = os.listdir(os.path.join(path, dirname))
            for filename in filenames:
                for line in open(os.path.join(path, dirname, filename)):
                    line = line.decode("latin-1").strip().split("|")
                    if len(line) != len(attribute_names):
                        print("Error! # of attributes is mismatch: %d(actual) != %d(expected)" \
                                % (len(line), len(attribute_names)) )
                        sys.exit(-1)
                    if (relation != "all") and (relation is not None):
                        if line[0] != relation:
                            continue
                    datum = OrderedDict()
                    for key, val in zip(attribute_names, line):
                        datum[key.decode("utf-8")] = val
                    data.append(datum)
            return data
        self.dataset = {}
        for relation in ["all", "Explicit", "Implicit", "AltLex", "EntRel", "NoRel"]:
            for section in xrange(0, 24):
                section = "%02d" % section
                self.dataset[relation + "." + section] = load_section(path, relation, section)
        
    def get(self, relation, section):
        return self.dataset[relation + "." + section]

    def dump(self, path_dir):
        # NOTE: Just do ``data = json.load(open(path))'' when loading.
        for split_rule in sections_dict.keys():
            print("Processing a split rule: %s" % split_rule)
            print("    -train: %s-%s" % (sections_dict[split_rule]["train"][0],
                                        sections_dict[split_rule]["train"][-1]))
            print("    -val: %s-%s" % (sections_dict[split_rule]["val"][0],
                                        sections_dict[split_rule]["val"][-1]))
            print("    -test: %s-%s" % (sections_dict[split_rule]["test"][0],
                                        sections_dict[split_rule]["test"][-1]))
            # create a sub directory
            path_subdir = os.path.join(path_dir,
                "train_%s_%s.val_%s_%s.test_%s_%s" %\
                    (sections_dict[split_rule]["train"][0],
                     sections_dict[split_rule]["train"][-1],
                     sections_dict[split_rule]["val"][0],
                     sections_dict[split_rule]["val"][-1],
                     sections_dict[split_rule]["test"][0],
                     sections_dict[split_rule]["test"][-1])
                    )
            if not os.path.exists(path_subdir):
                os.makedirs(path_subdir)
                print("    Created a directory: %s" % path_subdir)

            for relation in ["all", "Explicit", "Implicit", "AltLex", "EntRel", "NoRel"]:
                for split in ["train", "val", "test"]:
                    # merge
                    data_merged = []
                    for sec in sections_dict[split_rule][split]:
                        data_sec = self.get(relation, sec)
                        data_merged.extend(data_sec)
                    print("    RELATION=%s & SPLIT=%s: %d" % (relation, split, len(data_merged)))
                    # dump a json file
                    with open(os.path.join(path_subdir, "pdtb2.0.%s.%s.json" % (relation, split)), "w") as f:
                        json.dump(data_merged, f, indent=2)
                    # dump text files corresponding to attribute types
                    f_arg1 = open(os.path.join(path_subdir, "pdtb2.0.%s.Arg1_RawText.%s.txt" % \
                            (relation, split)), "w")
                    f_arg2 = open(os.path.join(path_subdir, "pdtb2.0.%s.Arg2_RawText.%s.txt" % \
                            (relation, split)), "w")
                    f_connraw = open(os.path.join(path_subdir, "pdtb2.0.%s.Connective_RawText.%s.txt" % \
                            (relation, split)), "w")
                    f_conninf = open(os.path.join(path_subdir, "pdtb2.0.%s.Conn1.%s.txt" % \
                            (relation, split)), "w")
                    f_conncls = open(os.path.join(path_subdir, "pdtb2.0.%s.ConnHeadSemClass1.%s.txt" % \
                            (relation, split)), "w")
                    for datum in data_merged:
                        f_arg1.write("%s\n" % datum[u"Arg1_RawText"].encode("utf-8"))
                        f_arg2.write("%s\n" % datum[u"Arg2_RawText"].encode("utf-8"))
                        f_connraw.write("%s\n" % datum[u"Connective_RawText"].encode("utf-8"))
                        f_conninf.write("%s\n" % datum[u"Conn1"].encode("utf-8"))
                        f_conncls.write("%s\n" % datum[u"ConnHeadSemClass1"].encode("utf-8"))
                    f_arg1.flush()
                    f_arg2.flush()
                    f_connraw.flush()
                    f_conninf.flush()
                    f_conncls.flush()
                    f_arg1.close()
                    f_arg2.close()
                    f_connraw.close()
                    f_conninf.close()
                    f_conncls.close()
    
