"""
Complete example to train a model
The task is to segment blood vessels in fundus images.
"""
from torch.utils.data import DataLoader
import dataclasses as dc
import dataclasses as dc
import simple_parsing as sp
import torch as T
import torch.optim as optim
import torchvision.transforms as tvt

from simplepytorch.datasets import DRIVE
from simplepytorch import trainlib as TL


def dataset_standardize_shape(img_labels_mask):
    img, labels, mask = img_labels_mask
    _rs = T.nn.functional.interpolate
    rv = (
        _rs(img.unsqueeze(0), (512, 512), mode='bilinear',
            align_corners=False).squeeze(0),
        _rs(labels.unsqueeze(0).unsqueeze(0), (512, 512), mode='nearest',
            ).squeeze(0).squeeze(0),
        _rs(mask.unsqueeze(0).unsqueeze(0), (512, 512), mode='nearest'
            ).squeeze(0).squeeze(0))
    return rv


def dataset_transforms(label_name='1st_manual'):
    assert label_name in {'1st_manual', '2nd_manual'}
    return tvt.Compose([
        DRIVE.as_tensor([label_name, 'mask'], return_numpy_array=True),
        DRIVE.crop_background,
        #  ... in future we could experiment with additional preprocessing steps
        lambda x: x/255,
        DRIVE.to_tensors(device='cpu'),
        dataset_standardize_shape
        # resize to (512,512) probly not necessary, but could be done using T.nn.functional.interpolate.with mode='nearest' for labels and 'bilinear' for image
    ])

class MyCELoss(T.nn.BCELoss):
    def forward(self, yhat, y, region_of_interest_mask):
        # the chosen dataset offers extra data that we will ignore for example...
        return super().forward(yhat, y.unsqueeze_(1).float())  # blindly ignore region of interest


@dc.dataclass
class CmdLineOptions:
    experiment_id:str
    epochs:int = 3
    device:str = 'cuda:0' if T.cuda.is_available() else "cpu"



def get_train_config(args:CmdLineOptions):
    DSET_TRAIN = DRIVE(use_train_set=True, getitem_transform=dataset_transforms('1st_manual'))
    DSET_VAL = DRIVE(use_train_set=False, getitem_transform=dataset_transforms('1st_manual'))
    MODEL = T.hub.load(
        'mateuszbuda/brain-segmentation-pytorch', 'unet',
        in_channels=3, out_channels=1, init_features=4, pretrained=False).to(args.device)

    CFG = TL.TrainConfig(
        experiment_id=args.experiment_id,
        model=MODEL,
        optimizer=optim.Adam(MODEL.parameters(), lr=0.005, betas=(.85, .9)),
        train_dset=DSET_TRAIN,
        val_dset=DSET_VAL,
        test_dset=None,
        train_loader=DataLoader(DSET_TRAIN, batch_size=5, num_workers=4, shuffle=True),
        val_loader=DataLoader(DSET_VAL, batch_size=5, num_workers=4, shuffle=True),
        test_loader=None,
        loss_fn=MyCELoss(),
        device=args.device,
        epochs=args.epochs,
        result_factory = lambda: TL.SegmentationResult(
            model_final_layer=T.nn.Softmax2d(), metrics=('dice', 'loss', 'confusion_matrix'))
    )
    return CFG


def plot_before_after(prediction_before_training, prediction_after_training, label, epochs):
    # visualize a segmentation
    with T.no_grad():
        from matplotlib import pyplot as plt
        fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
        ax1.imshow(prediction_before_training.squeeze().cpu().numpy())
        ax1.set_title('Before Training')
        ax2.imshow(prediction_after_training.squeeze().cpu().numpy())
        ax2.set_title(f'After Training {epochs} epochs')
        ax3.imshow(label)
        ax3.set_title('Ground truth')
        plt.show(block=True)


def main():
    # get commandline args
    par = sp.ArgumentParser()
    par.add_arguments(CmdLineOptions, dest='args')
    args = par.parse_args().args

    # load a training configuration
    cfg = get_train_config(args)

    # train model and visualize before/after results
    prediction_before_training = cfg.model(
        cfg.val_dset[0][0].unsqueeze(0).to(cfg.device))
    #
    cfg.train(cfg)  # stores results to file
    #
    prediction_after_training = cfg.model(
        cfg.val_dset[0][0].unsqueeze(0).to(cfg.device))
    label = cfg.val_dset[0][1].squeeze().cpu().numpy()

    plot_before_after(prediction_before_training, prediction_after_training, label, epochs=cfg.epochs)


if __name__ == "__main__":
    # Can run like this:
    # python examples/simple_example.py  --experiment_id test_experiment

    main()
