from typing import Any, List
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 DonutDatasetStream
from utils.donut_model_pl_stream import DonutModelPLModuleStream
from utils.callbacks import PushToHubCallback
import warnings
from datasets import load_dataset, interleave_datasets
from torchdata.datapipes.iter import IterableWrapper
import json


class CustomWrapperIterator(IterableWrapper):
    def __init__(self, iterable, deepcopy=True, total_len=None):
        super().__init__(iterable, deepcopy)
        self.total_len = total_len

    def __len__(self):
        if self.total_len:
            return self.total_len
        return super().__len__()

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 = []

    ### PROCESS FUNC START ###

    def add_tokens(list_of_tokens: List[str]):
        """
        Add special tokens to tokenizer and resize the token embeddings of the decoder
        """
        newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)
        if newly_added_num > 0:
            model.decoder.resize_token_embeddings(len(processor.tokenizer))
            added_tokens.extend(list_of_tokens)

    def json2token(obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
        """
        Convert an ordered JSON object into a token sequence
        """
        if type(obj) == dict:
            if len(obj) == 1 and "text_sequence" in obj:
                return obj["text_sequence"]
            else:
                output = ""
                if sort_json_key:
                    keys = sorted(obj.keys(), reverse=True)
                else:
                    keys = obj.keys()
                for k in keys:
                    if update_special_tokens_for_json_key:
                        add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
                    output += (
                        fr"<s_{k}>"
                        + json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
                        + fr"</s_{k}>"
                    )
                return output
        elif type(obj) == list:
            return r"<sep/>".join(
                [json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
            )
        else:
            obj = str(obj)
            if f"<{obj}/>" in added_tokens:
                obj = f"<{obj}/>"  # for categorical special tokens
            return obj

    def process(row, split):
        task_start_token, prompt_end_token = "<s_cord-v2>", "<s_cord-v2>"
        ground_truth = json.loads(row["ground_truth"])
        if "gt_parses" in ground_truth:  # when multiple ground truths are available, e.g., docvqa
            assert isinstance(ground_truth["gt_parses"], list)
            gt_jsons = ground_truth["gt_parses"]
        else:
            assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
            gt_jsons = [ground_truth["gt_parse"]]

        gt_token_sequences = (
            [
                json2token(
                    gt_json,
                    update_special_tokens_for_json_key=split == "train",
                    sort_json_key=False,
                )
                + processor.tokenizer.eos_token
                for gt_json in gt_jsons  # load json from list of json
            ]
        )

        add_tokens([task_start_token, prompt_end_token])
        prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(prompt_end_token)

        # change if not 3 channels
        if row['image'].mode != "RGB":
            row['image'] = row['image'].convert("RGB")

        # inputs
        pixel_values = processor(row["image"], random_padding=split == "train", return_tensors="pt").pixel_values
        pixel_values = pixel_values.squeeze()

        # targets
        input_ids = processor.tokenizer(
            gt_token_sequences,
            add_special_tokens=False,
            max_length=config.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )["input_ids"].squeeze(0)

        labels = input_ids.clone()
        labels[labels == processor.tokenizer.pad_token_id] = -100  # model doesn't need to predict pad token
        return {"pixel_values": pixel_values, "labels": labels, 'target_sequence': gt_token_sequences }

    def proces_train(row):
        return process(row, 'train')
    
    def proces_val(row):
        return process(row, 'validation')




    ### PROCESS FUNC END ###

    # train_dataset_process = DonutDatasetStream(
    #                     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_process = DonutDatasetStream(
    #                     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
    #                 )
    
    dataset = load_dataset(config.dataset_path, streaming=True)
    val_dataset = dataset.pop('validation') 
    train_dataset = interleave_datasets(list(dataset.values()))
    train_length = sum(split.num_examples for split in dataset[list(dataset.keys())[0]].info.splits.values() if split.name != 'validation')
    val_length = list(val_dataset.info.splits.values())[-1].num_examples


    train_dataset = train_dataset.map(proces_train, remove_columns = ['image', 'ground_truth'])
    val_dataset = val_dataset.map(proces_val, remove_columns = ['image', 'ground_truth'])

    # train_dataset = train_dataset.with_format('torch')
    # val_dataset = val_dataset.with_format('torch')

    train_dataset = CustomWrapperIterator(train_dataset, total_len=train_length)
    val_dataset = CustomWrapperIterator(val_dataset, total_len=val_length)

    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, num_workers=0)
    val_dataloader = DataLoader(val_dataset, batch_size=1, num_workers=0)

    login(hug_token, True)

    model_module = DonutModelPLModuleStream(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)

    if not os.path.exists(config.checkpoint_path):
        os.mkdir(config.checkpoint_path)

    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", "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")

    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)