#
# (c) FFRI Security, Inc., 2020-2021 / Author: FFRI Security, Inc.
#

from enum import Enum
from statistics import mean
from typing import Any, DefaultDict, Dict, List, Optional, Tuple, Union

import lief
import numpy as np
from bs4 import BeautifulSoup

from .exceptions import NotSupported
from .feature_extractor import FeatureExtractor
from .utils import (
    make_defaultdict_from_dict_elem,
    make_onehot_dict_from_bitflag,
    make_onehot_from_str_keys,
    vectorize_selected_features,
    vectorize_with_feature_hasher,
    ver_str_to_int,
)

class DosHeaderFeatureExtractor(FeatureExtractor):
    feature_name = "dos_header"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

class RichHeaderFeatureExtractor(FeatureExtractor):
    feature_name = "rich_header"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    def extract_raw_features(self, raw_json: dict) -> dict:
        rich_header = make_defaultdict_from_dict_elem(
            raw_json, self.feature_name
        )

        if rich_header["entries"]:
            # list of (comp.id, count)
            entries: Optional[List[Tuple[str, int]]] = [
                (
                    f'{entry["id"]:04x}{entry["build_id"]:04x}',
                    entry["count"],
                )
                for entry in rich_header["entries"]
            ]
        else:
            entries = None

        return {
            "key": rich_header["key"],
            "entries": entries,
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = ["key", "entries"]
        post_process_funcs = {
            "entries": lambda x: vectorize_with_feature_hasher(x, 50),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class HeaderFeatureExtractor(FeatureExtractor):
    feature_name = "header"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def machine_to_onehot(machine: str) -> Dict[str, int]:
        encoded_data = {
            "INVALID": 0,
            "UNKNOWN": 0,
            "AMD64": 0,
            "ARM": 0,
            "ARM64": 0,
            "I386": 0,
            "OTHER": 0,
        }
        if machine in encoded_data.keys():
            encoded_data[machine] = 1
        else:
            encoded_data["OTHER"] = 1
        return encoded_data

    @staticmethod
    def characteristics_to_onehot(chracteristics: int) -> Dict[str, int]:
        return make_onehot_dict_from_bitflag(
            [
                "RELOCS_STRIPPED",
                "EXECUTABLE_IMAGE",
                "LINE_NUMS_STRIPPED",
                "LOCAL_SYMS_STRIPPED",
                "AGGRESSIVE_WS_TRIM",
                "LARGE_ADDRESS_AWARE",
                "BYTES_REVERSED_LO",
                "CHARA_32BIT_MACHINE",
                "DEBUG_STRIPPED",
                "REMOVABLE_RUN_FROM_SWAP",
                "NET_RUN_FROM_SWAP",
                "SYSTEM",
                "DLL",
                "UP_SYSTEM_ONLY",
                "BYTES_REVERSED_HI",
            ],
            chracteristics,
            lief.PE.HEADER_CHARACTERISTICS,
        )

    def extract_raw_features(
        self, raw_json: dict
    ) -> Dict[str, Union[int, Dict[str, int]]]:
        header = raw_json[self.feature_name]
        return {
            "signature": header["signature"],
            "machine": self.machine_to_onehot(header["machine"]),
            "numberof_sections": header["numberof_sections"],
            "time_date_stamp": header["time_date_stamp"],
            "numberof_symbols": header["numberof_symbols"],
            "pointerto_symbol_table": header["pointerto_symbol_table"],
            "sizeof_optional_header": header["sizeof_optional_header"],
            "characteristics": self.characteristics_to_onehot(
                header["characteristics"]
            ),
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "signature",
            "machine",
            "numberof_sections",
            "time_date_stamp",
            "numberof_symbols",
            "time_date_stamp",
            "numberof_symbols",
            "pointerto_symbol_table",
            "sizeof_optional_header",
            "characteristics",
        ]
        post_process_funcs = {
            "machine": lambda x: list(x.values()),
            "characteristics": lambda x: list(x.values()),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

# NOTE: https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#dll-characteristics
class DllCharacteristics(Enum):
    IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA = 0x0020
    IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE = 0x0040
    IMAGE_DLLCHARACTERISTICS_FORCE_INTEGRITY = 0x0080
    IMAGE_DLLCHARACTERISTICS_NX_COMPAT = 0x0100
    IMAGE_DLLCHARACTERISTICS_NO_ISOLATION = 0x0200
    IMAGE_DLLCHARACTERISTICS_NO_SEH = 0x0400
    IMAGE_DLLCHARACTERISTICS_NO_BIND = 0x0800
    IMAGE_DLLCHARACTERISTICS_APPCONTAINER = 0x1000
    IMAGE_DLLCHARACTERISTICS_WDM_DRIVER = 0x2000
    IMAGE_DLLCHARACTERISTICS_GUARD_CF = 0x4000
    IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE = 0x8000

class OptionalHeaderFeatureExtractor(FeatureExtractor):
    feature_name = "optional_header"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def subsystem_to_onehot(subsystem: str) -> Dict[str, int]:
        return make_onehot_from_str_keys(
            [
                "UNKNOWN",
                "NATIVE",
                "WINDOWS_GUI",
                "WINDOWS_CUI",
                "OS2_CUI",
                "POSIX_CUI",
                "NATIVE_WINDOWS",
                "WINDOWS_CE_GUI",
                "EFI_APPLICATION",
                "EFI_BOOT_SERVICE_DRIVER",
                "EFI_RUNTIME_DRIVER",
                "EFI_ROM",
                "XBOX",
                "WINDOWS_BOOT_APPLICATION",
                "Out of range",
            ],
            subsystem,
        )

    @staticmethod
    def dll_characteristics_to_onehot(
        dll_characteristics: int,
    ) -> Dict[str, int]:
        encoded_data = {
            str(k): int((dll_characteristics & k.value) != 0)
            for k in DllCharacteristics
        }
        return encoded_data

    def extract_raw_features(self, raw_json: dict) -> dict:
        optional_header = make_defaultdict_from_dict_elem(
            raw_json, self.feature_name
        )
        return {
            "magic": int(optional_header["magic"] == "PE32"),
            "major_linker_version": optional_header["major_linker_version"],
            "minor_linker_version": optional_header["minor_linker_version"],
            "sizeof_code": optional_header["sizeof_code"],
            "sizeof_initialized_data": optional_header[
                "sizeof_initialized_data"
            ],
            "sizeof_uninitialized_data": optional_header[
                "sizeof_uninitialized_data"
            ],
            "addressof_entrypoint": optional_header["addressof_entrypoint"],
            "baseof_code": optional_header["baseof_code"],
            "baseof_data": optional_header["baseof_data"],
            "imagebase": optional_header["imagebase"],
            "section_alignment": optional_header["section_alignment"],
            "file_alignment": optional_header["file_alignment"],
            "major_operating_system_version": optional_header[
                "major_operating_system_version"
            ],
            "minor_operating_system_version": optional_header[
                "minor_operating_system_version"
            ],
            "major_image_version": optional_header["major_image_version"],
            "minor_image_version": optional_header["minor_image_version"],
            "major_subsystem_version": optional_header[
                "major_subsystem_version"
            ],
            "minor_subsystem_version": optional_header[
                "minor_subsystem_version"
            ],
            "win32_version_value": optional_header["win32_version_value"],
            "sizeof_image": optional_header["sizeof_image"],
            "sizeof_headers": optional_header["sizeof_headers"],
            "checksum": optional_header["checksum"],
            "sizeof_stack_reserve": optional_header["sizeof_stack_reserve"],
            "sizeof_stack_commit": optional_header["sizeof_stack_commit"],
            "sizeof_heap_reserve": optional_header["sizeof_heap_reserve"],
            "sizeof_heap_commit": optional_header["sizeof_heap_commit"],
            "loader_flags": optional_header["loader_flags"],
            "numberof_rva_and_size": optional_header["numberof_rva_and_size"],
            "subsystem": self.subsystem_to_onehot(optional_header["subsystem"]),
            "dll_characteristics": self.dll_characteristics_to_onehot(
                optional_header["dll_characteristics"]
            ),
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "major_linker_version",
            "minor_linker_version",
            "sizeof_code",
            "sizeof_initialized_data",
            "sizeof_uninitialized_data",
            "addressof_entrypoint",
            "baseof_code",
            "baseof_data",
            "imagebase",
            "section_alignment",
            "file_alignment",
            "major_operating_system_version",
            "minor_operating_system_version",
            "major_image_version",
            "minor_image_version",
            "major_subsystem_version",
            "minor_subsystem_version",
            "win32_version_value",
            "sizeof_image",
            "sizeof_headers",
            "checksum",
            "sizeof_stack_reserve",
            "sizeof_stack_commit",
            "sizeof_heap_reserve",
            "sizeof_heap_commit",
            "loader_flags",
            "numberof_rva_and_size",
            "subsystem",
            "dll_characteristics",
        ]
        post_process_funcs = {
            "subsystem": lambda x: list(x.values()),
            "dll_characteristics": lambda x: list(x.values()),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class DataDirectoriesFeatureExtractor(FeatureExtractor):
    feature_name = "data_directories"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def pair_with_section_type(
        data_directories: Optional[List[dict]], record_name: str
    ) -> Optional[List[Tuple[str, int]]]:
        if data_directories:
            return [
                (data_directory["type"], data_directory[record_name])
                for data_directory in data_directories
            ]
        else:
            return None

    def extract_raw_features(self, raw_json: dict) -> dict:
        data_directories = (
            raw_json[self.feature_name]
            if self.feature_name in raw_json.keys()
            else None
        )
        if data_directories:
            section: Optional[List[Optional[str]]] = [
                data_directory["section"]
                if "section" in data_directory.keys()
                else None
                for data_directory in data_directories
            ]
        else:
            section = None

        return {
            "RVA": self.pair_with_section_type(data_directories, "RVA"),
            "size": self.pair_with_section_type(data_directories, "size"),
            "section": section,
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = ["RVA", "size", "section"]
        post_process_funcs = {
            "RVA": lambda x: vectorize_with_feature_hasher(x, 50),
            "size": lambda x: vectorize_with_feature_hasher(x, 50),
            "section": lambda x: vectorize_with_feature_hasher(
                list(filter(lambda y: y is not None, x)) if x else None, 20
            ),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class SectionsFeatureExtractor(FeatureExtractor):
    feature_name = "sections"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def pair_with_section_name(
        sections: Optional[List[dict]], record_name: str
    ) -> Optional[List[Tuple[str, int]]]:
        if sections:
            return [
                (section["name"], section[record_name]) for section in sections
            ]
        else:
            return None

    # TODO: Reconsider the way to vectorize section characteristics
    @staticmethod
    def vectorize_characteristics(
        sections: Optional[List[dict]],
    ) -> Dict[str, int]:
        encoded_data = {
            "TYPE_NO_PAD": 0,
            "CNT_CODE": 0,
            "CNT_INITIALIZED_DATA": 0,
            "CNT_UNINITIALIZED_DATA": 0,
            "LNK_OTHER": 0,
            "LNK_INFO": 0,
            "LNK_REMOVE": 0,
            "LNK_COMDAT": 0,
            "GPREL": 0,
            "MEM_PURGEABLE": 0,
            "MEM_16BIT": 0,
            "MEM_LOCKED": 0,
            "MEM_PRELOAD": 0,
            "ALIGN_1BYTES": 0,
            "ALIGN_2BYTES": 0,
            "ALIGN_4BYTES": 0,
            "ALIGN_8BYTES": 0,
            "ALIGN_16BYTES": 0,
            "ALIGN_32BYTES": 0,
            "ALIGN_64BYTES": 0,
            "ALIGN_128BYTES": 0,
            "ALIGN_256BYTES": 0,
            "ALIGN_512BYTES": 0,
            "ALIGN_1024BYTES": 0,
            "ALIGN_2048BYTES": 0,
            "ALIGN_4096BYTES": 0,
            "ALIGN_8192BYTES": 0,
            "LNK_NRELOC_OVFL": 0,
            "MEM_DISCARDABLE": 0,
            "MEM_NOT_CACHED": 0,
            "MEM_NOT_PAGED": 0,
            "MEM_SHARED": 0,
            "MEM_EXECUTE": 0,
            "MEM_READ": 0,
            "MEM_WRITE": 0,
        }
        if sections:
            for section in sections:
                for characteristic in section["characteristics"]:
                    encoded_data[characteristic] += 1
        return encoded_data

    # TODO: Reconsider the way to vectorize section types
    @staticmethod
    def vectorize_types(sections: Optional[List[dict]]) -> Dict[str, int]:
        encoded_data = {
            "TEXT": 0,
            "TLS_": 0,
            "IDATA": 0,
            "DATA": 0,
            "BSS": 0,
            "RESOURCE": 0,
            "RELOCATION": 0,
            "EXPORT": 0,
            "DEBUG": 0,
            "UNKNOWN": 0,
            "Out of range": 0,
        }
        if sections:
            for section in sections:
                for type_ in section["types"]:
                    encoded_data[type_] += 1
        return encoded_data

    def extract_raw_features(self, raw_json: dict) -> dict:
        sections = (
            raw_json[self.feature_name]
            if self.feature_name in raw_json.keys()
            else None
        )
        return {
            "pointerto_relocation": self.pair_with_section_name(
                sections, "pointerto_relocation"
            ),
            "pointerto_line_numbers": self.pair_with_section_name(
                sections, "pointerto_line_numbers"
            ),
            "numberof_relocations": self.pair_with_section_name(
                sections, "numberof_relocations"
            ),
            "numberof_line_numbers": self.pair_with_section_name(
                sections, "numberof_line_numbers"
            ),
            "entropy": self.pair_with_section_name(sections, "entropy"),
            "characteristics": self.vectorize_characteristics(sections),
            "types": self.vectorize_types(sections),
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "pointerto_relocation",
            "pointerto_line_numbers",
            "numberof_relocations",
            "numberof_line_numbers",
            "characteristics",
            "entropy",
            "types",
        ]
        post_process_funcs = {
            "pointerto_relocation": lambda x: vectorize_with_feature_hasher(
                x, 50
            ),
            "pointerto_line_numbers": lambda x: vectorize_with_feature_hasher(
                x, 50
            ),
            "numberof_relocations": lambda x: vectorize_with_feature_hasher(
                x, 50
            ),
            "numberof_line_numbers": lambda x: vectorize_with_feature_hasher(
                x, 50
            ),
            "entropy": lambda x: vectorize_with_feature_hasher(x, 50),
            "characteristics": lambda x: list(x.values()),
            "types": lambda x: list(x.values()),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class RelocationsFeatureExtractor(FeatureExtractor):
    feature_name = "relocations"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def flatten_relocation_entries(
        relocations: Optional[List[dict]],
    ) -> List[Tuple[int, str]]:
        flattened_entries = list()
        if relocations:
            for relocation in relocations:
                virtual_address = relocation["virtual_address"]
                for entry in relocation["entries"]:
                    flattened_entries.append(
                        ((virtual_address + entry["position"], entry["type"]))
                    )
        return flattened_entries

    def count_relocation_types(
        self, flattend_entries: Optional[List[Tuple[int, str]]]
    ) -> List[int]:
        buckets = {
            "ABSOLUTE": 0,
            "HIGH": 0,
            "LOW": 0,
            "HIGHLOW": 0,
            "HIGHADJ": 0,
            "MIPS_JMPADDR | ARM_MOV32A | ARM_MOV32 | RISCV_HI20": 0,
            "SECTION": 0,
            "REL | ARM_MOV32T | THUMB_MOV32 | RISCV_LOW12I": 0,
            "RISCV_LOW12S": 0,
            "DIR64": 0,
            "HIGH3ADJ": 0,
            "Out of range": 0,
            "MIPS_JMPADDR16 | IA64_IMM64": 0,
        }
        # NOTE: https://github.com/lief-project/LIEF/commit/38f48a6e12f6cd46a657a4f354686333e8df4140
        if self.ver == 2020:
            buckets.pop("MIPS_JMPADDR16 | IA64_IMM64")
            buckets["MIPS_JMPADDR16 | IA64_DIR64"] = 0

        if flattend_entries:
            for _, type_ in flattend_entries:
                buckets[type_] += 1
        return list(buckets.values())

    def extract_raw_features(self, raw_json: dict) -> dict:
        relocations = (
            raw_json[self.feature_name]
            if self.feature_name in raw_json.keys()
            else None
        )
        return {
            "flattened_entries": self.flatten_relocation_entries(relocations)
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = ["flattened_entries"]
        post_process_funcs = {
            "flattened_entries": lambda x: self.count_relocation_types(x),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class TlsFeatureExtractor(FeatureExtractor):
    feature_name = "tls"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def characteristics_to_onehot(
        characteristics: Optional[int],
    ) -> Dict[str, int]:
        if characteristics:
            characteristics_str: Optional[str] = str(
                lief.PE.SECTION_CHARACTERISTICS(characteristics & 0xF00000)
            )
            has_extra_bits = int((characteristics & 0xFF0FFFFF) != 0)
        else:
            characteristics_str = None
            has_extra_bits = 0
        encoded_data = make_onehot_from_str_keys(
            [
                "SECTION_CHARACTERISTICS.ALIGN_1BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_2BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_4BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_8BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_16BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_32BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_64BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_128BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_256BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_512BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_1024BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_2048BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_4096BYTES",
                "SECTION_CHARACTERISTICS.ALIGN_8192BYTES",
                "SECTION_CHARACTERISTICS.???",
            ],
            characteristics_str,
        )
        encoded_data["has_extra_bits"] = has_extra_bits
        return encoded_data

    @staticmethod
    def data_directory_to_onehot(
        data_directory: Optional[str],
    ) -> Dict[str, int]:
        return make_onehot_from_str_keys(
            [
                "EXPORT_TABLE",
                "IMPORT_TABLE",
                "RESOURCE_TABLE",
                "EXCEPTION_TABLE",
                "CERTIFICATE_TABLE",
                "BASE_RELOCATION_TABLE",
                "DEBUG",
                "ARCHITECTURE",
                "GLOBAL_PTR",
                "TLS_TABLE",
                "LOAD_CONFIG_TABLE",
                "BOUND_IMPORT",
                "IAT",
                "DELAY_IMPORT_DESCRIPTOR",
                "CLR_RUNTIME_HEADER",
            ],
            data_directory,
        )

    def extract_raw_features(self, raw_json: dict) -> dict:
        tls = make_defaultdict_from_dict_elem(raw_json, self.feature_name)
        return {
            "callbacks": int(
                bool(tls["callbacks"])
            ),  # NOTE: Tls callback functions exists or not
            "addressof_raw_data": int(bool(tls["addressof_raw_data"])),
            "addressof_index": tls["addressof_index"],
            "addressof_callbacks": tls["addressof_callbacks"],
            "sizoeof_zero_fill": tls["sizeof_zero_fill"],
            "characteristics": self.characteristics_to_onehot(
                tls["characteristics"]
            ),
            "data_directory": self.data_directory_to_onehot(
                tls["data_directory"]
            ),
            "section": tls["section"],
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "callbacks",
            "addressof_raw_data",
            "addressof_index",
            "addressof_callbacks",
            "sizoeof_zero_fill",
            "characteristics",
            "data_directory",
        ]
        post_process_funcs = {
            "characteristics": lambda x: list(x.values()),
            "data_directory": lambda x: list(x.values()),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class ExportFeatureExtractor(FeatureExtractor):
    feature_name = "export"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def make_export_apis(
        entries: Optional[List[dict]], name: Optional[str]
    ) -> Optional[List[str]]:
        if entries is None or name is None:
            return None
        export_apis: List[str] = list()
        for entry in entries:
            if "forward_information" in entry.keys():
                api_name = "{}:{}".format(
                    entry["forward_information"]["library"],
                    entry["forward_information"]["function"],
                )
            else:
                api_name = "{}:{}".format(name, entry["name"])
            export_apis.append(api_name)
        return export_apis

    @staticmethod
    def pair_with_export_api(
        entries: Optional[List[dict]],
        export_apis: Optional[List[str]],
        record_name: str,
    ) -> Optional[List[Tuple[str, int]]]:
        if export_apis is not None and entries is not None:
            return [
                (export_api, entry[record_name])
                for export_api, entry in zip(export_apis, entries)
            ]
        else:
            return None

    def extract_raw_features_from_entries(
        self, entries: Optional[List[dict]], export_apis: Optional[List[str]]
    ) -> dict:
        prefix = "entries"
        return {
            f"{prefix}_ordinal": self.pair_with_export_api(
                entries, export_apis, "ordinal"
            ),
            f"{prefix}_address": self.pair_with_export_api(
                entries, export_apis, "address"
            ),
            f"{prefix}_is_extern": self.pair_with_export_api(
                entries, export_apis, "is_extern"
            ),
        }

    def extract_raw_features(self, raw_json: dict) -> dict:
        export = make_defaultdict_from_dict_elem(raw_json, self.feature_name)
        entries = export["entries"]
        export_apis = self.make_export_apis(entries, export["name"])
        return {
            "export_flags": export["export_flags"],
            "timestamp": export["timestamp"],
            "major_version": export["major_version"],
            "minor_version": export["minor_version"],
            "ordinal_base": export["ordinal_base"],
            "name": export["name"],
            **self.extract_raw_features_from_entries(entries, export_apis),
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        selected_features = [
            "export_flags",
            "timestamp",
            "major_version",
            "minor_version",
            "ordinal_base",
            "entries_ordinal",
            "entries_address",
            "entries_is_extern",
        ]
        post_process_funcs = {
            "entries_ordinal": lambda x: vectorize_with_feature_hasher(x, 50),
            "entries_address": lambda x: vectorize_with_feature_hasher(x, 50),
            "entries_is_extern": lambda x: vectorize_with_feature_hasher(x, 50),
        }
        return vectorize_selected_features(
            raw_features,
            selected_features,
            post_process_funcs,
            self.feature_name,
        )

class DebugFeatureExtractor(FeatureExtractor):
    feature_name = "debug"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    # TODO: vectorize debug feature
    # def extract_raw_features(self, raw_json: dict) -> dict:
    #     debug = _make_defaultdict_from_dict(raw_json, self.feature_name)
    #     raw_features = dict()
    #     return raw_features

class ImportsFeatureExtractor(FeatureExtractor):
    feature_name = "imports"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def pair_with_import_dlls(
        imports: Optional[List[dict]], record_name: str
    ) -> Optional[List[Tuple[str, int]]]:
        if imports:
            return [
                (import_["name"], import_[record_name]) for import_ in imports
            ]
        else:
            return None

    @staticmethod
    def pair_with_import_apis(
        imports: Optional[List[dict]], record_name: str
    ) -> Optional[List[Tuple[str, int]]]:
        if imports is None:
            return None
        result: List[Tuple[str, int]] = list()
        for import_ in imports:
            dll_name = import_["name"]
            for entry in import_["entries"]:
                result.append(
                    (
                        f"{dll_name}:{ImportsFeatureExtractor.get_api_name(entry)}",
                        entry[record_name],
                    )
                )
        return result

    @staticmethod
    def get_api_name(entry: dict) -> str:
        if "name" in entry.keys():
            return entry["name"]
        else:
            return f'ord{entry["ordinal"]}'

    @staticmethod
    def flatten_api_entries(
        imports: Optional[List[dict]],
    ) -> Optional[List[str]]:
        if imports is None:
            return None
        return [
            f'{import_["name"]}:{ImportsFeatureExtractor.get_api_name(entry)}'
            for import_ in imports
            for entry in import_["entries"]
        ]

    @staticmethod
    def flatten_dll_entries(
        imports: Optional[List[dict]],
    ) -> Optional[List[str]]:
        if imports is None:
            return None
        return [import_["name"] for import_ in imports]

    def extract_raw_features(self, raw_json: dict) -> dict:
        imports = (
            raw_json[self.feature_name]
            if self.feature_name in raw_json.keys()
            else None
        )
        return {
            "dll_names": self.flatten_dll_entries(imports),
            "api_names": self.flatten_api_entries(imports),
            "forwarder_chain": self.pair_with_import_dlls(
                imports, "forwarder_chain"
            ),
            "timedatestamp": self.pair_with_import_dlls(
                imports, "timedatestamp"
            ),
            "import_address_table_rva": self.pair_with_import_dlls(
                imports, "import_address_table_rva"
            ),
            "iat_address": self.pair_with_import_apis(imports, "iat_address"),
            "data": self.pair_with_import_apis(imports, "data"),
            "hint": self.pair_with_import_apis(imports, "hint"),
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "dll_names",
            "api_names",
            "forwarder_chain",
            "timedatestamp",
            "import_address_table_rva",
            "iat_address",
            "data",
            "hint",
        ]
        post_process_funcs = {
            "dll_names": lambda x: vectorize_with_feature_hasher(x, 100),
            "api_names": lambda x: vectorize_with_feature_hasher(x, 500),
            "forwarder_chain": lambda x: vectorize_with_feature_hasher(x, 100),
            "timedatestamp": lambda x: vectorize_with_feature_hasher(x, 100),
            "import_address_table_rva": lambda x: vectorize_with_feature_hasher(
                x, 100
            ),
            "iat_address": lambda x: vectorize_with_feature_hasher(x, 100),
            "data": lambda x: vectorize_with_feature_hasher(x, 100),
            "hint": lambda x: vectorize_with_feature_hasher(x, 100),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

# NOTE: In LIEF 0.11 or later, dict structure of resources_tree is different from 0.10.1 one.
# TODO: When updating LIEF version 0.10.1 to 0.11, the extract_raw_features should be modified.
class ResourcesTreeFeatureExtractor(FeatureExtractor):
    feature_name = "resources_tree"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    def extract_raw_features(self, raw_json: dict) -> dict:
        resources_tree = make_defaultdict_from_dict_elem(
            raw_json, self.feature_name
        )
        return {
            "characteristics": resources_tree["characteristics"],
            "major_version": resources_tree["major_version"],
            "minor_version": resources_tree["minor_version"],
            "numberof_id_entries": resources_tree["numberof_id_entries"],
            "numberof_name_entries": resources_tree["numberof_name_entries"],
            "time_date_stamp": resources_tree["time_date_stamp"],
        }

class ResourcesManagerFeatureExtractor(FeatureExtractor):
    feature_name = "resources_manager"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def fixed_version_file_flags_to_onehot(
        file_flags: Optional[int],
    ) -> Dict[str, int]:
        return make_onehot_dict_from_bitflag(
            [
                "DEBUG",
                "INFOINFERRED",
                "PATCHED",
                "PRERELEASE",
                "PRIVATEBUILD",
                "SPECIALBUILD",
            ],
            file_flags,
            lief.PE.FIXED_VERSION_FILE_FLAGS,
        )

    @staticmethod
    def fixed_version_os_to_onehot(version_os: Optional[str]) -> Dict[str, int]:
        return make_onehot_from_str_keys(
            [
                "DOS",
                "DOS_WINDOWS16",
                "DOS_WINDOWS32",
                "NT",
                "NT_WINDOWS32",
                "OS216",
                "OS216_PM16",
                "OS232",
                "OS232_PM32",
                "PM16",
                "PM32",
                "UNKNOWN",
                "WINDOWS16",
                "WINDOWS32",
                "Out of range",
            ],
            version_os,
        )

    @staticmethod
    def fixed_version_file_type_to_onehot(
        file_type: Optional[str],
    ) -> Dict[str, int]:
        return make_onehot_from_str_keys(
            [
                "APP",
                "DLL",
                "DRV",
                "FONT",
                "STATIC_LIB",
                "VXD",
                "UNKNOWN",
                "Out of range",
            ],
            file_type,
        )

    @staticmethod
    def fixed_version_file_subtype_to_onehot(
        file_subtype: Optional[str],
    ) -> Dict[str, int]:
        return make_onehot_from_str_keys(
            [
                "DRV_COMM",
                "DRV_DISPLAY",
                "DRV_INSTALLABLE",
                "DRV_KEYBOARD",
                "DRV_LANGUAGE",
                "DRV_MOUSE",
                "DRV_NETWORK",
                "DRV_PRINTER",
                "DRV_SOUND",
                "DRV_SYSTEM",
                "DRV_VERSIONED_PRINTER",
                "UNKNOWN",
                "Out of range",
            ],
            file_subtype,
        )

    @staticmethod
    def extract_raw_features_from_html(html: List[str]) -> dict:
        prefix = "html"

        def is_html(e: str) -> bool:
            try:
                return bool(BeautifulSoup(e, "html.parser").find())
            except Exception:
                return False

        number_of_non_htmls = sum(is_html(e) for e in html)
        return {
            f"{prefix}_avg_len": mean(len(e) for e in html) if html else 0,
            f"{prefix}_number_of_non_htmls": number_of_non_htmls if html else 0,
        }

    @staticmethod
    def extract_raw_features_from_version(
        version: DefaultDict[str, Any]
    ) -> dict:
        prefix = "version"
        return {
            f"{prefix}_type": version["type"],
            # NOTE: should be VS_VERSION_INFO, but other strings might come
            f"{prefix}_key": int(version["key"] == "VS_VERSION_INFO"),
        }

    def extract_raw_features_from_fixed_file_info(
        self, fixed_file_info: DefaultDict[str, Any]
    ) -> dict:
        prefix = "version"
        cprefix = "fixed_file_info"
        return {
            f"{prefix}_{cprefix}_signature": fixed_file_info["signature"],
            f"{prefix}_{cprefix}_struct_version": fixed_file_info[
                "struct_version"
            ],
            f"{prefix}_{cprefix}_file_version_MS": fixed_file_info[
                "file_version_MS"
            ],
            f"{prefix}_{cprefix}_file_version_LS": fixed_file_info[
                "file_version_LS"
            ],
            f"{prefix}_{cprefix}_file_flags_mask": self.fixed_version_file_flags_to_onehot(
                fixed_file_info["file_flags_mask"]
            ),
            f"{prefix}_{cprefix}_file_flags": self.fixed_version_file_flags_to_onehot(
                fixed_file_info["file_flags"]
            ),
            f"{prefix}_{cprefix}_file_os": self.fixed_version_os_to_onehot(
                fixed_file_info["file_os"]
            ),
            f"{prefix}_{cprefix}_file_type": self.fixed_version_file_type_to_onehot(
                fixed_file_info["file_type"]
            ),
            f"{prefix}_{cprefix}_file_subtype": self.fixed_version_file_subtype_to_onehot(
                fixed_file_info["file_subtype"]
            ),
            f"{prefix}_{cprefix}_file_date_MS": fixed_file_info["file_date_MS"],
            f"{prefix}_{cprefix}_file_date_LS": fixed_file_info["file_date_LS"],
        }

    def extract_raw_features_from_string_file_info(
        self, string_file_info: DefaultDict[str, Any]
    ) -> dict:
        prefix = "version"
        cprefix = "string_file_info"
        # TODO: convert langcode_items array to feature vector
        # langcode_items = _make_defaultdict_from_dict(string_file_info, "lang_code_item")

        return {
            f"{prefix}_{cprefix}_type": string_file_info["type"],
            f"{prefix}_{cprefix}_key": int(
                string_file_info["key"] == "StringFileInfo"
            ),
        }

    def extract_raw_features_from_var_file_info(
        self, var_file_info: DefaultDict[str, Any]
    ) -> dict:
        prefix = "version"
        cprefix = "var_file_info"

        return {
            f"{prefix}_{cprefix}_type": var_file_info["type"],
            f"{prefix}_{cprefix}_key": int(
                var_file_info["key"] == "VarFileInfo"
            ),
            # f"{prefix}_{cprefix}_translations": self.translations_to_featurevector(var_file_info["translations"])
        }

    def extract_raw_features(self, raw_json: dict) -> dict:
        resources_manager = make_defaultdict_from_dict_elem(
            raw_json, self.feature_name
        )
        version = make_defaultdict_from_dict_elem(resources_manager, "version")
        raw_features = {
            # NOTE: extracted but not converted to feature vector
            # TODO: should parse manifest XML file
            "manifest": resources_manager["manifest"],
            **self.extract_raw_features_from_version(version),
            **self.extract_raw_features_from_fixed_file_info(
                make_defaultdict_from_dict_elem(version, "fixed_file_info")
            ),
            **self.extract_raw_features_from_string_file_info(
                make_defaultdict_from_dict_elem(version, "string_file_info")
            ),
            **self.extract_raw_features_from_var_file_info(
                make_defaultdict_from_dict_elem(version, "var_file_info")
            ),
        }
        if self.ver >= 2021:
            raw_features.update(
                self.extract_raw_features_from_html(
                    resources_manager["html"]
                    if "html" in resources_manager.keys()
                    else []
                )
            )
            raw_features["has_string_table"] = (
                "string_table" in resources_manager.keys()
            )
            raw_features["has_icons"] = "icons" in resources_manager.keys()
            raw_features["has_dialogs"] = "dialogs" in resources_manager.keys()
            raw_features["has_accelerator"] = (
                "accelerator" in resources_manager.keys()
            )
        return raw_features

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "version_type",
            "version_key",
            "version_fixed_file_info_signature",
            "version_fixed_file_info_struct_version",
            "version_fixed_file_info_file_version_MS",
            "version_fixed_file_info_file_version_LS",
            "version_fixed_file_info_file_flags_mask",
            "version_fixed_file_info_file_flags",
            "version_fixed_file_info_file_os",
            "version_fixed_file_info_file_type",
            "version_fixed_file_info_file_subtype",
            "version_fixed_file_info_file_date_MS",
            "version_fixed_file_info_file_date_LS",
            "version_string_file_info_type",
            "version_string_file_info_key",
            "version_var_file_info_type",
            "version_var_file_info_key",
        ]
        if self.ver >= 2021:
            features_selected.append("html_avg_len")
            features_selected.append("html_number_of_non_htmls")
            features_selected.append("has_string_table")
            features_selected.append("has_icons")
            features_selected.append("has_dialogs")
            features_selected.append("has_accelerator")
        post_process_funcs = {
            "version_fixed_file_info_file_flags_mask": lambda x: list(
                x.values()
            ),
            "version_fixed_file_info_file_flags": lambda x: list(x.values()),
            "version_fixed_file_info_file_os": lambda x: list(x.values()),
            "version_fixed_file_info_file_type": lambda x: list(x.values()),
            "version_fixed_file_info_file_subtype": lambda x: list(x.values()),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class SignatureFeatureExtractor(FeatureExtractor):
    feature_name = "signature"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()
        if self.ver >= 2021:
            self.feature_name = "signatures"

    # TODO: should be implemented
    # @staticmethod
    # def crts_to_feature_vector(crts):
    #     pass

    def extract_raw_features_from_content_info(
        self, content_info: DefaultDict[str, Any]
    ) -> dict:
        prefix = "content_info"
        return {
            f"{prefix}_content_type": content_info["content_type"],
            f"{prefix}_digest_algorithm": content_info["digest_algorithm"],
            f"{prefix}_type": content_info["type"],
        }

    def extract_raw_features_from_signer_info(
        self, signer_info: DefaultDict[str, Any]
    ) -> dict:
        prefix = "signer_info"
        return {
            f"{prefix}_version": signer_info["version"],
            f"{prefix}_digest_algorithm": signer_info["digest_algorithm"],
            f"{prefix}_signature_algorithm": signer_info["signature_algorithm"],
        }

    def extract_raw_features_from_authenticated_attributres(
        self, authenticated_attributes: DefaultDict[str, Any]
    ) -> dict:
        prefix = "signer_info"
        cprefix = "authenticated_attributes"
        return {
            f"{prefix}_{cprefix}_content_type": authenticated_attributes[
                "content_type"
            ],
            f"{prefix}_{cprefix}_program_name": authenticated_attributes[
                "program_name"
            ],
            f"{prefix}_{cprefix}_url": authenticated_attributes["url"],
            f"{prefix}_{cprefix}_message_digest": authenticated_attributes[
                "message_digest"
            ],
        }

    def extract_raw_features(self, raw_json: dict) -> dict:
        if self.ver == 2020:
            signature = make_defaultdict_from_dict_elem(
                raw_json, self.feature_name
            )
            signer_info = make_defaultdict_from_dict_elem(
                signature, "signer_info"
            )
            return {
                "version": signature["version"],
                # NOTE: extracted but not converted to feature vector
                **self.extract_raw_features_from_content_info(
                    make_defaultdict_from_dict_elem(signature, "content_info")
                ),
                # NOTE: extracted but not converted to feature vector
                **self.extract_raw_features_from_signer_info(signer_info),
                # NOTE: extracted but not converted to feature vector
                **self.extract_raw_features_from_authenticated_attributres(
                    make_defaultdict_from_dict_elem(
                        signer_info, "authenticated_attributes"
                    )
                ),
                # NOTE: extracted but not converted to feature vector
                "signer_info_issuer": signer_info["issuer"],
                # TODO: make feature vector from certificates array
                "certificates": signature["certificates"],
            }
        elif self.ver == 2021:
            # TODO: extract other properties contained in raw_json["signatures"]
            return {
                "number_of_signatures": len(raw_json["signatures"])
                if "signatures" in raw_json.keys()
                else 0
            }
        else:
            raise NotSupported(self.ver, self.__class__.__name__)

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        if self.ver == 2020:
            features_selected = ["version"]
            return vectorize_selected_features(
                raw_features, features_selected, {}, self.feature_name
            )
        elif self.ver == 2021:
            features_selected = ["number_of_signatures"]
            return vectorize_selected_features(
                raw_features, features_selected, {}, self.feature_name
            )
        else:
            raise NotSupported(self.ver, self.__class__.__name__)

class SymbolsFeatureExtractor(FeatureExtractor):
    feature_name = "symbols"

    @staticmethod
    def make_base_type_hist() -> Dict[str, int]:
        return {
            "NULL": 0,
            "VOID": 0,
            "CHAR": 0,
            "SHORT": 0,
            "INT": 0,
            "LONG": 0,
            "FLOAT": 0,
            "DOUBLE": 0,
            "STRUCT": 0,
            "UNION": 0,
            "ENUM": 0,
            "MOE": 0,
            "BYTE": 0,
            "WORD": 0,
            "UINT": 0,
            "DWORD": 0,
            "Out of range": 0,
        }

    @staticmethod
    def make_complex_type_hist() -> Dict[str, int]:
        return {
            "NULL": 0,
            "POINTER": 0,
            "FUNCTION": 0,
            "ARRAY": 0,
            "COMPLEX_TYPE_SHIFT": 0,
            "Out of range": 0,
        }

    @staticmethod
    def make_storage_class_hist() -> Dict[str, int]:
        return {
            "END_OF_FUNCTION": 0,
            "NULL": 0,
            "AUTOMATIC": 0,
            "EXTERNAL": 0,
            "STATIC": 0,
            "REGISTER": 0,
            "EXTERNAL_DEF": 0,
            "LABEL": 0,
            "UNDEFINED_LABEL": 0,
            "MEMBER_OF_STRUCT": 0,
            "UNION_TAG": 0,
            "TYPE_DEFINITION": 0,
            "UDEFINED_STATIC": 0,
            "ENUM_TAG": 0,
            "MEMBER_OF_ENUM": 0,
            "REGISTER_PARAM": 0,
            "BIT_FIELD": 0,
            "BLOCK": 0,
            "FUNCTION": 0,
            "END_OF_STRUCT": 0,
            "FILE": 0,
            "SECTION": 0,
            "WEAK_EXTERNAL": 0,
            "CLR_TOKEN": 0,
            "Out of range": 0,
        }

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        if self.ver == 2020:
            raise NotSupported(self.ver, self.__class__.__name__)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def pair_with_symbol_name(
        symbols: List[dict], record_name: str
    ) -> List[Tuple[str, int]]:
        return [(symbol["name"], symbol[record_name]) for symbol in symbols]

    def extract_raw_features(self, raw_json: dict) -> dict:
        if self.feature_name in raw_json.keys():
            symbols = raw_json[self.feature_name]
        else:
            symbols = list()

        base_type_hist = self.make_base_type_hist()
        complex_type_hist = self.make_complex_type_hist()
        storage_class_hist = self.make_storage_class_hist()

        for symbol in symbols:
            base_type_hist[symbol["base_type"]] += 1
            complex_type_hist[symbol["complex_type"]] += 1
            storage_class_hist[symbol["storage_class"]] += 1

        return {
            "base_type": base_type_hist,
            "complex_type": complex_type_hist,
            "storage_class": storage_class_hist,
            "numberof_aux_symbols": self.pair_with_symbol_name(
                symbols, "numberof_aux_symbols"
            ),
            "section_number": self.pair_with_symbol_name(
                symbols, "section_number"
            ),
            "size": self.pair_with_symbol_name(symbols, "size"),
            "value": self.pair_with_symbol_name(symbols, "value"),
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "base_type",
            "complex_type",
            "storage_class",
            "numberof_aux_symbols",
            "section_number",
            "size",
            "value",
        ]
        post_process_funcs = {
            "base_type": lambda x: list(x.values()),
            "complex_type": lambda x: list(x.values()),
            "storage_class": lambda x: list(x.values()),
            "numberof_aux_symbols": lambda x: vectorize_with_feature_hasher(
                x, 50
            ),
            "section_number": lambda x: vectorize_with_feature_hasher(x, 50),
            "size": lambda x: vectorize_with_feature_hasher(x, 50),
            "value": lambda x: vectorize_with_feature_hasher(x, 50),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

# NOTE: https://docs.microsoft.com/en-us/windows/win32/api/heapapi/nf-heapapi-heapcreate#parameters
class ProcessHeapFlags(Enum):
    HEAP_CREATE_ENABLE_EXECUTE = 0x00040000
    HEAP_GENERATE_EXCEPTIONS = 0x00000004
    HEAP_NO_SERIALIZE = 0x00000001

class LoadConfigurationFeatureExtractor(FeatureExtractor):
    feature_name: str = "load_configuration"

    def __init__(self, ver: str) -> None:
        self.ver = ver_str_to_int(ver)
        super(FeatureExtractor, self).__init__()

    @staticmethod
    def version_to_onehot(ver: Optional[str]) -> Dict[str, int]:
        return make_onehot_from_str_keys(
            [
                "UNKNOWN",
                "SEH",
                "WIN_8_1",
                "WIN10_0_9879",
                "WIN10_0_14286",
                "WIN10_0_14383",
                "WIN10_0_14901",
                "WIN10_0_15002",
                "WIN10_0_16237",
            ],
            ver,
        )

    @staticmethod
    def process_heap_flags_to_onehot(phf: Optional[int]) -> Dict[str, int]:
        encoded_data = {str(k): 0 for k in ProcessHeapFlags}
        if phf:
            for k in ProcessHeapFlags:
                if (phf & k.value) != 0:
                    encoded_data[str(k)] = 1
        return encoded_data

    def extract_raw_features_from_code_integrity(
        self, code_integrity: DefaultDict[str, Any]
    ) -> dict:
        prefix = "code_integrity"
        return {
            f"{prefix}_flags": code_integrity["flags"],
            f"{prefix}_catalog": code_integrity["catalog"],
            f"{prefix}_catalog_offset": code_integrity["catalog_offset"],
            f"{prefix}_reserved": code_integrity["reserved"],
        }

    def extract_raw_features(self, raw_json: dict) -> dict:
        load_configuration = make_defaultdict_from_dict_elem(
            raw_json, self.feature_name
        )

        return {
            "version": self.version_to_onehot(load_configuration["version"]),
            "characteristics": load_configuration["characteristics"],
            "timedatestamp": load_configuration["timedatestamp"],
            "major_version": load_configuration["major_version"],
            "minor_version": load_configuration["minor_version"],
            "global_flags_clear": load_configuration["global_flags_clear"],
            "global_flags_set": load_configuration["global_flags_set"],
            "critical_section_default_timeout": load_configuration[
                "critical_section_default_timeout"
            ],
            "decommit_free_block_threshold": load_configuration[
                "decommit_free_block_threshold"
            ],
            "decommit_total_free_threshold": load_configuration[
                "decommit_total_free_threshold"
            ],
            "lock_prefix_table": load_configuration["lock_prefix_table"],
            "maximum_allocation_size": load_configuration[
                "maximum_allocation_size"
            ],
            "virtual_memory_threshold": load_configuration[
                "virtual_memory_threshold"
            ],
            "process_affinity_mask": load_configuration[
                "process_affinity_mask"
            ],
            "process_heap_flags": self.process_heap_flags_to_onehot(
                load_configuration["process_heap_flags"]
            ),
            "csd_version": load_configuration["csd_version"],
            "reserved1": load_configuration["reserved1"],
            "security_cookie": load_configuration["security_cookie"],
            "se_handler_table": load_configuration["se_handler_table"],
            "se_handler_count": load_configuration["se_handler_count"],
            "guard_cf_check_function_pointer": load_configuration[
                "guard_cf_check_function_pointer"
            ],
            "guard_cf_dispatch_function_pointer": load_configuration[
                "guard_cf_dispatch_function_pointer"
            ],
            "guard_cf_function_table": load_configuration[
                "guard_cf_function_table"
            ],
            "guard_cf_function_count": load_configuration[
                "guard_cf_function_count"
            ],
            "guard_flags": load_configuration["guard_flags"],
            **self.extract_raw_features_from_code_integrity(
                make_defaultdict_from_dict_elem(
                    load_configuration, "code_integrity"
                )
            ),
            "guard_address_taken_iat_entry_table": load_configuration[
                "guard_address_taken_iat_entry_table"
            ],
            "guard_address_taken_iat_entry_count": load_configuration[
                "guard_address_taken_iat_entry_count"
            ],
            "guard_long_jump_target_table": load_configuration[
                "guard_long_jump_target_table"
            ],
            "guard_long_jump_target_count": load_configuration[
                "guard_long_jump_target_count"
            ],
            "dynamic_value_reloc_table": load_configuration[
                "dynamic_value_reloc_table"
            ],
            "hybrid_metadata_pointer": load_configuration[
                "hybrid_metadata_pointer"
            ],
            "guard_rf_failure_routine": load_configuration[
                "guard_rf_failure_routine"
            ],
            "guard_rf_failure_routine_function_pointer": load_configuration[
                "guard_rf_failure_routine_function_pointer"
            ],
            "dynamic_value_reloctable_offset": load_configuration[
                "dynamic_value_reloctable_offset"
            ],
            "dynamic_value_reloctable_section": load_configuration[
                "dynamic_value_reloctable_section"
            ],
            "reserved2": load_configuration["reserved2"],
            "guard_rf_verify_stackpointer_function_pointer": load_configuration[
                "guard_rf_verify_stackpointer_function_pointer"
            ],
            "hotpatch_table_offset": load_configuration[
                "hotpatch_table_offset"
            ],
            "reserved3": load_configuration["reserved3"],
            "addressof_unicode_string": load_configuration[
                "addressof_unicode_string"
            ],
        }

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        features_selected = [
            "version",
            "characteristics",
            "timedatestamp",
            "major_version",
            "minor_version",
            "minor_version",
            "global_flags_clear",
            "global_flags_set",
            "critical_section_default_timeout",
            "decommit_free_block_threshold",
            "decommit_total_free_threshold",
            "lock_prefix_table",
            "maximum_allocation_size",
            "virtual_memory_threshold",
            "process_affinity_mask",
            "process_heap_flags",
            "csd_version",
            "reserved1",
            "security_cookie",
            "se_handler_table",
            "se_handler_count",
            "guard_cf_check_function_pointer",
            "guard_cf_dispatch_function_pointer",
            "guard_cf_function_table",
            "guard_cf_function_count",
            "guard_flags",
            "code_integrity_flags",
            "code_integrity_catalog",
            "code_integrity_catalog_offset",
            "code_integrity_reserved",
            "guard_address_taken_iat_entry_table",
            "guard_address_taken_iat_entry_count",
            "guard_long_jump_target_table",
            "guard_long_jump_target_count",
            "dynamic_value_reloc_table",
            "hybrid_metadata_pointer",
            "guard_rf_failure_routine",
            "guard_rf_failure_routine_function_pointer",
            "dynamic_value_reloctable_offset",
            "dynamic_value_reloctable_section",
            "reserved2",
            "guard_rf_verify_stackpointer_function_pointer",
            "hotpatch_table_offset",
            "reserved3",
            "addressof_unicode_string",
        ]
        post_process_funcs = {
            "version": lambda x: list(x.values()),
            "process_heap_flags": lambda x: list(x.values()),
        }
        return vectorize_selected_features(
            raw_features,
            features_selected,
            post_process_funcs,
            self.feature_name,
        )

class LiefFeatureExtractor(FeatureExtractor):
    feature_name = "lief"

    def __init__(self, ver: str) -> None:
        super(FeatureExtractor, self).__init__()
        self.ver = ver_str_to_int(ver)
        self.extractors = [
            DosHeaderFeatureExtractor(ver),
            RichHeaderFeatureExtractor(ver),
            HeaderFeatureExtractor(ver),
            OptionalHeaderFeatureExtractor(ver),
            DataDirectoriesFeatureExtractor(ver),
            SectionsFeatureExtractor(ver),
            RelocationsFeatureExtractor(ver),
            TlsFeatureExtractor(ver),
            ExportFeatureExtractor(ver),
            # DebugFeatureExtractor(ver), # DebugFeature is currently not supported
            ImportsFeatureExtractor(ver),
            ResourcesTreeFeatureExtractor(ver),
            ResourcesManagerFeatureExtractor(ver),
            SignatureFeatureExtractor(ver),
            LoadConfigurationFeatureExtractor(ver),
        ]
        if self.ver >= 2021:
            self.extractors.append(SymbolsFeatureExtractor(ver))

    def extract_raw_features(self, raw_json: dict) -> dict:
        raw_features = {
            extractor.feature_name: extractor.extract_raw_features(raw_json)
            for extractor in self.extractors
        }
        # NOTE: top-level defined features
        raw_features["entrypoint"] = raw_json["entrypoint"]
        raw_features["virtual_size"] = raw_json["virtual_size"]
        return raw_features

    def vectorize_features(
        self, raw_features: dict
    ) -> Tuple[List[str], np.ndarray]:
        columns: List[str] = ["entrypoint", "virtual_size"]
        vectors: List[np.ndarray] = [
            raw_features["entrypoint"],
            raw_features["virtual_size"],
        ]
        for extractor in self.extractors:
            column, vector = extractor.vectorize_features(
                raw_features[extractor.feature_name]
            )
            columns += column
            vectors.append(vector)
        return columns, np.hstack(vectors)


'''
Created on Mar 20, 2014

@author: flurin, quentin
'''
import numpy as np
import scipy.sparse as sps
import DNL
import os.path
import warnings

class Assignment:
    '''
    This class computes or loads the assigment matrices. If they have already been calculated, then they area
    simply loaded. This part can take some time if the matrices must be calculated from scratch.
    '''
    def __init__(self, estim_param, network):
        '''
        Constructor
        '''
        # create precomputed directory if it does not exist
        if not os.path.exists(estim_param.path_precomp):
            os.makedirs(estim_param.path_precomp)

        '''-----------------------------------------------------------------------------------------------------------------
                            Loads the link flow assignment matrix if is exists, otherwise calculates it
        -----------------------------------------------------------------------------------------------------------------'''
        if os.path.isfile(estim_param.path_precomputed_asg + "_linkFlowAssgn_data.csv") and estim_param.forceRecompAssgnment is False:
            print "Loading pre-computed flow assignment matrix",
            data = np.loadtxt(estim_param.path_precomputed_asg + "_linkFlowAssgn_data.csv", delimiter=",")
            indices = np.loadtxt(estim_param.path_precomputed_asg + "_linkFlowAssgn_indices.csv", dtype='int_', delimiter=",")
            indptr = np.loadtxt(estim_param.path_precomputed_asg + "_linkFlowAssgn_indptr.csv", dtype='int_', delimiter=",")
            self.linkFlowAssgn = sps.csr_matrix((data, indices, indptr)).todense()
        else:
            print "Computing flow assignment matrix",
            self.linkFlowAssgn = np.asmatrix(DNL.compute_assg_mat(network, estim_param))
            linkFlowAssgn_csr = sps.csr_matrix(self.linkFlowAssgn)
            np.savetxt(estim_param.path_precomputed_asg + "_linkFlowAssgn_data.csv", linkFlowAssgn_csr.data, delimiter=",")
            np.savetxt(estim_param.path_precomputed_asg + "_linkFlowAssgn_indices.csv", linkFlowAssgn_csr.indices, fmt='%i', delimiter=",")
            np.savetxt(estim_param.path_precomputed_asg + "_linkFlowAssgn_indptr.csv", linkFlowAssgn_csr.indptr, fmt='%i', delimiter=",")

        self.linkFlowAssgn_prime = DNL.build_flow_assg_mat(self.linkFlowAssgn, network, estim_param, network.edges_ASE_dict)
        self.subRouteAssgn_prime = DNL.build_ODflow_assg_mat(self.linkFlowAssgn, network, estim_param, network.subroutes_VS_dict)
        self.accAssgn_prime = DNL.build_flow_assg_mat(self.linkFlowAssgn, network, estim_param, network.edges_TINF_dict)
        estim_param.print_incr_runtime()

        '''-----------------------------------------------------------------------------------------------------------------
                         Loads the occupation flow assignment matrix if is exists, otherwise calculates it
        -----------------------------------------------------------------------------------------------------------------'''
        if os.path.isfile(estim_param.path_precomputed_asg + "_acc_data.csv") and estim_param.forceRecompAssgnment is False:
            print "Loading pre-computed accumulation assignment matrix",
            data = np.loadtxt(estim_param.path_precomputed_asg + "_acc_data.csv", delimiter=",")
            indices = np.loadtxt(estim_param.path_precomputed_asg + "_acc_indices.csv", dtype='int_', delimiter=",")
            indptr = np.loadtxt(estim_param.path_precomputed_asg + "_acc_indptr.csv", dtype='int_', delimiter=",")
            self.acc = sps.csr_matrix((data, indices, indptr)).todense()
            #print "Occupation assignment matrix accumulation loaded."
        else:
            print "Computing accumulation assignment matrix",
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")  # Prevent some annoying numerical warnings from appearing
                self.acc = np.asmatrix(DNL.compute_assg_mat_accumulation(network, estim_param))
            acc_csr = sps.csr_matrix(self.acc)
            np.savetxt(estim_param.path_precomputed_asg + "_acc_data.csv", acc_csr.data, delimiter=",")
            np.savetxt(estim_param.path_precomputed_asg + "_acc_indices.csv", acc_csr.indices, fmt='%i', delimiter=",")
            np.savetxt(estim_param.path_precomputed_asg + "_acc_indptr.csv", acc_csr.indptr, fmt='%i', delimiter=",")
        estim_param.print_incr_runtime()


<filename>antenna_pattern/rectify_timing_offsets.py
"""
Methods for rectifying for overflows and timing offsets in GNU Radio samples.

Described in https://www.la1k.no/?p=3152.
"""

from combine_samples_and_angles import read_gnuradio_header_element, load_gnuradio_samples
from frequency_analysis import get_center_frequency_and_samplerate
import os
import numpy as np
from scipy.interpolate import interp1d

def get_timestamps_from_header(gnuradio_hdr_file):
    """
    Parse GNU Radio header file for timestamp tags.

    Parameters
    ----------
    gnuradio_hdr_file: str
        Path to GNU Radio header file

    Returns
    -------
    timestamps: ndarray
        List over timestamps
    num_items: ndarray
        List over number of items for which the timestamps are valid
    """

    handle = open(gnuradio_hdr_file)

    #get number of header elements contained in header file
    _, element_size = read_gnuradio_header_element(handle)
    handle.seek(0, 0)

    header_size = os.path.getsize(gnuradio_hdr_file)
    num_elements = header_size/element_size

    #get all header information
    timestamps = np.zeros(num_elements)
    samplerates = np.zeros(num_elements)
    num_items = np.zeros(num_elements)
    num_read = 0

    for i in np.arange(0, num_elements):
        info, element_size = read_gnuradio_header_element(handle)
        if element_size == 0:
            break

        num_read += element_size
        handle.seek(num_read, 0)

        timestamps[i] = info['rx_time']
        num_items[i] = info['nitems']

    return timestamps, num_items

def copy_first_part_of_header(filename_in, filename_out):
    """
    Copy first part of header file to a new header file.

    Parameters
    ----------
    filename_in: str
        Input filename. Input header is assumed to be filename_in.hdr.
    filename_out: str
        Output filename. Output header is assumed to be filename_out.hdr.
    """
    header_file_out = open(filename_out + '.hdr', 'w')
    header_file_in = open(filename_in + '.hdr')
    _, header_size = read_gnuradio_header_element(header_file_in)
    header_file_in.seek(0)
    header_data = header_file_in.read(header_size)

    header_file_out.write(header_data)

def fill_overflows(filename_in, filename_out):
    """
    Takes in a GNU Radio meta file sink file and fills places with overflows
    with a number of samples corresponding to the missing samples during the
    overflow in order to rectify somewhat for the timing issues.

    Parameters
    ----------
    filename_in: str
        Input filename
    filename_out: str
        Output filename for writing the filled samples. Will also produce a
        header file, containing the first header element of the original file.
    """

    timestamps, num_items = get_timestamps_from_header(filename_in + '.hdr')
    f_c, f_s = get_center_frequency_and_samplerate(filename_in)
    _, data = load_gnuradio_samples(filename_in, return_full_timestamps=False)

    #calculate missing samples
    expected_samples = np.diff(timestamps)*f_s
    missing_samples = expected_samples - num_items[0:-1]
    missing_samples = np.round(missing_samples).astype(int)

    #find indices where there is overflow problems
    problem_samples = np.arange(0, len(missing_samples))[missing_samples > 0]
    sample_counter = np.cumsum(num_items).astype(int)

    #fill missing samples, write to new file
    out_file = open(filename_out, 'w')
    prev_problem = 0
    for problem_sample in problem_samples:
        #contiguous part of the data
        out_file.write(data[sample_counter[prev_problem]:sample_counter[problem_sample]].tobytes())

        #add a number of samples corresponding to the number of missing samples, interpolated between the current and next sample
        interpolated_samples = np.repeat(interp1d([0, 1], [data[problem_sample], data[problem_sample+1]])(0.5), missing_samples[problem_sample]).astype(np.complex64)

        out_file.write(interpolated_samples.tobytes())

        prev_problem = problem_sample

    #write out last part of the data
    out_file.write(data[sample_counter[problem_sample]:].tobytes())

    #write header
    copy_first_part_of_header(filename_in, filename_out)

import sys
if __name__ == "__main__":
    if len(sys.argv) <= 2:
        print("Rectifies for overflows in GNU Radio meta file sink files by filling in the missing samples.\n")
        print("Usage: " + sys.argv[0] + " input_filename output_filename")
    else:
        fill_overflows(sys.argv[1], sys.argv[2])


<reponame>code-cullison/pyaspect
import os
import glob
import copy
import datetime
import importlib

import numpy as np
import pandas as pd

from scipy import signal

from pyaspect.specfemio.headers import RecordHeader
from pyaspect.specfemio.read import _read_headers
from pyaspect.specfemio.headers import CMTSolutionHeader
from pyaspect.specfemio.headers import StationHeader

#TODO this is the actual record. I need the header, then make record.py
class Record(RecordHeader):

    class _TraceData(object):

        class _XYZ(object):

            def __init__(self,ex,ny,z):

                self.ex = ex
                self.ny = ny
                self.z  = z

            def __str__(self):
                out_str  = f'Component E/X:\n{self.ex}\n\n'
                out_str += f'Component N/Y:\n{self.ny}\n\n'
                out_str += f'Component Z:\n{self.z}'
                return out_str

            def __repr__(self):
                out_str  = f'Component E/X:\n{self.ex.__repr__()}\n\n'
                out_str += f'Component N/Y:\n{self.ny.__repr__()}\n\n'
                out_str += f'Component Z:\n{self.z.__repr__()}'
                return out_str

        def __init__(self,data_df):

            self.df_x = data_df['comp_EX']
            self.df_y = data_df['comp_NY']
            self.df_z = data_df['comp_Z']

        def __getitem__(self,islice):
            return self._XYZ(self.df_x.loc[islice],self.df_y.loc[islice],self.df_z.loc[islice])

        def __str__(self):
            out_str  = f'Component E/X:\n{self.df_x}\n\n'
            out_str += f'Component N/Y:\n{self.df_y}\n\n'
            out_str += f'Component Z:\n{self.df_z}'
            return out_str

        def __repr__(self):
            out_str  = f'Component E/X:\n{self.df_x.__repr__()}\n\n'
            out_str += f'Component N/Y:\n{self.df_y.__repr__()}\n\n'
            out_str += f'Component Z:\n{self.df_z.__repr__()}'
            return out_str

    @classmethod
    def from_header_and_data(cls,rheader,data_df):
        '''
        Alternative constructor StationHeader

        This constructor takes two requried arguments, a RecordHeader
        and a DataFrame with the data and indexed the same as
        the solutions and stations dataframes in the Record header

        :pyaspect:specfemio:headers:RecordHeader
        :pandas:DataFrame

        * ``rheader``
        * ``data_df``
        '''

        if not isinstance(rheader,RecordHeader):
            raise TypeError(f'arg: \'rheader\' must be of type RecordHeader')

        if not isinstance(data_df,pd.DataFrame):
            raise TypeError(f'arg: \'data_df\' must be of type pandas.DataFrame')

        return Record('None',dtype='b',rheader=rheader,data_df=data_df)

    def __init__(self,header_fqp,dtype='b',rheader=None,data_df=None):

        if not isinstance(rheader,RecordHeader):
            rheader = _read_headers(header_fqp)
        else:
            if not isinstance(data_df,pd.DataFrame):
                raise ArgumentError('data_df must be provided')

        super(Record,self).__init__(name=rheader.name,
                                    solutions_h=rheader.get_solutions_header_list(),
                                    stations_h=rheader.get_stations_header_list(),
                                    proj_id=rheader.proj_id,
                                    rid=rheader.rid,
                                    iter_id=rheader.iter_id,
                                    is_reciprocal=rheader.is_reciprocal)

        self['header_path'] = os.path.dirname(header_fqp)

        if not isinstance(data_df,pd.DataFrame):
            self._load_data(dtype=dtype)
        else:
            self['data_df'] = data_df

        if not self._check_dfs_are_ok():
            raise Exception('pd.DataFrame indices do not match')

    def __str__(self):
        out_str  = f'{super(Record, self).__str__()}\n\n'
        out_str += f'Data:\n {self.data_df}'
        return out_str

    def __repr__(self):
        out_str  = f'{super(Record, self).__repr__()}\n\n'
        out_str += f'Data:\n {self.data_df.__repr__()}'
        return out_str

    def __getitem__(self, kslice):

        if not isinstance(kslice, str):
            dslice = super(Record, self)._get_df_slice_index(kslice,self.data_df,is_stations=True)
            c_data_df = self.data_df.reset_index()[dslice]
            c_rheader = super(Record, self).__getitem__(kslice)
            return Record(rheader=c_rheader,data_df=c_data_df)
        else:
            return super(Record, self).__getitem__(kslice)

    def _read_specfem_bin_trace(self,fpath,dtype=np.float32):
        return np.fromfile(fpath, dtype=dtype)

    def _load_data(self,dtype='b',sl=slice(None,None,None),scale=1.0,rfunc=None):

        if dtype != 'b' and _rfunc == None:
            raise Exception('can only read binary type data for the time being')

        #FIXME: add read ascii
        read_func = self._read_specfem_bin_trace
        if rfunc != None:
            read_func = rfunc

        l_data = []
        for eidx, edf in self.stations_df.groupby(level='eid'):
            for sidx, sdf in edf.groupby(level='sid'):
                for tidx, tdf in sdf.groupby(level='trid'):
                    for gidx, gdf in tdf.groupby(level='gid'):
                        fp_prefix = gdf.loc[(eidx,sidx,tidx,gidx),"data_fqdn"]
                        fp = os.path.join(self.header_path,fp_prefix)
                        match_fp = fp + '.*X[XYZEN].sem*'
                        data_dict = {'eid':eidx,'sid':sidx,'trid':tidx,'gid':gidx}
                        for filepath in glob.glob(match_fp):
                            comp = filepath.split('.')[-2][-1]
                            if comp == 'X' or comp == 'E':
                                data_dict['comp_EX'] = scale*read_func(filepath)
                            elif comp == 'Y' or comp == 'N':
                                data_dict['comp_NY'] = scale*read_func(filepath)
                            elif comp == 'Z':
                                data_dict['comp_Z'] = scale*read_func(filepath)
                            else:
                                raise Exception(f'Could not find component: "{comp}"')

                        l_data.append(data_dict)

        self['data_df'] = pd.DataFrame.from_records(l_data, index=self['default_stat_midx'])

    def _check_dfs_are_ok(self):

        data_idx = self.data_df.index
        solu_idx = self.solutions_df.index
        stat_idx = self.stations_df.index
        stat_2_solu_idx = self.stations_df.reset_index().set_index(['eid','sid']).index

        return all(data_idx == stat_idx) and all(solu_idx == stat_2_solu_idx.unique())

    @property
    def data(self):
        return self._TraceData(self.data_df)

    @property
    def data_df(self):
        return self['data_df']

    @property
    def component_names(self):
        return ['comp_EX','comp_NY','comp_Z']

def make_recip_cmt_record_from_rgf(recip_record,rgf_cmt_data_df,mt_dict):

    solu_df = recip_record.solutions_df
    stat_df = recip_record.stations_df
    l_recip_cmtsolutions = []
    l_recip_cmtstations  = []
    proj_id = recip_record.proj_id
    for eidx, edf in rgf_cmt_data_df.groupby(level='eid'):
        eid    = eidx
        mt     = mt_dict[eid]
        print(f'mt[{eid}]:\n{mt.m6_up_south_east()}')
        for tidx, tdf in edf.groupby(level='trid'):
            date      = datetime.datetime.now()
            lon_xc    = solu_df.loc[(eidx,0),'lon_xc']
            lat_yc    = solu_df.loc[(eidx,0),'lat_yc']
            depth     = solu_df.loc[(eidx,0),'depth']
            elevation = 0.
            network   = stat_df.loc[(tidx,0,eidx,0),'network']
            stat_header = StationHeader(name=f'Reciprocal-Station:{tidx}',
                                        lat_yc=lat_yc,
                                        lon_xc=lon_xc,
                                        depth=depth,
                                        elevation=elevation,
                                        network=network,
                                        proj_id=proj_id,
                                        eid=eid,
                                        sid=0,
                                        trid=tidx,
                                        gid=0)
            l_recip_cmtstations.append(stat_header)
            cmt_lon_xc    = stat_df.loc[(tidx,0,eidx,0),'lon_xc']
            cmt_lat_yc    = stat_df.loc[(tidx,0,eidx,0),'lat_yc']
            cmt_depth     = stat_df.loc[(tidx,0,eidx,0),'depth']

        cmt_header = CMTSolutionHeader(ename=f'Reciprocal-CMT:{eid}',
                                       lat_yc=cmt_lat_yc,
                                       lon_xc=cmt_lon_xc,
                                       depth=cmt_depth,
                                       tshift=0,
                                       date=date,
                                       hdur=0,
                                       mt=mt,
                                       proj_id=proj_id,
                                       eid=eid,
                                       sid=0)

        l_recip_cmtsolutions.append(cmt_header)

    record_h = RecordHeader(name=f'Reciprocal of:{recip_record.name}',
                            solutions_h=l_recip_cmtsolutions,
                            stations_h=l_recip_cmtstations,
                            proj_id=proj_id,
                            rid=recip_record.rid,
                            iter_id=recip_record.iter_id,
                            is_reciprocal=False)

    return Record.from_header_and_data(record_h,rgf_cmt_data_df)

def make_cmt_data_from_rgf(rgf_df,mt_dict,force_stf,cmt_stf):

    comp_dict = {'comp_EX':0,'comp_NY':1,'comp_Z':2}

    rgf_events = list(rgf_df.index.get_level_values('eid').unique())
    #print(f'rgf_events: {rgf_events}')
    mt_events  = list(mt_dict.keys())
    #print(f'mt_events: {mt_events}')
    #print(f'all: {rgf_events == mt_events}')

    if not rgf_events == mt_events:
        raise Exception('RGF-events do not match MomentTensors-events')

    l_recip_cmt_traces = []
    for eidx, edf in rgf_df.groupby(level='eid'):

        mw = mt_dict[eidx].magnitude
        m0 = mt_dict[eidx].moment
        mt_arr = mt_dict[eidx].m6_up_south_east()

        wzz =  mt_arr[0] #mrr
        wyy =  mt_arr[1] #mtt
        wxx =  mt_arr[2] #mpp
        wyz = -mt_arr[3] #mrt
        wxz =  mt_arr[4] #mrp
        wxy = -mt_arr[5] #mtp

        for tidx, tdf in edf.groupby(level='trid'):
            d_recip_cmt = {'eid':eidx,'sid':0,'trid':tidx,'gid':0}
            for comp_key in comp_dict.keys():
                ic = comp_dict[comp_key]

                composite_trace  = wxx*1*rgf_df.loc[(eidx,tidx, 0,ic, 0),'data'] #Matrix: Mee
                composite_trace += wyy*1*rgf_df.loc[(eidx,tidx, 1,ic, 1),'data'] #Matrix: Mnn
                composite_trace += wzz*1*rgf_df.loc[(eidx,tidx, 2,ic, 2),'data'] #Matrix: Mzz

                #Matrix: M1/Mxy
                composite_trace += wxy*1*rgf_df.loc[(eidx,tidx, 1,ic, 0),'data']
                composite_trace += wxy*1*rgf_df.loc[(eidx,tidx, 0,ic, 1),'data']

                #Matrix: M2/Mxz
                composite_trace += wxz*1*rgf_df.loc[(eidx,tidx, 0,ic, 2),'data']
                composite_trace += wxz*1*rgf_df.loc[(eidx,tidx, 2,ic, 0),'data']

                #Matrix: M3/Myz
                composite_trace += wyz*1*rgf_df.loc[(eidx,tidx, 1,ic, 2),'data']
                composite_trace += wyz*1*rgf_df.loc[(eidx,tidx, 2,ic, 1),'data']

                #deconvolve and then convolved
                deconv = 1.0/force_stf[0]
                scaled_trace = deconv*np.convolve(composite_trace,cmt_stf)[:len(cmt_stf)]

                # convert back to single precision
                scaled_trace = scaled_trace.astype(np.float32)

                #add trace index to trace map dictionary
                d_recip_cmt[comp_key] = scaled_trace

            # add the mapped dictionary to a list so that it can be converted to a pd.DataFrame
            l_recip_cmt_traces.append(d_recip_cmt)

    return pd.DataFrame.from_records(l_recip_cmt_traces, index=('eid','sid','trid','gid'))

def make_rgf_data(record,fl,fh,fs):

    comp_dict = {'comp_EX':0,'comp_NY':1,'comp_Z':2}
    coord_dict = {0:'lon_xc',1:'lat_yc',2:'depth'}
    sos = signal.butter(3, [fl,fh], 'bp', fs=fs, output='sos')

    l_rgf_traces = []
    m_df = pd.merge(record.stations_df,record.data_df,on=['eid','sid','trid','gid'])
    for eidx, edf in m_df.groupby(level='eid'):
        for sidx, sdf in edf.groupby(level='sid'):
            for tidx, tdf in sdf.groupby(level='trid'):
                for comp_key in comp_dict.keys():
                    ie = tidx
                    ig = eidx
                    fi = sidx
                    for di in range(3):
                        rgf_dict = {'eid':tidx,'trid':eidx,'fid':sidx}
                        rgf_dict['cid'] = comp_dict[comp_key]
                        coord_key = coord_dict[di]
                        rgf_dict['did'] = di
                        ip1 = di+1     #coord + h
                        im1 = ip1 + 3  #coord - h
                        if di == 2:
                            tm1 = ip1
                            ip1 = im1
                            im1 = tm1
                        rgf_dict['data'] = calulate_spacial_derivative(m_df,
                                                                       eidx,
                                                                       sidx,
                                                                       tidx,
                                                                       ip1,
                                                                       im1,
                                                                       sos,
                                                                       comp_key,
                                                                       coord_key)

                        rgf_dict['component'] = comp_key
                        l_rgf_traces.append(rgf_dict)

    return pd.DataFrame.from_records(l_rgf_traces, index=('eid','trid','cid','fid','did'))

# make_rgf_data is poor choice of name, but some research code needs that
# function name
def calc_dataframe_rgf(record,fl,fh,fs):
    return make_rgf_data(record,fl,fh,fs)

def calulate_spacial_derivative(tdf,eidx,sidx,tidx,g_p1,g_m1,sos,comp_key,coord_key):
    gidx_0  = pd.IndexSlice[eidx,sidx,tidx,0]
    gidx_p1 = pd.IndexSlice[eidx,sidx,tidx,g_p1]
    gidx_m1 = pd.IndexSlice[eidx,sidx,tidx,g_m1]
    df_0    = tdf.loc[gidx_0]
    df_p1   = tdf.loc[gidx_p1]
    df_m1   = tdf.loc[gidx_m1]
    data_p1 = signal.sosfilt(sos, df_p1[comp_key].astype(np.float64))
    data_m1 = signal.sosfilt(sos, df_m1[comp_key].astype(np.float64))
    c_p1    = df_p1[coord_key]
    c_m1    = df_m1[coord_key]
    c_0     = df_0[coord_key]
    delta   = 0.5*(c_p1 - c_m1)
    h       = 2.0*np.abs(delta)
    c       = c_m1 + delta

    assert h != 0
    assert c_0-c == 0

    h_scale  = 1/h
    mt_trace = h_scale*(data_p1 - data_m1)

    return mt_trace

def calc_series_composite_recip_cmt_trace(eid,trid,mt_arr,rgf_df,force_stf,cmt_stf):

    comp_dict = {'comp_EX':0,'comp_NY':1,'comp_Z':2}

    wzz =  mt_arr[0] #mrr
    wyy =  mt_arr[1] #mtt
    wxx =  mt_arr[2] #mpp
    wyz = -mt_arr[3] #mrt
    wxz =  mt_arr[4] #mrp
    wxy = -mt_arr[5] #mtp

    cmt_trace_dict = {'eid':eid, 'trid':trid}
    for comp_key in comp_dict.keys():
        ic = comp_dict[comp_key]

        composite_trace  = wxx*rgf_df.loc[(eid,trid,0,ic, 0),'data'].copy() #Matrix: Mee
        composite_trace += wyy*rgf_df.loc[(eid,trid,1,ic, 1),'data'] #Matrix: Mnn
        composite_trace += wzz*rgf_df.loc[(eid,trid,2,ic, 2),'data'] #Matrix: Mzz

        #Matrix: M1/Mxy
        composite_trace += wxy*rgf_df.loc[(eid,trid,1,ic, 0),'data']
        composite_trace += wxy*rgf_df.loc[(eid,trid,0,ic, 1),'data']

        #Matrix: M2/Mxz
        composite_trace += wxz*rgf_df.loc[(eid,trid,0,ic, 2),'data']
        composite_trace += wxz*rgf_df.loc[(eid,trid,2,ic, 0),'data']

        #Matrix: M3/Myz
        composite_trace += wyz*rgf_df.loc[(eid,trid,1,ic, 2),'data']
        composite_trace += wyz*rgf_df.loc[(eid,trid,2,ic, 1),'data']

        #deconvolve and then convolved
        deconv = 1.0/force_stf[0]
        scaled_trace = deconv*np.convolve(composite_trace.astype(np.float64),cmt_stf.astype(np.float64))[:len(cmt_stf)]

        # convert back to single precision
        cmt_trace_dict[comp_key] = scaled_trace.astype(np.float32)

    return pd.Series(cmt_trace_dict)

def calc_dataframe_composite_recipt_cmt_traces_for_one_event(eid,mt,rgf_df,force_stf,cmt_stf):

    mt_arr = mt.m6_up_south_east()

    edf = None
    ntr = rgf_df.index.get_level_values('trid').nunique()
    for tidx in range(ntr):

        tseries = calc_series_composite_recip_cmt_trace(eid,tidx,mt_arr,rgf_df,force_stf,cmt_stf)
        edf = pd.concat([edf,tseries.to_frame().T])

    edf.set_index(['eid','trid'],inplace=True)
    return edf

def calc_dataframe_composite_recipt_cmt_for_all_events(rgf_df,mt_dict,force_stf,cmt_stf):

    comp_dict = {'comp_EX':0,'comp_NY':1,'comp_Z':2}

    rgf_events = list(rgf_df.index.get_level_values('eid').unique())
    mt_events  = list(mt_dict.keys())

    if not rgf_events == mt_events:
        raise Exception('RGF-events do not match MomentTensors-events')

    rdf = None
    ne = rgf_df.index.get_level_values('eid').nunique()
    for eidx in range(ne):

        df = calc_dataframe_composite_recipt_cmt_traces_for_one_event(eidx,mt_dict[eidx],rgf_df,force_stf,cmt_stf)
        rdf = pd.concat([rdf,df])

    return rdf