from termcolor import colored
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.transforms import ColorJitter, ToPILImage, RandomGrayscale
import torch.optim.lr_scheduler as lr_scheduler
import os
import numpy as np


from utils_data.image_transforms import ArrayToTensor, PILToNumpy, ToTensor, RandomBlur
from training.actors.batch_processing import GLUNetBatchPreprocessing
from training.actors.self_supervised_actor import GLUNetBasedActor
from training.trainers.matching_trainer import MatchingTrainer
from utils_data.loaders import Loader
from admin.multigpu import MultiGPU
from models.GLUNet.GLU_Net import glunet_vgg16
from training.losses.basic_losses import EPE
from datasets.MegaDepth.megadepth import MegaDepthDataset
from datasets.synthetic_warp_dataset import WarpingDataset
from utils_data.geometric_transformation_sampling.synthetic_warps_sampling import SynthecticAffHomoTPSTransfo
from training.losses.multiscale_loss import MultiScaleFlow
from datasets.object_augmented_dataset.synthetic_object_augmentation_for_pairs_multiple_ob import RandomAffine
from datasets.object_augmented_dataset import MSCOCO, AugmentedImagePairsDatasetMultipleObjects


def run(settings):
    settings.description = 'Default train settings for DiMP with ResNet18 as backbone.'
    settings.data_mode = 'local'
    settings.batch_size = 16
    settings.n_threads = 8
    settings.multi_gpu = True
    settings.print_interval = 500
    settings.lr = 0.0001
    settings.scheduler_steps = [100, 120, 130]
    settings.n_epochs = 150

    # training parameters
    # loss applied in non-black target prime regions (to account for warping). It is very important, if applied to
    # black regions, weird behavior. If valid mask, doesn't learn interpolation in non-visibile regions.
    settings.compute_mask_zero_borders = True
    settings.apply_mask = False  # valid visible matches, we apply mask_zero_borders instead
    settings.nbr_plot_images = 4
    settings.dataset_callback_fn = 'sample_new_items'  # use to resample image pair at each epoch

    # synthetic transfo parameters
    settings.ori_dataset_img_size = 750  # size of the images, on which to apply the warp
    settings.crop_size = 520  # size of the crop, which is also the final training size
    settings.parametrize_with_gaussian = False
    settings.transformation_types = ['hom', 'tps', 'afftps']
    settings.random_t = 0.25
    settings.random_s = 0.45
    settings.random_alpha = np.pi / 12
    settings.random_t_tps_for_afftps = 0.15
    settings.random_t_hom = 0.5
    settings.random_t_tps = 0.5

    # object dataset
    settings.number_of_objects = 1
    settings.min_object_area = 1300

    # 1. Define training and validation datasets
    '''
    Training dataset: Here, we generate random flow field, that are used to warp individual megadepth images.
    After central cropping, it leads to a pair of images, related with a known ground-truth synthetic flow field.
    Here, we additionally introduce independently moving objects on top and modify the synthetic flow accordingly.
    '''

    # image dataset: single megadepth images, need to output images at resolution settings.ori_dataset_img_size.
    # the synthetic flow will be generated at this resolution as well. Then a second image will be created by warpin
    # the original image with the sampled flow. Then both images along with the flow will be centered crop, resulting
    # in a final image size equal to settings.crop_size
    megadepth_cfg = {'scene_info_path': os.path.join(settings.env.megadepth_training, 'scene_info'),
                     'train_num_per_scene': 300, 'val_num_per_scene': 25,  'pad_to_same_shape': False,
                     'two_views': False, 'output_image_size': [settings.ori_dataset_img_size,
                                                               settings.ori_dataset_img_size]}
    base_image_training_dataset = MegaDepthDataset(root=settings.env.megadepth_training, cfg=megadepth_cfg,
                                                   split='train')

    # synthetic warping parameters to generate the synthetic flow
    synthetic_flow_generator = SynthecticAffHomoTPSTransfo(
        size_output_flow=settings.ori_dataset_img_size, random_t=settings.random_t, random_s=settings.random_s,
        use_cuda=False,
        random_alpha=settings.random_alpha, random_t_tps_for_afftps=settings.random_t_tps_for_afftps,
        random_t_hom=settings.random_t_hom, random_t_tps=settings.random_t_tps,
        transformation_types=settings.transformation_types,
        parametrize_with_gaussian=settings.parametrize_with_gaussian)

    # actual training dataset is pairs of images, warped with the flow generated by the synthetic flow generator.
    # the final image pair is then cropped in the center.
    color_aug_transforms = transforms.Compose([ToPILImage(), ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6,
                                                                         hue=0.2), RandomGrayscale(p=0.2),
                                               PILToNumpy(),
                                               RandomBlur(sigma=(0.2, 2.0), kernel_size=(3, 7), probability=0.5)])
    training_dataset = WarpingDataset(original_image_dataset=base_image_training_dataset,
                                      synthetic_flow_generator=synthetic_flow_generator,
                                      compute_mask_zero_borders=settings.compute_mask_zero_borders,
                                      crop_size=settings.crop_size, output_size=settings.crop_size,
                                      source_image_transform=color_aug_transforms)

    # object dataset
    fg_tform = RandomAffine(p_flip=0.0, max_rotation=30.0,
                            max_shear=0, max_ar_factor=0.,
                            max_scale=0.3, pad_amount=0)
    coco_dataset_train = MSCOCO(root=settings.env.coco, split='train', version='2014', min_area=settings.min_object_area)

    # FINAL TRAINING DATASET: we then adds the object on the dataset
    source_img_transforms = transforms.Compose([ArrayToTensor(get_float=False)])
    target_img_transforms = transforms.Compose([ArrayToTensor(get_float=False)])
    flow_transform = transforms.Compose([ArrayToTensor()])
    co_transform = None
    training_dataset = AugmentedImagePairsDatasetMultipleObjects(
        foreground_image_dataset=coco_dataset_train, background_image_dataset=training_dataset,
        foreground_transform=fg_tform, source_image_transform=source_img_transforms,
        target_image_transform=target_img_transforms, flow_transform=flow_transform, co_transform=co_transform,
        compute_mask_zero_borders=settings.compute_mask_zero_borders, number_of_objects=settings.number_of_objects)

    # validation dataet
    base_image_val_dataset = MegaDepthDataset(root=settings.env.megadepth_training, cfg=megadepth_cfg, split='val')

    val_dataset = WarpingDataset(original_image_dataset=base_image_val_dataset,
                                 synthetic_flow_generator=synthetic_flow_generator,
                                 compute_mask_zero_borders=settings.compute_mask_zero_borders,
                                 crop_size=settings.crop_size, output_size=settings.crop_size,
                                 source_image_transform=color_aug_transforms)
    val_dataset = AugmentedImagePairsDatasetMultipleObjects(
        foreground_image_dataset=coco_dataset_train, background_image_dataset=val_dataset,
        foreground_transform=fg_tform, source_image_transform=source_img_transforms,
        target_image_transform=target_img_transforms, flow_transform=flow_transform, co_transform=co_transform,
        compute_mask_zero_borders=settings.compute_mask_zero_borders, number_of_objects=settings.number_of_objects)

    # 2. Define dataloaders
    train_loader = Loader('train', training_dataset, batch_size=settings.batch_size, shuffle=True,
                          drop_last=False, training=True, num_workers=settings.n_threads)

    val_loader = Loader('val', val_dataset, batch_size=settings.batch_size, shuffle=False,
                        epoch_interval=1.0, training=False, num_workers=settings.n_threads)

    # 3. Define model
    model = glunet_vgg16(global_corr_type='feature_corr_layer', normalize='relu_l2norm',
                         normalize_features=True, cyclic_consistency=True,
                         local_corr_type='feature_corr_layer', give_flow_to_refinement_module=False,
                         local_decoder_type='OpticalFlowEstimator',
                         global_decoder_type='CMDTop',
                         use_interp_instead_of_deconv=False)  # in original GLUNet, we set it to False
    # but better results are obtained with using simple bilinear interpolation instead of deconvolutions.
    print(colored('==> ', 'blue') + 'model created.')

    # Wrap the network for multi GPU training
    if settings.multi_gpu:
        model = MultiGPU(model)

    # 4. Define batch_processing
    batch_processing = GLUNetBatchPreprocessing(settings, apply_mask=settings.apply_mask,
                                                apply_mask_zero_borders=settings.compute_mask_zero_borders,
                                                sparse_ground_truth=False)

    # 5, Define loss module
    objective = EPE()
    weights_level_loss = [0.32, 0.08, 0.02, 0.01]
    loss_module_256 = MultiScaleFlow(level_weights=weights_level_loss[:2], loss_function=objective,
                                     downsample_gt_flow=True)
    loss_module = MultiScaleFlow(level_weights=weights_level_loss[2:], loss_function=objective, downsample_gt_flow=True)

    # 6. Define actor
    GLUNetActor = GLUNetBasedActor(model, objective=loss_module, objective_256=loss_module_256,
                                   batch_processing=batch_processing)

    # 7. Define Optimizer
    optimizer = \
        optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
                   lr=settings.lr,
                   weight_decay=0.0004)

    # 8. Define Scheduler
    scheduler = lr_scheduler.MultiStepLR(optimizer,
                                         milestones=settings.scheduler_steps,
                                         gamma=0.5)

    # 9. Define Trainer
    trainer = MatchingTrainer(GLUNetActor, [train_loader, val_loader], optimizer, settings, lr_scheduler=scheduler)

    trainer.train(settings.n_epochs, load_latest=True, fail_safe=True)







