import os
import random
from collections import defaultdict
from typing import List, Dict

import torch
import pandas as pd
from tqdm import tqdm
import xml.etree.ElementTree as ET

from LMTrace.models import MODEL_ARCH_SINGLE, MODEL_ARCH_SIAMESE
from LMTrace.utils import exclude_and_sample

TOKEN = "token"

OID = "oid"
SOID = "s_oid"
TOID = "t_oid"

ID = "id"
SID = 'sid'
TID = 'tid'

TEXT = "text"
STEXT = "s_text"
TTEXT = "t_text"

EMBD = "embd"

INPUT_ID = "input_ids"
TK_TYPE = "token_type_ids"
ATTEN_MASK = "attention_mask"

DEFAULT_SOURCE_FILE = "source_art.csv"
DEFAULT_TARGET_FILE = "target_art.csv"
DEFAULT_LINK_FILE = "link.csv"


class TraceExamples:
    """
    Manage the examples read from raw dataset

    examples:
    valid_examples = CodeSearchNetReader(data_dir).get_examples(type="valid", num_limit=valid_num, summary_only=True)
    valid_examples = Examples(valid_examples)
    valid_examples.update_features(model)
    valid_examples.update_embd(model)
    """

    def __init__(self, s_arts, t_arts, links):
        """
        Index the raw examples with numeric ids (sid and tid) and The origin id is named as s_oid,t_oid.
        :param raw_examples: A list of dictionary with keys: s_oid, s_text, t_oid, t_text
        """

        self.s_index, self.t_index = dict(), dict()  # artifact index (id->text)
        self.rs_index, self.rt_index = dict(), dict()  # reversed artifact index (oid->id)
        self.s2t, self.t2s = defaultdict(set), defaultdict(set)  # true links
        self.sid_cnt, self.tid_cnt, self.lk_cnt = 0, 0, 0

        for i, s_oid in enumerate(s_arts):
            self.s_index[i] = {SOID: s_oid, TOKEN: s_arts[s_oid]}
            self.rs_index[s_oid] = i
        for i, t_oid in enumerate(t_arts):
            self.t_index[i] = {TOID: t_oid, TOKEN: t_arts[t_oid]}
            self.rt_index[t_oid] = i

        for lk in links:
            s_oid, t_oid = lk
            if s_oid not in self.rs_index or t_oid not in self.rt_index:
                continue
            sid, tid = self.rs_index[s_oid], self.rt_index[t_oid]
            if tid not in self.s2t[sid]:
                self.lk_cnt += 1

            self.s2t[sid].add(tid)
            self.t2s[tid].add(sid)

    def __len__(self):
        return self.lk_cnt

    def build_feature_entry(self, sid, tid, label, tokenizer, model_arch, max_seq_length=256):
        s_tks = self.s_index[sid][TOKEN]
        t_tks = self.t_index[tid][TOKEN]
        pair_feature = tokenizer(
            text=s_tks,
            text_pair=t_tks,
            return_attention_mask=True,
            return_token_type_ids=True,
            add_special_tokens=True,
            padding="max_length",
            max_length=max_seq_length,
            truncation="longest_first",
        )

        entry = {
            SID: sid,
            TID: tid,
            'label': label
        }

        if model_arch == MODEL_ARCH_SINGLE:
            entry['input_ids'] = pair_feature[INPUT_ID]
            entry['attention_mask'] = pair_feature[ATTEN_MASK]
            entry['token_type_ids'] = pair_feature[TK_TYPE]
        elif model_arch == MODEL_ARCH_SIAMESE:
            entry['s_input_ids'] = self.s_index[sid][INPUT_ID]
            entry['s_attention_mask'] = self.s_index[sid][ATTEN_MASK]
            entry['t_input_ids'] = self.t_index[tid][INPUT_ID]
            entry['t_attention_mask'] = self.t_index[tid][ATTEN_MASK]
        else:
            raise Exception("")
        return entry

    def gen_trace_dataset(self, tokenizer, model_arch, max_seq_length, query_num=-1, with_label=True):
        """
        Generate trace challenge based on true links.
        :param tokenizer:
        :param model_arch:
        :return:
        """
        dataset = []
        for sid in self.s_index:
            if query_num > 0 and len(dataset) > query_num:
                break
            for tid in self.t_index:
                if with_label:
                    label = 1 if tid in self.s2t[sid] else 0
                else:
                    label = None
                dataset.append(
                    self.build_feature_entry(sid, tid, label, tokenizer, model_arch, max_seq_length=max_seq_length))
        return dataset

    def gen_training_dataset(self, tokenizer, model_arch, max_seq_length=256, resample_rate=1):
        """
        Generate training examples with random negative examples.
        The dataset is a balanced dataset with equal size of positive and negative examples.
        Create features for both single and siamese models.
        :return:
        """

        for index in [self.s_index, self.t_index]:
            for id in index:
                feature = tokenizer(text=index[id][TOKEN],
                                    return_attention_mask=True,
                                    return_token_type_ids=False,
                                    truncation="longest_first",
                                    max_length=max_seq_length,
                                    padding="max_length")
                index[id][INPUT_ID] = feature[INPUT_ID]
                index[id][ATTEN_MASK] = feature[ATTEN_MASK]

        dataset = []

        for sid in tqdm(self.s2t, desc="random_neg_sampling_dataset"):
            pos_tids = self.s2t[sid]
            for pos_tid in pos_tids:
                for i in range(resample_rate):
                    dataset.append(self.build_feature_entry(sid, pos_tid, 1, tokenizer, model_arch))
            neg_tids = exclude_and_sample(set(self.t_index.keys()), pos_tids, num=len(pos_tids) * resample_rate)
            for n_tid in neg_tids:
                dataset.append(self.build_feature_entry(sid, n_tid, 0, tokenizer, model_arch))
        return dataset

    def update_embd(self, model):
        """
        update the artifact embeddings for Siamese-LM model
        :param model: Siamese model
        :return:
        """
        with torch.no_grad():
            model.eval()
            for index in [self.s_index, self.t_index]:
                for id in tqdm(index, desc="update embedding"):
                    feature = index[id]
                    input_tensor = torch.tensor(feature[INPUT_ID]).view(1, -1).to(model.device)
                    mask_tensor = torch.tensor(feature[ATTEN_MASK]).view(1, -1).to(model.device)
                    embd = model.get_LM(input_tensor, mask_tensor)[0]
                    embd_cpu = embd.to('cpu')
                    index[id][EMBD] = embd_cpu

    def split_dataset(self, out_dir, weigths={'train': 8, 'valid': 1, 'test': 1}):
        """
        Split dataset into train, validation, test and write them into CSV format.
        Each source artifact are treated as a single query and the weights are applied on the source artifacts.


        :param trace_examples:
        :param out_dir:
        :return:
        """

        def write_arts(art_dict, file_path):
            df = pd.DataFrame(columns=[ID, TEXT])
            df[ID] = art_dict.keys()
            df[TEXT] = art_dict.values()
            df.to_csv(file_path)

        queries = list(self.s2t.keys())
        random.shuffle(queries)
        q_iter = iter(queries)
        total_w = sum(weigths.values())

        for slice_name in weigths:
            slice_dir = os.path.join(out_dir, slice_name)
            if not os.path.isdir(slice_dir):
                os.makedirs(slice_dir)
            n = self.lk_cnt * (weigths[slice_name] / total_w)
            cnt = 0
            sarts, tarts = dict(), dict()
            lk_s_col, lk_t_col = [], []
            while cnt < n:
                sid = next(q_iter, None)
                if not sid:
                    break
                for tid in self.s2t[sid]:
                    sart, tart = self.s_index[sid], self.t_index[tid]
                    sarts[sart[SOID]] = sart[TOKEN]
                    tarts[tart[TOID]] = tart[TOKEN]
                    lk_s_col.append(sart[SOID])
                    lk_t_col.append(tart[TOID])
                    cnt += 1

            write_arts(sarts, os.path.join(slice_dir, DEFAULT_SOURCE_FILE))
            write_arts(tarts, os.path.join(slice_dir, DEFAULT_TARGET_FILE))
            lk_df = pd.DataFrame(columns=[SID, TID])
            lk_df[SID] = lk_s_col
            lk_df[TID] = lk_t_col
            lk_df.to_csv(os.path.join(slice_dir, DEFAULT_LINK_FILE))


class DataReader:
    def __init__(self, name, ):
        self.name = name

    def get_examples(self) -> TraceExamples:
        raise NotImplementedError


class CSVReader(DataReader):
    """
    dataset reader for csv file. A datatset contains three files:
    1. source_art.csv: contain 'id' and 'text' column
    2. target_art.csv: contain 'id' and 'text' column
    3. link.csv: contain 'sid' and 'tid' column
    """

    def __init__(self, data_dir, dataset_name='CSVReader', sart=DEFAULT_SOURCE_FILE, tart=DEFAULT_TARGET_FILE,
                 lk=DEFAULT_LINK_FILE):
        super().__init__(dataset_name)
        self.data_dir = data_dir
        self.sart = os.path.join(data_dir, sart)
        self.tart = os.path.join(data_dir, tart)
        self.lk = os.path.join(data_dir, lk)

    def get_examples(self):
        sart_df = pd.read_csv(self.sart)
        tart_df = pd.read_csv(self.tart)
        lk_df = pd.read_csv(self.lk)
        parse_art = lambda data_df: {row[ID]: row[TEXT] for index, row in data_df.iterrows()}
        s_index, t_index = parse_art(sart_df), parse_art(tart_df)
        links = [(row[SID], row[TID]) for index, row in lk_df.iterrows()]
        return TraceExamples(s_index, t_index, links)


class PACISReader(DataReader):
    def __init__(self, dir_path):
        super().__init__("PACIS")
        self.dir_path = dir_path
        self.source = os.path.join(dir_path, "SSDFinal.xml")
        self.target = os.path.join(dir_path, "SRS.xml")
        self.link = os.path.join(dir_path, "TraceMatrixSSD2SRS.txt")

    def get_examples(self) -> TraceExamples:
        sarts = self.read_artifact(self.source)
        tarts = self.read_artifact(self.target)
        lk_df = pd.read_csv(self.link)
        links = [(row[0], row[1]) for index, row in lk_df.iterrows()]
        return TraceExamples(sarts, tarts, links)

    def read_artifact(self, path):
        arts = dict()
        root = ET.parse(path).getroot()
        for art in root.findall('artifact'):
            id = art.find('art_id').text
            title = art.find('art_title').text
            content = art.find('art_content').text
            title = title if title else ""
            content = content if content else ""
            arts[id] = title + "\n" + content
        return arts


if __name__ == "__main__":
    rd = PACISReader('../data/PACIS')
    examples = rd.get_examples()
    print(len(examples))
