# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import logging
from dataclasses import dataclass
from operator import itemgetter
from statistics import mean, median
from typing import Any, Optional

import faiss  # @manual=//faiss/python:pyfaiss_gpu
import numpy as np
from faiss.contrib.evaluation import (  # @manual=//faiss/contrib:faiss_contrib_gpu
    knn_intersection_measure,
)
from scipy.optimize import curve_fit

from .descriptors import DatasetDescriptor, IndexDescriptor
from .index import Index, IndexFromCodec, IndexFromFactory

logger = logging.getLogger(__name__)


def range_search_pr_curve(dist_ann: np.ndarray, metric_score: np.ndarray, gt_rsm: float):
    assert dist_ann.shape == metric_score.shape
    assert dist_ann.ndim == 1
    l = len(dist_ann)
    if l == 0:
        return {
            "dist_ann": [],
            "metric_score_sample": [],
            "cum_score": [],
            "precision": [],
            "recall": [],
            "unique_key": [],
        }
    sort_by_dist_ann = dist_ann.argsort()
    dist_ann = dist_ann[sort_by_dist_ann]
    metric_score = metric_score[sort_by_dist_ann]
    cum_score = np.cumsum(metric_score)
    precision = cum_score / np.arange(1, len(cum_score) + 1)
    recall = cum_score / gt_rsm
    unique_key = np.round(precision * 100) * 100 + np.round(recall * 100)
    tbl = np.vstack([dist_ann, metric_score, cum_score, precision, recall, unique_key])
    group_by_dist_max_cum_score = np.empty(l, bool)
    group_by_dist_max_cum_score[-1] = True
    group_by_dist_max_cum_score[:-1] = dist_ann[1:] != dist_ann[:-1]
    tbl = tbl[:, group_by_dist_max_cum_score]
    _, unique_key_idx = np.unique(tbl[5], return_index=True)
    dist_ann, metric_score, cum_score, precision, recall, unique_key = tbl[
        :, np.sort(unique_key_idx)
    ].tolist()
    return {
        "dist_ann": dist_ann,
        "metric_score_sample": metric_score,
        "cum_score": cum_score,
        "precision": precision,
        "recall": recall,
        "unique_key": unique_key,
    }


def optimizer(op, search, cost_metric, perf_metric):
    totex = op.num_experiments()
    rs = np.random.RandomState(123)
    if totex > 1:
        experiments = rs.permutation(totex - 2) + 1
        experiments = [0, totex - 1] + list(experiments)
    else:
        experiments = [0]

    print(f"total nb experiments {totex}, running {len(experiments)}")

    for cno in experiments:
        key = op.cno_to_key(cno)
        parameters = op.get_parameters(key)

        (max_perf, min_cost) = op.predict_bounds(key)
        if not op.is_pareto_optimal(max_perf, min_cost):
            logger.info(
                f"{cno=:4d} {str(parameters):50}: SKIP, {max_perf=:.3f} {min_cost=:.3f}",
            )
            continue

        logger.info(f"{cno=:4d} {str(parameters):50}: RUN")
        cost, perf = search(
            parameters,
            cost_metric,
            perf_metric,
        )
        logger.info(f"{cno=:4d} {str(parameters):50}: DONE, {cost=:.3f} {perf=:.3f}")
        op.add_operating_point(key, perf, cost)


def distance_ratio_measure(I, R, D_GT, metric):
    sum_of_R = np.sum(np.where(I >= 0, R, 0))
    sum_of_D_GT = np.sum(np.where(I >= 0, D_GT, 0))
    if metric == faiss.METRIC_INNER_PRODUCT:
        return (sum_of_R / sum_of_D_GT).item()
    elif metric == faiss.METRIC_L2:
        return (sum_of_D_GT / sum_of_R).item()
    else:
        raise RuntimeError(f"unknown metric {metric}")


# range_metric possible values:
#
# radius
#    [0..radius) -> 1
#    [radius..inf) -> 0
#
# [[radius1, score1], ...]
#    [0..radius1) -> score1
#    [radius1..radius2) -> score2
#
# [[radius1_from, radius1_to, score1], ...]
#    [radius1_from, radius1_to) -> score1,
#    [radius2_from, radius2_to) -> score2
def get_range_search_metric_function(range_metric, D, R):
    if D is not None:
        assert R is not None
        assert D.shape == R.shape
    if isinstance(range_metric, list):
        aradius, ascore, aradius_from, aradius_to = [], [], [], []
        radius_to = 0
        for rsd in range_metric:
            assert isinstance(rsd, list)
            if len(rsd) == 3:
                radius_from, radius_to, score = rsd
            elif len(rsd) == 2:
                radius_from = radius_to
                radius_to, score = rsd
            else:
                raise AssertionError(f"invalid range definition {rsd}")
            # radius_from and radius_to are compressed distances,
            # we need to convert them to real embedding distances.
            if D is not None:
                sample_idxs = np.argwhere((D <= radius_to) & (D > radius_from))
                assert len(sample_idxs) > 0
                real_radius = np.mean(R[sample_idxs]).item()
            else:
                real_radius = mean([radius_from, radius_to])
            logger.info(
                f"range_search_metric_function {radius_from=} {radius_to=} {real_radius=} {score=}"
            )
            aradius.append(real_radius)
            ascore.append(score)
            aradius_from.append(radius_from)
            aradius_to.append(radius_to)

        def sigmoid(x, a, b, c):
            return a / (1 + np.exp(b * x - c))

        cutoff = max(aradius)
        popt, _ = curve_fit(sigmoid, aradius, ascore, [1, 5, 5])

        for r in np.arange(0, cutoff + 0.05, 0.05):
            logger.info(f"range_search_metric_function {r=} {sigmoid(r, *popt)=}")

        assert isinstance(cutoff, float)
        return (
            cutoff,
            lambda x: np.where(x < cutoff, sigmoid(x, *popt), 0),
            popt.tolist(),
            list(zip(aradius, ascore, aradius_from, aradius_to, strict=True)),
        )
    else:
        # Assuming that the range_metric is a float,
        # so the range is [0..range_metric).
        # D is the result of a range_search with a radius of range_metric,
        # but both range_metric and D may be compressed distances.
        # We approximate the real embedding distance as max(R).
        if R is not None:
            real_range = np.max(R).item()
        else:
            real_range = range_metric
        logger.info(f"range_search_metric_function {range_metric=} {real_range=}")
        assert isinstance(real_range, float)
        return real_range * 2, lambda x: np.where(x < real_range, 1, 0), [], []


@dataclass
class Benchmark:
    training_vectors: Optional[DatasetDescriptor] = None
    database_vectors: Optional[DatasetDescriptor] = None
    query_vectors: Optional[DatasetDescriptor] = None
    index_descs: Optional[list[IndexDescriptor]] = None
    range_ref_index_desc: Optional[str] = None
    k: Optional[int] = None
    distance_metric: str = "L2"

    def __post_init__(self):
        if self.distance_metric == "IP":
            self.distance_metric_type = faiss.METRIC_INNER_PRODUCT
        elif self.distance_metric == "L2":
            self.distance_metric_type = faiss.METRIC_L2
        else:
            raise ValueError

    def set_io(self, benchmark_io):
        self.io = benchmark_io
        self.io.distance_metric = self.distance_metric
        self.io.distance_metric_type = self.distance_metric_type

    def get_index_desc(self, factory: str) -> Optional[IndexDescriptor]:
        for desc in self.index_descs:
            if desc.factory == factory:
                return desc
        return None

    def range_search_reference(self, index, parameters, range_metric):
        logger.info("range_search_reference: begin")
        if isinstance(range_metric, list):
            assert len(range_metric) > 0
            m_radius = (
                max(rm[-2] for rm in range_metric)
                if self.distance_metric_type == faiss.METRIC_L2
                else min(rm[-2] for rm in range_metric)
            )
        else:
            m_radius = range_metric

        lims, D, I, R, P = self.range_search(
            index,
            parameters,
            radius=m_radius,
        )
        flat = index.factory == "Flat"
        (
            gt_radius,
            range_search_metric_function,
            coefficients,
            coefficients_training_data,
        ) = get_range_search_metric_function(
            range_metric,
            D if not flat else None,
            R if not flat else None,
        )
        logger.info("range_search_reference: end")
        return (
            gt_radius,
            range_search_metric_function,
            coefficients,
            coefficients_training_data,
        )

    def estimate_range(self, index, parameters, range_scoring_radius):
        D, I, R, P = index.knn_search(
            parameters,
            self.query_vectors,
            self.k,
        )
        samples = []
        for i, j in np.argwhere(R < range_scoring_radius):
            samples.append((R[i, j].item(), D[i, j].item()))
        if len(samples) > 0:  # estimate range
            samples.sort(key=itemgetter(0))
            return median(r for _, r in samples[-3:])
        else:  # ensure at least one result
            i, j = np.argwhere(R.min() == R)[0]
            return D[i, j].item()

    def range_search(
        self,
        index: Index,
        search_parameters: Optional[dict[str, int]],
        radius: Optional[float] = None,
        gt_radius: Optional[float] = None,
    ):
        logger.info("range_search: begin")
        if radius is None:
            assert gt_radius is not None
            radius = (
                gt_radius
                if index.is_flat()
                else self.estimate_range(
                    index,
                    search_parameters,
                    gt_radius,
                )
            )
        logger.info(f"Radius={radius}")
        return index.range_search(
            search_parameters=search_parameters,
            query_vectors=self.query_vectors,
            radius=radius,
        )

    def range_ground_truth(self, gt_radius, range_search_metric_function):
        logger.info("range_ground_truth: begin")
        flat_desc = self.get_index_desc("Flat")
        lims, D, I, R, P = self.range_search(
            flat_desc.index,
            search_parameters=None,
            radius=gt_radius,
        )
        gt_rsm = np.sum(range_search_metric_function(R)).item()
        logger.info("range_ground_truth: end")
        return gt_rsm

    def range_search_benchmark(
        self,
        results: dict[str, Any],
        index: Index,
        metric_key: str,
        radius: float,
        gt_radius: float,
        range_search_metric_function,
        gt_rsm: float,
    ):
        logger.info(f"range_search_benchmark: begin {index.get_index_name()}")

        def experiment(parameters, cost_metric, perf_metric):
            nonlocal results
            key = index.get_range_search_name(
                search_parameters=parameters,
                query_vectors=self.query_vectors,
                radius=radius,
            )
            key += metric_key
            if key in results["experiments"]:
                metrics = results["experiments"][key]
            else:
                lims, D, I, R, P = self.range_search(
                    index,
                    parameters,
                    radius=radius,
                    gt_radius=gt_radius,
                )
                range_search_metric = range_search_metric_function(R)
                range_search_pr = range_search_pr_curve(D, range_search_metric, gt_rsm)
                range_score_sum = np.sum(range_search_metric).item()
                metrics = P | {
                    "range_score_sum": range_score_sum,
                    "range_score_max_recall": range_score_sum / gt_rsm,
                    "range_search_pr": range_search_pr,
                }
                results["experiments"][key] = metrics
            return metrics[cost_metric], metrics[perf_metric]

        for cost_metric in ["time"]:
            for perf_metric in ["range_score_max_recall"]:
                op = index.get_operating_points()
                optimizer(
                    op,
                    experiment,
                    cost_metric,
                    perf_metric,
                )
        logger.info("range_search_benchmark: end")
        return results

    def knn_ground_truth(self):
        logger.info("knn_ground_truth: begin")
        flat_desc = self.get_index_desc("Flat")
        self.gt_knn_D, self.gt_knn_I, _, _ = flat_desc.index.knn_search(
            search_parameters=None,
            query_vectors=self.query_vectors,
            k=self.k,
        )
        logger.info("knn_ground_truth: end")

    def knn_search_benchmark(self, results: dict[str, Any], index: Index):
        index_name = index.get_index_name()
        logger.info(f"knn_search_benchmark: begin {index_name}")

        def experiment(parameters, cost_metric, perf_metric):
            nonlocal results
            key = index.get_knn_search_name(
                parameters,
                self.query_vectors,
                self.k,
            )
            key += "knn"
            if key in results["experiments"]:
                metrics = results["experiments"][key]
            else:
                D, I, R, P = index.knn_search(parameters, self.query_vectors, self.k)
                metrics = P | {
                    "knn_intersection": knn_intersection_measure(I, self.gt_knn_I),
                    "distance_ratio": distance_ratio_measure(
                        I, R, self.gt_knn_D, self.distance_metric_type
                    ),
                }
                results["experiments"][key] = metrics
            return metrics[cost_metric], metrics[perf_metric]

        for cost_metric in ["time"]:
            for perf_metric in ["knn_intersection", "distance_ratio"]:
                op = index.get_operating_points()
                optimizer(
                    op,
                    experiment,
                    cost_metric,
                    perf_metric,
                )
        logger.info("knn_search_benchmark: end")
        return results

    def train(self, results):
        xq = self.io.get_dataset(self.query_vectors)
        self.d = xq.shape[1]
        if self.get_index_desc("Flat") is None:
            self.index_descs.append(IndexDescriptor(factory="Flat"))
        for index_desc in self.index_descs:
            if index_desc.factory is not None:
                index = IndexFromFactory(
                    d=self.d,
                    metric=self.distance_metric,
                    database_vectors=self.database_vectors,
                    search_params=index_desc.search_params,
                    construction_params=index_desc.construction_params,
                    factory=index_desc.factory,
                    training_vectors=self.training_vectors,
                )
                index.set_io(self.io)
                index.train()
                index_desc.index = index
                results["indices"][index.get_codec_name()] = {"code_size": index.get_code_size()}
            else:
                index = IndexFromCodec(
                    d=self.d,
                    metric=self.distance_metric,
                    database_vectors=self.database_vectors,
                    search_params=index_desc.search_params,
                    construction_params=index_desc.construction_params,
                    path=index_desc.path,
                    bucket=index_desc.bucket,
                )
                index.set_io(self.io)
                index_desc.index = index
                results["indices"][index.get_codec_name()] = {"code_size": index.get_code_size()}
        return results

    def benchmark(self, result_file=None):
        logger.info("begin evaluate")

        faiss.omp_set_num_threads(24)
        results = {"indices": {}, "experiments": {}}
        results = self.train(results)

        # knn search
        self.knn_ground_truth()
        for index_desc in self.index_descs:
            results = self.knn_search_benchmark(
                results=results,
                index=index_desc.index,
            )

        # range search
        if self.range_ref_index_desc is not None:
            index_desc = self.get_index_desc(self.range_ref_index_desc)
            if index_desc is None:
                raise ValueError(f"Unknown range index {self.range_ref_index_desc}")
            if index_desc.range_metrics is None:
                raise ValueError(f"Range index {index_desc.factory} has no radius_score")
            results["metrics"] = {}
            for metric_key, range_metric in index_desc.range_metrics.items():
                (
                    gt_radius,
                    range_search_metric_function,
                    coefficients,
                    coefficients_training_data,
                ) = self.range_search_reference(
                    index_desc.index, index_desc.search_params, range_metric
                )
                results["metrics"][metric_key] = {
                    "coefficients": coefficients,
                    "training_data": coefficients_training_data,
                }
                gt_rsm = self.range_ground_truth(gt_radius, range_search_metric_function)
                for index_desc in self.index_descs:
                    if not index_desc.index.supports_range_search():
                        continue
                    results = self.range_search_benchmark(
                        results=results,
                        index=index_desc.index,
                        metric_key=metric_key,
                        radius=index_desc.radius,
                        gt_radius=gt_radius,
                        range_search_metric_function=range_search_metric_function,
                        gt_rsm=gt_rsm,
                    )
        if result_file is not None:
            self.io.write_json(results, result_file, overwrite=True)
        logger.info("end evaluate")
        return results
