from collections.abc import Iterable, Sequence
from functools import singledispatch
from typing import List, Optional, Type, Union

import numpy
import pandas
from pandas.core.dtypes.inference import is_sequence
from pydantic import create_model
from pydantic.main import ModelMetaclass

from tktl.core.exceptions import TaktileSdkError
from tktl.core.serializers.base import CustomDeserializingModelT
from tktl.core.serializers.rest.schema import (
    get_array_model,
    get_custom_model,
    get_dataframe_model,
    get_flat_array_model,
    get_mapping_model,
    get_nested_sequence_model,
    get_series_model,
    get_single_value_model,
)
from tktl.core.serializers.rest.utils import pandas_to_python_type
from tktl.core.t import RestSchemaTypes


@singledispatch
def to_pydantic(
    value, name: str = RestSchemaTypes.SINGLE_VALUE.value, unique_id: str = None
) -> Type[CustomDeserializingModelT]:
    """
    Single dispatch function for serializing values of each of the supported kinds. Each
    of the dispatched functions' implementation is tailored for one of Taktile's supported
    input types. All the implementations return a custom Pydantic model that implements
    a `deserialize` method to be used when deserializing incoming data as a REST request

    Parameters
    ----------
    value : Any
        Input value to be serialized
    name : str
        Name of the model
    unique_id : str
        Unique identifier of the model

    Returns
    -------
    single_value_model : CustomDeserializingModelT
        The Pydantic-based model
    """

    single_value_model = get_single_value_model(value)
    single_value_model.__name__ = get_model_unique_name(name=name, unique_id=unique_id)
    return single_value_model


@to_pydantic.register
def _(
    value: ModelMetaclass,
    name: str = RestSchemaTypes.CUSTOM_MODEL.value,
    unique_id: str = None,
):
    model = get_custom_model(value)
    model.__name__ = get_model_unique_name(name=name, unique_id=unique_id)
    return model


@to_pydantic.register
def _(
    value: pandas.DataFrame,
    name: str = RestSchemaTypes.DATAFRAME.value,
    unique_id: str = None,
    nullable=True,
):
    type_map = {}
    example_map = {}

    for col in value.columns:
        var_type = pandas_to_python_type(value[col])
        if var_type is None:
            types = (type(v) for _, v in value[col].iteritems() if not pandas.isna(v))
            try:
                var_type = next(types)  # type of first item that is not None
            except StopIteration:
                raise TaktileSdkError(
                    f"Column {col} contains all null values and unable to infer its type with .dtype"
                )

        if nullable:
            var_type = Optional[var_type]

        type_map[col] = (var_type, None)

        # Make nxt_val (pandas.isna is recursive so need to check for iterable first)
        vals = (
            v
            for _, v in value[col].iteritems()
            if isinstance(v, Iterable) or not pandas.isna(v)
        )

        try:
            nxt_val = next(vals)
        except StopIteration:
            nxt_val = None  # python built-in version of NaN

        example_map[col] = nxt_val

    index_type = type(next((t for t in value.index)))
    type_map["index"] = (index_type, None)

    base_model = create_model(name, **type_map)  # noqa
    inner_model_unique_id = get_model_unique_name(
        name, unique_id=f"{unique_id}__dt_validate"
    )
    dataframe_model = get_dataframe_model(
        base_model=base_model, example=[example_map], unique_id=inner_model_unique_id
    )
    dataframe_model.__name__ = get_model_unique_name(name=name, unique_id=unique_id)
    return dataframe_model


@to_pydantic.register
def _(
    value: pandas.Series,
    name: str = RestSchemaTypes.SERIES.value,
    unique_id: str = None,
) -> Type[CustomDeserializingModelT]:
    nullable = True if value.isna().sum() > 0 else False
    type_name = "Outcome" if not value.name else value.name
    type_map = {}

    try:
        nxt_val = next(
            v
            for v in value.to_dict().values()
            if isinstance(v, Iterable) or not pandas.isna(v)
        )
    except StopIteration:
        nxt_val = None

    example_map = {type_name: [nxt_val]}

    types = (type(v) for k, v in value.to_dict().items() if not pandas.isna(v))
    var_type = next(types)  # type of first item that is not None
    if var_type == bool:
        var_type = Union[var_type, float]
    if nullable:
        var_type = Optional[var_type]
    type_map[type_name] = (List[var_type], None)
    base_model = create_model(name, **type_map)  # noqa
    series_model = get_series_model(
        series=value, base_model=base_model, example=example_map
    )
    series_model.__name__ = get_model_unique_name(name=name, unique_id=unique_id)
    return series_model


@to_pydantic.register
def _(
    value: numpy.ndarray, name: str = RestSchemaTypes.ARRAY.value, unique_id: str = None
) -> Type[CustomDeserializingModelT]:
    as_list = value.tolist()
    if len(value.shape) == 1:
        first_non_empty = next(v for v in as_list if not pandas.isna(v))
        array_model = get_flat_array_model(
            type(first_non_empty), example=[first_non_empty]
        )
        name = RestSchemaTypes.FLAT_ARRAY.value
        array_model.__name__ = name if not unique_id else f"{name}__{unique_id}"
    else:
        example = replace_nans_and_infs(value[0]).tolist()
        list_model = get_nested_sequence_model(as_list)
        array_model = get_array_model(list_model, example=[example])
        array_model.__name__ = get_model_unique_name(name=name, unique_id=unique_id)
    return array_model


@to_pydantic.register
def _(
    value: Sequence, name: str = RestSchemaTypes.SEQUENCE.value, unique_id: str = None
) -> Type[CustomDeserializingModelT]:
    if isinstance(value, list) or isinstance(value, tuple):
        if len(value) > 0:
            first = value[0]
            if isinstance(first, list):
                sequence_model = get_nested_sequence_model(first)
                sequence_model.__name__ = get_model_unique_name(
                    name=name, unique_id=unique_id
                )

            elif isinstance(first, dict):
                mapping_model = _dict_model(
                    value=first, name=RestSchemaTypes.DICT.value
                )
                sequence_model = get_mapping_model(base_model=mapping_model)
                sequence_model.__name__ = get_model_unique_name(
                    name=name, unique_id=unique_id
                )
            # more specific than numpy's is_sequence implementation
            elif not is_sequence(first):
                sequence_model = get_single_value_model(first)
                sequence_model.__name__ = get_model_unique_name(
                    name=name, unique_id=unique_id
                )
            else:
                raise
        else:
            raise ValueError("Input must be non-empty")
    else:
        raise
    return sequence_model


@to_pydantic.register
def _(
    value: dict, name: str = RestSchemaTypes.DICT.value, unique_id: str = None
) -> Type[CustomDeserializingModelT]:
    mapping_model = _dict_model(value=value, name=name)
    dict_model = get_mapping_model(mapping_model)
    dict_model.__name__ = get_model_unique_name(name=name, unique_id=unique_id)
    return dict_model


def _dict_model(value, name):
    types = {k: (type(v), v) for k, v in value.items()}
    return create_model(name, **types)


def _get_df_example(example_map):
    class Config:
        schema_extra = {"example": example_map}

    return Config


def get_model_unique_name(name: str, unique_id: str) -> str:
    return name if not unique_id else f"{name}__{unique_id}"


def replace_nans_and_infs(array) -> numpy.array:
    as_float = numpy.isnan(array.astype(numpy.float))
    as_inf = numpy.isinf(array.astype(numpy.float))
    replaced_nans = numpy.where(as_float, None, array)
    return numpy.where(as_inf, 0.0, replaced_nans)
