import copy
import os
from functools import lru_cache
from collections import OrderedDict

import torch
from portal.portal_io import cache_file

from mmf.common import Sample
from mmf.datasets import BaseDataset


class TextCommonClsDataset(BaseDataset):
    def __init__(self, config, dataset_type, dataset_name, data_folder=None, *args, **kwargs):
        super().__init__(dataset_name, config, dataset_type, *args, **kwargs)
        self._data = []
        self._cols = {}
        self._ext_cols = []
        if hasattr(config, "ext_cols"):
            self._ext_cols = config.ext_cols

        if hasattr(config, "num_cols"):
            self._num_cols = config.num_cols

        data_path = getattr(config.data_path, dataset_type)
        for file in data_path:
            if file.startswith("oss://"):
                real_data_path = cache_file(file, root=config.data_dir)
            else:
                real_data_path = os.path.join(config.data_dir, file)

            with open(real_data_path) as f:
                header = f.readline().strip().split("\t")
                for idx, col in enumerate(header):
                    self._cols[col] = idx
                for line in f:
                    self._data.append(line.strip().split("\t"))

    def __len__(self):
        return len(self._data)

    def __getitem__(self, idx):
        return self.load_item(idx)

    @lru_cache(maxsize=32)
    def load_item(self, idx):
        sample = Sample()
        item = dict(zip(self._cols, self._data[idx]))
        query_item = self.text_processor(item)
        label = self.label_processor(item)["label"]
        sample.tokens = query_item["text"]
        sample.mask = torch.ne(sample.tokens, self.text_processor.processor.PAD_INDEX)
        sample.len = torch.IntTensor([query_item["length"]])
        sample.targets = torch.tensor(label, dtype=torch.int64)

        for num_col in self._num_cols:
            sample[num_col] = torch.FloatTensor([float(x) for x in item[num_col].split(" ")])

        for col in self._ext_cols:
            if col in item:
                setattr(sample, col, item[col])
            else:
                setattr(sample, col, None)
        return sample

    def format_for_prediction(self, report):
        scores_list = torch.softmax(report.scores, dim=1).cpu().numpy().tolist()
        answers = report.scores.argmax(dim=1)

        predictions = answers.cpu().numpy().tolist()
        result = OrderedDict()

        for col in self._ext_cols:
            if col in report:
                result[col] = report[col]

        if "logits" in report:
            result["logits"] = report.logits.cpu().numpy().tolist()

        result["predict"] = predictions
        result["scores"] = scores_list

        current_cols = copy.deepcopy(self._ext_cols)
        current_cols.append("predict")
        current_cols.append("scores")
        current_cols.append("logits")

        items = []
        for idx in range(len(result["predict"])):
            item = {
                key: result[key][idx]
                for key in current_cols
            }

            items.append(item)

        return items


class TextBertClsDataset(TextCommonClsDataset):
    def __init__(self, config, dataset_type, dataset_name, data_folder=None, *args, **kwargs):
        super().__init__(config, dataset_type, dataset_name, data_folder, *args, **kwargs)

    def __getitem__(self, idx):
        return self.load_item(idx)

    @lru_cache(maxsize=32)
    def load_item(self, idx):
        sample = Sample()
        item = dict(zip(self._cols, self._data[idx]))
        query_item = self.text_processor(item)
        label = self.label_processor(item)["label"]
        sample.input_ids = query_item["input_ids"]
        sample.input_mask = query_item["input_mask"]
        sample.token_type_ids = query_item["segment_ids"]
        sample.input_len = torch.IntTensor([len(query_item["text"])])
        sample.targets = torch.tensor(label, dtype=torch.int64)

        for col in self._ext_cols:
            setattr(sample, col, item[col])

        return sample

import copy
import os
from functools import lru_cache
from collections import OrderedDict

import torch
from portal.portal_io import cache_file

from mmf.common import Sample
from mmf.datasets import BaseDataset


class TextCommonClsDataset(BaseDataset):
    def __init__(self, config, dataset_type, dataset_name, data_folder=None, *args, **kwargs):
        super().__init__(dataset_name, config, dataset_type, *args, **kwargs)
        self._data = []
        self._cols = {}
        self._ext_cols = []
        if hasattr(config, "ext_cols"):
            self._ext_cols = config.ext_cols

        if hasattr(config, "num_cols"):
            self._num_cols = config.num_cols

        data_path = getattr(config.data_path, dataset_type)
        for file in data_path:
            if file.startswith("oss://"):
                real_data_path = cache_file(file, root=config.data_dir)
            else:
                real_data_path = os.path.join(config.data_dir, file)

            with open(real_data_path) as f:
                header = f.readline().strip().split("\t")
                for idx, col in enumerate(header):
                    self._cols[col] = idx
                for line in f:
                    self._data.append(line.strip().split("\t"))

    def __len__(self):
        return len(self._data)

    def __getitem__(self, idx):
        return self.load_item(idx)

    @lru_cache(maxsize=32)
    def load_item(self, idx):
        sample = Sample()
        item = dict(zip(self._cols, self._data[idx]))
        query_item = self.text_processor(item)
        label = self.label_processor(item)["label"]
        sample.tokens = query_item["text"]
        sample.mask = torch.ne(sample.tokens, self.text_processor.processor.PAD_INDEX)
        sample.len = torch.IntTensor([query_item["length"]])
        sample.targets = torch.tensor(label, dtype=torch.int64)

        for num_col in self._num_cols:
            sample[num_col] = torch.FloatTensor([float(x) for x in item[num_col].split(" ")])

        for col in self._ext_cols:
            if col in item:
                setattr(sample, col, item[col])
            else:
                setattr(sample, col, None)
        return sample

    def format_for_prediction(self, report):
        scores_list = torch.softmax(report.scores, dim=1).cpu().numpy().tolist()
        answers = report.scores.argmax(dim=1)

        predictions = answers.cpu().numpy().tolist()
        result = OrderedDict()

        for col in self._ext_cols:
            if col in report:
                result[col] = report[col]

        if "logits" in report:
            result["logits"] = report.logits.cpu().numpy().tolist()

        result["predict"] = predictions
        result["scores"] = scores_list

        current_cols = copy.deepcopy(self._ext_cols)
        current_cols.append("predict")
        current_cols.append("scores")
        current_cols.append("logits")

        items = []
        for idx in range(len(result["predict"])):
            item = {
                key: result[key][idx]
                for key in current_cols
            }

            items.append(item)

        return items


class TextBertClsDataset(TextCommonClsDataset):
    def __init__(self, config, dataset_type, dataset_name, data_folder=None, *args, **kwargs):
        super().__init__(config, dataset_type, dataset_name, data_folder, *args, **kwargs)

    def __getitem__(self, idx):
        return self.load_item(idx)

    @lru_cache(maxsize=32)
    def load_item(self, idx):
        sample = Sample()
        item = dict(zip(self._cols, self._data[idx]))
        query_item = self.text_processor(item)
        label = self.label_processor(item)["label"]
        sample.input_ids = query_item["input_ids"]
        sample.input_mask = query_item["input_mask"]
        sample.token_type_ids = query_item["segment_ids"]
        sample.input_len = torch.IntTensor([len(query_item["text"])])
        sample.targets = torch.tensor(label, dtype=torch.int64)

        for col in self._ext_cols:
            setattr(sample, col, item[col])

        return sample
