from transformers import VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel
import torch
from torch.utils.data import DataLoader
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint
import pytorch_lightning as pl
import os
from huggingface_hub import login
import argparse
from sconf import Config
from utils.checkpoint import CustomCheckpointIO
from utils.donut_dataset_stream import DonutDataset
from utils.donut_model_pl import DonutModelPLModule
from utils.callbacks import PushToHubCallback
import warnings
from datasets import load_dataset




def main(config, hug_token):

    config_vision = VisionEncoderDecoderConfig.from_pretrained(
        config.pretrained_model_path)
    config_vision.encoder.image_size = config.image_size
    config_vision.decoder.max_length = config.max_length

    processor = DonutProcessor.from_pretrained(config.start_model_path)
    model = VisionEncoderDecoderModel.from_pretrained(
        config.pretrained_model_path, config=config_vision)

    processor.image_processor.size = config.image_size[::-1]
    processor.image_processor.do_align_long_axis = False

    added_tokens = []

    dataset = load_dataset(config.dataset_path)
    dataset.train_test_split(test_size=0.1)
    print(dataset)

    # train_dataset = DonutDataset(
    #                     dataset, 
    #                     processor=processor, 
    #                     model=model, 
    #                     max_length=config.max_length,
    #                     split="train", 
    #                     task_start_token="<s_cord-v2>", 
    #                     prompt_end_token="<s_cord-v2>",
    #                     added_tokens=added_tokens,
    #                     sort_json_key=False,  # cord dataset is preprocessed, so no need for this
    #                 )

    # val_dataset = DonutDataset(
    #                     dataset, 
    #                     processor=processor, 
    #                     model=model, 
    #                     max_length=config.max_length,
    #                     split="validation", 
    #                     task_start_token="<s_cord-v2>", 
    #                     prompt_end_token="<s_cord-v2>",
    #                     added_tokens=added_tokens,
    #                     sort_json_key=False,  # cord dataset is preprocessed, so no need for this
    #                 )

    # model.config.pad_token_id = processor.tokenizer.pad_token_id
    # model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]

    # train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=1)
    # val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1)

    # login(hug_token, True)

    # model_module = DonutModelPLModule(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader)
    
    # wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name)

    # checkpoint_callback = ModelCheckpoint(
    #     monitor="val_metric",
    #     dirpath=config.checkpoint_path,
    #     filename="artifacts",
    #     save_top_k=1,
    #     save_last=False,
    #     mode="min",
    # )

    # custom_ckpt = CustomCheckpointIO()

    # trainer = pl.Trainer(
    #     accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
    #     devices=1,
    #     max_epochs=config.train_config.max_epochs,
    #     val_check_interval=config.train_config.val_check_interval,
    #     check_val_every_n_epoch=config.train_config.check_val_every_n_epoch,
    #     gradient_clip_val=config.train_config.gradient_clip_val,
    #     precision=16, # we'll use mixed precision
    #     plugins=custom_ckpt,
    #     num_sanity_val_steps=0,
    #     logger=wandb_logger,
    #     callbacks=[PushToHubCallback(output_model_path=config.output_model_path), checkpoint_callback],
    # )

    # trainer.fit(model_module)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, required=True)
    args, left_argv = parser.parse_known_args()
    config = Config(args.config)
    config.argv_update(left_argv)

    hug_token = os.environ.get("HUG_TOKEN", None)

    if not torch.cuda.is_available():
        warnings.warn("You don't have cuda available, training might be taking long time or impossible")

    if not hug_token:
        raise Exception("You need to set up HUG_TOKEN in enviroments to push output model to hub")
    main(config, hug_token)