<reponame>dreinq/DeepQ
import multiprocessing as mp
from logging import Logger
from typing import Tuple

import torch
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from MCQ.datasets import SiftLike, Enumerate
from MCQ.envs import Env
from MCQ.utils import Saver
from MCQ.utils.runtime import Timer
from MCQ.metrics import QuantizationError, Reconstruct
from MCQ import Consts

class Eval:
    def __init__(self, evalOther: bool, savePath: str, dataset: SiftLike, env: Env, model: nn.Module, logger: Logger = None):
        if not evalOther:
            self._model = model.cuda()
            self._env = env
            Saver.Load(savePath, model=self._model)
            self._model.eval()
            self._dataset = dataset
        else:
            pass
        torch.autograd.set_grad_enabled(False)

        self._obsMean = None
        self._obsStd = None

        self._logger = logger or Consts.Logger

    def Encode(self, dataset, icm = False, C = None):
        dataloader = DataLoader(dataset, batch_size=100000, shuffle=False, num_workers=mp.cpu_count())
        B = list()
        ticker = Timer()
        for x in tqdm(dataloader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            x = x.cuda()
            if self._env.DoNormalizationOnObs:
                newX = (x - self._obsMean) / self._obsStd
            b = self._model.Encode(newX, icm=True, C=C, shift=self._obsMean, scale=self._obsStd)
            B.append(b.detach().cpu())
        interval, _ = ticker.Tick()
        self._logger.info("Encode %d samples for %.2f seconds, %.2e s/sample, %.2e samples/s", len(dataset), interval, interval / len(dataset), len(dataset) / interval)
        B = torch.cat(B, 0)
        return B

    def EncodeFast(self, dataset, icm=True, C = None):
        dataloader = DataLoader(dataset, batch_size=100000, shuffle=False, num_workers=mp.cpu_count())
        B = list()
        icmC = C.reshape(-1, C.shape[-1])
        cUnary = (icmC ** 2).sum(-1)
        cPair = 2 * (icmC @ icmC.t())
        ticker = Timer()
        for x in tqdm(dataloader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            x = x.cuda()
            if self._env.DoNormalizationOnObs:
                newX = (x - self._obsMean) / self._obsStd
            b = self._model.EncodeFast(newX, icm=True, realX=x, C=icmC, cUnary=cUnary, cPair=cPair)
            B.append(b.detach().cpu())
        interval, _ = ticker.Tick()
        self._logger.info("Encode %d samples for %.2f seconds, %.2e s/sample, %.2e samples/s", len(dataset), interval, interval / len(dataset), len(dataset) / interval)
        B = torch.cat(B, 0)
        return B

    def GetCodebook(self, X, icm=False):
        if self._env.DoNormalizationOnObs:
            newX = X - self._obsMean.to(X.device)
            newX /= self._obsStd.to(X.device)
        B = self._model.Encode(newX)
        C, qError = self._env.Eval(X, B)
        if not icm:
            return C
        B = self._model.Encode(newX, icm=True, C=C, shift=self._obsMean, scale=self._obsStd)
        newC, newQError = self._env.Eval(X, B)
        if newQError.mean() > qError.mean():
            return C
        return newC

    @staticmethod
    def Retrieval(queries, C, B):
        M, _, _ = C.shape
        # if M == 1:
        #     return Eval._retrievalVQ(queryLoader, C, B)

        topK: Tuple[torch.Tensor, torch.Tensor] = None

        N = len(queries)
        ix = torch.arange(len(queries))[:, None]

        baseLoader = DataLoader(Enumerate(B), batch_size=1000, shuffle=False, num_workers=0)
        for i, b in tqdm(baseLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            quantized = Reconstruct(C, b)
            # [N, 100]
            dist = ((queries[:, None, ...] - quantized) ** 2).sum(-1)
            i = i.to(b.device).expand_as(dist)
            if topK is not None:
                # [N, 200]
                preparedI = torch.cat([i, topK[0]], -1)
                # [N, 200]
                preparedD = torch.cat([dist, topK[1]], -1)
            else:
                preparedI = i
                preparedD = dist
            # _, [N, 100]
            _, iy = torch.topk(preparedD, k=100, dim=-1, largest=False, sorted=True)
            topK = (preparedI[ix.expand_as(iy), iy], preparedD[ix.expand_as(iy), iy])
        return topK[0]
        # # [NB, D]
        # quantized = Reconstruct(C, B)
        # results = []
        # for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
        #     q = q.cuda()
        #     # [NQ, NB]
        #     dist = ((q[:, None, ...] - quantized) ** 2).sum(-1)
        #     _, indices = torch.topk(dist, k=100, dim=-1, largest=False, sorted=True)
        #     results.append(indices.detach())
        # return torch.cat(results, 0)

    @staticmethod
    def RetrievalSlow(queries, C, B):
        M, _, _ = C.shape
        queryLoader = DataLoader(queries, batch_size=1, num_workers=0, shuffle=False)
        if M == 1:
            return Eval._retrievalVQ(queryLoader, C, B)
        # [NB, D]
        quantized = Reconstruct(C, B)
        results = []
        for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            q = q.cuda()
            # [NQ, NB]
            dist = ((q[:, None, ...] - quantized) ** 2).sum(-1)
            _, indices = torch.topk(dist, k=100, dim=-1, largest=False, sorted=True)
            results.append(indices.detach())
        return torch.cat(results, 0)

    @staticmethod
    def _retrievalVQ(queryLoader, C, B):
        # [NB, 1] -> [NB]
        B = B.squeeze()
        results = []
        slot = B.shape[0] // C.shape[1]
        for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            q = q.cuda()
            # [N, K] <- [N, 1, D], [1, K, D]
            dist = ((q[:, None, ...] - C) ** 2).sum(-1)
            # [N, NB]
            allDist = dist[:, B]
            # [N, slot]
            _, indices = torch.topk(allDist, k=slot, dim=-1, largest=False, sorted=True)
            results.append(indices.detach())
            # [NQ, 100]
        return torch.cat(results, 0)

    @staticmethod
    def Recall(results, groundtruths):
        # [N, 100] = [N, 100] == [N, 1]
        isin = results == groundtruths
        recalls = []
        for i in range(1, isin.shape[-1] + 1):
            recalls.append(isin[:, :i].any(-1).float().mean())
        return torch.Tensor(recalls)

    def Test(self, C=None, B=None):
        sift = self._dataset
        if self._env.DoNormalizationOnObs:
            self._obsMean = sift.data.mean(0).cuda()
            self._obsStd = sift.data.std(0).cuda()
        if C is None:
            sift.Train()
            C = self.GetCodebook(sift.data, icm=True).cuda()
        if B is None:
            sift.Encode(device="cpu")
            B = self.Encode(sift, icm=True, C=C)
            self._logger.info("Quantization error in base: %.8e", QuantizationError(sift.data.cuda(), C.cuda(), B.cuda()).mean())
        sift.Query(device="cuda")
        # queryLoader = DataLoader(sift, batch_size=1, shuffle=False, num_workers=mp.cpu_count())
        results = self.Retrieval(sift.data, C.cuda(), B.cuda())
        sift.Gt(device="cuda")
        recalls = self.Recall(results, sift.data[:, :1]) * 100
        self._logger.info("R @ 1: %.2f%%", recalls[0])
        self._logger.info("R @ 10: %.2f%%", recalls[9])
        self._logger.info("R @ 100: %.2f%%", recalls[99])

import multiprocessing as mp
from logging import Logger
from typing import Tuple

import torch
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from MCQ.datasets import SiftLike, Enumerate
from MCQ.envs import Env
from MCQ.utils import Saver
from MCQ.utils.runtime import Timer
from MCQ.metrics import QuantizationError, Reconstruct
from MCQ import Consts

class Eval:
    def __init__(self, evalOther: bool, savePath: str, dataset: SiftLike, env: Env, model: nn.Module, logger: Logger = None):
        if not evalOther:
            self._model = model.cuda()
            self._env = env
            Saver.Load(savePath, model=self._model)
            self._model.eval()
            self._dataset = dataset
        else:
            pass
        torch.autograd.set_grad_enabled(False)

        self._obsMean = None
        self._obsStd = None

        self._logger = logger or Consts.Logger

    def Encode(self, dataset, icm = False, C = None):
        dataloader = DataLoader(dataset, batch_size=100000, shuffle=False, num_workers=mp.cpu_count())
        B = list()
        ticker = Timer()
        for x in tqdm(dataloader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            x = x.cuda()
            if self._env.DoNormalizationOnObs:
                newX = (x - self._obsMean) / self._obsStd
            b = self._model.Encode(newX, icm=True, C=C, shift=self._obsMean, scale=self._obsStd)
            B.append(b.detach().cpu())
        interval, _ = ticker.Tick()
        self._logger.info("Encode %d samples for %.2f seconds, %.2e s/sample, %.2e samples/s", len(dataset), interval, interval / len(dataset), len(dataset) / interval)
        B = torch.cat(B, 0)
        return B

    def EncodeFast(self, dataset, icm=True, C = None):
        dataloader = DataLoader(dataset, batch_size=100000, shuffle=False, num_workers=mp.cpu_count())
        B = list()
        icmC = C.reshape(-1, C.shape[-1])
        cUnary = (icmC ** 2).sum(-1)
        cPair = 2 * (icmC @ icmC.t())
        ticker = Timer()
        for x in tqdm(dataloader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            x = x.cuda()
            if self._env.DoNormalizationOnObs:
                newX = (x - self._obsMean) / self._obsStd
            b = self._model.EncodeFast(newX, icm=True, realX=x, C=icmC, cUnary=cUnary, cPair=cPair)
            B.append(b.detach().cpu())
        interval, _ = ticker.Tick()
        self._logger.info("Encode %d samples for %.2f seconds, %.2e s/sample, %.2e samples/s", len(dataset), interval, interval / len(dataset), len(dataset) / interval)
        B = torch.cat(B, 0)
        return B

    def GetCodebook(self, X, icm=False):
        if self._env.DoNormalizationOnObs:
            newX = X - self._obsMean.to(X.device)
            newX /= self._obsStd.to(X.device)
        B = self._model.Encode(newX)
        C, qError = self._env.Eval(X, B)
        if not icm:
            return C
        B = self._model.Encode(newX, icm=True, C=C, shift=self._obsMean, scale=self._obsStd)
        newC, newQError = self._env.Eval(X, B)
        if newQError.mean() > qError.mean():
            return C
        return newC

    @staticmethod
    def Retrieval(queries, C, B):
        M, _, _ = C.shape
        # if M == 1:
        #     return Eval._retrievalVQ(queryLoader, C, B)

        topK: Tuple[torch.Tensor, torch.Tensor] = None

        N = len(queries)
        ix = torch.arange(len(queries))[:, None]

        baseLoader = DataLoader(Enumerate(B), batch_size=1000, shuffle=False, num_workers=0)
        for i, b in tqdm(baseLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            quantized = Reconstruct(C, b)
            # [N, 100]
            dist = ((queries[:, None, ...] - quantized) ** 2).sum(-1)
            i = i.to(b.device).expand_as(dist)
            if topK is not None:
                # [N, 200]
                preparedI = torch.cat([i, topK[0]], -1)
                # [N, 200]
                preparedD = torch.cat([dist, topK[1]], -1)
            else:
                preparedI = i
                preparedD = dist
            # _, [N, 100]
            _, iy = torch.topk(preparedD, k=100, dim=-1, largest=False, sorted=True)
            topK = (preparedI[ix.expand_as(iy), iy], preparedD[ix.expand_as(iy), iy])
        return topK[0]
        # # [NB, D]
        # quantized = Reconstruct(C, B)
        # results = []
        # for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
        #     q = q.cuda()
        #     # [NQ, NB]
        #     dist = ((q[:, None, ...] - quantized) ** 2).sum(-1)
        #     _, indices = torch.topk(dist, k=100, dim=-1, largest=False, sorted=True)
        #     results.append(indices.detach())
        # return torch.cat(results, 0)

    @staticmethod
    def RetrievalSlow(queries, C, B):
        M, _, _ = C.shape
        queryLoader = DataLoader(queries, batch_size=1, num_workers=0, shuffle=False)
        if M == 1:
            return Eval._retrievalVQ(queryLoader, C, B)
        # [NB, D]
        quantized = Reconstruct(C, B)
        results = []
        for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            q = q.cuda()
            # [NQ, NB]
            dist = ((q[:, None, ...] - quantized) ** 2).sum(-1)
            _, indices = torch.topk(dist, k=100, dim=-1, largest=False, sorted=True)
            results.append(indices.detach())
        return torch.cat(results, 0)

    @staticmethod
    def _retrievalVQ(queryLoader, C, B):
        # [NB, 1] -> [NB]
        B = B.squeeze()
        results = []
        slot = B.shape[0] // C.shape[1]
        for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
            q = q.cuda()
            # [N, K] <- [N, 1, D], [1, K, D]
            dist = ((q[:, None, ...] - C) ** 2).sum(-1)
            # [N, NB]
            allDist = dist[:, B]
            # [N, slot]
            _, indices = torch.topk(allDist, k=slot, dim=-1, largest=False, sorted=True)
            results.append(indices.detach())
            # [NQ, 100]
        return torch.cat(results, 0)

    @staticmethod
    def Recall(results, groundtruths):
        # [N, 100] = [N, 100] == [N, 1]
        isin = results == groundtruths
        recalls = []
        for i in range(1, isin.shape[-1] + 1):
            recalls.append(isin[:, :i].any(-1).float().mean())
        return torch.Tensor(recalls)

    def Test(self, C=None, B=None):
        sift = self._dataset
        if self._env.DoNormalizationOnObs:
            self._obsMean = sift.data.mean(0).cuda()
            self._obsStd = sift.data.std(0).cuda()
        if C is None:
            sift.Train()
            C = self.GetCodebook(sift.data, icm=True).cuda()
        if B is None:
            sift.Encode(device="cpu")
            B = self.Encode(sift, icm=True, C=C)
            self._logger.info("Quantization error in base: %.8e", QuantizationError(sift.data.cuda(), C.cuda(), B.cuda()).mean())
        sift.Query(device="cuda")
        # queryLoader = DataLoader(sift, batch_size=1, shuffle=False, num_workers=mp.cpu_count())
        results = self.Retrieval(sift.data, C.cuda(), B.cuda())
        sift.Gt(device="cuda")
        recalls = self.Recall(results, sift.data[:, :1]) * 100
        self._logger.info("R @ 1: %.2f%%", recalls[0])
        self._logger.info("R @ 10: %.2f%%", recalls[9])
        self._logger.info("R @ 100: %.2f%%", recalls[99])
# if M == 1: #     return Eval._retrievalVQ(queryLoader, C, B) # [N, 100] # [N, 200] # [N, 200] # _, [N, 100] # # [NB, D] # quantized = Reconstruct(C, B) # results = [] # for q in tqdm(queryLoader, ncols=40, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"): #     q = q.cuda() #     # [NQ, NB] #     dist = ((q[:, None, ...] - quantized) ** 2).sum(-1) #     _, indices = torch.topk(dist, k=100, dim=-1, largest=False, sorted=True) #     results.append(indices.detach()) # return torch.cat(results, 0) # [NB, D] # [NQ, NB] # [NB, 1] -> [NB] # [N, K] <- [N, 1, D], [1, K, D] # [N, NB] # [N, slot] # [NQ, 100] # [N, 100] = [N, 100] == [N, 1] # queryLoader = DataLoader(sift, batch_size=1, shuffle=False, num_workers=mp.cpu_count())