__all__ = ["ModelAdapter"]

from icevision.imports import *
from icevision.metrics import *
from icevision.engines.lightning.lightning_model_adapter import LightningModelAdapter
from torch.nn import CrossEntropyLoss
from icevision.models.fastai import unet


class ModelAdapter(LightningModelAdapter, ABC):
    """Lightning module specialized for EfficientDet, with metrics support.

    The methods `forward`, `training_step`, `validation_step`, `validation_epoch_end`
    are already overriden.

    # Arguments
        model: The pytorch model to use.
        metrics: `Sequence` of metrics to use.

    # Returns
        A `LightningModule`.
    """

    def __init__(self, model: nn.Module, metrics: List[Metric] = None):
        super().__init__(metrics=metrics)
        self.model = model
        self.loss_func = CrossEntropyLoss()

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def training_step(self, batch, batch_idx):
        (xb, yb), _ = batch
        preds = self(xb)

        loss = self.compute_loss(preds, yb)

        self.log("train_loss", loss)

        return loss

    def compute_loss(self, preds, yb):
        return self.loss_func(preds, yb)

    def validation_step(self, batch, batch_idx):
        self._shared_eval(batch=batch, loss_log_key="val")

    def _shared_eval(self, batch, loss_log_key):
        (xb, yb), records = batch

        preds = self(xb)
        loss = self.compute_loss(preds, yb)

        preds = self.convert_raw_predictions(
            batch=xb,
            raw_preds=preds,
            records=records,
        )

        self.accumulate_metrics(preds)

        self.log(f"{loss_log_key}_loss", loss)

    def convert_raw_predictions(self, batch, raw_preds, records):
        return unet.convert_raw_predictions(
            batch=batch,
            raw_preds=raw_preds,
            records=records,
        )

    def validation_epoch_end(self, outs):
        self.finalize_metrics()

    def test_step(self, batch, batch_idx):
        self._shared_eval(batch=batch, loss_log_key="test")

    def test_epoch_end(self, outs):
        self.finalize_metrics()
