import enum
from collections.abc import Callable, Hashable
from typing import (
    Any,
    Literal,
    NotRequired,
    Required,
    TypeAlias,
    TypedDict,
)

import numpy as np
from pandas import DataFrame
from sentence_transformers import SentenceTransformer
from spacy.language import Language as SpacyModel
from spacy.tokens.doc import Doc as SpacyDoc
from spacy.tokens.token import Token as SpacyToken
from torch import Tensor

__all__ = [
    'SentenceTransformer',
    'SpacyModel',
    'SpacyDoc',
    'SpacyToken',
    'Tensor',
]


# ** logging
class LoggingLevels(enum.IntEnum):
    DEBUG = 10
    INFO = 20
    WARNING = 30
    ERROR = 40
    CRITICAL = 50


# ** models
class LanguageModels(enum.StrEnum):
    SENTENCE_TRANSFORMER = enum.auto()
    SPACY = enum.auto()


class ONNXExecutionProvider(enum.StrEnum):
    CPU = 'CPUExecutionProvider'


class STFRModelTypes(enum.StrEnum):
    ALL_MPNET_BASE_V2 = 'all-mpnet-base-v2'
    ALL_DISTILROBERTA_V1 = 'all-distilroberta-v1'
    ALL_MINI_LM_L12_V2 = 'all-MiniLM-L12-v2'
    ALL_MINI_LM_L6_V2 = 'all-MiniLM-L6-v2'


class SpacyModelTypes(enum.StrEnum):
    DE_CORE_NEWS_SM = 'de_core_news_sm'
    DE_CORE_NEWS_MD = 'de_core_news_md'
    DE_CORE_NEWS_LG = 'de_core_news_lg'
    DE_DEP_NEWS_TRF = 'de_dep_news_trf'


class STFRQuantFilenames(enum.StrEnum):
    ONNX_Q_UINT8 = 'onnx/model_quint8_avx2.onnx'


class STFRModelArgs(TypedDict):
    provider: NotRequired[ONNXExecutionProvider]
    file_name: NotRequired[STFRQuantFilenames]
    export: NotRequired[bool]


Model: TypeAlias = SentenceTransformer | SpacyModel
ModelLoaderFunc: TypeAlias = Callable[..., Model]


class ModelLoaderInfo(TypedDict):
    func: ModelLoaderFunc
    kwargs: dict[str, Any]


ModelLoaderMap: TypeAlias = dict[LanguageModels, ModelLoaderInfo]


class STFRBackends(enum.StrEnum):
    TORCH = enum.auto()
    ONNX = enum.auto()
    OPENVINO = enum.auto()


class STFRDeviceTypes(enum.StrEnum):
    CPU = enum.auto()
    GPU = enum.auto()


# ** pipelines
ResultHandling: TypeAlias = list[tuple[bool, str | None]]


class EntryPoints(enum.StrEnum):
    TIMELINE = 'TIMELINE'
    TIMELINE_POST = 'TIMELINE_POSTPROCESSING'
    TIMELINE_CANDS = 'TIMELINE_CANDIDATES'
    TIMELINE_TK_GRAPH_RESCALED = 'TIMELINE_TK_GRAPH_RESCALED'
    TK_GRAPH_POST = 'TK-GRAPH_POSTPROCESSING'
    TK_GRAPH_ANALYSIS = 'TK-GRAPH_ANALYSIS'
    TK_GRAPH_ANALYSIS_RESCALED = 'TK-GRAPH_ANALYSIS_RESCALED'
    TOKEN_ANALYSIS = 'TOKEN_ANALYSIS'


# ** datasets
PandasIndex: TypeAlias = int | np.int64
ObjectID: TypeAlias = int
Embedding: TypeAlias = SpacyDoc | Tensor
DataFrameTLFiltered: TypeAlias = DataFrame

# ** graphs
NodeTitle: TypeAlias = str
EdgeWeight: TypeAlias = int
CytoExportFileTypes: TypeAlias = Literal[
    'JPEG',
    'PDF',
    'PNG',
    'PS',
    'SVG',
]
CytoExportPageSizes: TypeAlias = Literal[
    'A0',
    'A1',
    'A2',
    'A3',
    'A4',
    'A5',
    'Auto',
    'Legal',
    'Letter',
    'Tabloid',
]
CytoLayouts: TypeAlias = Literal[
    'attribute-circle',
    'attribute-grid',
    'attributes-layout',
    'circular',
    'cose',
    'degree-circle',
    'force-directed',
    'force-directed-cl',
    'fruchterman-rheingold',
    'grid',
    'hierarchical',
    'isom',
    'kamada-kawai',
    'stacked-node-layout',
]
CytoLayoutProperties: TypeAlias = dict[str, float | bool]
CytoNodeID: TypeAlias = int


class NodeData(TypedDict):
    id: NodeTitle
    label: NodeTitle


class EdgeData(TypedDict):
    source: NodeTitle
    target: NodeTitle
    weight: EdgeWeight


class WeightData(TypedDict):
    min: EdgeWeight
    max: EdgeWeight


class CytoscapePosition(TypedDict):
    x: int
    y: int


class CytoscapeData(TypedDict, total=False):
    data: Required[EdgeData | NodeData]
    position: CytoscapePosition
    grabbable: bool
    locked: bool
    classes: str


# ** timeline
TimelineCandidates: TypeAlias = dict[ObjectID, tuple[tuple[PandasIndex, ...], ...]]


# ** Dash (Dashboard)
HTMLTable: TypeAlias = list[dict[Hashable, Any]]
HTMLColumns: TypeAlias = list[dict[str, str]]
