src/protocol/fpsi_adapter.py

from __future__ import annotations

import json
import math
import os
import random
import subprocess
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any

import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import normalize

from src import config
from src.protocol.elsh_params import recommended_params


@dataclass
class MatchResult:
    person: str
    filename: str
    hamming_distance: int
    matched: bool
    score: float


@dataclass
class ProtocolSummary:
    mode: str
    dim: int
    delta: int
    L: int
    k: int
    p_delta: float
    communication_mb: float
    time_s: float
    gao_feasible: bool
    tar: float
    far: float
    accuracy: float
    gallery_size: int
    query_count: int
    binary_density: float
    notes: list[str]


class FuzzyPSIAdapter:
    def __init__(self, mode: str | None = None) -> None:
        self.mode = mode or config.DEFAULT_MODE
        if self.mode not in config.SUPPORTED_MODES:
            self.mode = "simulation"

    def binarize(self, features: np.ndarray, target_dim: int) -> tuple[np.ndarray, float]:
        max_components = min(features.shape[0], features.shape[1])
        use_pca = (
            target_dim < features.shape[1]
            and target_dim <= max_components
            and features.shape[0] <= config.PCA_FIT_SAMPLE_LIMIT
        )
        if use_pca:
            pca = PCA(n_components=target_dim, random_state=42)
            projected = pca.fit_transform(features)
            variance = float(pca.explained_variance_ratio_.sum())
        elif target_dim < features.shape[1]:
            projected = features[:, :target_dim]
            variance = min(1.0, float(target_dim / features.shape[1]))
        else:
            projected = features[:, :target_dim]
            variance = 1.0
        projected = normalize(projected)
        binary_codes = (projected > 0).astype(np.uint8)
        return binary_codes, variance

    def calibrate_threshold(self, binary_codes: np.ndarray, labels: np.ndarray) -> dict[str, float | int]:
        rng = np.random.default_rng(42)
        label_to_idx: dict[Any, list[int]] = {}
        for i, label in enumerate(labels):
            label_to_idx.setdefault(label, []).append(i)

        multi_labels = [label for label, idx in label_to_idx.items() if len(idx) >= 2]
        unique_labels = np.array(list(label_to_idx.keys()))

        genuine_dists: list[int] = []
        impostor_dists: list[int] = []
        for label in multi_labels[:300]:
            indices = label_to_idx[label]
            for i in range(min(len(indices) - 1, 3)):
                genuine_dists.append(int(np.sum(binary_codes[indices[i]] != binary_codes[indices[i + 1]])))

        for _ in range(max(len(genuine_dists) * 2, 32)):
            l1, l2 = rng.choice(unique_labels, 2, replace=False)
            i1 = rng.choice(label_to_idx[l1])
            i2 = rng.choice(label_to_idx[l2])
            impostor_dists.append(int(np.sum(binary_codes[i1] != binary_codes[i2])))

        genuine = np.array(genuine_dists if genuine_dists else [0])
        impostor = np.array(impostor_dists if impostor_dists else [binary_codes.shape[1]])

        best_delta = 0
        best_acc = 0.0
        best_tar = 0.0
        best_far = 1.0
        for delta in range(binary_codes.shape[1]):
            tar = float(np.mean(genuine <= delta))
            far = float(np.mean(impostor <= delta))
            acc = (tar + (1.0 - far)) / 2.0
            if acc > best_acc:
                best_delta = delta
                best_acc = acc
                best_tar = tar
                best_far = far

        adjusted_delta = max(0, int(best_delta + config.MATCH_THRESHOLD_MARGIN))
        return {
            "delta": adjusted_delta,
            "tar": best_tar,
            "far": best_far,
            "accuracy": best_acc,
            "genuine_mean": float(genuine.mean()),
            "impostor_mean": float(impostor.mean()),
        }

    def gao_feasible(self, dim: int, delta: int) -> bool:
        return dim > 8 * delta + 8

    def select_l(self, dim: int, delta: int) -> tuple[int, int, float]:
        k, p_delta, l_min = recommended_params(dim, delta, config.DEFAULT_SECURITY_LAMBDA)
        return l_min, k, p_delta

    def query_against_gallery(
        self,
        query_feature: np.ndarray,
        gallery_features: np.ndarray,
        gallery_people: np.ndarray,
        gallery_filenames: np.ndarray,
        dim: int,
        calibration_features: np.ndarray | None = None,
        calibration_people: np.ndarray | None = None,
    ) -> tuple[MatchResult, ProtocolSummary, dict[str, Any]]:
        calibration_features = gallery_features if calibration_features is None else calibration_features
        calibration_people = gallery_people if calibration_people is None else calibration_people

        binarization_pool = np.vstack([query_feature.reshape(1, -1), calibration_features])
        pool_binary, variance = self.binarize(binarization_pool, dim)
        query_binary = pool_binary[0]
        calibration_binary = pool_binary[1:]

        gallery_pool = np.vstack([query_feature.reshape(1, -1), gallery_features])
        gallery_binary = self.binarize(gallery_pool, dim)[0][1:]

        labels = np.concatenate([np.array(["__query__"], dtype=object), calibration_people.astype(object)])
        threshold_stats = self.calibrate_threshold(pool_binary, labels)
        delta = int(threshold_stats["delta"])
        L, k, p_delta = self.select_l(dim, delta)
        distances = np.sum(gallery_binary != query_binary, axis=1)
        best_idx = int(np.argmin(distances))
        best_dist = int(distances[best_idx])
        matched = best_dist <= delta
        density = float(np.mean(query_binary))
        score = 1.0 - (best_dist / max(dim, 1))

        if self.mode == "full" and config.FULL_PROTOCOL_ENABLED:
            communication_mb, protocol_time, full_notes = self._run_full_protocol(gallery_binary, query_binary, dim, delta, L)
            notes = ["full protocol mode"] + full_notes
        else:
            communication_mb, protocol_time = self._simulate_protocol_metrics(gallery_binary, query_binary, dim, delta, L)
            notes = ["simulation mode", "full protocol disabled or unavailable"]

        summary = ProtocolSummary(
            mode=self.mode if self.mode == "simulation" or config.FULL_PROTOCOL_ENABLED else "simulation",
            dim=dim,
            delta=delta,
            L=L,
            k=k,
            p_delta=p_delta,
            communication_mb=communication_mb,
            time_s=protocol_time,
            gao_feasible=self.gao_feasible(dim, delta),
            tar=float(threshold_stats["tar"]),
            far=float(threshold_stats["far"]),
            accuracy=float(threshold_stats["accuracy"]),
            gallery_size=int(len(gallery_features)),
            query_count=1,
            binary_density=density,
            notes=notes,
        )
        match = MatchResult(
            person=str(gallery_people[best_idx]),
            filename=str(gallery_filenames[best_idx]),
            hamming_distance=best_dist,
            matched=matched,
            score=score,
        )
        details = {
            "variance_explained": variance,
            "query_binary": query_binary.tolist(),
            "best_index": best_idx,
            "top5_distances": [int(x) for x in np.sort(distances)[:5]],
            "calibration_size": int(len(calibration_binary)),
        }
        return match, summary, details

    def _simulate_protocol_metrics(
        self,
        gallery_binary: np.ndarray,
        query_binary: np.ndarray,
        dim: int,
        delta: int,
        L: int,
    ) -> tuple[float, float]:
        gallery_size = len(gallery_binary)
        hit_ratio = max(0.02, min(0.35, (delta + 1) / max(dim, 1) * 6.0))
        candidate_count = max(1, int(gallery_size * hit_ratio))
        communication_mb = (L * dim + candidate_count * 16 + 500000) / (1024.0 * 1024.0)
        base_time = 0.05 + candidate_count * 0.0012 + L * 0.001
        return float(communication_mb), float(base_time)

    def _run_full_protocol(
        self,
        gallery_binary: np.ndarray,
        query_binary: np.ndarray,
        dim: int,
        delta: int,
        L: int,
    ) -> tuple[float, float, list[str]]:
        sender_path, receiver_path = self._write_binary_pair(gallery_binary, query_binary, dim)
        build_dir = config.FULL_PROTOCOL_BUILD_DIR
        receiver_bin = build_dir / "fpsi_receiver"
        sender_bin = build_dir / "fpsi_sender"
        if not receiver_bin.exists() or not sender_bin.exists():
            communication_mb, protocol_time = self._simulate_protocol_metrics(gallery_binary, query_binary, dim, delta, L)
            return communication_mb, protocol_time, ["native binaries missing; fell back to simulated metrics"]

        port = 26000 + random.randint(0, 900)
        recv_cmd = [str(receiver_bin), str(port), str(len(query_binary.reshape(1, -1))), str(dim), str(delta), str(L)]
        send_cmd = [str(sender_bin), "127.0.0.1", str(port), str(len(gallery_binary)), str(dim), str(delta), str(L)]

        recv_proc = subprocess.Popen(recv_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        time.sleep(1.0)
        send_proc = subprocess.Popen(send_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        send_out, send_err = send_proc.communicate(timeout=300)
        recv_out, recv_err = recv_proc.communicate(timeout=300)
        comm_mb = 0.0
        total_time = 0.0
        for line in recv_out.splitlines():
            if "Total:" in line:
                parts = line.split()
                for i, token in enumerate(parts):
                    if token.endswith("s,"):
                        total_time = float(token[:-2])
                    elif token == "MB":
                        comm_mb = float(parts[i - 1])
        notes = []
        if send_err.strip():
            notes.append(f"sender stderr: {send_err.strip()[:200]}")
        if recv_err.strip():
            notes.append(f"receiver stderr: {recv_err.strip()[:200]}")
        if comm_mb == 0.0 and total_time == 0.0:
            sim_comm, sim_time = self._simulate_protocol_metrics(gallery_binary, query_binary, dim, delta, L)
            return sim_comm, sim_time, notes + ["native run returned no parsable totals; using simulated metrics"]
        return comm_mb, total_time, notes

    def _write_binary_pair(self, gallery_binary: np.ndarray, query_binary: np.ndarray, dim: int) -> tuple[Path, Path]:
        runtime_dir = config.OUTPUT_DIR / "runtime_inputs"
        runtime_dir.mkdir(parents=True, exist_ok=True)
        sender_path = runtime_dir / f"sender_d{dim}.bin"
        receiver_path = runtime_dir / f"receiver_d{dim}.bin"

        self._write_binary_dataset(sender_path, gallery_binary)
        self._write_binary_dataset(receiver_path, query_binary.reshape(1, -1))
        return sender_path, receiver_path

    @staticmethod
    def _write_binary_dataset(path: Path, vectors: np.ndarray) -> None:
        vectors = np.asarray(vectors, dtype=np.uint8)
        n, d = vectors.shape
        with open(path, "wb") as fh:
            fh.write(int(n).to_bytes(4, "little"))
            fh.write(int(d).to_bytes(4, "little"))
            fh.write(vectors.tobytes())

    def export_summary(self, match: MatchResult, summary: ProtocolSummary, details: dict[str, Any]) -> str:
        payload = {
            "match": asdict(match),
            "summary": asdict(summary),
            "details": details,
        }
        return json.dumps(payload, indent=2)