import json
from pathlib import Path
from data import Vocab
from train_target import train_target_wants_text
from model import ModelArgs, MusicalPositionEmbedTransformer
import dataclasses
import torch
from model_accelerated import MusicalPositionEmbedTransformerAccelerated
from util import add_layer_sequence, remove_module_prefix
from embedder import Embedder
from classifier_model import SimpleClassifier
import logging


def load_config(path):
    with open(path, "r") as f:
        return json.load(f)


def construct_model(config, unset_enable_flash=True):
    vocab = Vocab.from_config(config)

    args = ModelArgs.from_config(vocab, config)
    model = MusicalPositionEmbedTransformer(
        vocab,
        dataclasses.replace(
            args,
            cache=True,
            enable_flash=False if unset_enable_flash else args.enable_flash,
        ),
        (
            Embedder(load_pretrained_weights=False)
            if train_target_wants_text(config["train_target"])
            else None
        ),
    )

    return vocab, model


def model_state_dict_from_checkpoint(state_dict):
    return add_layer_sequence(remove_module_prefix(state_dict["model_state_dict"]))


def strip_checkpoint(checkpoint_path, output_path):
    state_dict = torch.load(
        checkpoint_path,
        map_location=torch.device("cpu"),
    )
    state_dict = model_state_dict_from_checkpoint(state_dict)

    torch.save(state_dict, output_path)


def strip_encoder(model_path, output_path):
    state_dict = torch.load(
        model_path,
        map_location=torch.device("cpu"),
    )
    encoder_state_dict = {
        ".".join(k.split(".")[1:]): v
        for k, v in state_dict.items()
        if k.startswith("encoder")
    }
    torch.save(encoder_state_dict, output_path)


def load_quantize_model(config, model_path, from_checkpoint=False):
    vocab, model = construct_model(config)

    state_dict = torch.load(
        model_path,
        map_location=torch.device("cpu"),
    )

    if from_checkpoint:
        state_dict = model_state_dict_from_checkpoint(state_dict)

    model.load_state_dict(state_dict)

    model.eval()

    torch.ao.quantization.quantize_dynamic(model, dtype=torch.qint8, inplace=True)

    model.transformer.jit()

    return vocab, model


def load_cuda_model(config, model_path, from_checkpoint=False):
    vocab, model = construct_model(config, unset_enable_flash=False)

    state_dict = torch.load(
        model_path,
        map_location=torch.device("cpu"),
    )

    if from_checkpoint:
        state_dict = model_state_dict_from_checkpoint(state_dict)

    model.load_state_dict(state_dict)

    model = model.cuda()
    model.eval()

    return vocab, model


def load_accelerated_model(config, path):
    vocab = Vocab.from_config(config)
    params = ModelArgs.from_config(vocab, config)

    encoder_path = Path.joinpath(path, "encoder.pt")
    if encoder_path.exists():
        encoder_state_dict = torch.load(encoder_path, map_location=torch.device("cpu"))
    else:
        logging.warn("No encoder checkpoint found, using model checkpoint")
        state_dict = torch.load(
            Path.joinpath(path, "model.pt"),
            map_location=torch.device("cpu"),
        )
        encoder_state_dict = {
            ".".join(k.split(".")[1:]): v
            for k, v in state_dict.items()
            if k.startswith("encoder")
        }
        del state_dict

    embedder = Embedder(load_pretrained_weights=False)
    embedder.eval()
    embedder.load_state_dict(encoder_state_dict)
    del encoder_state_dict
    embedder = embedder.cuda()

    return vocab, MusicalPositionEmbedTransformerAccelerated(
        params,
        Path.joinpath(path, "model.trt"),
        Path.joinpath(path, "model_one_step.trt"),
        embedder,
    )


def load_critic(model, critic_path):
    critic = SimpleClassifier(model.params.dim)
    critic.load_state_dict(torch.load(critic_path))
    return critic


def load_by_path(path):
    path = Path(path)
    config = load_config(Path.joinpath(path, "config.json"))

    accelerated = bool(config["inference"]["accelerated"])
    use_critic = bool(config["inference"]["critic"])

    if accelerated:
        vocab, model = load_accelerated_model(config, path)
    else:
        vocab, model = load_quantize_model(config, Path.joinpath(path, "model.pt"))

    if use_critic:
        critic = load_critic(model, Path.joinpath(path, "critic.pt"))
    else:
        critic = None

    return vocab, model, critic


if __name__ == "__main__":
    import sys
    from main import TrainState

    if sys.argv[1] == "strip":
        strip_checkpoint(sys.argv[2], sys.argv[3])
    elif sys.argv[1] == "stripencoder":
        strip_encoder(sys.argv[2], sys.argv[3])