import argparse
import json
import os
import pickle

import numpy as np
import torch

from .config import FLAGS
from .torch_model import Generator


class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self


def load_checkpoint(filepath, device):
    assert os.path.isfile(filepath)
    print("Loading '{}'".format(filepath))
    checkpoint_dict = torch.load(filepath, map_location=device)
    print("Complete.")
    return checkpoint_dict


def convert_to_haiku(a, h, device):
    generator = Generator(h).to(device)
    state_dict_g = load_checkpoint(a.checkpoint_file, device)
    generator.load_state_dict(state_dict_g["generator"])
    generator.eval()
    generator.remove_weight_norm()
    hk_map = {}
    for a, b in generator.state_dict().items():
        print(a, b.shape)
        if a.startswith("conv_pre"):
            a = "generator/~/conv1_d"
        elif a.startswith("conv_post"):
            a = "generator/~/conv1_d_1"
        elif a.startswith("ups."):
            ii = a.split(".")[1]
            a = f"generator/~/ups_{ii}"
        elif a.startswith("resblocks."):
            _, x, y, z, _ = a.split(".")
            ver = h.resblock
            a = f"generator/~/res_block{ver}_{x}/~/{y}_{z}"
        print(a, b.shape)
        if a not in hk_map:
            hk_map[a] = {}
        if len(b.shape) == 1:
            hk_map[a]["b"] = b.numpy()
        else:
            if "ups" in a:
                hk_map[a]["w"] = np.rot90(b.numpy(), k=1, axes=(0, 2))
            elif "conv" in a:
                hk_map[a]["w"] = np.swapaxes(b.numpy(), 0, 2)
            else:
                hk_map[a]["w"] = b.numpy()

    FLAGS.ckpt_dir.mkdir(parents=True, exist_ok=True)
    with open(FLAGS.ckpt_dir / "hk_hifi.pickle", "wb") as f:
        pickle.dump(hk_map, f)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint-file", required=True)
    parser.add_argument("--config-file", required=True)
    a = parser.parse_args()

    config_file = a.config_file
    with open(config_file) as f:
        data = f.read()

    json_config = json.loads(data)
    h = AttrDict(json_config)

    device = torch.device("cpu")
    convert_to_haiku(a, h, device)


if __name__ == "__main__":
    main()