from copy import deepcopy
from typing import Any, Sequence, Tuple, List, Callable, Dict
import torch
import argparse
import pdb
from cassle.utils.lars import LARSWrapper
from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
from cassle.utils.scheduler import ConstantWarmupCosineAnnealingLR
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
from functools import partial
from cassle.model.cbam_resnet import resnet18_cbam
from torch import nn
import torch.nn.functional as F
from cassle.losses.barlow import barlow_loss_func
import numpy as np
from tqdm import *
from torchvision.models import resnet18, resnet50


def base_cbam_wrapper5(Method=object):
    class BaseCbamWrapper(Method):
        def __init__(self, **kwargs) -> None:
            super().__init__(**kwargs)
            self.output_dim = kwargs["output_dim"]
            self.checkpoint_dir = kwargs["checkpoint_dir"]
            self.automatic_optimization = False
            self.cifar = kwargs["cifar"]
            self.pretrained_model = kwargs["pretrained_model"]
            self.encoder_name = kwargs["encoder"]
            self.current_task_idx = kwargs["task_idx"]
            self.proj_hidden_dim = kwargs["proj_hidden_dim"]
            self.output_dim = kwargs["output_dim"]
            self.num_tasks = kwargs["num_tasks"]
            self.addition_images = kwargs["addition_images"]
            self.method = kwargs["sample_method"]
            self.distill_lam = kwargs["distill_lam"]
            self.indices= []
            self.indices_distances = []
            
            if self.current_task_idx > 0:
                #distill projector
                self.distill_proj_hidden_dim = 2048
                self.distill_projector = nn.Sequential(
                    nn.Linear(self.features_dim, self.distill_proj_hidden_dim),
                    nn.BatchNorm1d(self.distill_proj_hidden_dim),
                    nn.ReLU(),
                    nn.Linear(self.distill_proj_hidden_dim, self.features_dim),
                    nn.ReLU(),)

                # load cbam_resnet for task 1~t
                print(f"Loading previous task checkpoint {self.pretrained_model}...")
                state_dict = torch.load(self.pretrained_model, map_location="cpu")["state_dict"]
                self.load_state_dict(state_dict, strict=False)
                


        def on_train_start(self):
            super().on_train_start()
            if self.current_task_idx > 0:
                # save frozen encoder for distillation
                self.frozen_encoder = deepcopy(self.encoder)
                self.frozen_projector = deepcopy(self.projector)
                for pg in self.frozen_encoder.parameters():
                    pg.requires_grad = False
                for pg in self.frozen_projector.parameters():
                    pg.requires_grad = False

        
        def distill_forward(self, x1, x2):
            x1 = self.distill_projector(x1)
            x2 = self.distill_projector(x2)
            return x1, x2
                
        @torch.no_grad()
        def frozen_forward(self, X):
            feats = self.frozen_encoder(X)
            return feats
        
        def calculate_distill_loss(self, old_fea, fea):
            return torch.mean(torch.norm(old_fea - fea, p=2, dim=1))

        
        def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
            opt = self.optimizers(use_pl_optimizer=True)
            _, (X1, X2), _ = batch[f"task{self.current_task_idx}"]
            loss = 0
            barlow_loss = 0

            out = super().training_step(batch, batch_idx)
            barlow_loss += out["loss"]
            loss += barlow_loss
            
            if self.current_task_idx > 0:
                distill_loss = 0
                f1, f2 = out["feats"]
                frozen_feats1, frozen_feats2 = self.frozen_forward(X1), self.frozen_forward(X2)

                # feature_distill_loss PFR
                f_f1, f_f2 = self.distill_forward(f1, f2)

                feature_distill_loss = self.distill_lam*(self.calculate_distill_loss(frozen_feats1, f_f1) + self.calculate_distill_loss(frozen_feats2, f_f2))
                distill_loss += feature_distill_loss


                loss += distill_loss

                self.log_dict({"barlow_loss": barlow_loss, "distill_loss": distill_loss}, on_epoch=True, sync_dist=True, prog_bar=True)

                out.update({"loss": loss})
            else:
                self.log_dict({"barlow_loss": barlow_loss}, on_epoch=True, sync_dist=True, prog_bar=True)
            opt.zero_grad()
            self.manual_backward(loss)
            opt.step()
            # self.ema_update()
            return out

        def on_train_epoch_end(self):
            super().on_train_epoch_end()
            sch = self.lr_schedulers()
            sch.step()


        @property
        def learnable_params(self) -> List[dict]:
            """Adds projector parameters to parent's learnable parameters.

            Returns:
                List[dict]: list of learnable parameters.
            """
            params = [
            {"name": "encoder", "params": self.encoder.parameters()},
            {
                "name": "classifier",
                "params": self.classifier.parameters(),
                "lr": self.classifier_lr,
                "weight_decay": 0,
            },
            {"params": self.projector.parameters()}
        ]
            if self.current_task_idx > 0:
                params.append({"name": "distill_projector", "params": self.distill_projector.parameters()})
            return params
    
            
    return BaseCbamWrapper
