import numpy as np
import pandas as pd
from src.gena_lm.utils import (concatenate_encodings,
                               get_service_token_encodings,
                               symmetric_pad_and_truncate_context)
from torch.utils.data import Dataset


class DeepSeaDataset(Dataset):
    def __init__(self, datapath: str, tokenizer, max_seq_len: int = 512):
        self.data = pd.read_csv(datapath, sep=",", header=None, dtype={i: np.int8 for i in range(1, 920)})
        self.data["targets"] = list(self.data.iloc[:, 1:].values)
        self.data = self.data[[0, "targets"]]
        self.data.columns = ["seq", "targets"]

        self.max_seq_len = max_seq_len
        self.tokenizer = tokenizer

        self.service_token_encodings = get_service_token_encodings(self.tokenizer)

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

    def __getitem__(self, idx):
        sample = self.data["seq"][idx]
        targets = self.data["targets"][idx]

        features = self.tokenizer(
            sample, add_special_tokens=False, padding=False, truncation=False, return_tensors="np"
        )
        _, _, features, padding = symmetric_pad_and_truncate_context(
            left_encoding=None,
            right_encoding=None,
            mid_encoding=features,
            n_service_tokens=2,
            max_seq_len=self.max_seq_len,
            PAD_id=self.tokenizer.pad_token_id,
        )
        features = concatenate_encodings(
            [
                self.service_token_encodings["CLS"],
                features,
                self.service_token_encodings["SEP"],
                padding,
            ]
        )

        # return features generated by tokenizer and labels
        for fn in features:
            features[fn] = features[fn][0]
        features["labels"] = targets.astype(np.float32)
        return features
