from __future__ import annotations

from typing import Iterator, List

import numpy as np
import pandas as pd

from blazematch._rust import compute_feature_matrix
from blazematch.config import FieldComparison, Metric, NUMERIC_METRICS
from blazematch.preprocess import apply_preprocess


class FeatureComputer:
    """Computes similarity features for candidate pairs using Rust acceleration.

    Uses a single Rust call (compute_feature_matrix) for all comparisons at once,
    avoiding per-comparison Python→Rust boundary overhead.
    """

    def __init__(
        self,
        df_left: pd.DataFrame,
        df_right: pd.DataFrame,
        comparisons: List[FieldComparison],
    ):
        self._df_left = df_left
        self._df_right = df_right
        self._comparisons = comparisons

    @property
    def feature_names(self) -> List[str]:
        return [c.output_name for c in self._comparisons]

    def compute(self, pairs: pd.DataFrame) -> np.ndarray:
        """Compute features for candidate pairs via single Rust call.

        Args:
            pairs: DataFrame with columns [left_idx, right_idx]

        Returns:
            (N_pairs, N_features) numpy array of float64
        """
        left_indices = pairs["left_idx"].values
        right_indices = pairs["right_idx"].values
        n_pairs = len(pairs)

        # Separate string and numeric comparisons, tracking original positions
        string_cols_left = []
        string_cols_right = []
        string_metrics = []
        string_positions = []  # original index in comparisons list

        numeric_cols_left = []
        numeric_cols_right = []
        numeric_metrics = []
        numeric_positions = []

        for idx, comp in enumerate(self._comparisons):
            # Use numpy indexing instead of pandas .iloc for performance
            left_arr = self._df_left[comp.left_field].values[left_indices]
            right_arr = self._df_right[comp.right_field].values[right_indices]
            left_vals = pd.Series(left_arr)
            right_vals = pd.Series(right_arr)

            # Apply preprocessing if configured
            if comp.preprocess:
                left_vals = apply_preprocess(left_vals, comp.preprocess)
                right_vals = apply_preprocess(right_vals, comp.preprocess)

            if comp.metric in NUMERIC_METRICS:
                if comp.metric == Metric.DATE_DISTANCE:
                    fmt = comp.date_format or "%Y-%m-%d"
                    left_days = pd.to_datetime(left_vals, format=fmt, errors="coerce")
                    right_days = pd.to_datetime(right_vals, format=fmt, errors="coerce")
                    ref = pd.Timestamp("1970-01-01")
                    numeric_cols_left.append(
                        ((left_days - ref).dt.days).fillna(0).to_numpy(dtype=np.float64)
                    )
                    numeric_cols_right.append(
                        ((right_days - ref).dt.days).fillna(0).to_numpy(dtype=np.float64)
                    )
                else:
                    numeric_cols_left.append(
                        pd.to_numeric(left_vals, errors="coerce").fillna(0.0).to_numpy(dtype=np.float64)
                    )
                    numeric_cols_right.append(
                        pd.to_numeric(right_vals, errors="coerce").fillna(0.0).to_numpy(dtype=np.float64)
                    )
                numeric_metrics.append(comp.metric.value)
                numeric_positions.append(idx)
            else:
                string_cols_left.append(
                    left_vals.astype(str).fillna("").tolist()
                )
                string_cols_right.append(
                    right_vals.astype(str).fillna("").tolist()
                )
                string_metrics.append(comp.metric.value)
                string_positions.append(idx)

        # Stack numeric columns into 2D numpy arrays (n_cols, n_pairs) for
        # zero-copy transfer to Rust, avoiding expensive .tolist() conversion.
        if numeric_cols_left:
            numeric_left_2d = np.ascontiguousarray(
                np.vstack(numeric_cols_left), dtype=np.float64
            )
            numeric_right_2d = np.ascontiguousarray(
                np.vstack(numeric_cols_right), dtype=np.float64
            )
        else:
            numeric_left_2d = np.empty((0, n_pairs), dtype=np.float64)
            numeric_right_2d = np.empty((0, n_pairs), dtype=np.float64)

        # Single Rust call: returns flat array with [string_features..., numeric_features...]
        n_features_total = len(string_positions) + len(numeric_positions)
        flat_result = compute_feature_matrix(
            string_cols_left,
            string_cols_right,
            string_metrics,
            numeric_left_2d,
            numeric_right_2d,
            numeric_metrics,
            n_pairs,
        )

        # Reshape Rust output: columns are [string_0, string_1, ..., numeric_0, numeric_1, ...]
        rust_matrix = np.asarray(flat_result).reshape(n_pairs, n_features_total)

        # Reorder columns back to original comparison order
        n_comparisons = len(self._comparisons)
        result = np.empty((n_pairs, n_comparisons), dtype=np.float64)

        # String features are first in Rust output
        for rust_col, orig_col in enumerate(string_positions):
            result[:, orig_col] = rust_matrix[:, rust_col]

        # Numeric features follow
        n_string = len(string_positions)
        for rust_col_offset, orig_col in enumerate(numeric_positions):
            result[:, orig_col] = rust_matrix[:, n_string + rust_col_offset]

        return result

    def compute_chunked(
        self, pairs_iter: Iterator[pd.DataFrame]
    ) -> Iterator[np.ndarray]:
        """Process pairs in chunks, yielding feature matrices."""
        for chunk in pairs_iter:
            yield self.compute(chunk)
