from typing import Iterable, List, Optional, Type, Union

import numpy as np
import pandas as pd
from lamin_utils import logger
from lamindb_setup.dev._docs import doc_args
from lnschema_core import Feature, FeatureSet, Modality, Registry, ids
from lnschema_core.types import FieldAttr, ListLike

from lamindb._utils import attach_func_to_class_method
from lamindb.dev.hashing import hash_set

from . import _TESTING
from ._feature import convert_numpy_dtype_to_lamin_feature_type
from ._registry import init_self_from_db
from ._save import bulk_create
from .dev._priors import NUMBER_TYPE


def get_related_name(features_type: Registry):
    candidates = [
        field.related_name
        for field in FeatureSet._meta.related_objects
        if field.related_model == features_type
    ]
    if not candidates:
        raise ValueError(
            f"Can't create feature sets from {features_type.__name__} because it's not"
            " related to it!\nYou need to create a link model between FeatureSet and"
            " your Registry in your custom schema.\nTo do so, add a"
            " line:\nfeature_sets = models.ManyToMany(FeatureSet,"
            " related_name='mythings')\n"
        )
    return candidates[0]


def validate_features(features: List[Registry]) -> Registry:
    """Validate and return feature type."""
    try:
        if len(features) == 0:
            raise ValueError("Provide list of features with at least one element")
    except TypeError:
        raise ValueError("Please pass a ListLike of features, not a single feature")
    if not hasattr(features, "__getitem__"):
        raise TypeError("features has to be list-like")
    if not isinstance(features[0], Registry):
        raise TypeError(
            "features has to store feature records! use .from_values() otherwise"
        )
    feature_types = set([feature.__class__ for feature in features])
    if len(feature_types) > 1:
        raise TypeError("feature_set can only contain a single type")
    for feature in features:
        if feature._state.adding:
            raise ValueError("Can only construct feature sets from validated features")
    return next(iter(feature_types))  # return value in set of cardinality 1


def __init__(self, *args, **kwargs):
    if len(args) == len(self._meta.concrete_fields):
        super(FeatureSet, self).__init__(*args, **kwargs)
        return None
    # now we proceed with the user-facing constructor
    if len(args) > 1:
        raise ValueError("Only one non-keyword arg allowed: features")
    features: Iterable[Registry] = kwargs.pop("features") if len(args) == 0 else args[0]
    type: Optional[Union[type, str]] = kwargs.pop("type") if "type" in kwargs else None
    modality: Optional[Modality] = (
        kwargs.pop("modality") if "modality" in kwargs else None
    )
    name: Optional[str] = kwargs.pop("name") if "name" in kwargs else None
    if len(kwargs) > 0:
        raise ValueError(
            "Only features, type, modality, name are valid keyword arguments"
        )
    # now code
    features_registry = validate_features(features)
    if type is None:
        type = None if features_registry == Feature else NUMBER_TYPE
    n_features = len(features)
    features_hash = hash_set({feature.id for feature in features})
    feature_set = FeatureSet.filter(hash=features_hash).one_or_none()
    if feature_set is not None:
        logger.success(f"loaded: {feature_set}")
        init_self_from_db(self, feature_set)
        return None
    else:
        hash = features_hash
    self._features = (get_related_name(features_registry), features)
    if modality is not None:
        if not isinstance(modality, Modality):
            raise TypeError("modality needs to be of type Modality")
        if modality._state.adding:
            raise ValueError(
                "unvalidated modality, save to registry if you're sure it is correct:"
                f" {modality}"
            )
    super(FeatureSet, self).__init__(
        id=ids.base62_20(),
        name=name,
        type=get_type_str(type),
        n=n_features,
        modality=modality,
        registry=features_registry.__get_name_with_schema__(),
        hash=hash,
    )


@doc_args(FeatureSet.save.__doc__)
def save(self, *args, **kwargs) -> None:
    """{}"""
    super(FeatureSet, self).save(*args, **kwargs)
    if hasattr(self, "_features"):
        related_name, records = self._features
        # if values are stored in their dedicated table, we can bulk_create
        if related_name != "features":
            bulk_create(records)
        else:
            # otherwise, we currently need to save one by one
            for record in records:
                record.save()
        getattr(self, related_name).set(records)


def get_type_str(type: Optional[Union[Type, str]]) -> Optional[str]:
    if type is not None:
        type_str = type.__name__ if not isinstance(type, str) else type
    else:
        type_str = None
    if type == "int" or type == "float":
        type_str = NUMBER_TYPE
    return type_str


@classmethod  # type:ignore
@doc_args(FeatureSet.from_values.__doc__)
def from_values(
    cls,
    values: ListLike,
    field: FieldAttr = Feature.name,
    type: Optional[Union[Type, str]] = None,
    name: Optional[str] = None,
    modality: Optional[Modality] = None,
    **kwargs,
) -> Optional["FeatureSet"]:
    """{}"""
    if not isinstance(field, FieldAttr):
        raise TypeError(
            "Argument `field` must be a Registry field, e.g., `Feature.name`"
        )
    if len(values) == 0:
        raise ValueError("Provide a list of at least one value")
    registry = field.field.model
    if registry != Feature and type is None:
        type = NUMBER_TYPE
        logger.debug(f"setting feature set to '{NUMBER_TYPE}'")
    validated = registry.validate(values, field=field)
    if validated.sum() == 0:
        logger.warning("no validated features, skip creating feature set")
        return None
    validated_values = np.array(values)[validated]
    validated_features = registry.from_values(validated_values, field=field, **kwargs)
    feature_set = FeatureSet(
        features=validated_features,
        name=name,
        modality=modality,
        type=get_type_str(type),
    )
    return feature_set


@classmethod  # type:ignore
@doc_args(FeatureSet.from_df.__doc__)
def from_df(
    cls,
    df: "pd.DataFrame",
    field: FieldAttr = Feature.name,
    name: Optional[str] = None,
    modality: Optional[Modality] = None,
) -> Optional["FeatureSet"]:
    """{}"""
    registry = field.field.model
    validated = registry.validate(df.columns, field=field)
    if validated.sum() == 0:
        logger.warning("no validated features, skip creating feature set")
        return None
    if registry == Feature:
        validated_features = Feature.from_df(df.loc[:, validated])
        feature_set = FeatureSet(
            validated_features, name=name, type=None, modality=modality
        )
    else:
        dtypes = [col.dtype for (_, col) in df.loc[:, validated].items()]
        if len(set(dtypes)) != 1:
            raise ValueError(f"data types are heterogeneous: {set(dtypes)}")
        type = convert_numpy_dtype_to_lamin_feature_type(dtypes[0])
        validated_features = registry.from_values(df.columns[validated], field=field)
        feature_set = FeatureSet(
            features=validated_features,
            name=name,
            modality=modality,
            type=get_type_str(type),
        )
    return feature_set


METHOD_NAMES = [
    "__init__",
    "from_values",
    "from_df",
    "save",
]

if _TESTING:
    from inspect import signature

    SIGS = {
        name: signature(getattr(FeatureSet, name))
        for name in METHOD_NAMES
        if name != "__init__"
    }

for name in METHOD_NAMES:
    attach_func_to_class_method(name, FeatureSet, globals())
