import logging
from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass, field
from typing import (
    Any,
    Callable,
    Collection,
    Dict,
    List,
    Mapping,
    Optional,
    Tuple,
    Type,
    Union,
)

from whylogs.core.datatypes import DataType, StandardTypeMapper, TypeMapper
from whylogs.core.metrics.metrics import Metric, MetricConfig
from whylogs.core.resolvers import NO_FI_RESOLVER, MetricSpec, ResolverSpec
from whylogs.core.schema import DeclarativeSchema
from whylogs.core.segmentation_partition import SegmentationPartition
from whylogs.core.stubs import pd
from whylogs.core.validators.validator import Validator
from whylogs.experimental.core.metrics.udf_metric import (
    _reset_metric_udfs,
    generate_udf_resolvers,
)
from whylogs.experimental.core.validators.validator import generate_validators

logger = logging.getLogger(__name__)


@dataclass
class UdfSpec:
    """
    Defines UDFs to apply to matching input columns.

    For UDFs matched by column_name(s), the function is passed a dictionary or
    dataframe with the named columns available (the UDF will not be called unless
    all the named columns are available). The output column name is the key in
    the udfs dictionary.

    For UDFs matched by column_type, the function is passed the value or Pandas
    series. The output column name is the key in the udfs dictionary prefixed
    by the input column name.

    You must specify exactly one of column_names or column_type.
    """

    column_names: Optional[List[str]] = None
    udfs: Dict[str, Callable[[Any], Any]] = field(
        default_factory=dict
    )  # new column name -> callable to compute new column value
    column_type: Optional[DataType] = None

    def __post_init__(self):
        if self.column_type is not None:
            if self.column_names:
                raise ValueError("UdfSpec cannot specify both column_names and column_type")
        elif self.column_names is None:
            raise ValueError("UdfSpec must specify column_names or column_type")
        elif len(self.column_names) == 0 or not all([isinstance(x, str) for x in self.column_names]):
            raise ValueError("UdfSpec column_names must be a non-empty list of strings")


def _apply_udfs_on_row(
    values: Union[List, Dict[str, List]], udfs: Dict, new_columns: Dict[str, Any], input_cols: Collection[str]
) -> None:
    for new_col, udf in udfs.items():
        if new_col in input_cols:
            continue

        try:
            new_columns[new_col] = udf(values)[0]
        except Exception:  # noqa
            new_columns[new_col] = None
            logger.exception(f"Evaluating UDF {new_col} failed")


def _apply_udfs_on_dataframe(
    pandas: pd.DataFrame, udfs: Dict, new_df: pd.DataFrame, input_cols: Collection[str]
) -> None:
    for new_col, udf in udfs.items():
        if new_col in input_cols:
            continue

        try:
            new_df[new_col] = pd.Series(udf(pandas))
        except Exception as e:  # noqa
            new_df[new_col] = pd.Series([None])
            logger.exception(f"Evaluating UDF {new_col} failed on columns {pandas.keys()} with error {e}")


def _apply_type_udfs(pandas: pd.Series, udfs: Dict, new_df: pd.DataFrame, input_cols: Collection[str]) -> None:
    for new_col, udf in udfs.items():
        if new_col in input_cols:
            continue

        try:
            new_df[new_col] = pd.Series(udf(pandas))
        except Exception as e:  # noqa
            new_df[new_col] = pd.Series([None])
            logger.exception(f"Evaluating UDF {new_col} failed on column {new_col} with error {e}")


class UdfSchema(DeclarativeSchema):
    """
    Subclass of DeclarativeSchema that runs the UDFs specified in udf_specs to
    create new columns before resolving metrics.
    """

    def __init__(
        self,
        resolvers: Optional[List[ResolverSpec]] = None,
        types: Optional[Dict[str, Any]] = None,
        default_config: Optional[MetricConfig] = None,
        type_mapper: Optional[TypeMapper] = None,
        cache_size: int = 1024,
        schema_based_automerge: bool = False,
        segments: Optional[Dict[str, SegmentationPartition]] = None,
        validators: Optional[Dict[str, List[Validator]]] = None,
        udf_specs: Optional[List[UdfSpec]] = None,
    ) -> None:
        super().__init__(
            resolvers=resolvers,
            types=types,
            default_config=default_config,
            type_mapper=type_mapper,
            cache_size=cache_size,
            schema_based_automerge=schema_based_automerge,
            segments=segments,
            validators=validators,
        )
        udf_specs = udf_specs if udf_specs else []
        self.multicolumn_udfs = [spec for spec in udf_specs if spec.column_names]
        self.type_udfs = defaultdict(list)
        for spec in udf_specs:
            if spec.column_type:
                self.type_udfs[spec.column_type].append(spec)

    def copy(self) -> "UdfSchema":
        copy = super().copy()
        copy.multicolumn_udfs = deepcopy(self.multicolumn_udfs)
        copy.type_udfs = deepcopy(self.type_udfs)
        return copy

    def _run_udfs_on_row(
        self, row: Mapping[str, Any], new_columns: Dict[str, Any], input_cols: Collection[str]
    ) -> None:
        for spec in self.multicolumn_udfs:
            if spec.column_names and set(spec.column_names).issubset(set(row.keys())):
                inputs = {col: [row[col]] for col in spec.column_names}
                _apply_udfs_on_row(inputs, spec.udfs, new_columns, input_cols)

        for column, value in row.items():
            why_type = type(self.type_mapper(type(value)))
            for spec in self.type_udfs[why_type]:
                udfs = {f"{column}.{key}": spec.udfs[key] for key in spec.udfs.keys()}
                _apply_udfs_on_row([value], udfs, new_columns, input_cols)

    def _run_udfs_on_dataframe(self, pandas: pd.DataFrame, new_df: pd.DataFrame, input_cols: Collection[str]) -> None:
        for spec in self.multicolumn_udfs:
            if spec.column_names and set(spec.column_names).issubset(set(pandas.keys())):
                _apply_udfs_on_dataframe(pandas[spec.column_names], spec.udfs, new_df, input_cols)

        for column, dtype in pandas.dtypes.items():
            why_type = type(self.type_mapper(dtype))
            for spec in self.type_udfs[why_type]:
                udfs = {f"{column}.{key}": spec.udfs[key] for key in spec.udfs.keys()}
                _apply_type_udfs(pandas[column], udfs, new_df, input_cols)

    def _run_udfs(
        self, pandas: Optional[pd.DataFrame] = None, row: Optional[Dict[str, Any]] = None
    ) -> Tuple[Optional[pd.DataFrame], Optional[Mapping[str, Any]]]:
        new_columns = deepcopy(row) if row else None
        new_df = pd.DataFrame() if pandas is not None else None
        if row is not None:
            self._run_udfs_on_row(row, new_columns, row.keys())  # type: ignore

        if pandas is not None:
            self._run_udfs_on_dataframe(pandas, new_df, pandas.keys())
            new_df = pd.concat([pandas, new_df], axis=1)

        return new_df, new_columns

    def apply_udfs(
        self, pandas: Optional[pd.DataFrame] = None, row: Optional[Dict[str, Any]] = None
    ) -> Tuple[Optional[pd.DataFrame], Optional[Mapping[str, Any]]]:
        return self._run_udfs(pandas, row)


_multicolumn_udfs: Dict[str, List[UdfSpec]] = defaultdict(list)
_resolver_specs: Dict[str, List[ResolverSpec]] = defaultdict(list)


def _reset_udfs(reset_metric_udfs: bool = True) -> None:
    if reset_metric_udfs:
        _reset_metric_udfs()

    global _multicolumn_udfs, _resolver_specs
    _multicolumn_udfs = defaultdict(list)
    _resolver_specs = defaultdict(list)


def register_dataset_udf(
    col_names: List[str],
    udf_name: Optional[str] = None,
    metrics: Optional[List[MetricSpec]] = None,
    namespace: Optional[str] = None,
    schema_name: str = "",
    anti_metrics: Optional[List[Metric]] = None,
) -> Callable[[Any], Any]:
    """
    Decorator to easily configure UDFs for your data set. Decorate your UDF
    functions, then call generate_udf_dataset_schema() to create a UdfSchema
    that includes the UDFs configured by your decorator parameters. The decorated
    function will automatically be a UDF in the UdfSchema.

    Specify udf_name to give the output of the UDF a name. udf_name
    defautls to the name of the decorated function. Note that all lambdas are
    named "lambda", so omitting udf_name on more than one lambda will result
    in name collisions. If you pass a namespace, it will be prepended to the UDF name.
    Specifying schema_name will register the UDF in a particular schema. If omitted,
    it will be registered to the defualt schema.

    If any metrics are passed via the metrics argument, they will be attached
    to the column produced by the UDF via the schema returned by generate_udf_dataset_schema().
    If metrics is None, the UDF output column will get the metrics determined by
    the other resolvers passed to generate_udf_dataset_schema(), or the STANDARD_UDF_RESOLVER
    by default. Any anti_metrics will be excluded from the metrics attached to the UDF output.
    """

    def decorator_register(func):
        global _multicolumn_udfs, _resolver_specs
        name = udf_name or func.__name__
        name = f"{namespace}.{name}" if namespace else name
        _multicolumn_udfs[schema_name].append(UdfSpec(col_names, {name: func}))
        if metrics:
            _resolver_specs[schema_name].append(ResolverSpec(name, None, deepcopy(metrics)))
        if anti_metrics:
            _resolver_specs[schema_name].append(ResolverSpec(name, None, [MetricSpec(m) for m in anti_metrics], True))

        return func

    return decorator_register


def unregister_udf(udf_name: str, namespace: Optional[str] = None, schema_name: str = "") -> None:
    global _multicolumn_udfs, _resolver_specs
    name = f"{namespace}.{udf_name}" if namespace else udf_name
    if schema_name not in _multicolumn_udfs:
        logger.warn(f"Can't unregister UDF {name} from non-existant schema {schema_name}")
        return

    found = False
    for spec in _multicolumn_udfs[schema_name]:
        if name in spec.udfs:
            found = True
            del spec.udfs[name]
    if not found:
        logger.warn(f"UDF {name} could not be found for unregistering")
    _resolver_specs[schema_name] = list(filter(lambda x: x.column_name != name, _resolver_specs[schema_name]))


def register_type_udf(
    col_type: Type,
    udf_name: Optional[str] = None,
    namespace: Optional[str] = None,
    schema_name: str = "",
    type_mapper: Optional[TypeMapper] = None,
) -> Callable[[Any], Any]:
    """
    Decorator to easily configure UDFs for your data set. Decorate your UDF
    functions, then call generate_udf_dataset_schema() to create a UdfSchema
    that includes the UDFs configured by your decorator parameters. The decorated
    function will automatically be a UDF in the UdfSchema.

    The registered function will be applied to any columns of the specified type.
    Specify udf_name to give the output of the UDF a name. udf_name
    defautls to the name of the decorated function. The output column name is the
    UDF name prefixed with the input column name. Note that all lambdas are
    named "lambda", so omitting udf_name on more than one lambda will result
    in name collisions. If you pass a namespace, it will be prepended to the UDF name.
    Specifying schema_name will register the UDF in a particular schema. If omitted,
    it will be registered to the defualt schema.

    """
    try:
        if not issubclass(col_type, DataType):
            type_mapper = type_mapper or StandardTypeMapper()
            col_type = type(type_mapper(col_type))
    except:  # noqa
        raise ValueError("Column type must be a DataType or Python type mapable to one")

    def decorator_register(func):
        global _multicolumn_udfs, _resolver_specs
        name = udf_name or func.__name__
        name = f"{namespace}.{name}" if namespace else name
        _multicolumn_udfs[schema_name].append(UdfSpec(None, {name: func}, col_type))

        return func

    return decorator_register


def generate_udf_specs(
    other_udf_specs: Optional[List[UdfSpec]] = None,
    schema_name: Union[str, List[str]] = "",
    include_default_schema: bool = True,
) -> List[UdfSpec]:
    """
    Generates a list UdfSpecs that implement the UDFs specified by the
    @register_dataset_udf, @register_type_udf, and @register_metric_udf
    decorators. You can provide a list of other_udf_specs to include in
    addition to those UDFs registered via the decorator.

    For example:

    @register_dataset_udf(col_names=["col1"])
    def add5(x):
        return x + 5

    schema = UdfSchema(STANDARD_RESOLVER, udf_specs=generate_udf_specs())
    why.log(data, schema=schema)

    This will attach a UDF to column "col1" that will generate a new column
    named "add5" containing the values in "col1" incremented by 5. Since these
    are appended to the STANDARD_UDF_RESOLVER, the default metrics are also tracked
    for every column.
    """
    specs = list(other_udf_specs) if other_udf_specs else []
    schema_name = schema_name if isinstance(schema_name, list) else [schema_name]
    if include_default_schema and "" not in schema_name:
        schema_name = [""] + schema_name

    for name in schema_name:
        specs += _multicolumn_udfs[name]
    return specs


DEFAULT_UDF_SCHEMA_RESOLVER = NO_FI_RESOLVER


def udf_schema(
    other_udf_specs: Optional[List[UdfSpec]] = None,
    resolvers: Optional[List[ResolverSpec]] = None,
    types: Optional[Dict[str, Any]] = None,
    default_config: Optional[MetricConfig] = None,
    type_mapper: Optional[TypeMapper] = None,
    cache_size: int = 1024,
    schema_based_automerge: bool = False,
    segments: Optional[Dict[str, SegmentationPartition]] = None,
    validators: Optional[Dict[str, List[Validator]]] = None,
    schema_name: Union[str, List[str]] = "",
    include_default_schema: bool = True,
) -> UdfSchema:
    """
    Returns a UdfSchema that implements any registered UDFs, along with any
    other_udf_specs or resolvers passed in.
    """
    resolver_specs = list(resolvers if resolvers is not None else DEFAULT_UDF_SCHEMA_RESOLVER)
    schema_names = schema_name if isinstance(schema_name, list) else [schema_name]
    if include_default_schema and "" not in schema_names:
        schema_names = [""] + schema_names

    for name in schema_names:
        resolver_specs += _resolver_specs[name]
        validators = generate_validators(validators, name, include_default_schema=True)

    resolver_specs += generate_udf_resolvers(schema_name, include_default_schema)
    return UdfSchema(
        resolver_specs,
        types,
        default_config,
        type_mapper,
        cache_size,
        schema_based_automerge,
        segments,
        validators,
        generate_udf_specs(other_udf_specs, schema_name, include_default_schema),
    )
