# src/fmatch/core/blocking_metrics.py
# -*- coding: utf-8 -*-
"""
Core calculation logic for analyzing and scoring blocking strategies.
"""

# ── Standard library imports ──────────────────────────────────────────────────
import logging
import math
import multiprocessing as mp
from typing import Optional

# ── Third-party imports ───────────────────────────────────────────────────────
import numpy as np
import pandas as pd
import psutil


# ── Setup ─────────────────────────────────────────────────────────────────────
log = logging.getLogger(__name__)

try:
    CPU_COUNT = mp.cpu_count() or 1
except:
    CPU_COUNT = 1

PSUTIL_AVAILABLE = psutil is not None

_BYTES_PER_PAIR = 100  # Assumed memory usage per candidate pair

# ── Core Metric Calculation Functions ─────────────────────────────────────────


def _calculate_gini(values: np.ndarray) -> float:
    """Calculate Gini coefficient for a numpy array of values."""
    if len(values) == 0:
        return 0.0

    sorted_values = np.sort(values)
    n = len(values)
    cumsum = np.cumsum(sorted_values)

    if cumsum[-1] == 0:
        return 0.0

    # Gini formula
    return (2 * np.sum((np.arange(1, n + 1) * sorted_values))) / (n * cumsum[-1]) - (
        n + 1
    ) / n


def _calculate_entropy(block_keys: pd.Series) -> float:
    """Calculate Shannon entropy for a series of blocking keys."""
    if block_keys.empty:
        return 0.0

    counts = block_keys.value_counts()
    total = counts.sum()
    if total == 0:
        return 0.0

    probs = counts / total
    probs = probs[probs > 0]  # Avoid log2(0)
    if probs.empty:
        return 0.0

    return -np.sum(probs * np.log2(probs))


def calculate_raw_signals(
    src_keys: pd.Series, key_len: int, mode: str, ref_keys: Optional[pd.Series] = None
) -> Optional[dict]:
    """Calculates all raw signals from engine.py (moved here)."""
    # This logic is extracted from your engine.py file
    from fmatch.core.engine import (
        BlockingMode,
    )  # Local import to avoid circular dependency

    if isinstance(mode, str):
        mode = BlockingMode(mode)

    N_src = len(src_keys)
    if N_src < 2:
        return None

    src_counts = src_keys.value_counts(dropna=False)
    blank_ratio_src = src_counts.get("", 0) / N_src
    # Loosened from 0.30 to 0.60: Accept columns with 40%+ coverage
    # Real-world domain/email fields often have 30-60% blanks but still massively reduce comparisons
    if blank_ratio_src > 0.60:
        return None

    if mode == BlockingMode.MATCH:
        if ref_keys is None or ref_keys.empty:
            return None
        N_ref = len(ref_keys)
        if N_ref == 0:
            return None
        ref_counts = ref_keys.value_counts(dropna=False)
        blank_ratio_ref = ref_counts.get("", 0) / N_ref
        # Loosened from 0.30 to 0.60: Accept columns with 40%+ coverage
        if blank_ratio_ref > 0.60:
            return None

    meaningful_src_counts = src_counts.drop("", errors="ignore")
    if meaningful_src_counts.empty:
        return {
            "RR": 1.0,
            "H": 0.0,
            "G": 0.0,
            "P1": 0.0,
            "S1": 1.0,
            "k": float(key_len),
        }

    # Entropy (H) - using a normalized version
    num_distinct = len(meaningful_src_counts)
    raw_H = _calculate_entropy(src_keys[src_keys != ""])
    H = raw_H / math.log2(num_distinct) if num_distinct > 1 else 0.0

    # Gini (G)
    G = _calculate_gini(meaningful_src_counts.values)
    # Largest Block (P1)
    P1 = meaningful_src_counts.max() / N_src
    # Singletons (S1)
    S1 = (meaningful_src_counts == 1).sum() / len(meaningful_src_counts)

    # Reduction Ratio (RR)
    RR = 0.0
    if mode == BlockingMode.DEDUPE:
        max_pairs = N_src * (N_src - 1) / 2.0
        pairs_after = (src_counts * (src_counts - 1) // 2).sum()
        if max_pairs > 0:
            RR = 1.0 - (pairs_after / max_pairs)
    elif mode == BlockingMode.MATCH and ref_keys is not None:
        N_ref = len(ref_keys)
        max_pairs = float(N_src * N_ref)
        ref_counts = ref_keys.value_counts(dropna=False)
        common_blocks = src_counts.index.intersection(ref_counts.index)
        pairs_after = (src_counts[common_blocks] * ref_counts[common_blocks]).sum()
        if max_pairs > 0:
            RR = 1.0 - (pairs_after / max_pairs)

    return {"RR": RR, "H": H, "G": G, "P1": P1, "S1": S1, "k": float(key_len)}


def normalize_signals(raw_signals: dict, preproc_variant: str) -> dict:
    """Normalizes raw signals from engine.py (moved here)."""
    # Extract raw values with safe defaults
    H = raw_signals.get("H", 0.0)
    G = raw_signals.get("G", 0.0)
    RR = raw_signals.get("RR", 0.0)
    P1 = raw_signals.get("P1", 0.0)
    S1 = raw_signals.get("S1", 0.0)
    k_raw = raw_signals.get("k", 2.0)

    normalized = {}

    # Entropy - keep as is
    normalized["Hn"] = H

    # Gini - inverted (lower is better)
    normalized["Gn"] = 1.0 - G

    # Reduction ratio - square root for better sensitivity
    normalized["RRn"] = math.sqrt(RR)

    # P1 (dominant block) - using 0.3 threshold from enhancement
    normalized["P1n"] = max(0, 1.0 - (P1 / 0.3)) if P1 < 0.3 else 0.0

    # S1 (singleton ratio) - NEW gradual scoring instead of cliff-edge
    # This gives partial credit for singleton ratios up to 0.7
    normalized["S1n"] = 1.0 - min(S1, 0.7) / 0.7

    # Key length normalization with phonetic variant handling
    is_phonetic = any(preproc_variant.startswith(pv) for pv in {"soundex", "metaphone"})
    if is_phonetic:
        normalized["kn"] = 1.0  # Phonetic variants have fixed output length
    else:
        min_k, max_k = 2.0, 6.0
        k_clamped = max(min_k, min(k_raw, max_k))
        normalized["kn"] = 1.0 - ((k_clamped - min_k) / (max_k - min_k))

    # Add algorithm fitness if present (for multi-algorithm support)
    if "af" in raw_signals:
        normalized["af"] = raw_signals["af"]

    return normalized


def quality_score(normalized_signals: dict, weights: Optional[dict] = None) -> float:
    """Calculates the quality score from normalized signals."""
    if weights is None:
        # Use 'balanced' weights as a default if none are provided
        weights = {
            "rr": 0.09,
            "h": 0.18,
            "g": 0.27,
            "p1": 0.225,
            "s1": 0.09,
            "k": 0.045,
            "af": 0.10,
        }

    score = (
        weights["rr"] * normalized_signals.get("RRn", 0.0)
        + weights["h"] * (normalized_signals.get("Hn", 0.0) ** 2)
        + weights["g"] * normalized_signals.get("Gn", 0.0)
        + weights["p1"] * normalized_signals.get("P1n", 0.0)
        + weights["s1"] * normalized_signals.get("S1n", 0.0)
        + weights["k"] * normalized_signals.get("kn", 0.0)
    )
    if "af" in normalized_signals and "af" in weights:
        score += weights["af"] * normalized_signals.get("af", 0.0)

    if normalized_signals.get("RRn", 0.0) > 0.999:
        score *= 0.1

    return score


# ── Metric Calculator Classes ────────────────────────────────────────────────


class ImpactMetricsCalculator:
    """Calculates real-world impact metrics from blocking strategy."""

    def __init__(self, blocking_result, source_df, ref_df, mode):
        self.result = blocking_result
        self.source_df = source_df
        self.ref_df = ref_df
        self.mode = mode

    def calculate(self):
        """Calculate comprehensive impact metrics using real data."""
        source_rows = len(self.source_df) if self.source_df is not None else 0
        ref_rows = len(self.ref_df) if self.ref_df is not None else source_rows

        if self.mode == "duplicate":
            comparisons_without = source_rows * (source_rows - 1) // 2
        else:
            comparisons_without = source_rows * ref_rows

        reduction_ratio = self.result.get("raw_metrics", {}).get("RR", 0)
        comparisons_with = int(comparisons_without * (1 - reduction_ratio))

        cores = min(CPU_COUNT, 8)
        comparisons_per_second = 10000 * cores

        runtime_without = comparisons_without / comparisons_per_second
        runtime_with = comparisons_with / comparisons_per_second

        peak_memory_bytes = min(comparisons_with, 1000000) * _BYTES_PER_PAIR
        peak_memory_mb = peak_memory_bytes / (1024 * 1024)

        available_memory_mb = (
            psutil.virtual_memory().available / (1024 * 1024)
            if PSUTIL_AVAILABLE
            else 8192
        )

        return {
            "comparisons_without": comparisons_without,  # ADDED
            "comparisons_with": comparisons_with,  # ADDED
            "comparisons_without_str": self._format_number(comparisons_without),
            "comparisons_with_str": self._format_number(comparisons_with),
            "reduction_percentage": reduction_ratio * 100,
            "runtime_without": runtime_without,  # ADDED
            "runtime_with": runtime_with,  # ADDED
            "runtime_without_str": self._format_time(runtime_without),
            "runtime_with_str": self._format_time(runtime_with),
            "time_saved_str": self._format_time(runtime_without - runtime_with),
            "memory_usage_mb": int(peak_memory_mb),
            "memory_percentage": (peak_memory_mb / available_memory_mb) * 100,
            "cpu_cores_utilized": cores,
            "comparisons_per_second": comparisons_per_second,
        }

    def _format_number(self, num):
        if num >= 1e9:
            return f"{num/1e9:.1f}B"
        if num >= 1e6:
            return f"{num/1e6:.1f}M"
        if num >= 1e3:
            return f"{num/1e3:.0f}K"
        return str(int(num))

    def _format_time(self, seconds):
        if seconds < 60:
            return f"{int(seconds)}s"
        if seconds < 3600:
            return f"{int(seconds/60)}m {int(seconds%60)}s"
        hours = int(seconds / 3600)
        minutes = int((seconds % 3600) / 60)
        return f"{hours}h {minutes}m"


class DistributionAnalyzer:
    """Analyzes actual block size distribution from the blocking strategy."""

    def __init__(self, blocking_result, source_df, ref_df, mode):
        self.result = blocking_result
        self.source_df = source_df
        self.ref_df = ref_df
        self.mode = mode

    def analyze(self):
        """Analyze actual block distribution from the data."""
        if self.source_df is None:
            return self._empty_distribution_result()

        try:
            from fmatch.core.engine import build_key

            # --- NEW LOGIC: Generate the block keys on the fly ---
            block_col = self.result.get("source_block_col")
            key_len = self.result.get("block_key_len")
            preproc = self.result.get("preproc")

            if not all([block_col, key_len, preproc]):
                log.warning("DistributionAnalyzer: Missing strategy info in result.")
                return self._empty_distribution_result()

            if block_col not in self.source_df.columns:
                log.warning(
                    f"DistributionAnalyzer: Blocking column '{block_col}' not in DataFrame."
                )
                return self._empty_distribution_result()

            block_key_series, _ = build_key(
                self.source_df[block_col],
                preproc,
                key_len,
                "default",  # Ruleset name is not critical for this analysis
            )
            # --- END NEW LOGIC ---

            block_counts = block_key_series.value_counts()
            block_counts = block_counts[block_counts.index != ""]

            if block_counts.empty:
                return self._empty_distribution_result()

            total_blocks = len(block_counts)
            block_sizes = block_counts.values

            hist, bins = np.histogram(block_sizes, bins=min(50, max(1, total_blocks)))
            histogram_data = [(int(bins[i]), int(hist[i])) for i in range(len(hist))]

            percentiles = {
                "p50": int(np.percentile(block_sizes, 50)),
                "p90": int(np.percentile(block_sizes, 90)),
                "p99": int(np.percentile(block_sizes, 99)),
            }

            largest_size = block_counts.iloc[0]
            total_records = len(self.source_df)

            gini = _calculate_gini(block_sizes)
            health = "extreme" if gini > 0.7 else "skewed" if gini > 0.5 else "balanced"

            problem_blocks = []
            threshold = total_records * 0.1
            for key, size in block_counts.head(5).items():
                if size > threshold:
                    problem_blocks.append(
                        {
                            "key": str(key)[:20],
                            "size": int(size),
                            "percentage": (size / total_records) * 100,
                        }
                    )

            return {
                "total_blocks": total_blocks,
                "largest_block_size": int(largest_size),
                "largest_block_percentage": (largest_size / total_records) * 100,
                "percentiles": percentiles,
                "distribution_health": health,
                "block_histogram": histogram_data,
                "problem_blocks": problem_blocks,
            }
        except Exception as e:
            log.error(f"Error during distribution analysis: {e}", exc_info=True)
            return self._empty_distribution_result()

    def _empty_distribution_result(self):
        return {
            "total_blocks": 0,
            "largest_block_size": 0,
            "largest_block_percentage": 0,
            "percentiles": {"p50": 0, "p90": 0, "p99": 0},
            "distribution_health": "unknown",
            "block_histogram": [],
            "problem_blocks": [],
        }
