from copy import deepcopy
from torch import nn
import torch as th
import pytorch_lightning as pl


from .models import CLIPClassifier
from .convnext_meta import build_covnext


class EMA(nn.Module):
    """Model Exponential Moving Average V2 from timm"""

    def __init__(self, model: nn.Module, decay: float = 0.9999):
        super(EMA, self).__init__()
        # make a copy of the model for accumulating moving average of weights
        self.module = deepcopy(model)
        self.module.eval()
        self.decay = decay

    def _update(self, model: nn.Module, update_fn):
        with th.no_grad():
            for ema_v, model_v in zip(
                self.module.state_dict().values(), model.state_dict().values()
            ):
                ema_v.copy_(update_fn(ema_v, model_v))

    def update(self, model):
        self._update(
            model, update_fn=lambda e, m: self.decay * e + (1.0 - self.decay) * m
        )

    def set(self, model):
        self._update(model, update_fn=lambda e, m: m)


class MosquitoClassifier(pl.LightningModule):
    def __init__(
        self,
        n_classes: int = 6,
        model_name: str = "ViT-L-14",
        dataset: str = None,
        head_version: int = 0,
        use_ema: bool = False,
    ):
        super().__init__()
        if dataset == "imagenet":
            self.cls = build_covnext(model_name, n_classes)

        else:
            self.cls = CLIPClassifier(n_classes, model_name, dataset, head_version)

        self.use_ema = use_ema
        if use_ema:
            self.ema = EMA(self.cls, decay=0.995)

    def forward(self, x: th.Tensor) -> th.Tensor:
        if self.use_ema and not self.training:
            print("Using EMA...")
            return nn.Softmax(dim=1)(self.ema.module(x))
        return nn.Softmax(dim=1)(self.cls(x))