# fmatch/core/sampling.py
"""
Enhanced sampling module for intelligent data analysis.
Provides reservoir sampling, adaptive sampling, and early accept optimization.
"""

import random
import math
from typing import List, Any, Optional, Tuple, Dict
import pandas as pd
import logging

log = logging.getLogger(__name__)

# Constants moved here for centralization
DEFAULT_SAMPLE_SIZE = 1000  # Increased from 500
MAX_SAMPLE_SIZE = 4000
TARGET_STANDARD_ERROR = 0.02
MIN_SAMPLE_FOR_ANALYSIS = 200


class ReservoirSampler:
    """Reservoir sampling for unbiased sampling from streams."""

    def __init__(self, k: int = 1000, random_state: Optional[int] = 42):
        self.k = k
        self.reservoir: List[Any] = []
        self.count = 0
        if random_state is not None:
            random.seed(random_state)

    def add(self, item: Any) -> None:
        """Add item using Vitter's Algorithm R."""
        self.count += 1
        if len(self.reservoir) < self.k:
            self.reservoir.append(item)
        else:
            # Use random.random() to avoid integer overflow
            if random.random() < self.k / self.count:
                self.reservoir[random.randrange(self.k)] = item

    def get_sample(self) -> List[Any]:
        """Return the current reservoir sample."""
        return list(self.reservoir)

    def get_sample_series(self, name: Optional[str] = None) -> pd.Series:
        """Return sample as pandas Series."""
        return pd.Series(self.reservoir, name=name)


class AdaptiveSampler:
    """Sample until standard error converges."""

    def __init__(
        self,
        target_se: float = TARGET_STANDARD_ERROR,
        max_samples: int = MAX_SAMPLE_SIZE,
    ):
        self.target_se = target_se
        self.max_samples = max_samples
        self.sample_sizes = [200, 500, 1000, 2000, 3000, 4000]
        self._cached_metrics: Dict[int, Any] = {}

    def sample_until_stable(
        self, series: pd.Series, metric_func: callable
    ) -> Tuple[List, float, float, Dict]:
        """Sample until metric converges or max samples reached."""
        samples = []
        series_clean = series.dropna()

        if len(series_clean) < MIN_SAMPLE_FOR_ANALYSIS:
            # Use all data if too small
            return (
                list(series_clean),
                0.0,
                0.0,
                {"converged": True, "reason": "small_dataset"},
            )

        for n in self.sample_sizes:
            if n > len(series_clean):
                n = len(series_clean)

            # Only sample what we need
            needed = n - len(samples)
            if needed > 0:
                new_samples = series_clean.sample(n=needed, random_state=42)
                samples.extend(new_samples.tolist())

            # Calculate metric and standard error
            metric_value = metric_func(samples)
            se = self._calculate_standard_error(metric_value, len(samples))

            # Cache for potential reuse
            self._cached_metrics[n] = {"metric": metric_value, "se": se}

            # Check convergence
            if se <= self.target_se or n >= self.max_samples or n >= len(series_clean):
                return (
                    samples,
                    metric_value,
                    se,
                    {
                        "converged": se <= self.target_se,
                        "samples_used": n,
                        "se": se,
                        "metric": metric_value,
                    },
                )

        return (
            samples,
            metric_value,
            se,
            {"converged": False, "samples_used": len(samples)},
        )

    def _calculate_standard_error(self, p: float, n: int) -> float:
        """Calculate standard error for a proportion."""
        if n == 0:
            return 1.0

        # Clamp p to avoid math domain errors
        p = max(0.001, min(0.999, p))

        # Standard error for proportion
        se = math.sqrt(p * (1 - p) / n)
        return se


class EarlyAcceptChecker:
    """Check if we can skip further sampling due to very clean data."""

    @staticmethod
    def should_early_accept(
        initial_metrics: Dict[str, float], sample_size: int
    ) -> Tuple[bool, str]:
        """Determine if metrics are good enough to skip further sampling."""
        if sample_size < MIN_SAMPLE_FOR_ANALYSIS:
            return False, "insufficient_samples"

        # Very clean data thresholds
        if (
            initial_metrics.get("overlap_ratio", 0) > 0.9
            and initial_metrics.get("variant_rate", 1) < 0.05
            and initial_metrics.get("null_rate", 1) < 0.01
        ):
            return True, "very_clean_data"

        # Extremely messy data (also skip sampling)
        if initial_metrics.get("overlap_ratio", 1) < 0.01:
            return True, "extremely_low_overlap"

        return False, "continue_sampling"


def stratified_domain_sample(series: pd.Series, target_size: int = 1000) -> pd.Series:
    """Sample domains with TLD stratification to avoid bias."""
    from fmatch.core.utils import normalize_domain

    # Extract TLD
    def get_tld(domain):
        normalized = normalize_domain(domain)
        if not normalized:
            return "unknown"
        parts = normalized.split(".")
        return parts[-1] if parts else "unknown"

    # Group by TLD
    series_clean = series.dropna()
    if series_clean.empty:
        return pd.Series([], dtype=series.dtype)

    tld_groups = series_clean.groupby(series_clean.map(get_tld))

    # Calculate proportional sample sizes
    samples = []
    for tld, group in tld_groups:
        # Proportional allocation with minimum of 3 per group
        group_size = len(group)
        total_size = len(series_clean)
        group_target = max(3, int(target_size * group_size / total_size))
        group_sample = group.sample(n=min(group_target, group_size), random_state=42)
        samples.extend(group_sample.tolist())

    # Return as series, trimmed to exact target size
    return pd.Series(samples[:target_size], dtype=series.dtype)
