# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

from __future__ import annotations

import logging
from typing import Any, Callable, Dict, Literal, Tuple

import numpy as np
import numpy.typing as npt
import pandas as pd
from balance import util as balance_util
from balance.adjustment import trim_weights
from balance.csv_utils import to_csv_with_defaults
from balance.sample_class import Sample
from balance.stats_and_plots import (
    general_stats,
    impact_of_weights_on_outcome,
    weighted_comparisons_plots,
    weighted_comparisons_stats,
    weighted_stats,
    weights_stats,
)
from balance.typing import FilePathOrBuffer
from balance.util import find_items_index_in_list, get_items_from_list_via_indices
from balance.utils.input_validation import _assert_type
from IPython.lib.display import FileLink
from plotly.graph_objs import Figure

logger: "logging.Logger" = logging.getLogger(__package__)


class BalanceDF:
    """
    Wrapper class around a Sample which provides additional balance-specific functionality

    This class encapsulates a pandas DataFrame along with a Sample object reference,
    providing methods for statistical analysis, plotting, and data transformation.
    """

    _model_matrix: Any = None

    def __init__(
        self: "BalanceDF",
        df: pd.DataFrame,
        sample: Sample,
        name: Literal["outcomes", "weights", "covars"],
    ) -> None:
        """A basic init method used by BalanceDFOutcomes,BalanceDFCovars, and BalanceDFWeights

        Args:
            self (BalanceDF): The object that is initiated.
            df (pd.DataFrame): a df from a sample object.
            sample (Sample): A sample object to be stored as reference.
            name (Literal["outcomes", "weights", "covars"]): The type of object that will be created. In practice, used for "outcomes", "weights" and "covars".
        """
        # NOTE: double underscore helps to add friction so that users do not change these objects.
        #       see details here: https://stackoverflow.com/a/1301369/256662
        # TODO: when refactoring the object class model, re-evaluate if we want to keep such objects protected or not.
        self.__sample = sample
        self.__df = df
        self.__name = name

    def __str__(self: "BalanceDF") -> str:
        name = self.__name
        sample = object.__repr__(self._sample)
        df = self.df.__repr__()
        return f"{name} from {sample}:\n{df}"

    def __repr__(self: "BalanceDF") -> str:
        return (
            f"({self.__class__.__module__}.{self.__class__.__qualname__})\n"
            f"{self.__str__()}"
        )

    #  Private API
    @staticmethod
    def _check_if_not_BalanceDF(
        BalanceDF_class_obj: "BalanceDF", object_name: str = "sample_BalanceDF"
    ) -> None:
        """Check if an object is BalanceDF, if not then it raises ValueError

        Args:
            BalanceDF_class_obj (BalanceDF): Object to check.
            object_name (str, optional): Object name (to use when raising the ValueError). Defaults to "sample_BalanceDF".

        Returns: None.

        Raises:
            ValueError: if BalanceDF_class_obj is not BalanceDF object.
        """
        if not isinstance(BalanceDF_class_obj, BalanceDF):
            raise ValueError(
                f"{object_name} must be balancedf_class.BalanceDF, is {type(BalanceDF_class_obj)}"
            )

    @property
    def _sample(self: "BalanceDF") -> "Sample":
        """Access __sample internal object.

        Args:
            self (BalanceDF): Object

        Returns:
            Sample: __sample
        """
        return self.__sample

    @property
    def _weights(
        self: "BalanceDF",
    ) -> pd.DataFrame | None:
        """Access the weight_column in __sample.

        Args:
            self (BalanceDF): Object

        Returns:
            pd.DataFrame | None: The weights (with no column name)
        """
        w = self._sample.weight_column
        return w.rename(None)

    # NOTE: only in the case of BalanceDFOutcomes can it result in a None value.
    def _BalanceDF_child_from_linked_samples(
        self: "BalanceDF",
    ) -> Dict[
        str,
        "BalanceDF"
        | "BalanceDFCovars"
        | "BalanceDFWeights"
        | "BalanceDFOutcomes"
        | None,
    ]:
        """Returns a dict with self and the same type of BalanceDF_child when created from the linked samples.

        For example, if this function is called from a BalanceDFCovars (originally created using `Sample.covars()`),
        that was invoked by a Sample with a target then the return dict will have the keys 'self' and 'target',
        with the BalanceDFCovars of the self and that of the target.

        If the object has nothing but self, then it will be a dict with only one key:value pair (of self).

        Args:
            self (BalanceDF): Object (used in practice only with children of BalanceDF).


        Returns:
            Dict[str, Union["BalanceDFCovars", "BalanceDFWeights", Union["BalanceDFOutcomes", None]],]:
            A dict mapping the link relationship to the result.
                First item is self, and it just returns it without using method on it.
                The other items are based on the objects in _links. E.g.: it can be 'target'
                and 'unadjusted', and it will return them after running the same BalanceDF child creation method on them.

        Examples:
        .. code-block:: python

                from balance.sample_class import Sample
                import pandas as pd

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")


                # keys depends on which samples are linked to the object:
                list(s1.covars()._BalanceDF_child_from_linked_samples().keys())  # ['self']
                list(s3.covars()._BalanceDF_child_from_linked_samples().keys())  # ['self', 'target']
                list(s3_null.covars()._BalanceDF_child_from_linked_samples().keys())  # ['self', 'target', 'unadjusted']

                # Indeed, all are of the same BalanceDF child type:
                s3.covars()._BalanceDF_child_from_linked_samples()
                # {'self': (balance.balancedf_class.BalanceDFCovars)
                # covars from <balance.sample_class.Sample object at 0x7f4392ea61c0>:
                #     a   b  c
                # 0  1 -42  x
                # 1  2   8  y
                # 2  3   2  z
                # 3  1 -42  v,
                # 'target': (balance.balancedf_class.BalanceDFCovars)
                # covars from <balance.sample_class.Sample object at 0x7f43958fbd90>:
                #     a  b  c
                # 0  1  4  x
                # 1  2  6  y
                # 2  3  8  z}

                s3_null.covars()._BalanceDF_child_from_linked_samples()
                # {'self': (balance.balancedf_class.BalanceDFCovars)
                # covars from <balance.sample_class.Sample object at 0x7f4392ea60d0>:
                #     a   b  c
                # 0  1 -42  x
                # 1  2   8  y
                # 2  3   2  z
                # 3  1 -42  v,
                # 'target': (balance.balancedf_class.BalanceDFCovars)
                # covars from <balance.sample_class.Sample object at 0x7f43958fbd90>:
                #     a  b  c
                # 0  1  4  x
                # 1  2  6  y
                # 2  3  8  z,
                # 'unadjusted': (balance.balancedf_class.BalanceDFCovars)
                # covars from <balance.sample_class.Sample object at 0x7f4392ea61c0>:
                #     a   b  c
                # 0  1 -42  x
                # 1  2   8  y
                # 2  3   2  z
                # 3  1 -42  v}

                the_dict = s3_null.covars()._BalanceDF_child_from_linked_samples()
                [v.__class__ for (k,v) in the_dict.items()]
                [balance.balancedf_class.BalanceDFCovars,
                balance.balancedf_class.BalanceDFCovars,
                balance.balancedf_class.BalanceDFCovars]


                # This also works for outcomes (returns None if there is none):
                s3.outcomes()._BalanceDF_child_from_linked_samples()
                # {'self': (balance.balancedf_class.BalanceDFOutcomes)
                #  outcomes from <balance.sample_class.Sample object at 0x7f4392ea61c0>:
                #      o
                #  0   7
                #  1   8
                #  2   9
                #  3  10,
                #  'target': None}

                # And also works for weights:
                s3.weights()._BalanceDF_child_from_linked_samples()
                # {'self': (balance.balancedf_class.BalanceDFWeights)
                #  weights from <balance.sample_class.Sample object at 0x7f4392ea61c0>:
                #       w
                #  0  0.5
                #  1  2.0
                #  2  1.0
                #  3  1.0,
                #  'target': (balance.balancedf_class.BalanceDFWeights)
                #  weights from <balance.sample_class.Sample object at 0x7f43958fbd90>:
                #       w
                #  0  0.5
                #  1  1.0
                #  2  2.0}
        """
        # NOTE: this assumes that the .__name is the same as the creation method (i.e.: .covars(), .weights(), .outcomes())
        BalanceDF_child_method = self.__name
        d: Dict[
            str,
            "BalanceDF"
            | "BalanceDFCovars"
            | "BalanceDFWeights"
            | "BalanceDFOutcomes"
            | None,
        ] = {"self": self}
        d.update(
            {
                k: getattr(v, BalanceDF_child_method)()
                for k, v in self._sample._links.items()
            }
        )
        return d

    def _call_on_linked(
        self: "BalanceDF",
        method: str,
        exclude: tuple[str, ...] = (),
        *args: Any,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Call a given method on the linked DFs of the BalanceDF object.
        Returns the result as a pandas DataFrame where the source column
        indicates where the result came from

        Args:
            self (BalanceDF): Object.
            method (str): A name of a method to call (e.g.: "mean", "std", etc.).
                Can also be a name of an attribute that is a DataFrame (e.g.: 'df')
            exclude (Tuple[str], optional): A tuple of strings which indicates which source should be excluded from the output. Defaults to ().
                E.g.: "self", "target".

        Returns:
            pd.DataFrame: A pandas DataFrame where the source column
                indicates where the result came from. E.g.: 'self', 'target', 'unadjusted'.
                And the columns are based on the method called. E.g.: using mean will give the
                per column mean, after applying `model_matrix` to the df from each object.

        Examples:
        .. code-block:: python

                from balance.sample_class import Sample
                import pandas as pd

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)

                print(s3.covars()._call_on_linked("mean").round(3))
                    #             a       b   c[v]   c[x]   c[y]   c[z]
                    # source
                    # self    1.889 -10.000  0.222  0.111  0.444  0.222
                    # target  2.429   6.857    NaN  0.143  0.286  0.571

                print(s3.covars()._call_on_linked("df").round(3))
                    #         a   b  c
                    # source
                    # self    1 -42  x
                    # self    2   8  y
                    # self    3   2  z
                    # self    1 -42  v
                    # target  1   4  x
                    # target  2   6  y
                    # target  3   8  z
        """
        output = []
        for k, v in self._BalanceDF_child_from_linked_samples().items():
            if v is not None and k not in exclude:
                v_att_method = getattr(v, method)
                if callable(v_att_method):
                    v_att_method = v_att_method(
                        *args, on_linked_samples=False, **kwargs
                    )
                output.append(v_att_method.assign(source=k).set_index("source"))

        return pd.concat(output)

        # return pd.concat(
        #     getattr(v, method)(on_linked_samples=False, *args, **kwargs)
        #     .assign(source=k)
        #     .set_index("source")
        #     if callable(getattr(v, method))
        #     else getattr(v, method)(on_linked_samples=False, *args, **kwargs)
        #     .assign(source=k)
        #     .set_index("source")
        #     for k, v in self._BalanceDF_child_from_linked_samples().items()
        #     if v is not None and k not in exclude
        # )

    def model_matrix(
        self: "BalanceDF", formula: str | list[str] | None = None
    ) -> pd.DataFrame:
        """Return a model_matrix version of the df inside the BalanceDF object using balance_util.model_matrix

        This can be used to turn all character columns into a one hot encoding columns.

        Args:
            self (BalanceDF): Object
            formula (str | list[str] | None, optional): Optional formula string (or list of
                formula strings) to pass to :func:`balance_util.model_matrix`. When
                provided, the model matrix is computed on demand for the formula and
                not cached on the object. Defaults to None.

        Returns:
            pd.DataFrame: The output from :func:`balance_util.model_matrix`

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                print(s1.covars().df)
                    # a   b  c
                    # 0  1 -42  x
                    # 1  2   8  y
                    # 2  3   2  z
                    # 3  1 -42  v

                print(s1.covars().model_matrix())
                    #      a     b  c[v]  c[x]  c[y]  c[z]
                    # 0  1.0 -42.0   0.0   1.0   0.0   0.0
                    # 1  2.0   8.0   0.0   0.0   1.0   0.0
                    # 2  3.0   2.0   0.0   0.0   0.0   1.0
                    # 3  1.0 -42.0   1.0   0.0   0.0   0.0

                print(s1.covars().model_matrix(formula="a + b"))
                    #      a     b
                    # 0  1.0 -42.0
                    # 1  2.0   8.0
                    # 2  3.0   2.0
                    # 3  1.0 -42.0
        """
        if formula is None:
            if not hasattr(self, "_model_matrix") or self._model_matrix is None:
                self._model_matrix = balance_util.model_matrix(
                    self.df, add_na=True, return_type="one"
                )["model_matrix"]
            return self._model_matrix

        return _assert_type(
            balance_util.model_matrix(
                self.df, add_na=True, return_type="one", formula=formula
            )["model_matrix"],
            pd.DataFrame,
        )

    def _descriptive_stats(
        self: "BalanceDF",
        stat: Literal["mean", "std", "var_of_mean", "ci_of_mean", "..."] = "mean",
        weighted: bool = True,
        numeric_only: bool = False,
        add_na: bool = True,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """
        Calls a given method from :func:`weighted_stats.descriptive_stats` on 'self'.
        This function knows how to extract the df and the weights from a BalanceDF object.

        Args:
            self (BalanceDF): An object to run stats on.
            stat (Literal["mean", "std", "var_of_mean", "ci_of_mean", "..."], optional): Defaults to "mean".
            weighted (bool, optional): Defaults to True.
            numeric_only (bool, optional): Defaults to False.
            add_na (bool, optional): Defaults to True.
            **kwargs: extra args to pass to descriptive_stats

        Returns:
            pd.DataFrame: Returns pd.DataFrame of the output (based on stat argument), for each of the columns in df.
        """
        if numeric_only:
            df = self.df.select_dtypes(include=[np.number])
        else:
            df = self.model_matrix()

        weights = (
            self._weights.values if (weighted and self._weights is not None) else None
        )
        wdf = weighted_stats.descriptive_stats(
            df,
            weights,
            stat,
            weighted=weighted,
            # Using numeric_only=True since we know that df is screened in this function
            # To only include numeric variables. So this saves descriptive_stats from
            # running model_matrix again.
            numeric_only=True,
            add_na=add_na,
            **kwargs,
        )
        return wdf

    def to_download(self: "BalanceDF", tempdir: str | None = None) -> FileLink:
        """Creates a downloadable link of the DataFrame, with ids, of the BalanceDF object.

        File name starts with tmp_balance_out_, and some random file name (using :func:`uuid.uuid4`).

        Args:
            self (BalanceDF): Object.
            tempdir (Optional[str], optional): Defaults to None (which then uses a temporary folder using :func:`tempfile.gettempdir`).

        Returns:
            FileLink: Embedding a local file link in an IPython session, based on path. Using :func:FileLink.

        Examples:
        .. code-block:: python

            import pandas as pd
            import tempfile
            from IPython.lib.display import FileLink
            from balance.sample_class import Sample
            sample = Sample.from_frame(
                pd.DataFrame(
                    {"id": ["1", "2"], "x": [0, 1], "weight": [1.0, 2.0]}
                ),
                id_column="id",
                weight_column="weight",
                standardize_types=False,
            )
            isinstance(sample.covars().to_download(tempdir=tempfile.gettempdir()), FileLink)
            # True
        """
        return balance_util._to_download(self._df_with_ids(), tempdir)

    #  Public API
    #  All these functions operate on multiple samples
    @property
    def df(self: "BalanceDF") -> pd.DataFrame:
        """
        Get the df of the BalanceDF object.

        The df is stored in the BalanceDF.__df object, that is set during the __init__ of the object.

        Args:
            self (BalanceDF): The object.

        Returns:
            pd.DataFrame: The df (this is __df, with no weights) from the BalanceDF object.

        Examples:
        .. code-block:: python

            import pandas as pd
            from balance.sample_class import Sample
            sample = Sample.from_frame(
                pd.DataFrame(
                    {"id": ["1", "2"], "x": [0, 1], "weight": [1.0, 2.0]}
                ),
                id_column="id",
                weight_column="weight",
                standardize_types=False,
            )
            sample.covars().df.columns.tolist()
            # ['x']
        """
        return self.__df

    def names(self: "BalanceDF") -> list[str]:
        """Returns the column names of the DataFrame (df) inside a BalanceDF object.

        Args:
            self (BalanceDF): The object.

        Returns:
            List: Of column names.

        Examples:
        .. code-block:: python

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s1.covars().names()
                # ['a', 'b', 'c']
                s1.weights().names()
                # ['w']
                s1.outcomes().names()
                # ['o']
        """
        return list(self.df.columns.values)

    def plot(
        self: "BalanceDF",
        on_linked_samples: bool = True,
        **kwargs: Any,
    ) -> list[Any] | npt.NDArray[Any] | dict[str, Figure] | None:
        """Plots the variables in the df of the BalanceDF object.

        See :func:`weighted_comparisons_plots.plot_dist` for details of various arguments that can be passed.
        The default plotting engine is plotly, but seaborn can be used for static plots.

        This function is inherited as is when invoking BalanceDFCovars.plot, but some modifications are made when
        preparing the data for BalanceDFOutcomes.plot and BalanceDFWeights.plot.

        Args:
            self (BalanceDF): Object (used in the plots as "sample" or "self")
            on_linked_samples (bool, optional): Determines if the linked samples should be included in the plot.
                Defaults to True.
            **kwargs: passed to :func:`weighted_comparisons_plots.plot_dist`.

        Returns:
            Union[Union[List, np.ndarray], Dict[str, Figure], None]:
                If library="plotly" then returns a dictionary containing plots if return_dict_of_figures is True. None otherwise.
                If library="seaborn" then returns None, unless return_axes is True. Then either a list or an np.array of matplotlib axis.

        Examples:
        .. code-block:: python

                import numpy as np
                import pandas as pd
                from numpy import random
                from balance.sample_class import Sample

                random.seed(96483)

                df = pd.DataFrame({
                    "id": range(100),
                    'v1': random.random_integers(11111, 11114, size=100).astype(str),
                    'v2': random.normal(size = 100),
                    'v3': random.uniform(size = 100),
                    "w": pd.Series(np.ones(99).tolist() + [1000]),
                }).sort_values(by=['v2'])

                s1 = Sample.from_frame(df,
                    id_column="id",
                    weight_column="w",
                )

                s2 = Sample.from_frame(
                    df.assign(w = pd.Series(np.ones(100))),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")
                s3_null.set_weights(random.random(size = 100) + 0.5)

                s3_null.covars().plot()
                s3_null.covars().plot(library = "seaborn")

                # Controlling the limits of the y axis using lim:
                s3_null.covars().plot(ylim = (0,1))
                s3_null.covars().plot(library = "seaborn",ylim = (0,1), dist_type = "hist")

                # Returning plotly qq plots:
                s3_null.covars().plot(dist_type = "qq")
        """
        if on_linked_samples:
            dfs_to_add = self._BalanceDF_child_from_linked_samples()
        else:
            dfs_to_add = {"self": self}

        # Create a list of dicts, each dict representing a dataframe and weights
        # Notice that we skip cases in which there is no data (i.e.: v is None)
        # None values are skipped in both dfs and names
        dfs = [
            {"df": v.df, "weight": v._weights}
            for k, v in dfs_to_add.items()
            if (v is not None)
        ]
        names = [k for k, v in dfs_to_add.items() if (v is not None)]

        # re-order dfs and names
        # NOTE: "target", if exists, is placed at the end of the dict so that comparative plotting functions,
        # can use it as the reference distribution for comparison.
        indices_of_ordered_names = find_items_index_in_list(
            names, ["unadjusted", "self", "adjusted", "target"]
        )
        dfs = get_items_from_list_via_indices(dfs, indices_of_ordered_names)
        names = get_items_from_list_via_indices(names, indices_of_ordered_names)

        return weighted_comparisons_plots.plot_dist(dfs, names=names, **kwargs)

    #  NOTE: The following functions use the _call_on_linked method
    #        to return information about the characteristics of linked Samples
    def mean(
        self: "BalanceDF", on_linked_samples: bool = True, **kwargs: Any
    ) -> pd.DataFrame:
        """Calculates a weighted mean on the df of the BalanceDF object.

        Args:
            self (BalanceDF): Object.
            on_linked_samples (bool, optional): Should the calculation be on self AND the linked samples objects? Defaults to True.
                If True, then uses :func:`_call_on_linked` with method "mean".
                If False, then uses :func:`_descriptive_stats` with method "mean".

        Returns:
            pd.DataFrame:
                With row per object: self if on_linked_samples=False, and self and others (e.g.: target and unadjusted) if True.
                Columns are for each of the columns in the relevant df (after applying :func:`model_matrix`)

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")

                print(s3_null.covars().mean())

                #                 a          b      c[v]      c[x]      c[y]      c[z]
                # source
                # self        1.888889 -10.000000  0.222222  0.111111  0.444444  0.222222
                # target      2.428571   6.857143       NaN  0.142857  0.285714  0.571429
                # unadjusted  1.888889 -10.000000  0.222222  0.111111  0.444444  0.222222
        """
        if on_linked_samples:
            return self._call_on_linked("mean", **kwargs)
        else:
            return self._descriptive_stats("mean", **kwargs)

    def std(
        self: "BalanceDF", on_linked_samples: bool = True, **kwargs: Any
    ) -> pd.DataFrame:
        """Calculates a weighted std on the df of the BalanceDF object.

        Args:
            self (BalanceDF): Object.
            on_linked_samples (bool, optional): Should the calculation be on self AND the linked samples objects? Defaults to True.
                If True, then uses :func:`_call_on_linked` with method "std".
                If False, then uses :func:`_descriptive_stats` with method "std".

        Returns:
            pd.DataFrame:
                With row per object: self if on_linked_samples=False, and self and others (e.g.: target and unadjusted) if True.
                Columns are for each of the columns in the relevant df (after applying :func:`model_matrix`)

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")

                print(s3_null.covars().std())

                    #                 a          b  c[v]      c[x]      c[y]      c[z]
                    # source
                    # self        0.886405  27.354812   0.5  0.377964  0.597614  0.500000
                    # target      0.963624   1.927248   NaN  0.462910  0.597614  0.654654
                    # unadjusted  0.886405  27.354812   0.5  0.377964  0.597614  0.500000
        """
        if on_linked_samples:
            return self._call_on_linked("std", **kwargs)
        else:
            return self._descriptive_stats("std", **kwargs)

    def var_of_mean(
        self: "BalanceDF", on_linked_samples: bool = True, **kwargs: Any
    ) -> pd.DataFrame:
        """Calculates a variance of the weighted mean on the df of the BalanceDF object.

        Args:
            self (BalanceDF): Object.
            on_linked_samples (bool, optional): Should the calculation be on self AND the linked samples objects? Defaults to True.
                If True, then uses :func:`_call_on_linked` with method "var_of_mean".
                If False, then uses :func:`_descriptive_stats` with method "var_of_mean".

        Returns:
            pd.DataFrame:
                With row per object: self if on_linked_samples=False, and self and others (e.g.: target and unadjusted) if True.
                Columns are for each of the columns in the relevant df (after applying :func:`model_matrix`)

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample
                from balance.stats_and_plots.weighted_stats import var_of_weighted_mean

                var_of_weighted_mean(pd.Series((1, 2, 3, 1)), pd.Series((0.5, 2, 1, 1)))
                    # 0    0.112178
                    # dtype: float64
                # This shows we got the first cell of 'a' as expected.

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")

                print(s3_null.covars().var_of_mean())
                    #                    a           b      c[v]      c[x]      c[y]      c[z]
                    # source
                    # self        0.112178  134.320988  0.042676  0.013413  0.082914  0.042676
                    # target      0.163265    0.653061       NaN  0.023324  0.069971  0.093294
                    # unadjusted  0.112178  134.320988  0.042676  0.013413  0.082914  0.042676
        """
        if on_linked_samples:
            return self._call_on_linked("var_of_mean", **kwargs)
        else:
            return self._descriptive_stats("var_of_mean", **kwargs)

    def ci_of_mean(
        self: "BalanceDF", on_linked_samples: bool = True, **kwargs: Any
    ) -> pd.DataFrame:
        """Calculates a confidence intervals of the weighted mean on the df of the BalanceDF object.

        Args:
            self (BalanceDF): Object.
            on_linked_samples (bool, optional): Should the calculation be on self AND the linked samples objects? Defaults to True.
                If True, then uses :func:`_call_on_linked` with method "ci_of_mean".
                If False, then uses :func:`_descriptive_stats` with method "ci_of_mean".
            kwargs: we can pass ci_of_mean arguments. E.g.: conf_level and round_ndigits.

        Returns:
            pd.DataFrame:
                With row per object: self if on_linked_samples=False, and self and others (e.g.: target and unadjusted) if True.
                Columns are for each of the columns in the relevant df (after applying :func:`model_matrix`)

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample
                from balance.stats_and_plots.weighted_stats import ci_of_weighted_mean

                ci_of_weighted_mean(pd.Series((1, 2, 3, 1)), pd.Series((0.5, 2, 1, 1)), round_ndigits = 3)
                # 0    (1.232, 2.545)
                # dtype: object
                # This shows we got the first cell of 'a' as expected.

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")

                print(s3_null.covars().ci_of_mean(round_ndigits = 3).T)
                    # source               self           target         unadjusted
                    # a          (1.232, 2.545)   (1.637, 3.221)     (1.232, 2.545)
                    # b       (-32.715, 12.715)   (5.273, 8.441)  (-32.715, 12.715)
                    # c[v]      (-0.183, 0.627)              NaN    (-0.183, 0.627)
                    # c[x]      (-0.116, 0.338)  (-0.156, 0.442)    (-0.116, 0.338)
                    # c[y]       (-0.12, 1.009)  (-0.233, 0.804)     (-0.12, 1.009)
                    # c[z]      (-0.183, 0.627)   (-0.027, 1.17)    (-0.183, 0.627)

                s3_2 = s1.set_target(s2)
                s3_null_2 = s3_2.adjust(method="null")
                print(s3_null_2.outcomes().ci_of_mean(round_ndigits = 3))
                    #                         o
                    # source
                    # self        (7.671, 9.44)
                    # unadjusted  (7.671, 9.44)
        """
        if on_linked_samples:
            return self._call_on_linked("ci_of_mean", **kwargs)
        else:
            return self._descriptive_stats("ci_of_mean", **kwargs)

    def mean_with_ci(
        self: "BalanceDF", round_ndigits: int = 3, on_linked_samples: bool = True
    ) -> pd.DataFrame:
        """
        Returns a table with means and confidence intervals (CIs) for all elements in the BalanceDF object.

        This method calculates the mean and CI for each column of the BalanceDF object using the BalanceDF.mean()
        and BalanceDF.ci_of_mean() methods, respectively. The resulting table contains (for each element such as self, target and adjust) two columns for each input
        column: one for the mean and one for the CI.

        Args:
            self (BalanceDF): The BalanceDF object.
            round_ndigits (int, optional): The number of decimal places to round the mean and CI to.
                Defaults to 3.
            on_linked_samples (bool, optional): A boolean indicating whether to include linked samples
                when calculating the mean. Defaults to True.

        Returns:
            pd.DataFrame: A table with two rows for each input column: one for the mean and one for the CI.
                The columns of the table are labeled with the names of the input columns.

        Examples:
        .. code-block:: python

                import numpy as np
                import pandas as pd

                from balance.sample_class import Sample

                s_o = Sample.from_frame(
                    pd.DataFrame({"o1": (7, 8, 9, 10), "o2": (7, 8, 9, np.nan), "id": (1, 2, 3, 4)}),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )

                t_o = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "o1": (7, 8, 9, 10, 11, 12, 13, 14),
                            "o2": (7, 8, 9, np.nan, np.nan, 12, 13, 14),
                            "id": (1, 2, 3, 4, 5, 6, 7, 8),
                        }
                    ),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )
                s_o2 = s_o.set_target(t_o)

                print(s_o2.outcomes().mean_with_ci())
                    # source            self  target             self           target
                    # _is_na_o2[False]  0.75   0.750   (0.326, 1.174)     (0.45, 1.05)
                    # _is_na_o2[True]   0.25   0.250  (-0.174, 0.674)    (-0.05, 0.55)
                    # o1                8.50  10.500   (7.404, 9.596)  (8.912, 12.088)
                    # o2                6.00   7.875   (2.535, 9.465)  (4.351, 11.399)
        """
        the_means = (
            self.mean(on_linked_samples=on_linked_samples).round(round_ndigits).T
        )
        the_cis = self.ci_of_mean(
            on_linked_samples=on_linked_samples, round_ndigits=round_ndigits
        ).T
        the_cis.columns = the_cis.columns.astype(str) + "_ci"
        return pd.concat([the_means, the_cis], axis=1)

    # NOTE: Summary could return also an str in case it is overridden in other children's methods.
    def summary(
        self: "BalanceDF", on_linked_samples: bool = True
    ) -> pd.DataFrame | str:
        """
        Returns a summary of the BalanceDF object.

        This method currently calculates the mean and confidence interval (CI) for each column of the object
        using the :func:`BalanceDF.mean_with_ci()` method. In the future, this method may be extended to include additional
        summary statistics.

        Args:
            self (BalanceDF): The BalanceDF object.
            on_linked_samples (bool, optional): A boolean indicating whether to include linked samples
                when calculating the mean and CI. Defaults to True.

        Returns:
            Union[pd.DataFrame, str]: A table with two rows for each input column: one for the mean and one for the CI.
                The columns of the table are labeled with the names of the input columns.

        Examples:
        .. code-block:: python

            import pandas as pd
            from balance.sample_class import Sample
            sample = Sample.from_frame(
                pd.DataFrame(
                    {"id": ["1", "2"], "x": [0, 1], "weight": [1.0, 2.0]}
                ),
                id_column="id",
                weight_column="weight",
                standardize_types=False,
            )
            sample.covars().summary().columns.tolist()
            # ['self', 'self_ci']
        """
        # TODO model matrix means to include categorical columns, fix model_matrix to accept DataFrame
        # TODO: include min/max/std/etc. show min/mean/max if there's a single column, just means if multiple (covars and outcomes)
        #       Doing so would either require to implement a min/max etc methods in BalanceDF and use them with _call_on_linked.
        #       Or, update _call_on_linked to deal with non functions, get 'df' from it, and apply the needed functions on it.
        # TODO add outcome variance ratio
        return self.mean_with_ci(on_linked_samples=on_linked_samples)

    def _get_df_and_weights(
        self: "BalanceDF",
        use_model_matrix: bool = True,
    ) -> Tuple[pd.DataFrame, npt.NDArray | None]:
        """Extract df values and weights from a BalanceDF object.

        Args:
            self (BalanceDF): Object
            use_model_matrix (bool, optional): If True, use :func:`model_matrix`.
                If False, use the raw df. Defaults to True.

        Returns:
            Tuple[pd.DataFrame, np.ndarray | None]:
                A pd.DataFrame output from running :func:`model_matrix` or using the raw df, and
                A np.ndarray of weights from :func:`_weights`, or just None (if there are no weights).
        """
        if use_model_matrix:
            df_values = self.model_matrix()
        else:
            df_values = self.df.copy()
        # get weights (like in BalanceDF._descriptive_stats)
        weights = self._weights.values if (self._weights is not None) else None
        return df_values, weights

    @staticmethod
    def _apply_comparison_stat_to_BalanceDF(
        comparison_func: Callable[..., pd.Series],
        sample_BalanceDF: "BalanceDF",
        target_BalanceDF: "BalanceDF",
        aggregate_by_main_covar: bool = False,
        use_model_matrix: bool = True,
        **kwargs: Any,
    ) -> pd.Series:
        """Generic helper to apply a weighted comparison statistic function to two BalanceDF objects.

        This helper function reduces code duplication across multiple comparison methods
        (asmd, kld, emd, cvmd, ks) by extracting the common pattern of:
        1. Validating inputs are BalanceDF objects
        2. Extracting df and weights from both objects
        3. Calling the comparison function with the extracted data

        When ``use_model_matrix`` is False, the raw DataFrames are combined, NA
        indicator columns are added via
        :func:`balance.util.add_na_indicator_to_combined`, and the result is
        split back so that both frames share consistent indicator columns.

        Args:
            comparison_func (Callable[..., pd.Series]): The comparison function from
                weighted_comparisons_stats to apply (e.g., asmd, kld, emd, cvmd, ks).
            sample_BalanceDF (BalanceDF): Sample object.
            target_BalanceDF (BalanceDF): Target object.
            aggregate_by_main_covar (bool, optional): Whether to aggregate by main covariate.
                Defaults to False. Passed to the comparison function.
            use_model_matrix (bool, optional): If True, use :func:`model_matrix` on the
                BalanceDF objects. If False, use raw covariates with NA indicators.
                Defaults to True.
            **kwargs: Additional keyword arguments to pass to the comparison function
                (e.g., std_type for asmd).

        Returns:
            pd.Series: The result from the comparison function.
        """
        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")

        sample_df_values, sample_weights = sample_BalanceDF._get_df_and_weights(
            use_model_matrix=use_model_matrix,
        )
        target_df_values, target_weights = target_BalanceDF._get_df_and_weights(
            use_model_matrix=use_model_matrix,
        )
        if not use_model_matrix:
            combined = balance_util.add_na_indicator_to_combined(
                pd.concat([sample_df_values, target_df_values], axis=0)
            )
            sample_n = sample_df_values.shape[0]
            sample_df_values = combined.iloc[:sample_n].copy()
            target_df_values = combined.iloc[sample_n:].copy()

        return comparison_func(
            sample_df_values,
            target_df_values,
            sample_weights,
            target_weights,
            aggregate_by_main_covar=aggregate_by_main_covar,
            **kwargs,
        )

    @staticmethod
    def _asmd_BalanceDF(
        sample_BalanceDF: "BalanceDF",
        target_BalanceDF: "BalanceDF",
        aggregate_by_main_covar: bool = False,
    ) -> pd.Series:
        """Run asmd on two BalanceDF objects

        Prepares the BalanceDF objects by passing them through :func:`_get_df_and_weights`, and
        then pass the df and weights from the two objects into :func:`weighted_comparisons_stats.asmd`.

        Note that this will works on the result of model_matrix (default behavior, no formula supplied),
        which is different than just the raw covars. E.g.: in case there are nulls (will produce an indicator column of that),
        as well as if there are categorical variables (transforming them using one hot encoding).

        Args:
            sample_df (BalanceDF): Object
            target_df (BalanceDF): Object
            aggregate_by_main_covar (bool, optional): See :func:`weighted_comparisons_stats.asmd`. Defaults to False.

        Returns:
            pd.Series: See :func:`weighted_comparisons_stats.asmd`

        Examples:
        .. code-block:: python

                from balance.balancedf_class import BalanceDF

                BalanceDF._asmd_BalanceDF(
                    Sample.from_frame(
                        pd.DataFrame(
                            {"id": (1, 2), "a": (1, 2), "b": (-1, 12), "weight": (1, 2)}
                        )
                    ).covars(),
                    Sample.from_frame(
                        pd.DataFrame(
                            {"id": (1, 2), "a": (3, 4), "b": (0, 42), "weight": (1, 2)}
                        )
                    ).covars(),
                )

                    # a             2.828427
                    # b             0.684659
                    # mean(asmd)    1.756543
                    # dtype: float64
        """
        return BalanceDF._apply_comparison_stat_to_BalanceDF(
            weighted_comparisons_stats.asmd,
            sample_BalanceDF,
            target_BalanceDF,
            aggregate_by_main_covar,
            std_type="target",
        )

    @staticmethod
    def _kld_BalanceDF(
        sample_BalanceDF: "BalanceDF",
        target_BalanceDF: "BalanceDF",
        aggregate_by_main_covar: bool = False,
    ) -> pd.Series:
        """Run KLD on two BalanceDF objects.

        Prepares the BalanceDF objects by using their raw df (with NA indicators), and
        then passes the df and weights from the two objects into :func:`weighted_comparisons_stats.kld`.

        Args:
            sample_BalanceDF (BalanceDF): Object
            target_BalanceDF (BalanceDF): Object
            aggregate_by_main_covar (bool, optional): See :func:`weighted_comparisons_stats.kld`. Defaults to False.

        Returns:
            pd.Series: See :func:`weighted_comparisons_stats.kld`.
        """
        return BalanceDF._apply_comparison_stat_to_BalanceDF(
            weighted_comparisons_stats.kld,
            sample_BalanceDF,
            target_BalanceDF,
            aggregate_by_main_covar,
            use_model_matrix=False,
        )

    @staticmethod
    def _emd_BalanceDF(
        sample_BalanceDF: "BalanceDF",
        target_BalanceDF: "BalanceDF",
        aggregate_by_main_covar: bool = False,
    ) -> pd.Series:
        """Run EMD on two BalanceDF objects.

        Prepares the BalanceDF objects by using their raw df (with NA indicators), and
        then passes the df and weights into :func:`weighted_comparisons_stats.emd`.

        Args:
            sample_BalanceDF (BalanceDF): Object.
            target_BalanceDF (BalanceDF): Object.
            aggregate_by_main_covar (bool, optional): See :func:`weighted_comparisons_stats.emd`. Defaults to False.

        Returns:
            pd.Series: See :func:`weighted_comparisons_stats.emd`.
        """
        return BalanceDF._apply_comparison_stat_to_BalanceDF(
            weighted_comparisons_stats.emd,
            sample_BalanceDF,
            target_BalanceDF,
            aggregate_by_main_covar,
            use_model_matrix=False,
        )

    @staticmethod
    def _cvmd_BalanceDF(
        sample_BalanceDF: "BalanceDF",
        target_BalanceDF: "BalanceDF",
        aggregate_by_main_covar: bool = False,
    ) -> pd.Series:
        """Run CVMD on two BalanceDF objects.

        Prepares the BalanceDF objects by using their raw df (with NA indicators), and
        then passes the df and weights into :func:`weighted_comparisons_stats.cvmd`.

        Args:
            sample_BalanceDF (BalanceDF): Object.
            target_BalanceDF (BalanceDF): Object.
            aggregate_by_main_covar (bool, optional): See :func:`weighted_comparisons_stats.cvmd`. Defaults to False.

        Returns:
            pd.Series: See :func:`weighted_comparisons_stats.cvmd`.
        """
        return BalanceDF._apply_comparison_stat_to_BalanceDF(
            weighted_comparisons_stats.cvmd,
            sample_BalanceDF,
            target_BalanceDF,
            aggregate_by_main_covar,
            use_model_matrix=False,
        )

    @staticmethod
    def _ks_BalanceDF(
        sample_BalanceDF: "BalanceDF",
        target_BalanceDF: "BalanceDF",
        aggregate_by_main_covar: bool = False,
    ) -> pd.Series:
        """Run KS on two BalanceDF objects.

        Prepares the BalanceDF objects by using their raw df (with NA indicators), and
        then passes the df and weights into :func:`weighted_comparisons_stats.ks`.

        Args:
            sample_BalanceDF (BalanceDF): Object.
            target_BalanceDF (BalanceDF): Object.
            aggregate_by_main_covar (bool, optional): See :func:`weighted_comparisons_stats.ks`. Defaults to False.

        Returns:
            pd.Series: See :func:`weighted_comparisons_stats.ks`.
        """
        return BalanceDF._apply_comparison_stat_to_BalanceDF(
            weighted_comparisons_stats.ks,
            sample_BalanceDF,
            target_BalanceDF,
            aggregate_by_main_covar,
            use_model_matrix=False,
        )

    def asmd(
        self: "BalanceDF",
        on_linked_samples: bool = True,
        target: "BalanceDF" | None = None,
        aggregate_by_main_covar: bool = False,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """ASMD is the absolute difference of the means of two groups (say, P and T), divided by some standard deviation (std).
        It can be std of P or of T, or of P and T.
        These are all variations on the absolute value of cohen's d (see: https://en.wikipedia.org/wiki/Effect_size#Cohen's_d).

        We can use asmd to compares multiple Samples (with and without adjustment) to a target population.

        Args:
            self (BalanceDF): Object from sample (with/without adjustment, but it needs some target)
            on_linked_samples (bool, optional): If to compare also to linked sample objects (specifically: unadjusted), or not. Defaults to True.
            target (Optional["BalanceDF"], optional): A BalanceDF (of the same type as the one used in self) to compare against.
                If None then it looks for a target in the self linked objects. Defaults to None.
            aggregate_by_main_covar (bool, optional): Defaults to False.
                If True, it will make sure to return the asmd DataFrame after averaging
                all the columns from using the one-hot encoding for categorical variables.
                See ::_aggregate_statistic_by_main_covar:: for more details.

        Raises:
            ValueError:
                If self has no target and no target is supplied.

        Returns:
            pd.DataFrame:
                If on_linked_samples is False, then only one row (index name depends on BalanceDF type, e.g.: covars), with asmd of self vs the target (depending if it's covars, or something else).
                If on_linked_samples is True, then two rows per source (self, unadjusted), each with the asmd compared to target, and a third row for the difference (self-unadjusted).

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                from copy import deepcopy


                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")

                s3_null_madeup_weights = deepcopy(s3_null)
                s3_null_madeup_weights.set_weights((1, 2, 3, 1))

                print(s3_null.covars().asmd().round(3))
                    #                     a      b  c[v]   c[x]   c[y]   c[z]  mean(asmd)
                    # source
                    # self               0.56  8.747   NaN  0.069  0.266  0.533       3.175
                    # unadjusted         0.56  8.747   NaN  0.069  0.266  0.533       3.175
                    # unadjusted - self  0.00  0.000   NaN  0.000  0.000  0.000       0.000

                # show that on_linked_samples = False works:
                print(s3_null.covars().asmd(on_linked_samples = False).round(3))
                    #            a      b  c[v]   c[x]   c[y]   c[z]  mean(asmd)
                    # index
                    # covars  0.56  8.747   NaN  0.069  0.266  0.533       3.175

                # verify this also works when we have some weights
                print(s3_null_madeup_weights.covars().asmd())
                    #                           a         b  c[v]  ...      c[y]      c[z]  mean(asmd)
                    # source                                       ...
                    # self               0.296500  8.153742   NaN  ...  0.000000  0.218218    2.834932
                    # unadjusted         0.560055  8.746742   NaN  ...  0.265606  0.533422    3.174566
                    # unadjusted - self  0.263555  0.592999   NaN  ...  0.265606  0.315204    0.33963
        """
        target_from_self = self._BalanceDF_child_from_linked_samples().get("target")

        if target is None:
            target = target_from_self

        if target is None:
            raise ValueError(
                f"Sample {object.__str__(self._sample)} has no target set, or target has no {self.__name} to compare against."
            )
        elif on_linked_samples:
            return balance_util.row_pairwise_diffs(
                self._call_on_linked(
                    "asmd",
                    exclude=("target",),
                    target=target,
                    aggregate_by_main_covar=aggregate_by_main_covar,
                    **kwargs,
                )
            )
        else:
            out = (
                pd.DataFrame(
                    self._asmd_BalanceDF(self, target, aggregate_by_main_covar)
                )
                .transpose()
                .assign(index=(self.__name,))
                .set_index("index")
            )
            return out

    def kld(
        self: "BalanceDF",
        on_linked_samples: bool = True,
        target: "BalanceDF" | None = None,
        aggregate_by_main_covar: bool = False,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Calculate KL divergence (KLD) to measure distributional differences between samples and target.

        KLD is the Kullback-Leibler divergence, an asymmetric measure of how one probability distribution
        differs from a reference distribution. Unlike ASMD which measures standardized mean differences,
        KLD captures differences in the entire distribution shape.

        We can use kld to compare multiple Samples (with and without adjustment) to a target population.

        Args:
            self (BalanceDF): Object from sample (with/without adjustment, but it needs some target)
            on_linked_samples (bool, optional): If to compare also to linked sample objects (specifically: unadjusted), or not. Defaults to True.
                If True, then uses :func:`_call_on_linked` with method "kld".
                If False, then uses :func:`_kld_BalanceDF` directly.
            target (Optional["BalanceDF"], optional): A BalanceDF (of the same type as the one used in self) to compare against.
                If None then it looks for a target in the self linked objects. Defaults to None.
            aggregate_by_main_covar (bool, optional): Defaults to False.
                If True, it will return the KLD results after aggregating columns
                that share a main covariate name into a single metric. This applies
                to columns created from the same underlying covariate via formula
                transformations or interactions (e.g., ``"age"``, ``"log(age)"``,
                ``"age:squared"``), not specifically to categorical one-hot encoding.
                See :func:`_aggregate_statistic_by_main_covar` for more details.

        Raises:
            ValueError:
                If self has no target and no target is supplied.

        Returns:
            pd.DataFrame:
                If on_linked_samples is False, then only one row (index name depends on BalanceDF type, e.g.: covars), with kld of self vs the target.
                If on_linked_samples is True, then two rows per source (self, unadjusted), each with the kld compared to target, and a third row for the difference (self-unadjusted).

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                # Create a simple example with numeric and categorical variables
                sample = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (1, 2, 3, 4),
                            "c": ("x", "x", "y", "z"),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                target = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (1, 2, 2, 3),
                            "c": ("x", "y", "y", "z"),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                sample_with_target = sample.set_target(target)

                # Compare only self to target without linked samples
                print(sample_with_target.covars().kld(on_linked_samples=False).round(6))
                    #            a        c  mean(kld)
                    # index
                    # covars   0.0  0.173287   0.086643
        """

        target_from_self = self._BalanceDF_child_from_linked_samples().get("target")

        if target is None:
            target = target_from_self

        if target is None:
            raise ValueError(
                f"Sample {object.__str__(self._sample)} has no target set, or target has no {self.__name} to compare against."
            )
        elif on_linked_samples:
            return balance_util.row_pairwise_diffs(
                self._call_on_linked(
                    "kld",
                    exclude=("target",),
                    target=target,
                    aggregate_by_main_covar=aggregate_by_main_covar,
                    **kwargs,
                )
            )
        else:
            out = (
                pd.DataFrame(self._kld_BalanceDF(self, target, aggregate_by_main_covar))
                .transpose()
                .assign(index=(self.__name,))
                .set_index("index")
            )
            return out

    def emd(
        self: "BalanceDF",
        on_linked_samples: bool = True,
        target: "BalanceDF" | None = None,
        aggregate_by_main_covar: bool = False,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Calculate Earth Mover's Distance (EMD) to compare distributions to a target.

        See: https://en.wikipedia.org/wiki/Earth_mover%27s_distance

        Args:
            self (BalanceDF): Object from sample (with/without adjustment, but it needs some target).
            on_linked_samples (bool, optional): If to compare also to linked sample objects (specifically: unadjusted).
                If True, then uses :func:`_call_on_linked` with method "emd". Defaults to True.
            target (Optional["BalanceDF"], optional): A BalanceDF (of the same type as the one used in self) to compare against.
                If None then it looks for a target in the self linked objects. Defaults to None.
            aggregate_by_main_covar (bool, optional): Defaults to False.
                If True, it will make sure to return the emd DataFrame after averaging
                all the columns from using the one-hot encoding for categorical variables.
                See :func:`_aggregate_statistic_by_main_covar` for more details.

        Raises:
            ValueError: If self has no target and no target is supplied.

        Returns:
            pd.DataFrame:
                If on_linked_samples is False, then only one row (index name depends on BalanceDF type, e.g.: covars),
                with emd of self vs the target.
                If on_linked_samples is True, then two rows per source (self, unadjusted), each with the emd compared
                to target, and a third row for the difference (self-unadjusted).

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                sample = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (1, 2, 3, 4),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                target = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (2, 3, 4, 5),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                sample.set_target(target).covars().emd(on_linked_samples=False)
        """
        target_from_self = self._BalanceDF_child_from_linked_samples().get("target")

        if target is None:
            target = target_from_self

        if target is None:
            raise ValueError(
                f"Sample {object.__str__(self._sample)} has no target set, or target has no {self.__name} to compare against."
            )
        elif on_linked_samples:
            return balance_util.row_pairwise_diffs(
                self._call_on_linked(
                    "emd",
                    exclude=("target",),
                    target=target,
                    aggregate_by_main_covar=aggregate_by_main_covar,
                    **kwargs,
                )
            )
        else:
            out = (
                pd.DataFrame(self._emd_BalanceDF(self, target, aggregate_by_main_covar))
                .transpose()
                .assign(index=(self.__name,))
                .set_index("index")
            )
            return out

    def cvmd(
        self: "BalanceDF",
        on_linked_samples: bool = True,
        target: "BalanceDF" | None = None,
        aggregate_by_main_covar: bool = False,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Calculate Cramér-von Mises distance (CVMD) to compare distributions to a target.

        See: https://en.wikipedia.org/wiki/Cram%C3%A9r%E2%80%93von_Mises_criterion

        Args:
            self (BalanceDF): Object from sample (with/without adjustment, but it needs some target).
            on_linked_samples (bool, optional): If to compare also to linked sample objects (specifically: unadjusted).
                If True, then uses :func:`_call_on_linked` with method "cvmd". Defaults to True.
            target (Optional["BalanceDF"], optional): A BalanceDF (of the same type as the one used in self) to compare against.
                If None then it looks for a target in the self linked objects. Defaults to None.
            aggregate_by_main_covar (bool, optional): Defaults to False.
                If True, it will make sure to return the cvmd DataFrame after averaging
                all the columns from using the one-hot encoding for categorical variables.
                See :func:`_aggregate_statistic_by_main_covar` for more details.

        Raises:
            ValueError: If self has no target and no target is supplied.

        Returns:
            pd.DataFrame:
                If on_linked_samples is False, then only one row (index name depends on BalanceDF type, e.g.: covars),
                with cvmd of self vs the target.
                If on_linked_samples is True, then two rows per source (self, unadjusted), each with the cvmd compared
                to target, and a third row for the difference (self-unadjusted).

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                sample = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (1, 2, 3, 4),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                target = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (2, 3, 4, 5),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                sample.set_target(target).covars().cvmd(on_linked_samples=False)
        """
        target_from_self = self._BalanceDF_child_from_linked_samples().get("target")

        if target is None:
            target = target_from_self

        if target is None:
            raise ValueError(
                f"Sample {object.__str__(self._sample)} has no target set, or target has no {self.__name} to compare against."
            )
        elif on_linked_samples:
            return balance_util.row_pairwise_diffs(
                self._call_on_linked(
                    "cvmd",
                    exclude=("target",),
                    target=target,
                    aggregate_by_main_covar=aggregate_by_main_covar,
                    **kwargs,
                )
            )
        else:
            out = (
                pd.DataFrame(
                    self._cvmd_BalanceDF(self, target, aggregate_by_main_covar)
                )
                .transpose()
                .assign(index=(self.__name,))
                .set_index("index")
            )
            return out

    def ks(
        self: "BalanceDF",
        on_linked_samples: bool = True,
        target: "BalanceDF" | None = None,
        aggregate_by_main_covar: bool = False,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Calculate Kolmogorov-Smirnov (KS) statistic to compare distributions to a target.

        See: https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Smirnov_test

        Args:
            self (BalanceDF): Object from sample (with/without adjustment, but it needs some target).
            on_linked_samples (bool, optional): If to compare also to linked sample objects (specifically: unadjusted).
                If True, then uses :func:`_call_on_linked` with method "ks". Defaults to True.
            target (Optional["BalanceDF"], optional): A BalanceDF (of the same type as the one used in self) to compare against.
                If None then it looks for a target in the self linked objects. Defaults to None.
            aggregate_by_main_covar (bool, optional): Defaults to False.
                If True, it will make sure to return the ks DataFrame after averaging
                all the columns from using the one-hot encoding for categorical variables.
                See :func:`_aggregate_statistic_by_main_covar` for more details.

        Raises:
            ValueError: If self has no target and no target is supplied.

        Returns:
            pd.DataFrame:
                If on_linked_samples is False, then only one row (index name depends on BalanceDF type, e.g.: covars),
                with ks of self vs the target.
                If on_linked_samples is True, then two rows per source (self, unadjusted), each with the ks compared
                to target, and a third row for the difference (self-unadjusted).

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                sample = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (1, 2, 3, 4),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                target = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "a": (2, 3, 4, 5),
                            "w": (1, 1, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                sample.set_target(target).covars().ks(on_linked_samples=False)
        """
        target_from_self = self._BalanceDF_child_from_linked_samples().get("target")

        if target is None:
            target = target_from_self

        if target is None:
            raise ValueError(
                f"Sample {object.__str__(self._sample)} has no target set, or target has no {self.__name} to compare against."
            )
        elif on_linked_samples:
            return balance_util.row_pairwise_diffs(
                self._call_on_linked(
                    "ks",
                    exclude=("target",),
                    target=target,
                    aggregate_by_main_covar=aggregate_by_main_covar,
                    **kwargs,
                )
            )
        else:
            out = (
                pd.DataFrame(self._ks_BalanceDF(self, target, aggregate_by_main_covar))
                .transpose()
                .assign(index=(self.__name,))
                .set_index("index")
            )
            return out

    def asmd_improvement(
        self: "BalanceDF",
        unadjusted: "BalanceDF" | None = None,
        target: "BalanceDF" | None = None,
    ) -> np.float64:
        """Calculates the improvement in mean(asmd) from before to after applying some weight adjustment.

        See :func:`weighted_comparisons_stats.asmd_improvement` for details.

        Args:
            self (BalanceDF): BalanceDF (e.g.: of self after adjustment)
            unadjusted (Optional["BalanceDF"], optional): BalanceDF (e.g.: of self before adjustment). Defaults to None.
            target (Optional["BalanceDF"], optional): To compare against. Defaults to None.

        Raises:
            ValueError: If target is not linked in self and also not provided to the function.
            ValueError: If unadjusted is not linked in self and also not provided to the function.

        Returns:
            np.float64: The improvement is taking the (before_mean_asmd-after_mean_asmd)/before_mean_asmd.
                The asmd is calculated using :func:`asmd`.

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                from copy import deepcopy


                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                s2 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3),
                            "b": (4, 6, 8),
                            "id": (1, 2, 3),
                            "w": (0.5, 1, 2),
                            "c": ("x", "y", "z"),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")

                s3_null_madeup_weights = deepcopy(s3_null)
                s3_null_madeup_weights.set_weights((1, 2, 3, 1))

                s3_null.covars().asmd_improvement() # 0. since unadjusted is just a copy of self
                s3_null_madeup_weights.covars().asmd_improvement() # 0.10698596233975825

                asmd_df = s3_null_madeup_weights.covars().asmd()
                print(asmd_df["mean(asmd)"])
                    # source
                    # self                 2.834932
                    # unadjusted           3.174566
                    # unadjusted - self    0.339634
                    # Name: mean(asmd), dtype: float64
                (asmd_df["mean(asmd)"][1] - asmd_df["mean(asmd)"][0]) / asmd_df["mean(asmd)"][1]  # 0.10698596233975825
                # just like asmd_improvement
        """
        if unadjusted is None:
            unadjusted = self._BalanceDF_child_from_linked_samples().get("unadjusted")
        if unadjusted is None:
            raise ValueError(
                f"Sample {object.__repr__(self._sample)} has no unadjusted set or unadjusted has no {self.__name}."
            )

        if target is None:
            target = self._BalanceDF_child_from_linked_samples().get("target")
        if target is None:
            raise ValueError(
                f"Sample {object.__repr__(self._sample)} has no target set or target has no {self.__name}."
            )

        sample_before_df, sample_before_weights = unadjusted._get_df_and_weights()
        sample_after_df, sample_after_weights = self._get_df_and_weights()
        target_df, target_weights = target._get_df_and_weights()

        return weighted_comparisons_stats.asmd_improvement(
            sample_before=sample_before_df,
            sample_after=sample_after_df,
            target=target_df,
            sample_before_weights=sample_before_weights,
            sample_after_weights=sample_after_weights,
            target_weights=target_weights,
        )

    def _df_with_ids(self: "BalanceDF") -> pd.DataFrame:
        """Creates a DataFrame of the BalanceDF, with ids.

        Args:
            self (BalanceDF): Object.

        Returns:
            pd.DataFrame: DataFrame with id_column and then the df.
        """
        return pd.concat((self._sample.id_column, self.df), axis=1)

    def to_csv(
        self: "BalanceDF",
        path_or_buf: FilePathOrBuffer | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> str | None:
        """Write df with ids from BalanceDF to a comma-separated values (csv) file.

        Uses :func:`pd.DataFrame.to_csv`.

        If an 'index' argument is not provided then it defaults to False.

        Args:
            self (BalanceDF): Object.
            path_or_buf (Optional[FilePathOrBuffer], optional): location where to save the csv.

        Returns:
            Optional[str]: If path_or_buf is None, returns the resulting csv format as a string. Otherwise returns None.

        Examples:
        .. code-block:: python

            import pandas as pd
            from balance.sample_class import Sample
            sample = Sample.from_frame(
                pd.DataFrame(
                    {"id": ["1", "2"], "x": [0, 1], "weight": [1.0, 2.0]}
                ),
                id_column="id",
                weight_column="weight",
                standardize_types=False,
            )
            csv_text = sample.covars().to_csv()
            "id" in csv_text
            # True
        """
        return to_csv_with_defaults(self._df_with_ids(), path_or_buf, *args, **kwargs)


class BalanceDFOutcomes(BalanceDF):
    def __init__(self: "BalanceDFOutcomes", sample: Sample) -> None:
        """A factory function to create BalanceDFOutcomes

        This is used through :func:`Sample.outcomes`.
        It initiates a BalanceDFOutcomes object by passing the relevant arguments to
        :func:`BalanceDF.__init__`.

        Args:
            self (BalanceDFOutcomes): Object that is initiated.
            sample (Sample): Object
        """
        super().__init__(sample._outcome_columns, sample, name="outcomes")

    # TODO: add the `relative_to` argument (with options 'self' and 'target')
    #       this will also require to update _relative_response_rates a bit.
    def relative_response_rates(
        self: "BalanceDFOutcomes",
        target: bool | pd.DataFrame = False,
        per_column: bool = False,
    ) -> pd.DataFrame | None:
        """Produces a summary table of number of responses and proportion of completed responses.

        See :func:`general_stats.relative_response_rates`.

        Args:
            self (BalanceDFOutcomes): Object
            target (Union[bool, pd.DataFrame], optional): Defaults to False.
                Determines what is passed to df_target in :func:`general_stats.relative_response_rates`
                If False: passes None.
                If True: passes the df from the target of sample (notice, it's the df of target, NOT target.outcome().df).
                    So it means it will count only rows that are all notnull rows (so if the target has covars and outcomes,
                        both will need to be notnull to be counted).
                    If you want to control this in a more specific way, pass pd.DataFrame instead.
                If pd.DataFrame: passes it as is.
            per_column (bool, optional): Default is False. See :func:`general_stats.relative_response_rates`.

        Returns:
            Optional[pd.DataFrame]: A column per outcome, and two rows.
                One row with number of non-null observations, and
                A second row with the proportion of non-null observations.

                If 'target' is set to True but there is no target, the function returns None.

        Examples:
        .. code-block:: python

                import numpy as np
                import pandas as pd

                from balance.sample_class import Sample


                s_o = Sample.from_frame(
                    pd.DataFrame({"o1": (7, 8, 9, 10), "o2": (7, 8, 9, np.nan), "id": (1, 2, 3, 4)}),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )

                print(s_o.outcomes().relative_response_rates())
                    #       o1    o2
                    # n    4.0   3.0
                    # %  100.0  75.0

                s_o.outcomes().relative_response_rates(target = True)
                # None

                # compared with a larger target

                t_o = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "o1": (7, 8, 9, 10, 11, 12, 13, 14),
                            "o2": (7, 8, 9, np.nan, np.nan, 12, 13, 14),
                            "id": (1, 2, 3, 4, 5, 6, 7, 8),
                        }
                    ),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )
                s_o2 = s_o.set_target(t_o)

                print(s_o2.outcomes().relative_response_rates(True, per_column = True))
                    #     o1    o2
                    # n   4.0   3.0
                    # %  50.0  50.0

                df_target = pd.DataFrame(
                        {
                            "o1": (7, 8, 9, 10, 11, 12, 13, 14),
                            "o2": (7, 8, 9, np.nan, np.nan, 12, 13, 14),
                        }
                    )

                print(s_o2.outcomes().relative_response_rates(target = df_target, per_column = True))
                    #     o1    o2
                    # n   4.0   3.0
                    # %  50.0  50.0
        """
        if type(target) is bool:
            # Then: get target from self:
            if target:
                self_target = self._BalanceDF_child_from_linked_samples().get("target")
                if self_target is None:
                    logger.warning("Sample does not have target set")
                    return None
                else:
                    df_target = self_target.df
            else:
                df_target = None
        else:
            df_target = target

        return general_stats.relative_response_rates(
            self.df, df_target, per_column=per_column
        )

    def target_response_rates(self: "BalanceDFOutcomes") -> pd.DataFrame | None:
        """Calculates relative_response_rates for the target in a Sample object.

        See :func:`general_stats.relative_response_rates`.

        Args:
            self (BalanceDFOutcomes): Object (with/without a target set)

        Returns:
            Optional[pd.DataFrame]: None if the object doesn't have a target.
                If the object has a target, it returns the output of :func:`general_stats.relative_response_rates`.

        Examples:
        .. code-block:: python

                import numpy as np
                import pandas as pd

                from balance.sample_class import Sample


                s_o = Sample.from_frame(
                    pd.DataFrame({"o1": (7, 8, 9, 10), "o2": (7, 8, 9, np.nan), "id": (1, 2, 3, 4)}),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )

                t_o = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "o1": (7, 8, 9, 10, 11, 12, 13, 14),
                            "o2": (7, 8, 9, np.nan, 11, 12, 13, 14),
                            "id": (1, 2, 3, 4, 5, 6, 7, 8),
                        }
                    ),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )
                s_o = s_o.set_target(t_o)

                print(s_o.outcomes().target_response_rates())
                    #       o1    o2
                    # n    8.0   7.0
                    # %  100.0  87.5
        """
        self_target = self._BalanceDF_child_from_linked_samples().get("target")
        if self_target is None:
            logger.warning("Sample does not have target set")
            return None
        else:
            return general_stats.relative_response_rates(self_target.df)

    def weights_impact_on_outcome_ss(
        self: "BalanceDFOutcomes",
        method: str = "t_test",
        conf_level: float = 0.95,
        round_ndigits: int | None = 3,
        w0: npt.NDArray | pd.Series | None = None,
        w1: npt.NDArray | pd.Series | None = None,
    ) -> pd.DataFrame | None:
        """
        Compare weighted outcomes using paired tests on y*w0 vs y*w1.

        Args:
            self (BalanceDFOutcomes): Object.
            method (str, optional): Statistical test to use. Defaults to "t_test".
            conf_level (float, optional): Confidence level for the mean difference interval.
                Defaults to 0.95.
            round_ndigits (int | None, optional): Optional rounding for numeric outputs.
                Defaults to 3.
            w0: Baseline weights. Defaults to a vector of ones.
            w1: Alternative weights. Defaults to the BalanceDF weights.

        Returns:
            Optional[pd.DataFrame]: Outcome-by-statistic table or None if weights are missing.

        Examples:
        .. code-block:: python

                import pandas as pd

                from balance.sample_class import Sample

                sample = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "id": (1, 2, 3, 4),
                            "weight": (1.0, 2.0, 1.0, 2.0),
                            "outcome": (1.0, 2.0, 3.0, 4.0),
                        }
                    ),
                    id_column="id",
                    weight_column="weight",
                    outcome_columns=("outcome",),
                    standardize_types=False,
                )

                impact = sample.outcomes().weights_impact_on_outcome_ss(
                    method="t_test", round_ndigits=3
                )
                print(impact.to_string())

        .. code-block:: text

                    mean_yw0  mean_yw1  mean_diff  diff_ci_lower  diff_ci_upper  t_stat  p_value    n
            outcome
            outcome       2.5       4.0        1.5         -1.547          4.547   1.567    0.215  4.0
        """
        if w1 is None:
            if self._weights is None:
                logger.warning("No weights available for outcome impact analysis.")
                return None
            w1 = self._weights

        w1_values = (
            w1.to_numpy() if isinstance(w1, pd.Series) else np.asarray(w1, dtype=float)
        )
        if w0 is None:
            w0_values = np.ones_like(w1_values, dtype=float)
        else:
            w0_values = (
                w0.to_numpy()
                if isinstance(w0, pd.Series)
                else np.asarray(w0, dtype=float)
            )

        model_matrix = self.model_matrix()
        results = {}
        for column in model_matrix.columns:
            results[column] = impact_of_weights_on_outcome.weights_impact_on_outcome_ss(
                y=model_matrix[column].to_numpy(),
                w0=w0_values,
                w1=w1_values,
                method=method,
                conf_level=conf_level,
            )

        impact_df = pd.DataFrame(results).T
        impact_df.index.name = "outcome"
        if round_ndigits is not None:
            numeric_cols = impact_df.select_dtypes(include=["number"]).columns
            impact_df[numeric_cols] = impact_df[numeric_cols].round(round_ndigits)
        return impact_df

    # TODO: it's a question if summary should produce a printable output or a DataFrame.
    #       The BalanceDF.summary method only returns a DataFrame. So it's a question
    #       what is the best way to structure this more generally.
    def summary(
        self: "BalanceDFOutcomes",
        on_linked_samples: bool | None = None,
        weights_impact_method: str | None = "t_test",
        weights_impact_conf_level: float = 0.95,
    ) -> str:
        """Produces summary printable string of a BalanceDFOutcomes object.

        Args:
            self (BalanceDFOutcomes): Object.
            on_linked_samples (Optional[bool]): Ignored. Only here since summary overrides BalanceDF.summary.
            weights_impact_method (Optional[str]): If provided, include a paired test
                of y*w0 vs y*w1 for each outcome (default is "t_test").
            weights_impact_conf_level (float): Confidence level for the mean difference
                interval when weights_impact_method is provided. Defaults to 0.95.

        Returns:
            str: A printable string, with mean of outcome variables and response rates.

        Examples:
        .. code-block:: python

                import numpy as np
                import pandas as pd

                from balance.sample_class import Sample

                s_o = Sample.from_frame(
                    pd.DataFrame({"o1": (7, 8, 9, 10), "o2": (7, 8, 9, np.nan), "id": (1, 2, 3, 4)}),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )

                t_o = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "o1": (7, 8, 9, 10, 11, 12, 13, 14),
                            "o2": (7, 8, 9, np.nan, np.nan, 12, 13, 14),
                            "id": (1, 2, 3, 4, 5, 6, 7, 8),
                        }
                    ),
                    id_column="id",
                    outcome_columns=("o1", "o2"),
                )
                s_o2 = s_o.set_target(t_o)

                print(s_o.outcomes().summary())
                    # 2 outcomes: ['o1' 'o2']
                    # Mean outcomes (with 95% confidence intervals):
                    # source            self             self
                    # _is_na_o2[False]  0.75   (0.326, 1.174)
                    # _is_na_o2[True]   0.25  (-0.174, 0.674)
                    # o1                8.50   (7.404, 9.596)
                    # o2                6.00   (2.535, 9.465)

                    # Weights impact on outcomes (t_test):
                    #                   mean_yw0  mean_yw1  mean_diff  diff_ci_lower  diff_ci_upper  t_stat  p_value    n
                    # outcome
                    # _is_na_o2[False]      0.75      0.75        0.0            0.0            0.0     NaN      NaN  4.0
                    # _is_na_o2[True]       0.25      0.25        0.0            0.0            0.0     NaN      NaN  4.0
                    # o1                    8.50      8.50        0.0            0.0            0.0     NaN      NaN  4.0
                    # o2                    6.00      6.00        0.0            0.0            0.0     NaN      NaN  4.0

                    # Response rates (relative to number of respondents in sample):
                    #       o1    o2
                    # n    4.0   3.0
                    # %  100.0  75.0

                print(s_o2.outcomes().summary())
                    # 2 outcomes: ['o1' 'o2']
                    # Mean outcomes (with 95% confidence intervals):
                    # source            self  target             self           target
                    # _is_na_o2[False]  0.75   0.750   (0.326, 1.174)     (0.45, 1.05)
                    # _is_na_o2[True]   0.25   0.250  (-0.174, 0.674)    (-0.05, 0.55)
                    # o1                8.50  10.500   (7.404, 9.596)  (8.912, 12.088)
                    # o2                6.00   7.875   (2.535, 9.465)  (4.351, 11.399)

                    # Weights impact on outcomes (t_test):
                    #                   mean_yw0  mean_yw1  mean_diff  diff_ci_lower  diff_ci_upper  t_stat  p_value    n
                    # outcome
                    # _is_na_o2[False]      0.75      0.75        0.0            0.0            0.0     NaN      NaN  4.0
                    # _is_na_o2[True]       0.25      0.25        0.0            0.0            0.0     NaN      NaN  4.0
                    # o1                    8.50      8.50        0.0            0.0            0.0     NaN      NaN  4.0
                    # o2                    6.00      6.00        0.0            0.0            0.0     NaN      NaN  4.0

                    # Response rates (relative to number of respondents in sample):
                    #       o1    o2
                    # n    4.0   3.0
                    # %  100.0  75.0
                    # Response rates (relative to notnull rows in the target):
                    #            o1    o2
                    # n   4.000000   3.0
                    # %  66.666667  50.0
                    # Response rates (in the target):
                    #        o1    o2
                    # n    8.0   6.0
                    # %  100.0  75.0
        """

        mean_outcomes_with_ci = self.mean_with_ci()
        relative_response_rates = self.relative_response_rates()
        target_response_rates = self.target_response_rates()
        weights_impact_clause = ""
        if weights_impact_method is not None:
            weights_impact = self.weights_impact_on_outcome_ss(
                method=weights_impact_method,
                conf_level=weights_impact_conf_level,
            )
            if weights_impact is not None:
                weights_impact_clause = (
                    "Weights impact on outcomes "
                    f"({weights_impact_method}):\n"
                    f"{weights_impact.to_string(max_cols=None)}\n\n"
                )
        if target_response_rates is None:
            target_clause = ""
            relative_to_target_clause = ""
        else:
            relative_to_target_response_rates = self.relative_response_rates(
                target=True, per_column=False
            )
            relative_to_target_clause = f"Response rates (relative to notnull rows in the target):\n {relative_to_target_response_rates}"

            target_clause = f"Response rates (in the target):\n {target_response_rates}"

        n_outcomes = self.df.shape[1]
        list_outcomes = np.array(self.df.columns, dtype=object)
        mean_outcomes_with_ci = mean_outcomes_with_ci
        relative_response_rates = relative_response_rates
        target_clause = target_clause

        out = (
            f"{n_outcomes} outcomes: {list_outcomes}\n"
            f"Mean outcomes (with 95% confidence intervals):\n"
            # TODO: in the future consider if to add an argument to transpose (.T) the output, in case there are multiple outcomes.
            f"{mean_outcomes_with_ci.to_string(max_cols=None)}\n\n"
            f"{weights_impact_clause}"
            "Response rates (relative to number of respondents in sample):\n"
            f"{relative_response_rates}\n"
            f"{relative_to_target_clause}\n"
            f"{target_clause}\n"
        )

        return out


class BalanceDFCovars(BalanceDF):
    def __init__(self: "BalanceDFCovars", sample: Sample) -> None:
        """A factory function to create BalanceDFCovars

        This is used through :func:`Sample.covars`.
        It initiates a BalanceDFCovars object by passing the relevant arguments to
        :func:`BalanceDF.__init__`.

        Args:
            self (BalanceDFCovars): Object that is initiated.
            sample (Sample): Object
        """
        super().__init__(sample._covar_columns(), sample, name="covars")

    def from_frame(
        self: "BalanceDFCovars",
        df: pd.DataFrame,
        weights: pd.Series | None = None,
    ) -> "BalanceDFCovars":
        """A factory function to create a BalanceDFCovars from a df.

        Although generally the main way the object is created is through the __init__ method.
        This method is useful when you need to create a BalanceDFCovars object directly from a DataFrame.

        Args:
            self (BalanceDFCovars): Object
            df (pd.DataFrame): A df.
            weights (Optional[pd.Series], optional): _description_. Defaults to None.

        Returns:
            BalanceDFCovars: Object.

        Examples:
        .. code-block:: python

            import pandas as pd
            from balance.balancedf_class import BalanceDFCovars
            covars = BalanceDFCovars.from_frame(pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
            covars.df.columns.tolist()
            # ['index', 'a', 'b']
        """
        df = df.reset_index()
        concat_list: list[pd.DataFrame | pd.Series] = [
            df,
            pd.Series(np.arange(0, df.shape[0]), name="id"),
        ]
        if weights is not None:
            concat_list.append(weights)
        df = pd.concat(concat_list, axis=1)
        return Sample.from_frame(df, id_column="id").covars()


class BalanceDFWeights(BalanceDF):
    def __init__(self: "BalanceDFWeights", sample: Sample) -> None:
        """A factory function to create BalanceDFWeights

        This is used through :func:`Sample.weights`.
        It initiates a BalanceDFWeights object by passing the relevant arguments to
        :func:`BalanceDF.__init__`.

        Args:
            self (BalanceDFWeights): Object that is initiated.
            sample (Sample): Object
        """
        super().__init__(sample.weight_column.to_frame(), sample, name="weights")

    @property
    def df(self: "BalanceDFWeights") -> pd.DataFrame:
        """Return the current weight column as a DataFrame.

        Args:
            self (BalanceDFWeights): The BalanceDFWeights instance.

        Returns:
            pd.DataFrame: DataFrame containing the current weight column.
        """
        return self._sample.weight_column.to_frame()

    # TODO: maybe add better control if there are no weights for unadjusted or target (the current default shows them in the legend, but not in the figure)
    def plot(
        self: "BalanceDFWeights", on_linked_samples: bool = True, **kwargs: Any
    ) -> list[Any] | npt.NDArray[Any] | dict[str, Figure] | None:
        """Plots kde (kernal density estimation) of the weights in a BalanceDFWeights object using seaborn (as default).

        It's possible to use other plots using dist_type with arguments such as "hist", "kde" (default), "qq", and "ecdf".
        Look at :func:`plot_dist` for more details.

        Args:
            self (BalanceDFWeights): a BalanceDFOutcomes object, with a set of variables.
            on_linked_samples (bool, optional): Determines if the linked samples should be included in the plot.
                Defaults to True.

        Returns:
            Union[Union[List, np.ndarray], Dict[str, Figure], None]:
                If library="plotly" then returns a dictionary containing plots if return_dict_of_figures is True. None otherwise.
                If library="seaborn" then returns None, unless return_axes is True. Then either a list or an np.array of matplotlib axis.

        Examples:
        .. code-block:: python

                import numpy as np
                import pandas as pd
                from numpy import random
                from balance.sample_class import Sample

                random.seed(96483)

                df = pd.DataFrame({
                    "id": range(100),
                    'v1': random.random_integers(11111, 11114, size=100).astype(str),
                    'v2': random.normal(size = 100),
                    'v3': random.uniform(size = 100),
                    "w": pd.Series(np.ones(99).tolist() + [1000]),
                }).sort_values(by=['v2'])

                s1 = Sample.from_frame(df,
                    id_column="id",
                    weight_column="w",
                    outcome_columns=["v1", "v2"],
                )

                s2 = Sample.from_frame(
                    df.assign(w = pd.Series(np.ones(100))),
                    id_column="id",
                    weight_column="w",
                    outcome_columns=["v1", "v2"],
                )

                s3 = s1.set_target(s2)
                s3_null = s3.adjust(method="null")
                s3_null.set_weights(random.random(size = 100) + 0.5)

                # default: seaborn with dist_type = "kde"
                s3_null.weights().plot()
        """
        default_kwargs = {
            "weighted": False,
            "library": "seaborn",
            "dist_type": "kde",
            "numeric_n_values_threshold": -1,
        }
        default_kwargs.update(kwargs)
        return super().plot(on_linked_samples=on_linked_samples, **default_kwargs)

    def design_effect(self: "BalanceDFWeights") -> np.float64:
        """Calculates Kish's design effect (deff) on the BalanceDFWeights weights.

        Extract the first column to get a pd.Series of the weights.

        See :func:`weights_stats.design_effect` for details.

        Args:
            self (BalanceDFWeights): Object.

        Returns:
            np.float64: Deff.

        Examples:
        .. code-block:: python

            import pandas as pd
            from balance.sample_class import Sample
            sample = Sample.from_frame(
                pd.DataFrame(
                    {"id": ["1", "2"], "weight": [1.0, 2.0]}
                ),
                id_column="id",
                weight_column="weight",
                standardize_types=False,
            )
            round(sample.weights().design_effect(), 3)
            # 1.111
        """
        return weights_stats.design_effect(self.df.iloc[:, 0])

    # TODO: in the future, consider if this type of overriding is the best solution.
    #       to reconsider as part of a larger code refactoring.
    @property
    def _weights(self: "BalanceDFWeights") -> None:
        """A BalanceDFWeights has no weights (its df is that of the weights.)

        Args:
            self (BalanceDFWeights): Object.

        Returns:
            NoneType: None.
        """
        return None

    def trim(
        self: "BalanceDFWeights",
        ratio: float | int | None = None,
        percentile: float | None = None,
        keep_sum_of_weights: bool = True,
    ) -> None:
        """Trim weights in the sample object.

        Uses :func:`adjustments.trim_weights` for the weights trimming.

        Args:
            self (BalanceDFWeights): Object.
            ratio (Optional[Union[float, int]], optional): Maps to weight_trimming_mean_ratio. Defaults to None.
            percentile (Optional[float], optional): Maps to weight_trimming_percentile. Defaults to None.
            keep_sum_of_weights (bool, optional): Maps to weight_trimming_percentile. Defaults to True.

        Returns:
            None. This function updates the :func:`_sample` using :func:`set_weights`

        Examples:
        .. code-block:: python

            import pandas as pd
            from balance.sample_class import Sample
            sample = Sample.from_frame(
                pd.DataFrame(
                    {"id": ["1", "2"], "weight": [1.0, 100.0]}
                ),
                id_column="id",
                weight_column="weight",
                standardize_types=False,
            )
            sample.weights().trim(percentile=0.5, keep_sum_of_weights=False)
            sample.weights().df["weight"].max() <= 100.0
            # True
        """
        # TODO: verify which object exactly gets updated - and explain it here.
        self._sample.set_weights(
            trim_weights(
                self.df.iloc[:, 0],
                weight_trimming_mean_ratio=ratio,
                weight_trimming_percentile=percentile,
                keep_sum_of_weights=keep_sum_of_weights,
            )
        )

    def summary(
        self: "BalanceDFWeights", on_linked_samples: bool | None = None
    ) -> pd.DataFrame:
        """
        Generates a summary of a BalanceDFWeights object.

        This function provides a comprehensive overview of the BalanceDFWeights object
        by calculating and returning a range of weight diagnostics.

        Args:
            self (BalanceDFWeights): The BalanceDFWeights object to be summarized.
            on_linked_samples (Optional[bool], optional): This parameter is ignored. It is only included
                because summary overrides BalanceDF.summary. Defaults to None.

        Returns:
            pd.DataFrame: A DataFrame containing various weight diagnostics such as 'design_effect',
                'effective_sample_proportion', 'effective_sample_size', sum of weights, and basic summary statistics
                from describe, 'nonparametric_skew', and 'weighted_median_breakdown_point' among others.

        Note:
            The weights are normalized to sum to the sample size, n.

        Examples:
        .. code-block:: python

                import pandas as pd
                from balance.sample_class import Sample

                s1 = Sample.from_frame(
                    pd.DataFrame(
                        {
                            "a": (1, 2, 3, 1),
                            "b": (-42, 8, 2, -42),
                            "o": (7, 8, 9, 10),
                            "c": ("x", "y", "z", "v"),
                            "id": (1, 2, 3, 4),
                            "w": (0.5, 2, 1, 1),
                        }
                    ),
                    id_column="id",
                    weight_column="w",
                    outcome_columns="o",
                )

                print(s1.weights().summary().round(2))
                    #                                 var   val
                    # 0                     design_effect  1.23
                    # 1       effective_sample_proportion  0.81
                    # 2             effective_sample_size  3.24
                    # 3                               sum  4.50
                    # 4                    describe_count  4.00
                    # 5                     describe_mean  1.00
                    # 6                      describe_std  0.56
                    # 7                      describe_min  0.44
                    # 8                      describe_25%  0.78
                    # 9                      describe_50%  0.89
                    # 10                     describe_75%  1.11
                    # 11                     describe_max  1.78
                    # 12                    prop(w < 0.1)  0.00
                    # 13                    prop(w < 0.2)  0.00
                    # 14                  prop(w < 0.333)  0.00
                    # 15                    prop(w < 0.5)  0.25
                    # 16                      prop(w < 1)  0.75
                    # 17                     prop(w >= 1)  0.25
                    # 18                     prop(w >= 2)  0.00
                    # 19                     prop(w >= 3)  0.00
                    # 20                     prop(w >= 5)  0.00
                    # 21                    prop(w >= 10)  0.00
                    # 22               nonparametric_skew  0.20
                    # 23  weighted_median_breakdown_point  0.25
        """
        # ----------------------------------------------------
        # Diagnostics on the weights
        # ----------------------------------------------------
        the_weights = self.df.iloc[
            :, 0
        ]  # should be ['weight'], but this is more robust in case a user uses other names
        weights_diag_var = []
        weights_diag_value = []

        # adding design_effect and variations
        the_weights_de = weights_stats.design_effect(the_weights)
        weights_diag_var.extend(
            ["design_effect", "effective_sample_proportion", "effective_sample_size"]
        )
        weights_diag_value.extend(
            [
                the_weights_de,
                1.0 / the_weights_de,
                float(len(the_weights)) / the_weights_de,
            ]
        )

        # adding sum of weights, and then normalizing them to n (sample size)
        weights_diag_var.append("sum")
        weights_diag_value.append(the_weights.sum())

        the_weights = the_weights / the_weights.mean()  # normalize weights to sum to n.

        # adding basic summary statistics from describe:
        tmp_describe = the_weights.describe()
        weights_diag_var.extend(["describe_" + i for i in tmp_describe.index])
        weights_diag_value.extend(tmp_describe.to_list())
        # TODO: decide if we want more quantiles of the weights.

        # adding prop_above_and_below
        tmp_props = weights_stats.prop_above_and_below(the_weights)
        weights_diag_var.extend(
            tmp_props.index.to_list()  # pyre-ignore[16]: existing defaults make sure this output is pd.Series with relevant methods.
        )
        weights_diag_value.extend(
            tmp_props.to_list()  # pyre-ignore[16]: existing defaults make sure this output is pd.Series with relevant methods.
        )
        # TODO: decide if we want more numbers (e.g.: 2/3 and 3/2)

        # adding nonparametric_skew and weighted_median_breakdown_point
        weights_diag_var.append("nonparametric_skew")
        weights_diag_value.append(weights_stats.nonparametric_skew(the_weights))

        weights_diag_var.append("weighted_median_breakdown_point")
        weights_diag_value.append(
            weights_stats.weighted_median_breakdown_point(the_weights)
        )

        return pd.DataFrame(
            {
                # "metric": "weights_diagnostics",
                "var": weights_diag_var,
                "val": weights_diag_value,
            }
        )
