import logging
from abc import ABC, abstractmethod
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Dict, Iterable

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

from fed_distill.common.tester import AccuracyTester, get_batch_accuracy

logger = logging.getLogger(__name__)


class Trainer(ABC):
    @abstractmethod
    def train_for(self, num_epochs: int = 1) -> Dict[str, Any]:
        raise NotImplementedError()

    @property
    @abstractmethod
    def model(self) -> nn.Module:
        raise NotImplementedError()

    def train_epoch(self) -> Dict[str, Any]:
        return self.train_for(1)


@dataclass
class BasicTrainer(Trainer):
    net: nn.Module
    optimizer: torch.optim.Optimizer
    scheduler: torch.optim.lr_scheduler._LRScheduler
    criterion: nn.Module
    loader: DataLoader
    device: str = "cuda"

    def __post_init__(self,) -> None:
        self.train_metrics = {"train_losses": [], "train_accs": []}
        self.epoch = 1

    def _train_batch(self, images, labels) -> float:
        self.optimizer.zero_grad()
        prediction = self.net(images.to(self.device))
        loss = self.criterion(prediction, labels.to(self.device))
        loss.backward()
        self.optimizer.step()
        return float(loss)

    def train_for(self, num_epochs: int = 1) -> None:
        self.model.train()
        for _ in range(num_epochs):
            logger.info(f"Training epoch %i", self.epoch)
            epoch_loss = 0
            epoch_acc = 0
            num_batches = 0
            for images, labels in self.loader:
                epoch_loss += self._train_batch(images, labels)
                epoch_acc += get_batch_accuracy(self.net, images.to(self.device), labels.to(self.device))
                self.net.train()
                num_batches += 1
            epoch_loss /= num_batches
            epoch_acc /= num_batches
            self.train_metrics["train_losses"].append(epoch_loss)
            self.train_metrics["train_accs"].append(epoch_acc)
            logger.info("Training accuracy: %f", epoch_acc)
            logger.info("Training loss: %f", epoch_loss)
            self.scheduler.step()
            self.epoch += 1
        return self.train_metrics

    @property
    def model(self):
        return self.net


class AccuracyTrainer(Trainer):
    def __init__(
        self,
        base_trainer: Trainer,
        accuracy_tester: AccuracyTester,
        result_model: nn.Module,
        acc_limit: float = 1.0,
    ) -> None:
        self.base = base_trainer
        self.tester = accuracy_tester
        self.result_model = result_model
        self.test_accs = []
        self.best_model = None
        self.best_acc = 0.0
        self.acc_limit = acc_limit

    @property
    def model(self):
        self.result_model.load_state_dict(self.best_model)
        return self.result_model

    def train_for(self, num_epochs: int = 1) -> Dict[str, Any]:
        base_training_dict = {}
        for _ in range(num_epochs):
            base_training_dict = self.base.train_epoch()
            model = self.base.model
            accuracy = self.tester(model)
            self.test_accs.append(accuracy)
            logger.info("Testing accuracy: %f", accuracy)
            if accuracy > self.best_acc:
                self.best_acc = accuracy
                self.best_model = deepcopy(model.state_dict())
            logger.info("Best accuracy: %f", self.best_acc)
            if self.best_acc >= self.acc_limit:
                logger.info("Training stopping early due to convergence")
                break

        final_dict_copy = deepcopy(base_training_dict)
        final_dict_copy["test_accs"] = self.test_accs
        final_dict_copy["best_model"] = self.best_model
        return final_dict_copy
