"""Shared Dask utilities, constants, and reusable helpers."""

from __future__ import annotations

from moderndid.distributed._utils import (
    CHUNKED_SE_THRESHOLD,
    MEMMAP_THRESHOLD,
    SE_CHUNK_SIZE,
    auto_tune_partitions,
    chunked_vcov,
    sum_global_stats,
)

__all__ = [
    "CHUNKED_SE_THRESHOLD",
    "MEMMAP_THRESHOLD",
    "SE_CHUNK_SIZE",
    "auto_tune_partitions",
    "chunked_vcov",
    "sum_global_stats",
]


def is_dask_collection(data) -> bool:
    """Check if data is a Dask DataFrame.

    Parameters
    ----------
    data : object
        Input data to check.

    Returns
    -------
    bool
        True if data is a ``dask.dataframe.DataFrame``.
    """
    try:
        import dask.dataframe as dd

        return isinstance(data, dd.DataFrame)
    except ImportError:
        _type_name = type(data).__module__ + "." + type(data).__qualname__
        if "dask" in _type_name.lower():
            raise ImportError(
                f"Input data appears to be a Dask object ({_type_name}) but "
                "the dask extra is not installed. Install with: "
                "uv pip install 'moderndid[dask]'"
            ) from None
        return False


def validate_dask_input(ddf, required_cols):
    """Validate that a Dask DataFrame has the required columns.

    Parameters
    ----------
    ddf : dask.dataframe.DataFrame
        The Dask DataFrame to validate.
    required_cols : list of str
        Column names that must be present.

    Raises
    ------
    ValueError
        If any required columns are missing.
    """
    missing = [c for c in required_cols if c not in ddf.columns]
    if missing:
        raise ValueError(f"Columns not found in Dask DataFrame: {missing}")


def get_default_partitions(client):
    """Compute default partition count from total cluster threads.

    Uses the total number of threads across all workers so that every
    thread has at least one partition to work on, rather than defaulting
    to the number of workers (which leaves most threads idle).

    Parameters
    ----------
    client : distributed.Client
        Dask distributed client.

    Returns
    -------
    int
        Recommended number of partitions (total threads, minimum 1).
    """
    info = client.scheduler_info()
    workers = info.get("workers", {})
    total_threads = sum(w.get("nthreads", 1) for w in workers.values())
    return max(total_threads, 1)


def get_or_create_client(client=None):
    """Get an existing Dask client or create a local one.

    Parameters
    ----------
    client : distributed.Client or None
        An existing Dask distributed client. If None, attempts to get
        the current client or creates a new ``LocalCluster`` client.

    Returns
    -------
    distributed.Client
        A Dask distributed client.
    """
    from distributed import Client

    if client is not None:
        return client
    try:
        return Client.current()
    except ValueError:
        return Client()


def prepare_cohort_wide_pivot(
    client,
    dask_data,
    g,
    cells,
    time_col,
    group_col,
    id_col,
    y_col,
    covariate_cols,
    n_partitions,
    extra_cols=None,
):
    """Build a single wide-pivoted DataFrame for all cells in a cohort.

    Instead of performing one shuffle join per :math:`(g, t)` cell, this
    function collects every time period required by the cohort's cells,
    merges them into a single wide DataFrame (one row per unit, one
    ``_y_{period}`` column per time period), and does a single
    ``repartition().persist()`` cycle.

    Parameters
    ----------
    client : distributed.Client
        Dask distributed client.
    dask_data : dask.dataframe.DataFrame
        Persisted Dask DataFrame in long panel format.
    g : int or float
        Treatment cohort identifier.
    cells : list of tuple
        Cell specifications ``(counter, g, t, pret, post_treat, action)``
        for this cohort. Only ``action == "compute"`` cells are used.
    time_col, group_col, id_col, y_col : str
        Column names.
    covariate_cols : list of str or None
        Covariate column names.
    n_partitions : int
        Number of Dask partitions for the wide DataFrame.
    extra_cols : list of str or None, default None
        Additional columns to include in the base (e.g. partition column
        for DDD).

    Returns
    -------
    wide_dask : dask.dataframe.DataFrame or None
        Persisted wide DataFrame, or ``None`` if no data.
    n_wide : int
        Number of units in the wide DataFrame.
    """
    all_times = set()
    for _counter, _g, t_cell, pret, _pt, action in cells:
        if action != "compute" or pret is None:
            continue
        all_times.add(t_cell)
        all_times.add(pret)

    all_times = sorted(all_times)
    if not all_times:
        return None, 0

    group_filter = (dask_data[group_col] == 0) | (dask_data[group_col] == g)
    filtered = dask_data.loc[group_filter]

    base_cols = [id_col, group_col]
    if extra_cols:
        base_cols = base_cols + [c for c in extra_cols if c not in base_cols]
    if covariate_cols:
        base_cols = base_cols + [c for c in covariate_cols if c not in base_cols]

    base = filtered.loc[filtered[time_col] == all_times[0]][base_cols]

    period_dfs = []
    for tp in all_times:
        period_y = filtered.loc[filtered[time_col] == tp][[id_col, y_col]]
        period_y = period_y.rename(columns={y_col: f"_y_{tp}"})
        period_dfs.append(period_y)

    # Balanced binary merge tree: O(log N) graph depth instead of O(N)
    level = period_dfs
    while len(level) > 1:
        next_level = []
        for i in range(0, len(level), 2):
            if i + 1 < len(level):
                next_level.append(level[i].merge(level[i + 1], on=id_col, how="inner"))
            else:
                next_level.append(level[i])
        level = next_level
    wide_y = level[0]
    base = base.merge(wide_y, on=id_col, how="inner")

    from distributed import wait

    wide_dask = base.repartition(npartitions=n_partitions).persist()
    wait(wide_dask)

    n_wide = len(wide_dask)
    if n_wide == 0:
        return None, 0

    return wide_dask, n_wide
