from __future__ import annotations

from collections import Counter

import pandas as pd


NULL_SENTINEL = "__MISSING__"
OTHER_SENTINEL = "__OTHER__"


def safe_stringify_categories(
    series: pd.Series,
    null_sentinel: str = NULL_SENTINEL,
) -> pd.Series:
    """
    Safely convert any series, including pandas Categoricals, to strings.

    This avoids errors like:
    'Cannot setitem on a Categorical with a new category'
    """
    out = series.copy()

    # Convert categoricals to object/string first to avoid category assignment issues.
    if pd.api.types.is_categorical_dtype(out):
        out = out.astype("object")

    try:
        out = out.where(out.notna(), null_sentinel)
    except Exception:
        out = out.astype("object").fillna(null_sentinel)

    return out.astype(str)


def collapse_rare_categories(
    series: pd.Series,
    min_count: int = 10,
    max_categories: int | None = None,
    other_sentinel: str = OTHER_SENTINEL,
    null_sentinel: str = NULL_SENTINEL,
) -> pd.Series:
    """
    Convert a categorical-like series to strings and collapse rare labels.

    Useful before adversarial validation or lightweight encoding.
    """
    s = safe_stringify_categories(series, null_sentinel=null_sentinel)

    counts = Counter(s.tolist())

    allowed = {value for value, count in counts.items() if count >= min_count}

    if max_categories is not None and len(allowed) > max_categories:
        most_common = [value for value, _ in counts.most_common(max_categories)]
        allowed = set(most_common)

    return s.where(s.isin(allowed), other_sentinel)


def frequency_encode(
    series: pd.Series,
    normalize: bool = True,
    null_sentinel: str = NULL_SENTINEL,
) -> pd.Series:
    """
    Replace each category with its frequency.

    Helpful for tree-based models in adversarial validation.
    """
    s = safe_stringify_categories(series, null_sentinel=null_sentinel)
    freqs = s.value_counts(normalize=normalize, dropna=False).to_dict()
    return s.map(freqs).astype(float)


def lightweight_encode_frame(
    df: pd.DataFrame,
    min_category_count: int = 10,
    max_categories: int = 100,
    null_sentinel: str = NULL_SENTINEL,
    other_sentinel: str = OTHER_SENTINEL,
) -> pd.DataFrame:
    """
    Make a dataframe safe for quick modeling by:
    - leaving numeric/bool columns as-is (with median fill for numeric)
    - converting datetime to integer nanoseconds
    - converting categorical/object columns to frequency encodings
    """
    out = pd.DataFrame(index=df.index)

    for col in df.columns:
        s = df[col]

        try:
            if pd.api.types.is_bool_dtype(s):
                out[col] = s.astype("float").fillna(0.0)

            elif pd.api.types.is_numeric_dtype(s):
                numeric = pd.to_numeric(s, errors="coerce")
                fill_value = float(numeric.median()) if numeric.notna().any() else 0.0
                out[col] = numeric.fillna(fill_value)

            elif pd.api.types.is_datetime64_any_dtype(s):
                dt = pd.to_datetime(s, errors="coerce")
                # int64 nanoseconds, NaT -> sentinel
                values = dt.view("int64")
                values = pd.Series(values, index=s.index)
                nat_mask = dt.isna()
                values.loc[nat_mask] = -1
                out[col] = values.astype("float")

            else:
                collapsed = collapse_rare_categories(
                    s,
                    min_count=min_category_count,
                    max_categories=max_categories,
                    other_sentinel=other_sentinel,
                    null_sentinel=null_sentinel,
                )
                out[col] = frequency_encode(
                    collapsed,
                    normalize=True,
                    null_sentinel=null_sentinel,
                )

        except Exception:
            # Final fallback: stringify and frequency encode.
            safe = safe_stringify_categories(s, null_sentinel=null_sentinel)
            out[col] = frequency_encode(safe, normalize=True, null_sentinel=null_sentinel)

    return out