# fmatch/core/tuning.py
from dataclasses import dataclass
from typing import Literal, Callable, Optional, List, Tuple
import random
import time
import logging

log = logging.getLogger(__name__)


@dataclass
class BlockCostModel:
    """Stores the results of a micro-benchmark to decide processing strategy."""

    a_small: float
    b_cdist: float
    cutoff_ops: int

    def choose(self, n_pairs: int) -> Literal["small", "cdist"]:
        """
        Chooses between the simple loop ('small') or vectorized 'cdist'
        based on the number of pairs and the benchmarked costs.
        """
        # Add logging for explainability
        t_small = self.a_small * n_pairs
        t_cdist = self.b_cdist * (
            n_pairs**0.5
        )  # Note: cdist is more like O(n) not O(n^2) for n pairs

        choice = "small" if t_small <= t_cdist else "cdist"
        if n_pairs > 10000:  # Only log for larger blocks
            log.debug(
                f"Cost model: {n_pairs} pairs -> {choice} (small={t_small:.4f}s, cdist={t_cdist:.4f}s)"
            )

        return choice


def _benchmark_function(func: Callable, rounds: int = 3) -> float:
    """Return best time (s) over N rounds."""
    times = []
    for _ in range(rounds):
        t0 = time.perf_counter()
        func()
        times.append(time.perf_counter() - t0)
    return min(times)


def calibrate_cost_model(
    sample_pairs: int = 30_000,
    gen_pairs: Optional[Callable[[int], Tuple[List[str], List[str]]]] = None,
) -> BlockCostModel:
    """
    Micro-benchmarks the simple loop vs. vectorized cdist to find the
    optimal crossover point, returning a tuned cost model.
    """
    from rapidfuzz import process, fuzz

    # Generate random string pairs for benchmarking if none are provided
    if gen_pairs is None:
        alphabet = "abcdefghijklmnopqrstuvwxyz "

        def gen(n):
            mk = lambda: "".join(
                random.choice(alphabet) for _ in range(random.randint(5, 22))
            )
            return [mk() for _ in range(n)], [mk() for _ in range(n)]

        gen_pairs = gen

    A, B = gen_pairs(sample_pairs)

    # Benchmark the simple, iterative approach (theoretically O(n))
    def small_loop_task():
        s = 0
        for a in A:
            s += fuzz.WRatio(a, B[0])  # Simulate a small inner loop
        return s

    # Benchmark the vectorized cdist approach on a smaller batch (theoretically O(n^2))
    batch_size = min(4000, sample_pairs)

    def cdist_task():
        process.cdist(
            A[:batch_size], B[:batch_size], scorer=fuzz.WRatio, score_cutoff=0
        )

    # Time both functions
    time_small = _benchmark_function(small_loop_task)
    time_cdist = _benchmark_function(cdist_task)

    # Fit rough coefficients to the cost models
    # a_small is the cost per operation for the simple loop
    a = max(time_small / max(sample_pairs, 1), 1e-9)
    # b_cdist is the cost per operation for the vectorized approach
    b = max(time_cdist / max(batch_size * batch_size, 1), 1e-12)

    # Calculate the crossover point where a*n = b*n^2  -->  n = a/b
    cutoff = int(a / b)

    return BlockCostModel(a_small=a, b_cdist=b, cutoff_ops=cutoff)
