import os
import warnings
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
warnings.simplefilter("ignore")

import torch
from torch.utils.data import DataLoader
import torchvision
from torchvision import transforms
import lightning as L
from lightning.pytorch import seed_everything
import lightning.pytorch.callbacks as callbk
from lightning.pytorch.loggers import TensorBoardLogger
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import Loaders
import Models
import defs
import argparse


torch.set_float32_matmul_precision('high')
torch.backends.cudnn.deterministic = True

Mean = [0.485, 0.456, 0.406]
Std = [0.229, 0.224, 0.225]

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    seed_everything(seed, workers=True)
    torch.backends.cudnn.deterministic = True

if __name__ == "__main__":    
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--Log_Name", type=str, default="logs_debug", help="the name of the directory of the log chkp")
    parser.add_argument("-me", "--maxEpochs", type=int, default=None)
    parser.add_argument("-w", "--Workers", type=int, default=0)
    parser.add_argument("-b", "--BatchSize", type=int, default=32)
    parser.add_argument("-m", "--Model", type=str, required=True)
    
    args = parser.parse_args()
    
    setup_seed(2023)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    batchSize = args.BatchSize #17 #8, 32
    numWorkers = args.Workers
    
    ckpLossBest = callbk.ModelCheckpoint(monitor="val_loss", filename="RARP-{epoch}-{val_loss:.4f}", save_top_k=5, mode='min')
    
    original_crop = torch.nn.Sequential(
        transforms.CenterCrop((300, 420)),
        transforms.Resize((224, 224), antialias=True, interpolation=transforms.InterpolationMode.BICUBIC),
        transforms.Normalize(Mean, Std)
    ).to(device)
    
    TrainDINOTransforms = Loaders.RARP_MAE_Augmentation(
        Init_Resize=(300, 420),
        GloblaCropsScale = (0.4, 1),
        Size = 224, 
        device = device,
        mean = Mean,
        std = Std
    )
    
    trainDataset = torchvision.datasets.DatasetFolder(
        "./Dataset_video/",
        loader=defs.load_Img_RBG_tensor_norm,
        extensions="webp",
        transform=TrainDINOTransforms
    )
    
    valDataset = torchvision.datasets.DatasetFolder(
        "./Dataset_video_Val/",
        loader=defs.load_Img_RBG_tensor_norm,
        extensions="webp",
        transform=original_crop
    )
    
    Train_DataLoader = DataLoader(
        trainDataset, 
        batch_size=batchSize, 
        num_workers=numWorkers, 
        shuffle=True, 
        drop_last=True,
        pin_memory=True,
        persistent_workers=numWorkers>0,
        #prefetch_factor=1
    )
    
    Val_DataLoader = DataLoader(
        valDataset, 
        batch_size=batchSize, 
        num_workers=numWorkers, 
        shuffle=False, 
        pin_memory=True,
        persistent_workers=numWorkers>0
    )
    
    MaxEpochs = 150
        
    if args.maxEpochs is not None:
        MaxEpochs = args.maxEpochs
    
    Model = Models.RARP_MAE(backbone=args.Model, activation_fn=torch.nn.SiLU(True))
    
    print(f"Model Used: {type(Model).__name__}")
    LogFileName = f"{args.Log_Name}" 
    print("Train Phase")        
    trainer = L.Trainer(
        deterministic=True,
        accelerator='gpu', 
        devices=1,
        logger=TensorBoardLogger(save_dir=f"./{LogFileName}"),
        callbacks=[ckpLossBest, callbk.EarlyStopping(monitor="val_loss", mode="min", patience=5)],
        max_epochs=MaxEpochs,
    )
    
    trainer.fit(Model, train_dataloaders=Train_DataLoader, val_dataloaders=Val_DataLoader)
    