from torch.utils.data import Dataset, DataLoader
from sklearn.feature_extraction.text import HashingVectorizer
import pickle
import pandas as pd

from properties import *


class FeatureGenerator(object):
    SAMPLE_HOME_DIR_PATH = "/workspace/Project/Pattern_Recognition/data/KernelDriver/ProcessId"
    FEATURE_HOME_DIR_PATH = "/workspace/Project/Pattern_Recognition/data/feature"

    FEATURE_DIR_NAME_DICT = dict(
        clean="ProcessIdClean",
        clean_zero="ProcessIdCleanZero",
        clean_pippo="ProcessIdCleanPippo",
        virus_share_500="ProcessIdVirusShare500",
        virus_share_1000="ProcessIdVirusShare1000"
    )

    def __init__(self):
        pass

    @classmethod
    def get_sample_dir_path_list(cls, sample_name):
        sample_home_dir_path = FeatureGenerator.SAMPLE_HOME_DIR_PATH
        feature_dir_name_dict = FeatureGenerator.FEATURE_DIR_NAME_DICT
        entire_list = list()

        if sample_name.split("_")[1] == "entire":
            sample_name_list = None

            if sample_name == "clean_entire":
                sample_name_list = ["clean", "clean_zero", "clean_pippo"]
            elif sample_name == "virus_entire":
                sample_name_list = ["virus_share_500", "virus_share_1000"]

            for sample_name_value in sample_name_list:
                data_dir_path = os.path.join(sample_home_dir_path, feature_dir_name_dict[sample_name_value])
                entire_list += list(map(lambda sample_name: os.path.join(data_dir_path, sample_name), os.listdir(data_dir_path)))

        else:
            data_dir_path = os.path.join(sample_home_dir_path, feature_dir_name_dict[sample_name])
            entire_list += list(map(lambda sample_name: os.path.join(data_dir_path, sample_name), os.listdir(data_dir_path)))

        return entire_list

    @classmethod
    def get_sample_df(cls, sample_dir_path: str):
        def _split_value(row):
            try:
                return row.split("=")[1]
            except Exception as e:
                pass

        df = pd.DataFrame(columns=["Time", "Pid", "MethodName", "ProcessName"])

        try:
            for sample_file_name in os.listdir(sample_dir_path):
                sample_file_path = os.path.join(sample_dir_path, sample_file_name)

                df_sample = pd.read_csv(sample_file_path, sep="\,", names=["Time", "Pid", "MethodName", "ProcessName"])
                df_sample = df_sample.applymap(lambda row: _split_value(row))
                df = df.append(df_sample)
        except Exception as e:
            print(e)

        df = df.dropna()
        df = df.sort_values(by=['Time'])
        df["ProcessName"] = df["ProcessName"].map(lambda row: row.split("\\")[-1])

        return df

    @classmethod
    def train_val_test_split(cls, normal_data_list, abnormal_data_list, train_ratio=0.7, val_ratio=0.1, test_ratio=0.2, is_shuffle=True):
        if is_shuffle:
            random.shuffle(normal_data_list)
            random.shuffle(abnormal_data_list)

        normal_dataset_size = len(normal_data_list)
        abnormal_dataset_size = len(abnormal_data_list)

        # Normal split
        train_normal_dataset_size = int(normal_dataset_size * train_ratio)
        val_normal_dataset_size = int(normal_dataset_size * val_ratio)

        train_normal_dataset = normal_data_list[:train_normal_dataset_size]
        remains_normal_dataset = normal_data_list[train_normal_dataset_size:]
        val_normal_dataset = remains_normal_dataset[:val_normal_dataset_size]
        test_normal_dataset = remains_normal_dataset[val_normal_dataset_size:]

        val_normal_dataset_size = len(val_normal_dataset)
        test_normal_dataset_size = len(test_normal_dataset)

        # Abnormal split
        val_abnormal_dataset = abnormal_data_list[:val_normal_dataset_size]
        remains_normal_dataset = abnormal_data_list[val_normal_dataset_size:]
        test_abnormal_dataset = remains_normal_dataset[:test_normal_dataset_size]

        val_abnormal_dataset_size = len(val_abnormal_dataset)
        test_abnormal_dataset_size = len(test_abnormal_dataset)

        # Generate dataset
        train_X = train_normal_dataset
        val_X = val_normal_dataset + val_abnormal_dataset
        test_X = test_normal_dataset + test_abnormal_dataset

        train_y = np.zeros(shape=(train_normal_dataset_size,))
        val_y = np.concatenate(
            [np.zeros(shape=(val_normal_dataset_size,)), np.ones(shape=(val_abnormal_dataset_size,))]
        )
        test_y = np.concatenate(
            [np.zeros(shape=(test_normal_dataset_size,)), np.ones(shape=(test_abnormal_dataset_size,))]
        )



        print(f"Entire Normal data size : {normal_dataset_size}\n"
              f"Entire Abnormal data size : {abnormal_dataset_size}\n"
              f"===========================================\n"
              f"Train Normal dataset size : {train_normal_dataset_size}\n"
              f"Val Normal | Abnormal dataset size : {val_normal_dataset_size} | {val_abnormal_dataset_size}\n"
              f"Test Normal | Abnormal dataset size : {test_normal_dataset_size} | {test_abnormal_dataset_size}\n")

        return dict(
            train_X=train_X,
            train_y=train_y,
            val_X=val_X,
            val_y=val_y,
            test_X=test_X,
            test_y=test_y
        )


class FeatureSequenceDataset(Dataset):

    FEATURE_HOME_DIR_PATH = "/workspace/Project/Pattern_Recognition/data/feature"

    FEATURE_DIR_NAME_DICT = dict(
        clean="ProcessIdClean",
        clean_zero="ProcessIdCleanZero",
        clean_pippo="ProcessIdCleanPippo",
        virus_share_500="ProcessIdVirusShare500",
        virus_share_1000="ProcessIdVirusShare1000"
    )

    def __init__(self, feature_name, vectorizer, is_preprocessed=True, is_vectorized=False, is_shuffle=True):
        super(FeatureSequenceDataset, self).__init__()
        self.feature_name = feature_name
        self.vectorizer = vectorizer
        self.is_preprocessed = is_preprocessed
        self.is_vectorized = is_vectorized
        self.is_shuffle = is_shuffle

        # Init
        self.init()

    def __getitem__(self, idx):
        feature_file_path = self.feature_file_path_list[idx]
        feature_df_dict = self.load_feature_df_dict(sample_path=feature_file_path)

        vectorized_feature_dict = self.preprocess(feature_df_dict=feature_df_dict)

        return vectorized_feature_dict

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

    @classmethod
    def get_feature_file_path_list(self, feature_name):
        feature_home_dir_path = FeatureSequenceDataset.FEATURE_HOME_DIR_PATH
        feature_dir_name_dict = FeatureSequenceDataset.FEATURE_DIR_NAME_DICT

        data_dir_path = os.path.join(feature_home_dir_path, feature_dir_name_dict[feature_name])

        return list(map(lambda sample_name: os.path.join(data_dir_path, sample_name), os.listdir(data_dir_path)))

    @classmethod
    def load_feature_df_dict(self, sample_path):
        feature_df_dict = None

        with open(sample_path, "rb") as f:
            feature_df_dict = pickle.load(f)

        return feature_df_dict

    def preprocess(self, feature_df_dict):
        vectorized_feature_dict = dict()

        for feature_name, feature_df in feature_df_dict.items():
            feature_df = feature_df[5:]
            vectorized_feature_dict[feature_name] = dict()

            for time_step_name in feature_df.columns:
                vectorized_feature_dict[feature_name][time_step_name] = self.vectorizer.fit_transform(
                    feature_df[time_step_name]).toarray()

        return vectorized_feature_dict

    def init(self):
        if self.feature_name.split("_")[1] == "entire":
            entire_list = list()
            feature_name_list = None

            if self.feature_name == "clean_entire":
                feature_name_list = ["clean", "clean_zero", "clean_pippo"]
            elif self.feature_name == "virus_entire":
                feature_name_list = ["virus_share_500", "virus_share_1000"]

            for feature_name in feature_name_list:
                entire_list += FeatureSequenceDataset.get_feature_file_path_list(feature_name=feature_name)

            self.feature_file_path_list = entire_list
        else:
            self.feature_file_path_list = FeatureSequenceDataset.get_feature_file_path_list(
                feature_name=self.feature_name)

        if self.is_shuffle:
            random.shuffle(self.feature_file_path_list)


class FeatureSequenceDatasetController(object):

    def __init__(self, dataset):
        self.dataset = dataset

    def __getitem__(self, idx):
        #         return self.dataset[idx]
        return DataLoader(
            dataset=self.dataset[idx],
            num_workers=1,
            pin_memory=True,
            batch_size=32
        )

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


class SampleSequenceDataset(Dataset):

    def __init__(self, data_seq, vectorizer=HashingVectorizer(), window_size=32):
        super(SampleSequenceDataset, self).__init__()
        self.data_seq = data_seq
        self.vectorizer = vectorizer
        self.window_size = window_size

    def __getitem__(self, idx):
        x = self.data_seq[idx:idx + self.window_size]
        x = self.vectorizer.fit_transform(x).toarray()
        x = torch.from_numpy(x)

        return x

    def __len__(self):
        return len(self.data_seq) - self.window_size


class SampleSequenceConcatDataset(Dataset):

    def __init__(self, data_seq, vectorizer=HashingVectorizer(), window_size=32, concat_column_list=["Pid", "MethodName", "ProcessName"], is_normal=True):
        super(SampleSequenceConcatDataset, self).__init__()
        self.data_seq = data_seq
        self.vectorizer = vectorizer
        self.window_size = window_size
        self.concat_column_list = concat_column_list
        self.is_normal = is_normal

    def __getitem__(self, idx):
        x = self.data_seq[idx:idx + self.window_size]
        x = torch.from_numpy(x)
        y = torch.Tensor([0]) if self.is_normal else torch.Tensor([1])

        x = x.float()
        y = y.float()

        return x, y

    def __len__(self):
        return len(self.data_seq) - self.window_size


class SampleSequenceDatasetController(object):

    def __init__(self, sample_dir_path_list, y_list, dataset_hparams, dataloader_hparams, is_shuffle=True):
        self.sample_dir_path_list = sample_dir_path_list
        self.y_list = y_list
        self.dataset_hparams = dataset_hparams
        self.dataloader_hparams = dataloader_hparams

        if is_shuffle:
            self._shuffle_data_list()
    
    @property
    def random_sample_idx_list(self):
        sample_idx_list = [i for i in range(len(self.sample_dir_path_list))]
        random.shuffle(sample_idx_list)

        return sample_idx_list

    def _shuffle_data_list(self):
        s = np.arange(self.y_list.shape[0])
        np.random.shuffle(s)

        self.sample_dir_path_list = np.array(self.sample_dir_path_list)[s]
        self.sample_dir_path_list = self.sample_dir_path_list.tolist()
        self.y_list = self.y_list[s]

    def get_df(self, sample_idx):
        df = FeatureGenerator.get_sample_df(sample_dir_path=self.sample_dir_path_list[sample_idx])
        df = df.reset_index()

        return df

    def get_dataset(self, sample_idx, dataset_hparams=None):
        dataset_hparams = dataset_hparams if dataset_hparams else self.dataset_hparams
        df = self.get_df(sample_idx=sample_idx)

        return dict(
            pid_dataset=SampleSequenceDataset(data_seq=df["Pid"], **dataset_hparams),
            method_name_dataset=SampleSequenceDataset(data_seq=df["MethodName"], **dataset_hparams),
            process_name_dataset=SampleSequenceDataset(data_seq=df["ProcessName"], **dataset_hparams)
        )

    def get_dataloader(self, sample_idx, dataset_hparams=None, dataloader_hparams=None):
        dataset_hparams = dataset_hparams if dataset_hparams else self.dataset_hparams
        dataloader_hparams = dataloader_hparams if dataloader_hparams else self.dataloader_hparams
        dataset_dict = self.get_dataset(sample_idx=sample_idx, dataset_hparams=dataset_hparams)

        return dict(
            pid_dataloader=DataLoader(dataset_dict["pid_dataset"], **dataloader_hparams),
            method_name_dataloader=DataLoader(dataset_dict["method_name_dataset"], **dataloader_hparams),
            process_name_dataloader=DataLoader(dataset_dict["process_name_dataset"], **dataloader_hparams)
        )

    def get_concat_dataset(self, sample_idx, dataset_hparams=None):
        dataset_hparams = dataset_hparams if dataset_hparams else self.dataset_hparams
        vectorizer = dataset_hparams.vectorizer
        df = self.get_df(sample_idx=sample_idx)
        is_normal = True if self.y_list[sample_idx] == 0 else False

        vectorized_list = list()

        for column_name in dataset_hparams.concat_column_list:
            series = df[column_name].astype(str).to_numpy()
            vectorized_list.append(vectorizer.fit_transform(series).toarray())

        data_seq = np.sum(vectorized_list, axis=0)

        return SampleSequenceConcatDataset(data_seq=data_seq, is_normal=is_normal, **dataset_hparams)

    def get_concat_dataloader(self, sample_idx, dataset_hparams=None, dataloader_hparams=None):
        dataset_hparams = dataset_hparams if dataset_hparams else self.dataset_hparams
        dataloader_hparams = dataloader_hparams if dataloader_hparams else self.dataloader_hparams
        dataset = self.get_concat_dataset(sample_idx=sample_idx, dataset_hparams=dataset_hparams)

        return DataLoader(dataset, **dataloader_hparams)

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