"""Core distribution fitting engine for Spark."""

import logging
from functools import reduce
from typing import Callable, List, Optional, Tuple, Union

import numpy as np
import pyspark.sql.functions as F
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql.types import NumericType

from spark_bestfit.discrete_fitting import (
    compute_discrete_histogram,
    create_discrete_fitting_udf,
    create_discrete_sample_data,
)
from spark_bestfit.distributions import DiscreteDistributionRegistry, DistributionRegistry
from spark_bestfit.fitting import FITTING_SAMPLE_SIZE, compute_data_summary, create_fitting_udf
from spark_bestfit.histogram import HistogramComputer
from spark_bestfit.results import DistributionFitResult, FitResults
from spark_bestfit.utils import get_spark_session

logger = logging.getLogger(__name__)

# Re-export for convenience
DEFAULT_EXCLUDED_DISTRIBUTIONS: Tuple[str, ...] = tuple(DistributionRegistry.DEFAULT_EXCLUSIONS)


class DistributionFitter:
    """Modern Spark distribution fitting engine.

    Efficiently fits ~100 scipy.stats distributions to data using Spark's
    parallel processing capabilities. Uses broadcast variables and Pandas UDFs
    to avoid data collection and minimize serialization overhead.

    Example:
        >>> from pyspark.sql import SparkSession
        >>> from spark_bestfit import DistributionFitter
        >>>
        >>> # Create your own SparkSession
        >>> spark = SparkSession.builder.appName("my-app").getOrCreate()
        >>> df = spark.createDataFrame([(float(x),) for x in data], ['value'])
        >>>
        >>> # Simple usage
        >>> fitter = DistributionFitter(spark)
        >>> results = fitter.fit(df, column='value')
        >>> best = results.best(n=1)[0]
        >>> print(f"Best: {best.distribution} with SSE={best.sse}")
        >>>
        >>> # With custom parameters
        >>> fitter = DistributionFitter(spark, random_seed=123)
        >>> results = fitter.fit(df, 'value', bins=100, support_at_zero=True)
        >>>
        >>> # Plot the best fit
        >>> fitter.plot(best, df, 'value', title='Best Fit')
    """

    def __init__(
        self,
        spark: Optional[SparkSession] = None,
        excluded_distributions: Optional[Tuple[str, ...]] = None,
        random_seed: int = 42,
    ):
        """Initialize DistributionFitter.

        Args:
            spark: SparkSession. If None, uses the active session.
            excluded_distributions: Distributions to exclude from fitting.
                Defaults to DEFAULT_EXCLUDED_DISTRIBUTIONS (slow distributions).
            random_seed: Random seed for reproducible sampling.

        Raises:
            RuntimeError: If no SparkSession provided and no active session exists
        """
        self.spark: SparkSession = get_spark_session(spark)
        self.excluded_distributions = (
            excluded_distributions if excluded_distributions is not None else DEFAULT_EXCLUDED_DISTRIBUTIONS
        )
        self.random_seed = random_seed
        self._registry = DistributionRegistry()
        self._histogram_computer = HistogramComputer()

    def fit(
        self,
        df: DataFrame,
        column: Optional[str] = None,
        columns: Optional[List[str]] = None,
        bins: Union[int, Tuple[float, ...]] = 50,
        use_rice_rule: bool = True,
        support_at_zero: bool = False,
        max_distributions: Optional[int] = None,
        enable_sampling: bool = True,
        sample_fraction: Optional[float] = None,
        max_sample_size: int = 1_000_000,
        sample_threshold: int = 10_000_000,
        num_partitions: Optional[int] = None,
        progress_callback: Optional[Callable[[int, int, float], None]] = None,
    ) -> FitResults:
        """Fit distributions to data column(s).

        Args:
            df: Spark DataFrame containing data
            column: Name of single column to fit distributions to
            columns: List of column names for multi-column fitting
            bins: Number of histogram bins or tuple of bin edges
            use_rice_rule: Use Rice rule to auto-determine bin count
            support_at_zero: Only fit non-negative distributions
            max_distributions: Limit number of distributions (for testing)
            enable_sampling: Enable sampling for large datasets
            sample_fraction: Fraction to sample (None = auto-determine)
            max_sample_size: Maximum rows to sample when auto-determining
            sample_threshold: Row count above which sampling is applied
            num_partitions: Spark partitions (None = auto-determine)
            progress_callback: Optional callback for progress updates.
                Called with (completed_tasks, total_tasks, percent_complete).
                Callback is invoked from background thread - ensure thread-safety.

        Returns:
            FitResults object with fitted distributions

        Raises:
            ValueError: If column not found, DataFrame empty, or invalid params
            TypeError: If column is not numeric

        Example:
            >>> # Single column
            >>> results = fitter.fit(df, column='value')
            >>> results = fitter.fit(df, 'value', bins=100, support_at_zero=True)
            >>>
            >>> # Multi-column
            >>> results = fitter.fit(df, columns=['col1', 'col2', 'col3'])
            >>> best_col1 = results.for_column('col1').best(n=1)[0]
            >>> best_per_col = results.best_per_column(n=1)
        """
        # Validate column/columns parameters
        if column is None and columns is None:
            raise ValueError("Must provide either 'column' or 'columns' parameter")
        if column is not None and columns is not None:
            raise ValueError("Cannot provide both 'column' and 'columns' - use one or the other")

        # Normalize to list of columns
        target_columns = [column] if column is not None else columns

        # Input validation for all columns
        for col in target_columns:
            self._validate_inputs(df, col, max_distributions, bins, sample_fraction)

        # Get row count (single operation for all columns)
        row_count = df.count()
        if row_count == 0:
            raise ValueError("DataFrame is empty")
        logger.info(f"Row count: {row_count}")

        # Sample if needed (single operation for all columns)
        df_sample = self._apply_sampling(
            df, row_count, enable_sampling, sample_fraction, max_sample_size, sample_threshold
        )

        # Get distributions to fit (same for all columns)
        distributions = self._registry.get_distributions(
            support_at_zero=support_at_zero,
            additional_exclusions=list(self.excluded_distributions),
        )
        if max_distributions is not None and max_distributions > 0:
            distributions = distributions[:max_distributions]

        # Start progress tracking if callback provided
        tracker = None
        if progress_callback is not None:
            from spark_bestfit.progress import ProgressTracker

            tracker = ProgressTracker(self.spark, progress_callback)
            tracker.start()

        try:
            # Fit each column and collect results
            all_results_dfs = []
            for col in target_columns:
                logger.info(f"Fitting column '{col}'...")
                results_df = self._fit_single_column(
                    df_sample=df_sample,
                    column=col,
                    row_count=row_count,
                    bins=bins,
                    use_rice_rule=use_rice_rule,
                    distributions=distributions,
                    num_partitions=num_partitions,
                )
                all_results_dfs.append(results_df)

            # Union all results using reduce for cleaner query plan
            combined_df = reduce(DataFrame.union, all_results_dfs)

            combined_df = combined_df.cache()
            total_results = combined_df.count()
            logger.info(
                f"Total results: {total_results} ({len(target_columns)} columns × ~{len(distributions)} distributions)"
            )

            return FitResults(combined_df)
        finally:
            if tracker is not None:
                tracker.stop()

    def _fit_single_column(
        self,
        df_sample: DataFrame,
        column: str,
        row_count: int,
        bins: Union[int, Tuple[float, ...]],
        use_rice_rule: bool,
        distributions: List[str],
        num_partitions: Optional[int],
    ) -> DataFrame:
        """Fit distributions to a single column (internal method).

        Args:
            df_sample: Sampled DataFrame
            column: Column name
            row_count: Original row count (for histogram computation)
            bins: Number of histogram bins
            use_rice_rule: Use Rice rule for bin count
            distributions: List of distribution names to fit
            num_partitions: Number of Spark partitions

        Returns:
            Spark DataFrame with fit results for this column
        """
        # Compute histogram
        y_hist, x_hist = self._histogram_computer.compute_histogram(
            df_sample, column, bins=bins, use_rice_rule=use_rice_rule, approx_count=row_count
        )
        logger.info(f"  Histogram for '{column}': {len(x_hist)} bins")

        # Broadcast histogram
        histogram_bc = self.spark.sparkContext.broadcast((y_hist, x_hist))

        # Create fitting sample
        data_sample = self._create_fitting_sample(df_sample, column, row_count)
        data_sample_bc = self.spark.sparkContext.broadcast(data_sample)

        # Compute data summary for provenance (once per column)
        data_summary = compute_data_summary(data_sample)

        try:
            # Create DataFrame of distributions
            dist_df = self.spark.createDataFrame([(dist,) for dist in distributions], ["distribution_name"])

            # Determine partitioning
            n_partitions = num_partitions or self._calculate_partitions(len(distributions))
            dist_df = dist_df.repartition(n_partitions)

            # Apply fitting UDF
            fitting_udf = create_fitting_udf(
                histogram_bc, data_sample_bc, column_name=column, data_summary=data_summary
            )
            results_df = dist_df.select(fitting_udf(F.col("distribution_name")).alias("result")).select("result.*")

            # Filter failed fits
            results_df = results_df.filter(F.col("sse") < float(np.inf))

            num_results = results_df.count()
            logger.info(f"  Fit {num_results}/{len(distributions)} distributions for '{column}'")

            return results_df

        finally:
            histogram_bc.unpersist()
            data_sample_bc.unpersist()

    def plot(
        self,
        result: DistributionFitResult,
        df: DataFrame,
        column: str,
        bins: Union[int, Tuple[float, ...]] = 50,
        use_rice_rule: bool = True,
        title: str = "",
        xlabel: str = "Value",
        ylabel: str = "Density",
        figsize: Tuple[int, int] = (12, 8),
        dpi: int = 100,
        show_histogram: bool = True,
        histogram_alpha: float = 0.5,
        pdf_linewidth: int = 2,
        title_fontsize: int = 14,
        label_fontsize: int = 12,
        legend_fontsize: int = 10,
        grid_alpha: float = 0.3,
        save_path: Optional[str] = None,
        save_format: str = "png",
    ):
        """Plot fitted distribution against data histogram.

        Args:
            result: DistributionFitResult to plot
            df: DataFrame with data
            column: Column name
            bins: Number of histogram bins
            use_rice_rule: Use Rice rule for bins
            title: Plot title
            xlabel: X-axis label
            ylabel: Y-axis label
            figsize: Figure size (width, height)
            dpi: Dots per inch for saved figures
            show_histogram: Show data histogram
            histogram_alpha: Histogram transparency (0-1)
            pdf_linewidth: Line width for PDF curve
            title_fontsize: Title font size
            label_fontsize: Axis label font size
            legend_fontsize: Legend font size
            grid_alpha: Grid transparency (0-1)
            save_path: Path to save figure (optional)
            save_format: Save format (png, pdf, svg)

        Returns:
            Tuple of (figure, axis) from matplotlib

        Example:
            >>> fitter.plot(best, df, 'value', title='Best Fit')
            >>> fitter.plot(best, df, 'value', figsize=(16, 10), dpi=300)
        """
        from spark_bestfit.plotting import plot_distribution

        # Compute histogram for plotting
        y_hist, x_hist = self._histogram_computer.compute_histogram(df, column, bins=bins, use_rice_rule=use_rice_rule)

        return plot_distribution(
            result=result,
            y_hist=y_hist,
            x_hist=x_hist,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            figsize=figsize,
            dpi=dpi,
            show_histogram=show_histogram,
            histogram_alpha=histogram_alpha,
            pdf_linewidth=pdf_linewidth,
            title_fontsize=title_fontsize,
            label_fontsize=label_fontsize,
            legend_fontsize=legend_fontsize,
            grid_alpha=grid_alpha,
            save_path=save_path,
            save_format=save_format,
        )

    def plot_comparison(
        self,
        results: List[DistributionFitResult],
        df: DataFrame,
        column: str,
        bins: Union[int, Tuple[float, ...]] = 50,
        use_rice_rule: bool = True,
        title: str = "Distribution Comparison",
        xlabel: str = "Value",
        ylabel: str = "Density",
        figsize: Tuple[int, int] = (12, 8),
        dpi: int = 100,
        show_histogram: bool = True,
        histogram_alpha: float = 0.5,
        pdf_linewidth: int = 2,
        title_fontsize: int = 14,
        label_fontsize: int = 12,
        legend_fontsize: int = 10,
        grid_alpha: float = 0.3,
        save_path: Optional[str] = None,
        save_format: str = "png",
    ):
        """Plot multiple distributions for comparison.

        Args:
            results: List of DistributionFitResult objects
            df: DataFrame with data
            column: Column name
            bins: Number of histogram bins
            use_rice_rule: Use Rice rule for bins
            title: Plot title
            xlabel: X-axis label
            ylabel: Y-axis label
            figsize: Figure size (width, height)
            dpi: Dots per inch
            show_histogram: Show data histogram
            histogram_alpha: Histogram transparency
            pdf_linewidth: PDF line width
            title_fontsize: Title font size
            label_fontsize: Label font size
            legend_fontsize: Legend font size
            grid_alpha: Grid transparency
            save_path: Path to save figure
            save_format: Save format

        Returns:
            Tuple of (figure, axis)

        Example:
            >>> top_3 = results.best(n=3)
            >>> fitter.plot_comparison(top_3, df, 'value')
        """
        from spark_bestfit.plotting import plot_comparison

        y_hist, x_hist = self._histogram_computer.compute_histogram(df, column, bins=bins, use_rice_rule=use_rice_rule)

        return plot_comparison(
            results=results,
            y_hist=y_hist,
            x_hist=x_hist,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            figsize=figsize,
            dpi=dpi,
            show_histogram=show_histogram,
            histogram_alpha=histogram_alpha,
            pdf_linewidth=pdf_linewidth,
            title_fontsize=title_fontsize,
            label_fontsize=label_fontsize,
            legend_fontsize=legend_fontsize,
            grid_alpha=grid_alpha,
            save_path=save_path,
            save_format=save_format,
        )

    @staticmethod
    def _validate_inputs(
        df: DataFrame,
        column: str,
        max_distributions: Optional[int],
        bins: Union[int, Tuple[float, ...]],
        sample_fraction: Optional[float],
    ) -> None:
        """Validate inputs."""
        if max_distributions == 0:
            raise ValueError("max_distributions cannot be 0")

        if column not in df.columns:
            raise ValueError(f"Column '{column}' not found in DataFrame. Available columns: {df.columns}")

        col_type = df.schema[column].dataType
        if not isinstance(col_type, NumericType):
            raise TypeError(f"Column '{column}' must be numeric, got {col_type}")

        if isinstance(bins, int) and bins <= 0:
            raise ValueError(f"bins must be positive, got {bins}")

        if sample_fraction is not None and not 0.0 < sample_fraction <= 1.0:
            raise ValueError(f"sample_fraction must be in (0, 1], got {sample_fraction}")

    def _apply_sampling(
        self,
        df: DataFrame,
        row_count: int,
        enable_sampling: bool,
        sample_fraction: Optional[float],
        max_sample_size: int,
        sample_threshold: int,
    ) -> DataFrame:
        """Apply sampling if needed."""
        if not enable_sampling or row_count <= sample_threshold:
            return df

        if sample_fraction is not None:
            fraction = sample_fraction
        else:
            fraction = min(max_sample_size / row_count, 0.35)

        logger.info(f"Sampling {fraction * 100:.1f}% of data ({int(row_count * fraction)} rows)")
        return df.sample(fraction=fraction, seed=self.random_seed)

    def _create_fitting_sample(self, df: DataFrame, column: str, row_count: int) -> np.ndarray:
        """Create sample for scipy distribution fitting."""
        sample_size = min(FITTING_SAMPLE_SIZE, row_count)
        fraction = min(sample_size / row_count, 1.0)
        sample_df = df.select(column).sample(fraction=fraction, seed=self.random_seed)
        return sample_df.toPandas()[column].values

    def _calculate_partitions(self, num_distributions: int) -> int:
        """Calculate optimal partition count."""
        total_cores = self.spark.sparkContext.defaultParallelism
        return min(num_distributions, total_cores * 2)

    def plot_qq(
        self,
        result: DistributionFitResult,
        df: DataFrame,
        column: str,
        max_points: int = 1000,
        title: str = "",
        xlabel: str = "Theoretical Quantiles",
        ylabel: str = "Sample Quantiles",
        figsize: Tuple[int, int] = (10, 10),
        dpi: int = 100,
        marker: str = "o",
        marker_size: int = 30,
        marker_alpha: float = 0.6,
        marker_color: str = "steelblue",
        line_color: str = "red",
        line_style: str = "--",
        line_width: float = 1.5,
        title_fontsize: int = 14,
        label_fontsize: int = 12,
        grid_alpha: float = 0.3,
        save_path: Optional[str] = None,
        save_format: str = "png",
    ):
        """Create a Q-Q plot to assess goodness-of-fit.

        A Q-Q (quantile-quantile) plot compares sample quantiles against
        theoretical quantiles from the fitted distribution. Points falling
        close to the reference line indicate a good fit.

        Args:
            result: DistributionFitResult to plot
            df: DataFrame with data
            column: Column name
            max_points: Maximum data points to sample for plotting
            title: Plot title
            xlabel: X-axis label
            ylabel: Y-axis label
            figsize: Figure size (width, height)
            dpi: Dots per inch for saved figures
            marker: Marker style for data points
            marker_size: Size of markers
            marker_alpha: Marker transparency (0-1)
            marker_color: Color of markers
            line_color: Color of reference line
            line_style: Style of reference line
            line_width: Width of reference line
            title_fontsize: Title font size
            label_fontsize: Axis label font size
            grid_alpha: Grid transparency (0-1)
            save_path: Path to save figure (optional)
            save_format: Save format (png, pdf, svg)

        Returns:
            Tuple of (figure, axis) from matplotlib

        Example:
            >>> best = results.best(n=1)[0]
            >>> fitter.plot_qq(best, df, 'value', title='Q-Q Plot')
        """
        from spark_bestfit.plotting import plot_qq

        # Sample data for Q-Q plot using sample() instead of orderBy(rand())
        # sample() operates per-partition without shuffle, much faster for large datasets
        row_count = df.count()
        fraction = min(max_points * 3 / row_count, 1.0) if row_count > 0 else 1.0
        sample_df = df.select(column).sample(fraction=fraction, seed=self.random_seed).limit(max_points)
        data = sample_df.toPandas()[column].values

        return plot_qq(
            result=result,
            data=data,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            figsize=figsize,
            dpi=dpi,
            marker=marker,
            marker_size=marker_size,
            marker_alpha=marker_alpha,
            marker_color=marker_color,
            line_color=line_color,
            line_style=line_style,
            line_width=line_width,
            title_fontsize=title_fontsize,
            label_fontsize=label_fontsize,
            grid_alpha=grid_alpha,
            save_path=save_path,
            save_format=save_format,
        )

    def plot_pp(
        self,
        result: DistributionFitResult,
        df: DataFrame,
        column: str,
        max_points: int = 1000,
        title: str = "",
        xlabel: str = "Theoretical Probabilities",
        ylabel: str = "Sample Probabilities",
        figsize: Tuple[int, int] = (10, 10),
        dpi: int = 100,
        marker: str = "o",
        marker_size: int = 30,
        marker_alpha: float = 0.6,
        marker_color: str = "steelblue",
        line_color: str = "red",
        line_style: str = "--",
        line_width: float = 1.5,
        title_fontsize: int = 14,
        label_fontsize: int = 12,
        grid_alpha: float = 0.3,
        save_path: Optional[str] = None,
        save_format: str = "png",
    ):
        """
        Create a P-P plot to assess goodness-of-fit.

        A P-P (probability-probability) plot compares the empirical CDF of the
        sample data against the theoretical CDF of the fitted distribution.
        Points falling close to the reference line indicate a good fit,
        particularly in the center of the distribution.

        Args:
            result: DistributionFitResult to plot
            df: DataFrame with data
            column: Column name
            max_points: Maximum data points to sample for plotting
            title: Plot title
            xlabel: X-axis label
            ylabel: Y-axis label
            figsize: Figure size (width, height)
            dpi: Dots per inch for saved figures
            marker: Marker style for data points
            marker_size: Size of markers
            marker_alpha: Marker transparency (0-1)
            marker_color: Color of markers
            line_color: Color of reference line
            line_style: Style of reference line
            line_width: Width of reference line
            title_fontsize: Title font size
            label_fontsize: Axis label font size
            grid_alpha: Grid transparency (0-1)
            save_path: Path to save figure (optional)
            save_format: Save format (png, pdf, svg)

        Returns:
            Tuple of (figure, axis) from matplotlib

        Example:
            >>> best = results.best(n=1)[0]
            >>> fitter.plot_pp(best, df, 'value', title='P-P Plot')
        """
        from spark_bestfit.plotting import plot_pp

        # Sample data for P-P plot using sample() instead of orderBy(rand())
        # sample() operates per-partition without shuffle, much faster for large datasets
        row_count = df.count()
        fraction = min(max_points * 3 / row_count, 1.0) if row_count > 0 else 1.0
        sample_df = df.select(column).sample(fraction=fraction, seed=self.random_seed).limit(max_points)
        data = sample_df.toPandas()[column].values

        return plot_pp(
            result=result,
            data=data,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            figsize=figsize,
            dpi=dpi,
            marker=marker,
            marker_size=marker_size,
            marker_alpha=marker_alpha,
            marker_color=marker_color,
            line_color=line_color,
            line_style=line_style,
            line_width=line_width,
            title_fontsize=title_fontsize,
            label_fontsize=label_fontsize,
            grid_alpha=grid_alpha,
            save_path=save_path,
            save_format=save_format,
        )


# Re-export for convenience
DEFAULT_EXCLUDED_DISCRETE_DISTRIBUTIONS: Tuple[str, ...] = tuple(DiscreteDistributionRegistry.DEFAULT_EXCLUSIONS)


class DiscreteDistributionFitter:
    """Spark distribution fitting engine for discrete (count) data.

    Efficiently fits scipy.stats discrete distributions to integer data using
    Spark's parallel processing capabilities. Uses MLE optimization since
    scipy discrete distributions don't have a built-in fit() method.

    Metric Selection:
        For discrete distributions, **AIC is recommended** for model selection:
        - ``aic``: Proper model selection criterion with complexity penalty
        - ``bic``: Similar to AIC but stronger penalty for complex models
        - ``ks_statistic``: Valid for ranking, but p-values are not reliable
        - ``sse``: Simple comparison metric

        The K-S test assumes continuous distributions. For discrete data,
        the K-S statistic can rank fits, but p-values are conservative and
        should not be used for hypothesis testing.

    Example:
        >>> from pyspark.sql import SparkSession
        >>> from spark_bestfit import DiscreteDistributionFitter
        >>>
        >>> spark = SparkSession.builder.appName("my-app").getOrCreate()
        >>> df = spark.createDataFrame([(x,) for x in count_data], ['counts'])
        >>>
        >>> fitter = DiscreteDistributionFitter(spark)
        >>> results = fitter.fit(df, column='counts')
        >>>
        >>> # Use AIC for model selection (recommended)
        >>> best = results.best(n=1, metric='aic')[0]
        >>> print(f"Best: {best.distribution} (AIC={best.aic:.2f})")
    """

    def __init__(
        self,
        spark: Optional[SparkSession] = None,
        excluded_distributions: Optional[Tuple[str, ...]] = None,
        random_seed: int = 42,
    ):
        """Initialize DiscreteDistributionFitter.

        Args:
            spark: SparkSession. If None, uses the active session.
            excluded_distributions: Distributions to exclude from fitting.
                Defaults to DEFAULT_EXCLUDED_DISCRETE_DISTRIBUTIONS.
            random_seed: Random seed for reproducible sampling.

        Raises:
            RuntimeError: If no SparkSession provided and no active session exists
        """
        self.spark: SparkSession = get_spark_session(spark)
        self.excluded_distributions = (
            excluded_distributions if excluded_distributions is not None else DEFAULT_EXCLUDED_DISCRETE_DISTRIBUTIONS
        )
        self.random_seed = random_seed
        self._registry = DiscreteDistributionRegistry()

    def fit(
        self,
        df: DataFrame,
        column: Optional[str] = None,
        columns: Optional[List[str]] = None,
        max_distributions: Optional[int] = None,
        enable_sampling: bool = True,
        sample_fraction: Optional[float] = None,
        max_sample_size: int = 1_000_000,
        sample_threshold: int = 10_000_000,
        num_partitions: Optional[int] = None,
        progress_callback: Optional[Callable[[int, int, float], None]] = None,
    ) -> FitResults:
        """Fit discrete distributions to integer data column(s).

        Args:
            df: Spark DataFrame containing integer count data
            column: Name of single column to fit distributions to
            columns: List of column names for multi-column fitting
            max_distributions: Limit number of distributions (for testing)
            enable_sampling: Enable sampling for large datasets
            sample_fraction: Fraction to sample (None = auto-determine)
            max_sample_size: Maximum rows to sample when auto-determining
            sample_threshold: Row count above which sampling is applied
            num_partitions: Spark partitions (None = auto-determine)
            progress_callback: Optional callback for progress updates.
                Called with (completed_tasks, total_tasks, percent_complete).
                Callback is invoked from background thread - ensure thread-safety.

        Returns:
            FitResults object with fitted distributions

        Raises:
            ValueError: If column not found, DataFrame empty, or invalid params
            TypeError: If column is not numeric

        Example:
            >>> # Single column
            >>> results = fitter.fit(df, column='counts')
            >>> best = results.best(n=1, metric='aic')
            >>>
            >>> # Multi-column
            >>> results = fitter.fit(df, columns=['counts1', 'counts2'])
            >>> best_per_col = results.best_per_column(n=1, metric='aic')
        """
        # Validate column/columns parameters
        if column is None and columns is None:
            raise ValueError("Must provide either 'column' or 'columns' parameter")
        if column is not None and columns is not None:
            raise ValueError("Cannot provide both 'column' and 'columns' - use one or the other")

        # Normalize to list of columns
        target_columns = [column] if column is not None else columns

        # Input validation for all columns
        for col in target_columns:
            self._validate_inputs(df, col, max_distributions, sample_fraction)

        # Get row count (single operation for all columns)
        row_count = df.count()
        if row_count == 0:
            raise ValueError("DataFrame is empty")
        logger.info(f"Row count: {row_count}")

        # Sample if needed (single operation for all columns)
        df_sample = self._apply_sampling(
            df, row_count, enable_sampling, sample_fraction, max_sample_size, sample_threshold
        )

        # Get distributions to fit (same for all columns)
        distributions = self._registry.get_distributions(
            additional_exclusions=list(self.excluded_distributions),
        )
        if max_distributions is not None and max_distributions > 0:
            distributions = distributions[:max_distributions]

        # Start progress tracking if callback provided
        tracker = None
        if progress_callback is not None:
            from spark_bestfit.progress import ProgressTracker

            tracker = ProgressTracker(self.spark, progress_callback)
            tracker.start()

        try:
            # Fit each column and collect results
            all_results_dfs = []
            for col in target_columns:
                logger.info(f"Fitting discrete column '{col}'...")
                results_df = self._fit_single_column(
                    df_sample=df_sample,
                    column=col,
                    row_count=row_count,
                    distributions=distributions,
                    num_partitions=num_partitions,
                )
                all_results_dfs.append(results_df)

            # Union all results using reduce for cleaner query plan
            combined_df = reduce(DataFrame.union, all_results_dfs)

            combined_df = combined_df.cache()
            total_results = combined_df.count()
            logger.info(
                f"Total results: {total_results} ({len(target_columns)} columns × ~{len(distributions)} distributions)"
            )

            return FitResults(combined_df)
        finally:
            if tracker is not None:
                tracker.stop()

    def _fit_single_column(
        self,
        df_sample: DataFrame,
        column: str,
        row_count: int,
        distributions: List[str],
        num_partitions: Optional[int],
    ) -> DataFrame:
        """Fit discrete distributions to a single column (internal method).

        Args:
            df_sample: Sampled DataFrame
            column: Column name
            row_count: Original row count
            distributions: List of distribution names to fit
            num_partitions: Number of Spark partitions

        Returns:
            Spark DataFrame with fit results for this column
        """
        # Create integer data sample for fitting
        sample_size = min(FITTING_SAMPLE_SIZE, row_count)
        fraction = min(sample_size / row_count, 1.0)
        sample_df = df_sample.select(column).sample(fraction=fraction, seed=self.random_seed)
        data_sample = sample_df.toPandas()[column].values.astype(int)
        data_sample = create_discrete_sample_data(data_sample, sample_size=FITTING_SAMPLE_SIZE)
        logger.info(f"  Data sample for '{column}': {len(data_sample)} values")

        # Compute discrete histogram (PMF)
        x_values, empirical_pmf = compute_discrete_histogram(data_sample)
        logger.info(f"  PMF for '{column}': {len(x_values)} unique values (range: {x_values.min()}-{x_values.max()})")

        # Broadcast histogram and data
        histogram_bc = self.spark.sparkContext.broadcast((x_values, empirical_pmf))
        data_sample_bc = self.spark.sparkContext.broadcast(data_sample)

        # Compute data summary for provenance (once per column)
        data_summary = compute_data_summary(data_sample.astype(float))

        try:
            # Create DataFrame of distributions
            dist_df = self.spark.createDataFrame([(dist,) for dist in distributions], ["distribution_name"])

            # Determine partitioning
            n_partitions = num_partitions or self._calculate_partitions(len(distributions))
            dist_df = dist_df.repartition(n_partitions)

            # Apply discrete fitting UDF
            fitting_udf = create_discrete_fitting_udf(
                histogram_bc, data_sample_bc, column_name=column, data_summary=data_summary
            )
            results_df = dist_df.select(fitting_udf(F.col("distribution_name")).alias("result")).select("result.*")

            # Filter failed fits
            results_df = results_df.filter(F.col("sse") < float(np.inf))

            num_results = results_df.count()
            logger.info(f"  Fit {num_results}/{len(distributions)} distributions for '{column}'")

            return results_df

        finally:
            histogram_bc.unpersist()
            data_sample_bc.unpersist()

    @staticmethod
    def _validate_inputs(
        df: DataFrame,
        column: str,
        max_distributions: Optional[int],
        sample_fraction: Optional[float],
    ) -> None:
        """Validate inputs."""
        if max_distributions == 0:
            raise ValueError("max_distributions cannot be 0")

        if column not in df.columns:
            raise ValueError(f"Column '{column}' not found in DataFrame. Available columns: {df.columns}")

        col_type = df.schema[column].dataType
        if not isinstance(col_type, NumericType):
            raise TypeError(f"Column '{column}' must be numeric, got {col_type}")

        if sample_fraction is not None and not 0.0 < sample_fraction <= 1.0:
            raise ValueError(f"sample_fraction must be in (0, 1], got {sample_fraction}")

    def _apply_sampling(
        self,
        df: DataFrame,
        row_count: int,
        enable_sampling: bool,
        sample_fraction: Optional[float],
        max_sample_size: int,
        sample_threshold: int,
    ) -> DataFrame:
        """Apply sampling if needed."""
        if not enable_sampling or row_count <= sample_threshold:
            return df

        if sample_fraction is not None:
            fraction = sample_fraction
        else:
            fraction = min(max_sample_size / row_count, 0.35)

        logger.info(f"Sampling {fraction * 100:.1f}% of data ({int(row_count * fraction)} rows)")
        return df.sample(fraction=fraction, seed=self.random_seed)

    def _calculate_partitions(self, num_distributions: int) -> int:
        """Calculate optimal partition count."""
        total_cores = self.spark.sparkContext.defaultParallelism
        return min(num_distributions, total_cores * 2)

    def plot(
        self,
        result: DistributionFitResult,
        df: DataFrame,
        column: str,
        title: str = "",
        xlabel: str = "Value",
        ylabel: str = "Probability",
        figsize: Tuple[int, int] = (12, 8),
        dpi: int = 100,
        show_histogram: bool = True,
        histogram_alpha: float = 0.7,
        pmf_linewidth: int = 2,
        title_fontsize: int = 14,
        label_fontsize: int = 12,
        legend_fontsize: int = 10,
        grid_alpha: float = 0.3,
        save_path: Optional[str] = None,
        save_format: str = "png",
    ):
        """Plot fitted discrete distribution against data histogram.

        Args:
            result: DistributionFitResult to plot
            df: DataFrame with data
            column: Column name
            title: Plot title
            xlabel: X-axis label
            ylabel: Y-axis label
            figsize: Figure size (width, height)
            dpi: Dots per inch for saved figures
            show_histogram: Show data histogram
            histogram_alpha: Histogram transparency (0-1)
            pmf_linewidth: Line width for PMF curve
            title_fontsize: Title font size
            label_fontsize: Axis label font size
            legend_fontsize: Legend font size
            grid_alpha: Grid transparency (0-1)
            save_path: Path to save figure (optional)
            save_format: Save format (png, pdf, svg)

        Returns:
            Tuple of (figure, axis) from matplotlib
        """
        from spark_bestfit.plotting import plot_discrete_distribution

        # Get data sample
        sample_df = df.select(column).sample(fraction=min(10000 / df.count(), 1.0), seed=self.random_seed)
        data = sample_df.toPandas()[column].values.astype(int)

        return plot_discrete_distribution(
            result=result,
            data=data,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            figsize=figsize,
            dpi=dpi,
            show_histogram=show_histogram,
            histogram_alpha=histogram_alpha,
            pmf_linewidth=pmf_linewidth,
            title_fontsize=title_fontsize,
            label_fontsize=label_fontsize,
            legend_fontsize=legend_fontsize,
            grid_alpha=grid_alpha,
            save_path=save_path,
            save_format=save_format,
        )
