"""

 Copyright (c) 2022, salesforce.com, inc.

 All rights reserved.

 SPDX-License-Identifier: BSD-3-Clause

 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause

"""

import argparse

import os

import random

import numpy as np

import pickle

import torch

import torch.backends.cudnn as cudnn

import wandb

from torch.utils.data import DataLoader

from torchinfo import summary

from tqdm import tqdm

import model.lavis.tasks as tasks

from model.lavis.common.config import Config

from model.lavis.common.dist_utils import get_rank

from model.lavis.common.logger import setup_logger

from local_config import WANDB_ENTITY

from model.lavis.common.registry import registry

from model.lavis.common.utils import now

# imports modules for registration

from model.lavis.common.optims import (

   LinearWarmupCosineLRScheduler,

   LinearWarmupStepLRScheduler,

)

from model.lavis.datasets.builders import *

from model.lavis.models import *

from model.lavis.processors import *

from model.lavis.runners import *

from model.lavis.tasks import *

from model.lavis.data.ReportDataset import MIMIC_CXR_Dataset

from local_config import PATH_TO_MIMIC_CXR

def parse_args():

    parser = argparse.ArgumentParser(description="Training")

    parser.add_argument("--cfg-path", required=True, help="path to configuration file.")

    parser.add_argument("--local_rank", type=int, default=0, help="local rank for distributed training.")

    parser.add_argument(

        "--options",

        nargs="+",

        help="override some settings in the used config, the key-value pair "

             "in xxx=yyy format will be merged into config file (deprecate), "

             "change to --cfg-options instead.",

    )

    args = parser.parse_args()

    # if 'LOCAL_RANK' not in os.environ:

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

    return args

def setup_seeds(config):

    seed = config.run_cfg.seed + get_rank()

    random.seed(seed)

    np.random.seed(seed)

    torch.manual_seed(seed)

    cudnn.benchmark = False

    cudnn.deterministic = True

def get_runner_class(cfg):

    """

    Get runner class from config. Default to epoch-based runner.

    """

    runner_cls = registry.get_runner_class(cfg.run_cfg.get("runner", "runner_base"))

    return runner_cls

def main():

    registry.mapping['paths']['cache_root'] = '.'

    cfg = Config(parse_args())

    job_id = now()

    # init_distributed_mode(cfg)

    setup_seeds(cfg)

    # set after init_distributed_mode() to only log on master.

    setup_logger()

    wandb_run = wandb.init(

        project=cfg.run_cfg.project_name,

        entity=WANDB_ENTITY,

        name=cfg.run_cfg.run_name

    )

    cfg.pretty_print()

    task = tasks.setup_task(cfg)

    # my report dataset

    datasets = {}

    datasets['mimic_cxr'] = {}

    datasets['mimic_cxr']['train'] = MIMIC_CXR_Dataset(vis_processor=None, text_processor=None, vis_root=f"{PATH_TO_MIMIC_CXR}/mimic-cxr-jpg/2.0.0",

                                                       split="train", cfg=cfg, truncate=None)

    datasets['mimic_cxr']['train_val'] = MIMIC_CXR_Dataset(vis_processor=None, text_processor=None,

                                                           vis_root=f"{PATH_TO_MIMIC_CXR}/mimic-cxr-jpg/2.0.0", split="train", cfg=cfg,

                                                           truncate=1000)  # 1000

    datasets['mimic_cxr']['val'] = MIMIC_CXR_Dataset(vis_processor=None, text_processor=None, vis_root=f"{PATH_TO_MIMIC_CXR}/mimic-cxr-jpg/2.0.0",

                                                     split="validate", cfg=cfg, truncate=None)

    datasets['mimic_cxr']['test'] = MIMIC_CXR_Dataset(vis_processor=None, text_processor=None, vis_root=f"{PATH_TO_MIMIC_CXR}/mimic-cxr-jpg/2.0.0",

                                                      split="test", cfg=cfg, truncate=None)

    model = task.build_model(cfg)

    print(summary(model, input_size=None, device='cpu'))

    if not cfg.run_cfg.evaluate:

        ''' training code '''

        runner = RunnerBase(

            cfg=cfg, job_id=job_id, task=task, model=model, datasets=datasets

        )

        runner.train(wandb_run)

    else:

        ''' precompute Q-Former output embeddings for all images '''

        model.cuda()

        model.eval()

        dataloader = DataLoader(datasets['mimic_cxr']['test'], batch_size=256, shuffle=False, num_workers=cfg.run_cfg.num_workers)

        embeddings = {}

        for i, batch in enumerate(tqdm(dataloader)):

            qformer_embs, _ = model.forward_image(batch['image'].cuda())

            for j, id in enumerate(batch['image_id']):

                dicom = datasets['mimic_cxr']['test'].id_to_dicom[id.item()]

                embeddings[dicom] = qformer_embs[j].cpu().detach().numpy()

        # save embeddings

        with open(f"pretraining/embs/{cfg.run_cfg.run_name}_embeddings_test.pkl", "wb") as f:

            pickle.dump(embeddings, f)

        dataloader = DataLoader(datasets['mimic_cxr']['val'], batch_size=256, shuffle=False, num_workers=cfg.run_cfg.num_workers)

        embeddings = {}

        for i, batch in enumerate(tqdm(dataloader)):

            qformer_embs, _ = model.forward_image(batch['image'].cuda())

            for j, id in enumerate(batch['image_id']):

                dicom = datasets['mimic_cxr']['val'].id_to_dicom[id.item()]

                embeddings[dicom] = qformer_embs[j].cpu().detach().numpy()

        # save embeddings

        with open(f"pretraining/embs/{cfg.run_cfg.run_name}_embeddings_val.pkl", "wb") as f:

            pickle.dump(embeddings, f)

        dataloader = DataLoader(datasets['mimic_cxr']['train'], batch_size=256, shuffle=False, num_workers=cfg.run_cfg.num_workers)

        embeddings = {}

        for i, batch in enumerate(tqdm(dataloader)):

            qformer_embs, _ = model.forward_image(batch['image'].cuda())

            for j, id in enumerate(batch['image_id']):

                dicom = datasets['mimic_cxr']['train'].id_to_dicom[id.item()]

                embeddings[dicom] = qformer_embs[j].cpu().detach().numpy()

        # save embeddings

        with open(f"pretraining/embs/{cfg.run_cfg.run_name}_embeddings_train.pkl", "wb") as f:

            pickle.dump(embeddings, f)

if __name__ == "__main__":

    main()