import pickle
import random
from typing import List, Tuple

from accelerate import init_empty_weights
import torch

from model.lumina_bert import LuminaBertConfig, LuminaBertForImageGeneration
from xllmx.data.data_reader import read_general
from xllmx.data.item_processor import ItemProcessorBase
from xllmx.solvers.finetune import FinetuneSolverBase


class ItemProcessor(ItemProcessorBase):
    def process_item(self, data_item: dict, training_mode=False) -> Tuple[List, List]:
        assert training_mode

        assert "caps" in data_item and len(data_item["caps"]) >= 1
        caption = random.choice(list(data_item["caps"].values()))

        if "s3://" in data_item["file"]:
            pre_tokenized = pickle.load(read_general(data_item["file"]))
        else:
            with open(data_item["file"], "rb") as f:
                pre_tokenized = pickle.load(f)

        image_media = {"type": "<|image|>", "tokens": self.convert_chameleon_tokens(pre_tokenized["token"])}

        text_media = {"type": "<|text|>", "text": caption}

        return [text_media, image_media], [-1, -1]

    def convert_chameleon_tokens(self, tokens):

        image_start_ori = 8197  # "<racm3:break>"
        image_end_ori = 8196  # "<eoss>"
        end_of_line_ori = 8803  # "<reserved08799>"

        image_start = 0
        image_end = 1
        end_of_line = 2

        cache = None
        for i in range(len(tokens)):
            token_id = tokens[i]
            if token_id == image_start_ori:
                cache = [image_start_ori]
                for j in range(i + 1, len(tokens)):
                    cache.append(tokens[j])
                    if tokens[j] == image_end_ori:
                        break
                break

        assert cache is not None

        for i, token in enumerate(cache):
            if token == image_start_ori:
                assert i == 0
                cache[i] = image_start
            elif token == image_end_ori:
                assert i == len(cache) - 1
                cache[i] = image_end
            elif token == end_of_line_ori:
                cache[i] = end_of_line
            elif 8804 <= token < 9804:
                assert i in [1, 2]
                cache[i] = token - 8004 + 1000
            elif 4 <= token < 8196:
                cache[i] = token - 4 + 2000
            else:
                raise ValueError()

        return cache

    def predict_item_token_length(self, data_item: dict) -> int:
        if "len" in data_item:
            return data_item["len"]
        else:
            raise 1


class Solver(FinetuneSolverBase):
    @classmethod
    def get_args_parser(cls):
        parser = super().get_args_parser()
        # task-specific parameters
        parser.add_argument("--max_seq_len", default=4096, type=int, help="max token length")
        parser.add_argument("--model_size", type=str, default="7B", choices=["2B", "7B", "34B"])
        parser.add_argument("--z_loss_weight", type=float, default=1e-5)
        return parser

    def _model_func(
        self,
        init_from: str,
    ) -> (LuminaBertForImageGeneration, None):

        # Only instantiate the model on rank0
        # Other ranks will receive the model weights from rank0 during FSDP wrapping (through `sync_module_states`)
        # See https://github.com/pytorch/pytorch/issues/105840
        if self.global_rank == 0:
            model = LuminaBertForImageGeneration.from_pretrained(
                init_from,
                max_position_embeddings=self.args.max_seq_len,
                z_loss_weight=self.args.z_loss_weight,
                torch_dtype=torch.bfloat16,
                device_map="cpu",
                attn_implementation="flash_attention_2",
            )
        else:
            with init_empty_weights():
                config = LuminaBertConfig.from_pretrained(
                    init_from,
                    max_position_embeddings=self.args.max_seq_len,
                    z_loss_weight=self.args.z_loss_weight,
                    torch_dtype=torch.bfloat16,
                    attn_implementation="flash_attention_2",
                )
                model = LuminaBertForImageGeneration(config)

        return model, None

    def _item_processor_func(self) -> ItemProcessorBase:
        return ItemProcessor()

    def _make_and_save_starting_point(self, save_path: str) -> None:

        size_specific_kwargs = {
            "7B": {
                "hidden_size": 4096,
                "num_attention_heads": 32,
                "num_hidden_layers": 32,
            },
            "2B": {
                "hidden_size": 2304,
                "num_attention_heads": 16,
                "num_key_value_heads": 8,
                "num_hidden_layers": 24,
            },
        }[self.args.model_size]

        config = LuminaBertConfig(
            vocab_size=8192 + 2000,
            max_position_embeddings=self.args.max_seq_len,
            torch_dtype=torch.bfloat16,
            device_map="cpu",
            attn_implementation="flash_attention_2",
            **size_specific_kwargs
        )

        model = LuminaBertForImageGeneration(config)

        model.save_pretrained(save_path, max_shard_size="10GB")


if __name__ == "__main__":
    args = Solver.get_args_parser().parse_args()
    solver = Solver(args)
    solver.run()
