from __future__ import annotations

from typing import cast
from typing import get_args
from typing import Literal
from typing import Union

import numpy as np
import pyarrow as pa


np_aliases = Literal['numpy', 'ndarray', 'np']
pd_aliases = Literal['pandas', 'pd', 'dataframe', 'df']
pt_aliases = Literal['pytorch', 'torch', 'pt', 'pt']
tf_aliases = Literal['tensorflow', 'tf']
pa_aliases = Literal['pyarrow', 'arrow', 'pa']
jx_aliases = Literal['jax', 'jx']

OutputCategory = Literal['np', 'pd', 'pt', 'tf', 'pa', 'jx']
OutputFormat = Union[
    np_aliases, pd_aliases, pt_aliases, tf_aliases, pa_aliases, jx_aliases,
]

output_aliases = {
    'np': np_aliases, 'pd': pd_aliases, 'pt': pt_aliases,
    'tf': tf_aliases, 'pa': pa_aliases, 'jx': jx_aliases,
}


def check_output_format(
    output_format: OutputFormat, category: OutputCategory | None = None,
) -> bool:
    output_formats: list[str] = []

    if category:
        output_literal = output_aliases.get(category)
        if output_literal is None:
            raise ValueError(f'Unknown output format category: {category}')
        output_formats = list(get_args(output_literal))
    else:
        for output_literal in get_args(OutputFormat):
            output_formats.extend(list(get_args(output_literal)))

    return output_format in output_formats


def validate_output_format(
    output_format: OutputFormat, schema: pa.Schema, as_tensor: bool = False,
) -> None:
    if not check_output_format(output_format=output_format):
        msg = f'Unknown output_format: "{output_format}"\n'
        msg += 'Valid output formats:\n'
        for output_literal in output_aliases.values():
            outputs = ', '.join(x for x in get_args(output_literal))
            msg += f'  {outputs}\n'

    if as_tensor and check_output_format(output_format, category='pd'):
        raise ValueError(
            f'output_format "{output_format}" is not compatible '
            f'with "as_tensor=True"',
        )

    # For PyTorch/TF/Jax we need to check the schema contains only numeric types
    if is_numeric_format(output_format):
        pd_dtypes = [t.to_pandas_dtype() for t in schema.types]
        for name, dt in zip(schema.names, pd_dtypes):
            if not (isinstance(dt, np.dtype) and np.issubdtype(dt, np.number)):
                raise ValueError(
                    f'Column "{name}" with dtype "{dt}" is not compatible '
                    f'with output_format: "{output_format}"',
                )


def is_numeric_format(output_format: OutputFormat) -> bool:
    for category in ('pt', 'tf', 'jx'):
        category = cast(OutputCategory, category)
        if check_output_format(output_format, category=category):
            return True
    return False


def get_format_name(output_format: OutputFormat) -> str:
    for aliases in output_aliases.values():
        aliases_list = list(get_args(aliases))
        if output_format in aliases_list:
            return aliases_list[0]  # the first alias in the list is the format name
    return f'Unknown output format: "{output_format}"'
