"""Cluster bootstrap for did_multiplegt estimator."""

from typing import NamedTuple

import numpy as np

from .numba import compute_column_std, gather_bootstrap_indices


class BootstrapResult(NamedTuple):
    """Result from the cluster bootstrap.

    Attributes
    ----------
    effects_se : ndarray
        Standard errors for effect estimates.
    placebos_se : ndarray or None
        Standard errors for placebo estimates.
    ate_se : float or None
        Standard error for ATE estimate.
    """

    effects_se: np.ndarray
    placebos_se: np.ndarray | None
    ate_se: float | None


def cluster_bootstrap(
    data,
    config,
    compute_func,
    biters=999,
    random_state=None,
):
    """Compute cluster bootstrap standard errors.

    Parameters
    ----------
    data : polars.DataFrame
        The preprocessed panel data.
    config : DIDInterConfig
        Configuration object with estimation parameters.
    compute_func : callable
        Function to compute estimates on bootstrap sample.
    biters : int, default 999
        Number of bootstrap iterations.
    random_state : int or None, default None
        Seed for random number generation.

    Returns
    -------
    BootstrapResult
        NamedTuple containing standard errors for effects, placebos, and ATE.
    """
    rng = np.random.default_rng(random_state)

    bs_group = config.cluster if config.cluster else config.gname

    data_sorted = data.sort(bs_group)
    cluster_col = data_sorted[bs_group].to_numpy()

    unique_clusters, first_indices, counts = np.unique(cluster_col, return_index=True, return_counts=True)
    n_clusters = len(unique_clusters)

    cluster_starts = first_indices.astype(np.int64)
    cluster_counts = counts.astype(np.int64)

    n_effects = config.effects
    n_placebos = config.placebo

    bresults_effects = np.full((biters, n_effects), np.nan)
    bresults_ate = np.full(biters, np.nan) if not config.trends_lin else None
    bresults_placebos = np.full((biters, n_placebos), np.nan) if n_placebos > 0 else None

    for b in range(biters):
        sampled_ids = rng.integers(0, n_clusters, size=n_clusters)
        row_indices = gather_bootstrap_indices(sampled_ids, cluster_starts, cluster_counts)

        df_boot = data_sorted[row_indices.tolist()]
        df_boot = df_boot.sort([config.gname, config.tname])

        result = compute_func(df_boot, config)

        for i in range(min(n_effects, len(result["effects"]))):
            bresults_effects[b, i] = result["effects"][i]

        if bresults_ate is not None and result.get("ate") is not None:
            bresults_ate[b] = result["ate"]

        if bresults_placebos is not None and result.get("placebos") is not None:
            for i in range(min(n_placebos, len(result["placebos"]))):
                bresults_placebos[b, i] = result["placebos"][i]

    effects_se = compute_column_std(bresults_effects)

    placebos_se = None
    if bresults_placebos is not None:
        placebos_se = compute_column_std(bresults_placebos)

    ate_se = None
    if bresults_ate is not None:
        valid_ate = bresults_ate[~np.isnan(bresults_ate)]
        ate_se = np.std(valid_ate, ddof=1) if len(valid_ate) > 1 else np.nan

    return BootstrapResult(
        effects_se=effects_se,
        placebos_se=placebos_se,
        ate_se=ate_se,
    )
