import pdb, os

import random

import argparse

import numpy as np

import torch

from torch.utils.data import DataLoader

from PathBLIP.dataset import Quiltdataset, ImageTextContrastiveCollator

from lavis.models import load_model

from trainer import Trainer

# set random seed

seed = 42

random.seed(seed)

np.random.seed(seed)

torch.manual_seed(seed)

torch.cuda.manual_seed(seed)

os.environ['PYTHONASHSEED'] = str(seed)

os.environ['TOKENIZERS_PARALLELISM']='false'

# set cuda devices

# os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3'

# device = "cuda:0,1,2,3" if torch.cuda.is_available() else "cpu"

train_config = {

    'num_epochs': 20,

    'warmup': 0.1,

    'lr': 2e-5,

    'weight_decay': 1e-4,

    'eval_batch_size': 8,

    'eval_steps': 1000,

    'save_steps': 1000,

}

train_dataset = Quiltdataset("../BLIP/LAVIS-main/quilt.csv")

train_collate_fn = ImageTextContrastiveCollator()

train_dataloader = DataLoader(train_dataset,

    batch_size=8,

    collate_fn=train_collate_fn,

    shuffle=True,

    pin_memory=True,

    num_workers=4,

    drop_last=True

    )

val_dataset = Quiltdataset("../test_samples.csv")

val_collate_fn = ImageTextContrastiveCollator()

val_dataloader = DataLoader(val_dataset,

    batch_size=4,

    collate_fn=val_collate_fn,

    shuffle=False,

    pin_memory=True,

    num_workers=4,

    )

# parser = argparse.ArgumentParser(description='training')

# parser.add_argument('--local_rank', type=int, help='local rank for dist')

# args = parser.parse_args()

torch.distributed.init_process_group(backend='nccl')

# print(args.local_rank)

# world_size = torch.cuda.device_count()

local_rank = torch.distributed.get_rank()

# print(local_rank)

torch.cuda.set_device(local_rank)

# parser.add_argument("--local-rank", type=int)

# args = parser.parse_args()

# if 'LOCAL_RANK' not in os.environ:

#     os.environ['LOCAL_RANK'] = str(args.local_rank)

model = load_model("blip2", "pretrain", checkpoint="../BLIP/blip2_pretrained.pth")

    # model.load_state_dict(torch.load('./checkpoints/vision_text_pretrain/t5/epoch10.pth',map_location='cpu'),strict=False)

model.cuda()

model_save_path = f'../BLIP/LAVIS-main/checkpoints/VL'

trainer = Trainer()

trainer.train(

    model,

    train_dataset,

    val_dataset,

    local_rank,

    warmup_ratio=train_config['warmup'],

    epochs=train_config['num_epochs'],

    optimizer_params={'lr':train_config['lr']},

    output_path=model_save_path,

    weight_decay=train_config['weight_decay'],

    use_amp=True,

    )