# Copyright (c) 2021, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import Compose, ToTensor, Normalize

from nvflare.apis.dxo import from_shareable, DataKind, DXO
from nvflare.apis.executor import Executor
from nvflare.apis.fl_constant import ReturnCode
from nvflare.apis.fl_context import FLContext
from nvflare.apis.shareable import Shareable, make_reply
from nvflare.apis.signal import Signal
from nvflare.app_common.app_constant import AppConstants
from simple_network import SimpleNetwork


class Cifar10Validator(Executor):
    
    def __init__(self, validate_task_name=AppConstants.TASK_VALIDATION):
        super(Cifar10Validator, self).__init__()

        self._validate_task_name = validate_task_name

        # Setup the model
        self.model = SimpleNetwork()
        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        self.model.to(self.device)

        # Preparing the dataset for testing.
        transforms = Compose([
            ToTensor(),
            Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])
        self.test_data = CIFAR10(root='~/data', train=False, transform=transforms)
        self.test_loader = DataLoader(self.test_data, batch_size=4, shuffle=False)

    def execute(self, task_name: str, shareable: Shareable, fl_ctx: FLContext, abort_signal: Signal) -> Shareable:
        if task_name == self._validate_task_name:
            model_owner = "?"
            try:
                try:
                    dxo = from_shareable(shareable)
                except:
                    self.log_error(fl_ctx, "Error in extracting dxo from shareable.")
                    return make_reply(ReturnCode.BAD_TASK_DATA)

                # Ensure data_kind is weights.
                if not dxo.data_kind == DataKind.WEIGHTS:
                    self.log_exception(fl_ctx, f"DXO is of type {dxo.data_kind} but expected type WEIGHTS.")
                    return make_reply(ReturnCode.BAD_TASK_DATA)

                # Extract weights and ensure they are tensor.
                model_owner = shareable.get_header(AppConstants.MODEL_OWNER, "?")
                weights = {k: torch.as_tensor(v, device=self.device) for k, v in dxo.data.items()}

                # Get validation accuracy
                val_accuracy = self.do_validation(weights, abort_signal)
                if abort_signal.triggered:
                    return make_reply(ReturnCode.TASK_ABORTED)

                self.log_info(fl_ctx, f"Accuracy when validating {model_owner}'s model on"
                                      f" {fl_ctx.get_identity_name()}"f's data: {val_accuracy}')

                dxo = DXO(data_kind=DataKind.METRICS, data={'val_acc': val_accuracy})
                return dxo.to_shareable()
            except:
                self.log_exception(fl_ctx, f"Exception in validating model from {model_owner}")
                return make_reply(ReturnCode.EXECUTION_EXCEPTION)
        else:
            return make_reply(ReturnCode.TASK_UNKNOWN)

    def do_validation(self, weights, abort_signal):
        self.model.load_state_dict(weights)

        self.model.eval()

        correct = 0
        total = 0
        with torch.no_grad():
            for i, (images, labels) in enumerate(self.test_loader):
                if abort_signal.triggered:
                    return 0

                images, labels = images.to(self.device), labels.to(self.device)
                output = self.model(images)

                _, pred_label = torch.max(output, 1)

                correct += (pred_label == labels).sum().item()
                total += images.size()[0]

            metric = correct/float(total)

        return metric

[((38, 21, 38, 36), 'simple_network.SimpleNetwork', 'SimpleNetwork', ({}, {}), '()', False, 'from simple_network import SimpleNetwork\n'), ((47, 25, 47, 82), 'torchvision.datasets.CIFAR10', 'CIFAR10', (), '', False, 'from torchvision.datasets import CIFAR10\n'), ((48, 27, 48, 82), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((39, 46, 39, 71), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((39, 22, 39, 42), 'torch.device', 'torch.device', ({(39, 35, 39, 41): '"""cuda"""'}, {}), "('cuda')", False, 'import torch\n'), ((39, 77, 39, 96), 'torch.device', 'torch.device', ({(39, 90, 39, 95): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n'), ((83, 19, 83, 54), 'nvflare.apis.shareable.make_reply', 'make_reply', ({(83, 30, 83, 53): 'ReturnCode.TASK_UNKNOWN'}, {}), '(ReturnCode.TASK_UNKNOWN)', False, 'from nvflare.apis.shareable import Shareable, make_reply\n'), ((92, 13, 92, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((44, 12, 44, 22), 'torchvision.transforms.ToTensor', 'ToTensor', ({}, {}), '()', False, 'from torchvision.transforms import Compose, ToTensor, Normalize\n'), ((45, 12, 45, 55), 'torchvision.transforms.Normalize', 'Normalize', ({(45, 22, 45, 37): '(0.5, 0.5, 0.5)', (45, 39, 45, 54): '(0.5, 0.5, 0.5)'}, {}), '((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))', False, 'from torchvision.transforms import Compose, ToTensor, Normalize\n'), ((77, 22, 77, 85), 'nvflare.apis.dxo.DXO', 'DXO', (), '', False, 'from nvflare.apis.dxo import from_shareable, DataKind, DXO\n'), ((100, 32, 100, 52), 'torch.max', 'torch.max', ({(100, 42, 100, 48): 'output', (100, 50, 100, 51): '1'}, {}), '(output, 1)', False, 'import torch\n'), ((55, 26, 55, 51), 'nvflare.apis.dxo.from_shareable', 'from_shareable', ({(55, 41, 55, 50): 'shareable'}, {}), '(shareable)', False, 'from nvflare.apis.dxo import from_shareable, DataKind, DXO\n'), ((63, 27, 63, 63), 'nvflare.apis.shareable.make_reply', 'make_reply', ({(63, 38, 63, 62): 'ReturnCode.BAD_TASK_DATA'}, {}), '(ReturnCode.BAD_TASK_DATA)', False, 'from nvflare.apis.shareable import Shareable, make_reply\n'), ((67, 30, 67, 68), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((72, 27, 72, 62), 'nvflare.apis.shareable.make_reply', 'make_reply', ({(72, 38, 72, 61): 'ReturnCode.TASK_ABORTED'}, {}), '(ReturnCode.TASK_ABORTED)', False, 'from nvflare.apis.shareable import Shareable, make_reply\n'), ((81, 23, 81, 65), 'nvflare.apis.shareable.make_reply', 'make_reply', ({(81, 34, 81, 64): 'ReturnCode.EXECUTION_EXCEPTION'}, {}), '(ReturnCode.EXECUTION_EXCEPTION)', False, 'from nvflare.apis.shareable import Shareable, make_reply\n'), ((58, 27, 58, 63), 'nvflare.apis.shareable.make_reply', 'make_reply', ({(58, 38, 58, 62): 'ReturnCode.BAD_TASK_DATA'}, {}), '(ReturnCode.BAD_TASK_DATA)', False, 'from nvflare.apis.shareable import Shareable, make_reply\n')]