from collections import OrderedDict
import torch
import torch.nn as nn
from torch.utils.data import ConcatDataset
import random
from catalyst.dl.experiment import ConfigExperiment
from dataset import *
from augmentation import train_aug, valid_aug


class Experiment(ConfigExperiment):
    def _postprocess_model_for_stage(self, stage: str, model: nn.Module):

        import warnings
        warnings.filterwarnings("ignore")

        random.seed(2411)
        np.random.seed(2411)
        torch.manual_seed(2411)

        model_ = model
        if isinstance(model, torch.nn.DataParallel):
            model_ = model_.module

        if stage == "warmup":
            if hasattr(model_, 'freeze'):
                model_.freeze(model_)
                print("Freeze backbone model using freeze method !!!")
            else:
                for param in model_.parameters():
                    param.requires_grad = False

                for param in model_.get_classifier().parameters():
                    param.requires_grad = True
                print("Freeze backbone model !!!")

        else:
            if hasattr(model_, 'unfreeze'):
                model_.unfreeze(model_)
                print("Unfreeze backbone model using unfreeze method !!!")
            else:
                for param in model_.parameters():
                    param.requires_grad = True

                print("Unfreeze backbone model !!!")
        #
        # import apex
        # model_ = apex.parallel.convert_syncbn_model(model_)

        return model_

    def get_datasets(self, stage: str, **kwargs):
        datasets = OrderedDict()

        """
        image_key: 'id'
        label_key: 'attribute_ids'
        """

        image_size = kwargs.get("image_size", [224, 224])
        train_csv = kwargs.get('train_csv', None)
        valid_csv = kwargs.get('valid_csv', None)
        with_any = kwargs.get('with_any', True)
        dataset_type = kwargs.get('dataset_type', 'RSNADataset')
        image_type = kwargs.get('image_type', 'jpg')
        normalization = kwargs.get('normalization', True)
        root = kwargs.get('root', None)

        print(f"Image Size: {image_size}")

        if train_csv:
            transform = train_aug(image_size)
            if dataset_type == 'RSNADataset':
                train_set = RSNADataset(
                    csv_file=train_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform,
                    mode='train',
                    image_type=image_type
                )
            elif dataset_type == 'RSNAMultiWindowsDataset':
                train_set = RSNAMultiWindowsDataset(
                    csv_file=train_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform
                )
            elif dataset_type == 'RSNADicomDataset':
                train_set = RSNADicomDataset(
                    csv_file=train_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform
                )
            elif dataset_type == "RSNARandomWindowDataset":
                train_set = RSNARandomWindowDataset(
                    csv_file=train_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform
                )
            else:
                raise("No Dataset: {}".format(dataset_type))
            datasets["train"] = train_set

        if valid_csv:
            transform = valid_aug(image_size)
            if dataset_type == 'RSNADataset':
                valid_set = RSNADataset(
                    csv_file=valid_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform,
                    mode='valid',
                    image_type=image_type
                )
            elif dataset_type == 'RSNAMultiWindowsDataset':
                valid_set = RSNAMultiWindowsDataset(
                    csv_file=valid_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform
                )
            elif dataset_type == 'RSNADicomDataset':
                valid_set = RSNADicomDataset(
                    csv_file=valid_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform,
                    mode='valid'
                )
            elif dataset_type == "RSNARandomWindowDataset":
                valid_set = RSNARandomWindowDataset(
                    csv_file=valid_csv,
                    root=root,
                    with_any=with_any,
                    transform=transform,
                    mode='valid'
                )
            else:
                raise("No Dataset: {}".format(dataset_type))
            datasets["valid"] = valid_set

        return datasets