# default_exp core


#hide
from nbdev.showdoc import *
from fastcore.test import *


# export 
import os

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
import warnings
import torchvision
from torchvision.datasets import MNIST, ImageFolder
from torchvision.transforms import ToTensor, Resize, Compose, CenterCrop, Normalize
import pytorch_lightning as pl
# from pytorch_lightning.metrics.functional import classification, f1
from pytorch_lightning.loggers import TensorBoardLogger

import fastai.vision.augment
import fastai.vision.data
# from fastai.vision.data import ImageDataLoaders
# from fastai.vision.augment import Resize


#export
class ImageFolderDataModule(pl.LightningDataModule):

    def __init__(self, data_dir, batch_size, transform):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.transform = transform
#         Compose([
#             Resize(256, interpolation=2),
#             CenterCrop(224),
#             ToTensor(),
# #             TODO: check whether normalize is the same for imagenet and fractalDB
#             Normalize(mean=[0.485, 0.456, 0.406],
#                                  std=[0.229, 0.224, 0.225])
#         ])

    def prepare_data(self, stage=None):
        pass
    
    def setup(self, stage=None):
        data_dir = self.data_dir
        transform = self.transform
        
        self.dls = fastai.vision.data.ImageDataLoaders.from_folder(data_dir, item_tfms=fastai.vision.augment.Resize(224))
        self.trainset = ImageFolder(os.path.join(data_dir, 'train'), transform)
        self.valset = ImageFolder(os.path.join(data_dir, 'valid'), transform)

    def train_dataloader(self):
        return DataLoader(self.trainset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.valset, batch_size=self.batch_size, shuffle=False)
    
    def test_dataloader(self):
        pass


data_dir = 'Datasets/cifar10'

transform = Compose([
        Resize(256, interpolation=2),
        CenterCrop(224),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])


dm = ImageFolderDataModule(data_dir, 128, transform)
dm.setup()


for x,y in dm.train_dataloader():
    test_eq(type(x), torch.Tensor) 
    test_eq(type(y), torch.Tensor) 
    break


#export
class CNNModule(pl.LightningModule):
    def __init__(self, model=None, pretrained=False, freeze_extractor=False, log_level=10, num_classes=None, weight_path=None):
        super().__init__()
        self.num_classes = num_classes
        self.pretrained = pretrained
        self.freeze_extractor = freeze_extractor

        assert model is not None, 'Select model from torchvision'
        assert num_classes is not None, 'Must configure number of classes with num_classes'
        
        if not model.startswith('resnet'):
            warnings.warn('models other than resnet variants may need different setup for finetuning to work.')
            
        # Prepare model for finetuning
        if weight_path is not None:
            param = torch.load(weight_path)
            backbone = eval(f'torchvision.models.{model}(pretrained={False})')     
            backbone.load_state_dict(param)
        else:
            backbone = eval(f'torchvision.models.{model}(pretrained={pretrained})')        
                    
        num_filters = backbone.fc.in_features
        layers = list(backbone.children())[:-1]
        
        self.feature_extractor = torch.nn.Sequential(*layers)
        self.classifier = nn.Linear(num_filters, num_classes)

    def forward(self, x):
        if self.freeze_extractor:
            self.feature_extractor.eval()
            with torch.no_grad():
                representations = self.feature_extractor(x).flatten(1)
        else:
            representations = self.feature_extractor(x).flatten(1)

        y = self.classifier(representations)
        return y

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        
        outputs = self.calculate_metrics(y_hat=y_hat, y=y)
        return outputs
    
    def training_epoch_end(self, outputs):
        avg_metrics = {}
        for metric in outputs[0].keys():
            val = torch.stack([x[metric] for x in outputs]).mean()
            self.logger.experiment.add_scalar(f"{metric}/train", val, self.current_epoch)
            avg_metrics[metric] = val

#         epoch_dictionary = {'loss': avg_metrics['loss']}
        
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        
        outputs = self.calculate_metrics(y_hat=y_hat, y=y)
        return outputs
    
    def validation_epoch_end(self, outputs):
        avg_metrics = {}
        for metric in outputs[0].keys():
            val = torch.stack([x[metric] for x in outputs]).mean()
            self.logger.experiment.add_scalar(f"{metric}/validation", val, self.current_epoch)
            avg_metrics[metric] = val

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02, weight_decay=1e-04)
#     >    return torch.optim.SGF(self.parameters(), lr=self.lr, aldsfk'a)
    
    
    def calculate_metrics(self, y, y_hat):
        loss = F.cross_entropy(y_hat, y)
        y_pred = y_hat.argmax(dim=1)
        acc = classification.accuracy(y_pred, y)
        f1_score = f1(y_pred, y, self.num_classes)
        return {
            "loss":loss,
            "acc": acc,
            "f1": f1_score
        }
    
    def on_sanity_check_start(self):
        self.logger.disable()

    def on_sanity_check_end(self):
        self.logger.enable() 



modelname = 'resnet18'
logger = TensorBoardLogger('tb_logs', name=modelname)
trainer = pl.Trainer(gpus=1, checkpoint_callback=False, logger=logger, fast_dev_run=5)
model = CNNModule(modelname, pretrained=True, num_classes=len(dm.trainset.classes))
test_eq(trainer.fit(model, dm), 1)


weight_path = 'FractalDB-1000_resnet50_epoch90.pth'
modelname = 'resnet50'
logger = TensorBoardLogger('tb_logs', name=modelname)
trainer = pl.Trainer(gpus=1, checkpoint_callback=False, logger=logger, fast_dev_run=5)
model = CNNModule(modelname, pretrained=True, num_classes=len(dm.trainset.classes), weight_path=weight_path)
test_eq(trainer.fit(model, dm), 1)





[((183, 9, 183, 53), 'pytorch_lightning.loggers.TensorBoardLogger', 'TensorBoardLogger', (), '', False, 'from pytorch_lightning.loggers import TensorBoardLogger\n'), ((184, 10, 184, 86), 'pytorch_lightning.Trainer', 'pl.Trainer', (), '', True, 'import pytorch_lightning as pl\n'), ((191, 9, 191, 53), 'pytorch_lightning.loggers.TensorBoardLogger', 'TensorBoardLogger', (), '', False, 'from pytorch_lightning.loggers import TensorBoardLogger\n'), ((192, 10, 192, 86), 'pytorch_lightning.Trainer', 'pl.Trainer', (), '', True, 'import pytorch_lightning as pl\n'), ((59, 15, 59, 82), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((62, 15, 62, 81), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((71, 8, 71, 36), 'torchvision.transforms.Resize', 'Resize', (), '', False, 'from torchvision.transforms import ToTensor, Resize, Compose, CenterCrop, Normalize\n'), ((72, 8, 72, 23), 'torchvision.transforms.CenterCrop', 'CenterCrop', ({(72, 19, 72, 22): '224'}, {}), '(224)', False, 'from torchvision.transforms import ToTensor, Resize, Compose, CenterCrop, Normalize\n'), ((73, 8, 73, 18), 'torchvision.transforms.ToTensor', 'ToTensor', ({}, {}), '()', False, 'from torchvision.transforms import ToTensor, Resize, Compose, CenterCrop, Normalize\n'), ((74, 8, 75, 51), 'torchvision.transforms.Normalize', 'Normalize', (), '', False, 'from torchvision.transforms import ToTensor, Resize, Compose, CenterCrop, Normalize\n'), ((114, 33, 114, 61), 'torch.nn.Sequential', 'torch.nn.Sequential', ({(114, 53, 114, 60): '*layers'}, {}), '(*layers)', False, 'import torch\n'), ((115, 26, 115, 61), 'torch.nn.Linear', 'nn.Linear', ({(115, 36, 115, 47): 'num_filters', (115, 49, 115, 60): 'num_classes'}, {}), '(num_filters, num_classes)', False, 'from torch import nn\n'), ((164, 15, 164, 40), 'torch.nn.functional.cross_entropy', 'F.cross_entropy', ({(164, 31, 164, 36): 'y_hat', (164, 38, 164, 39): 'y'}, {}), '(y_hat, y)', True, 'from torch.nn import functional as F\n'), ((55, 36, 55, 67), 'os.path.join', 'os.path.join', ({(55, 49, 55, 57): 'data_dir', (55, 59, 55, 66): '"""train"""'}, {}), "(data_dir, 'train')", False, 'import os\n'), ((56, 34, 56, 65), 'os.path.join', 'os.path.join', ({(56, 47, 56, 55): 'data_dir', (56, 57, 56, 64): '"""valid"""'}, {}), "(data_dir, 'valid')", False, 'import os\n'), ((101, 12, 101, 111), 'warnings.warn', 'warnings.warn', ({(101, 26, 101, 110): '"""models other than resnet variants may need different setup for finetuning to work."""'}, {}), "(\n    'models other than resnet variants may need different setup for finetuning to work.'\n    )", False, 'import warnings\n'), ((105, 20, 105, 43), 'torch.load', 'torch.load', ({(105, 31, 105, 42): 'weight_path'}, {}), '(weight_path)', False, 'import torch\n'), ((120, 17, 120, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((138, 18, 138, 59), 'torch.stack', 'torch.stack', ({(138, 30, 138, 58): '[x[metric] for x in outputs]'}, {}), '([x[metric] for x in outputs])', False, 'import torch\n'), ((154, 18, 154, 59), 'torch.stack', 'torch.stack', ({(154, 30, 154, 58): '[x[metric] for x in outputs]'}, {}), '([x[metric] for x in outputs])', False, 'import torch\n')]