import torch
from torch import nn

class Finetune(nn.Module):
    def __init__(self, backbone, feat_dim, num_class, **kwargs):
        super().__init__()
        self.backbone = backbone
        self.feat_dim = feat_dim
        self.num_class = num_class
        self.classifier = nn.Linear(feat_dim, num_class)
        self.loss_fn = nn.CrossEntropyLoss(reduction='mean')
        self.device = kwargs['device']
        self.kwargs = kwargs
    
    def observe(self, data):
        x, y = data['image'], data['label']
        x = x.to(self.device)
        y = y.to(self.device)
        logit = self.classifier(self.backbone(x)['features'])    
        loss = self.loss_fn(logit, y)

        pred = torch.argmax(logit, dim=1)

        acc = torch.sum(pred == y).item()
        return pred, acc / x.size(0), loss

    def inference(self, data):
        x, y = data['image'], data['label']
        x = x.to(self.device)
        y = y.to(self.device)
        
        logit = self.classifier(self.backbone(x)['features'])  
        pred = torch.argmax(logit, dim=1)

        acc = torch.sum(pred == y).item()
        return pred, acc / x.size(0)

    def forward(self, x):
        return self.classifier(self.backbone(x)['features'])  
    
    def before_task(self, task_idx, buffer, train_loader, test_loaders):
        pass

    def after_task(self, task_idx, buffer, train_loader, test_loaders):
        pass
    
    def get_parameters(self, config):
        train_parameters = []
        train_parameters.append({"params": self.backbone.parameters()})
        train_parameters.append({"params": self.classifier.parameters()})
        return train_parameters