#!/usr/bin/env python

"""

Stochastic Weight Averaging (SWA)

Averaging Weights Leads to Wider Optima and Better Generalization

https://github.com/timgaripov/swa

"""

import torch

import models

from tqdm import tqdm

import glob

def moving_average(net1, net2, alpha=1.):

    for param1, param2 in zip(net1.parameters(), net2.parameters()):

        param1.data *= (1.0 - alpha)

        param1.data += param2.data * alpha

def _check_bn(module, flag):

    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):

        flag[0] = True

def check_bn(model):

    flag = [False]

    model.apply(lambda module: _check_bn(module, flag))

    return flag[0]

def reset_bn(module):

    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):

        module.running_mean = torch.zeros_like(module.running_mean)

        module.running_var = torch.ones_like(module.running_var)

def _get_momenta(module, momenta):

    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):

        momenta[module] = module.momentum

def _set_momenta(module, momenta):

    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):

        module.momentum = momenta[module]

def bn_update(loader, model):

    """

        BatchNorm buffers update (if any).

        Performs 1 epochs to estimate buffers average using train dataset.

        :param loader: train dataset loader for buffers average estimation.

        :param model: model being update

        :return: None

    """

    if not check_bn(model):

        return

    model.train()

    momenta = {}

    model.apply(reset_bn)

    model.apply(lambda module: _get_momenta(module, momenta))

    n = 0

    pbar = tqdm(loader, unit="images", unit_scale=loader.batch_size)

    for batch in pbar:

        input, targets = batch['images'], batch['targets']

        input = input.cuda()

        b = input.size(0)

        momentum = b / (n + b)

        for module in momenta.keys():

            module.momentum = momentum

        model(input)

        n += b

    model.apply(lambda module: _set_momenta(module, momenta))

if __name__ == '__main__':

    import argparse

    from pathlib import Path

    from torchvision.transforms import Compose

    from torch.utils.data import DataLoader

    from augmentation import valid_aug

    from dataset import SIIMDataset

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("--input", type=str, help='input directory')

    parser.add_argument("--output", type=str, default='swa_model.pth', help='output model file')

    parser.add_argument("--batch-size", type=int, default=16, help='batch size')

    args = parser.parse_args()

    # directory = Path(args.input)

    # files = [f for f in directory.iterdir() if f.suffix == ".pth"]

    files = glob.glob(args.input + "/stage1/checkpoints/stage1.*.pth")

    files += glob.glob(args.input + "/stage2/checkpoints/stage1.*.pth")

    assert(len(files) > 1)

    net = models.Unet(

        encoder_name="resnet34",

        activation='sigmoid',

        classes=1,

        # center=True

    )

    checkpoint = torch.load(files[0])

    net.load_state_dict(checkpoint['model_state_dict'])

    for i, f in enumerate(files[1:]):

        # net2 = model.load(f)

        net2 = models.Unet(

            encoder_name="resnet34",

            activation='sigmoid',

            classes=1,

            # center=True

        )

        checkpoint = torch.load(f)

        net2.load_state_dict(checkpoint['model_state_dict'])

        moving_average(net, net2, 1. / (i + 2))

    test_csv = './csv/train_0.csv'

    root = "/raid/data/kaggle/siim/siim256/"

    # img_size = 128

    batch_size = 16

    train_transform = valid_aug()

    train_dataset = SIIMDataset(

        csv_file=test_csv,

        root=root,

        transform=train_transform,

        mode='train'

    )

    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, drop_last=True)

    net.cuda()

    bn_update(train_dataloader, net)

    # models.save(net, args.output)

    torch.save({

        'model_state_dict': net.state_dict()

    }, args.output)