import sys

import torch
from torchvision import transforms as T
from pytorch_lightning.callbacks import ModelCheckpoint
import pytorch_lightning as pl
from argparse import ArgumentParser

import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage

from torch.utils.data import Dataset, DataLoader
from datasets.cityscapes import CityScapes
from datasets.noaaice import NOAAICE
from models.BiSeNetV2.bisenetv2 import BiSeNetV2
from models.losses import BiSeNetV2Loss
from configs.default import cfg

MODELS = {'bisenetv2' : BiSeNetV2 }

#torch.backends.cudnn.enabled = False

# TODO move to it's own file
from bisect import bisect_right
import torch
import random
import math

class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):
    def __init__(
        self,
        optimizer,
        milestones,
        gamma=0.1,
        warmup_factor=1.0 / 3,
        warmup_iters=500,
        warmup_method="linear",
        last_epoch=-1,
    ):  
        if not list(milestones) == sorted(milestones):
            raise ValueError(
                "Milestones should be a list of" " increasing integers. Got {}",
                milestones,
            )

        if warmup_method not in ("constant", "linear"):
            raise ValueError(
                "Only 'constant' or 'linear' warmup_method accepted"
                "got {}".format(warmup_method)
            )
        self.milestones = milestones
        self.gamma = gamma
        self.warmup_factor = warmup_factor
        self.warmup_iters = warmup_iters
        self.warmup_method = warmup_method
        super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        warmup_factor = 1 
        if self.last_epoch < self.warmup_iters:
            if self.warmup_method == "constant":
                warmup_factor = self.warmup_factor
            elif self.warmup_method == "linear":
                alpha = self.last_epoch / self.warmup_iters
                warmup_factor = self.warmup_factor * (1 - alpha) + alpha
        return [
            base_lr
            * warmup_factor
            * self.gamma ** bisect_right(self.milestones, self.last_epoch)
            for base_lr in self.base_lrs
        ]



class LightningModel(pl.LightningModule):

    def __init__(self, cfg):
        super(LightningModel, self).__init__()
        self.model = MODELS[cfg.MODEL.MODEL_NAME](cfg)
        self.cfg = cfg
        self.loss = BiSeNetV2Loss()

 
    def forward(self, x):
        return self.model(x)


    def prepare_data(self):
        #TODO put code here that needs to prepare/verify dataset
        return


    def train_dataloader(self):

        print('Loading dataloader')

        '''
        transform=T.Compose([
            T.RandomResizedCrop([0.375, 1.], [512, 1024]),
            T.RandomHorizontalFlip(),
            T.ColorJitter(
                brightness=0.4,
                contrast=0.4,
                saturation=0.4
            ),
        ])
        '''

        # Define our augmentation pipeline.
        seq = iaa.Sequential([
                #iaa.Resize((self.cfg.TRAIN.INPUT_HEIGHT, self.cfg.TRAIN.INPUT_WIDTH)), 
                #iaa.Crop(percent=(0.0, 0.625)),
                iaa.Fliplr(0.5),
                #iaa.Sharpen((0.0, 1.0)),       # sharpen the image
                iaa.Affine(rotate=(-90, 90)),  # rotate by -45 to 45 degrees (affects segmaps)
                #iaa.Dropout([0.05, 0.2]),      # drop 5% or 20% of all pixels
                #iaa.ElasticTransformation(alpha=50, sigma=5)  # apply water effect (affects segmaps)
            ]) #, random_order=True)

        # TODO get dataset from catalogue or something
        #dataset = CityScapes(self.cfg.DATASET.PATH, trans_func=seq)
        dataset = NOAAICE(self.cfg.DATASET.PATH, trans_func=seq)

        train_dataloader = DataLoader(dataset, batch_size=self.cfg.TRAIN.BATCH_SIZE,
                                        num_workers=cfg.SYSTEM.NUM_WORKERS)

        print('Dataloader loaded')
        return train_dataloader

    '''
    def val_dataloader(self):
        return None

    def test_dataloader(self):
        return None
    '''


    def training_step(self, batch, batch_idx):
        x, y = batch

        res = self.forward(x)

        loss = self.loss(res, y)

        return {'loss' : loss}
       

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.02)
        #return torch.optim.SGD(self.parameters(), lr=5e-2, momentum=0.9)
        scheduler = WarmupMultiStepLR(optimizer, [40, 70], warmup_iters=10, warmup_factor=0.01)
        return [optimizer], [scheduler]

def main(args):
    cfg.merge_from_file(args.config)
    model = LightningModel(cfg)

    # DEFAULTS used by the Trainer
    checkpoint_callback = ModelCheckpoint(
        filepath=cfg.MODEL.TRAIN_PATH + '/{epoch}-{loss:.2f}',
        save_top_k=-1,  # Save all models at rate of 'period'
        period=10,      # Save every 10 epochs
        verbose=True,
        monitor='loss',
        mode='min',
        save_weights_only=True,
        prefix=cfg.MODEL.TRAIN_PREFIX
    )   

    trainer = pl.Trainer(checkpoint_callback=checkpoint_callback, gpus=1, max_epochs=120)
    #trainer = pl.Trainer(gpus=1)
    trainer.fit(model)

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--config', default=None, help='config yaml file to use (updates default)')
    args = parser.parse_args()

    main(args)

