import re
from typing import Dict, List

import numpy as np
import numpy.typing as npt
from sklearn.model_selection import train_test_split

from mlc.constraints.relation_constraint import (
    BaseRelationConstraint,
    Constant,
    Count,
    Feature,
    Log,
    ManySum,
)
from mlc.datasets.dataset import (
    CsvDataSource,
    Dataset,
    DataSource,
    DefaultIndexSorter,
    DownloadFileDataSource,
    Hdf5DataSource,
    Splitter,
    Task,
)


class MalwareSplitter(Splitter):
    def get_splits(self, dataset: Dataset) -> Dict[str, npt.NDArray[np.int_]]:
        _, y = dataset.get_x_y()
        i = np.arange(len(y))
        i_train, i_test = i[:12308], i[12308:]
        i_val, i_test = train_test_split(
            i_test, test_size=0.5, random_state=42
        )
        return {"train": i_train, "val": i_val, "test": i_test}


def get_relation_constraints(
    metadata: DataSource,
) -> List[BaseRelationConstraint]:

    features = metadata.load_data()["feature"].to_list()

    def regex_features(features_l: List[str], pattern: str) -> List[Feature]:
        p = re.compile(pattern)
        out = list(filter(lambda el: p.match(el), features_l))
        return [Feature(o) for o in out]

    features_g1 = regex_features(features, "^pesection_\d+_name$")
    g1_right = [f == Constant(832) for f in features_g1]
    g1 = Feature("header_NumberOfSections") == Count(g1_right, inverse=True)

    g2 = Feature("header_FileAlignment") <= Feature("header_SectionAlignment")

    g3 = Log(Feature("header_FileAlignment")) / Log(Constant(2.0))
    g3 = g3 % Constant(1.0)
    g3 = g3 == Constant(0.0)

    g4 = ManySum(regex_features(features, "^imp_.*$")) <= Feature(
        "api_import_nb"
    )

    g5 = ManySum(regex_features(features, "^.*\.dll$")) <= Feature(
        "dll_import_nb"
    )

    g6 = ManySum(regex_features(features, "^freq_byte_\d+$")) == Constant(1)

    freq = regex_features(features, "^freq_byte_\d+$")
    log_freq = [
        (f * (Log(f, safe_value=Constant(0.0)) / Log(Constant(2.0))))
        for f in freq
    ]
    g7 = Feature("generic_fileEntropy") == (Constant(0) - ManySum(log_freq))

    return [g1, g2, g3, g4, g5, g6, g7]


def create_dataset() -> Dataset:
    data_source = DownloadFileDataSource(
        url="https://uniluxembourg-my.sharepoint.com/:u:/g/personal/"
        "thibault_simonetto_uni_lu/"
        "EWpXWGxq_WxBtCJ5o1mr0fcBGcJYFrYchZD8esq-91yfJw?download=1",
        file_data_source=Hdf5DataSource(
            path="./data/mlc/malware/malware.hdf5"
        ),
    )
    metadata_source = DownloadFileDataSource(
        url="https://uniluxembourg-my.sharepoint.com/:x:/g/personal/"
        "thibault_simonetto_uni_lu/"
        "Ef-mGuyoTZtBqnUe6JA9HK4BiqyxHrVYgO7ssJ0Qe6T_AQ?download=1",
        file_data_source=CsvDataSource(
            path="./data/mlc/malware/malware_metadata.csv"
        ),
    )
    tasks = [
        Task(
            name="malicious",
            task_type="classification",
            evaluation_metric="mcc",
        )
    ]
    sorter = DefaultIndexSorter()
    splitter = MalwareSplitter()
    relation_constraints = get_relation_constraints(metadata_source)

    malware = Dataset(
        name="malware",
        data_source=data_source,
        metadata_source=metadata_source,
        tasks=tasks,
        sorter=sorter,
        splitter=splitter,
        relation_constraints=relation_constraints,
    )
    return malware


datasets = [
    {
        "name": "malware",
        "fun_create": create_dataset,
    }
]
