import argparse
import json
import os
import re
from collections import defaultdict

import numpy as np
import torch


def load_json(path: str):
    with open(path) as f:
        return json.load(f)


def parse_shape_info(flat_dir: str):
    data = load_json(os.path.join(flat_dir, "shape.json"))
    flat_info = defaultdict(lambda: defaultdict(list))
    for k, shape in data.items():
        matched = re.match(r"decoder.layers.\d+", k)
        if matched is None:
            flat_key = "flat_param_0"
        else:
            flat_key = f"{matched[0]}.flat_param_0"
        flat_info[flat_key]["names"].append(k)
        flat_info[flat_key]["shapes"].append(shape)
        flat_info[flat_key]["numels"].append(int(np.prod(shape)))
    return flat_info


def convert(flat_dir: str, output_dir: str, part: int):
    flat_path = os.path.join(flat_dir, f"reshard-model_part-{part}-shard0.pt")
    output_path = os.path.join(output_dir, f"reshard-model_part-{part}.pt")
    flat_meta = load_json(os.path.join(flat_dir, "flat-meta.json"))
    flat_sd = torch.load(flat_path)
    print(f"Loaded flat state dict from {flat_path}")
    output_sd = {}
    for flat_key, param_meta in flat_meta.items():
        flat_param = flat_sd["model"][flat_key]
        assert (
            sum(param_meta["numels"]) == flat_param.numel()
        ), f'flat {flat_key} {flat_param.numel()} vs {sum(param_meta["numels"])}'
        for name, shape, param in zip(
            param_meta["names"], param_meta["shapes"], flat_param.split(param_meta["numels"])
        ):
            output_sd[name] = param.view(shape)

    torch.save(output_sd, output_path)
    print(f"Saved unflat state dict to {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("flat_dir")
    parser.add_argument("output_dir")
    parser.add_argument("part", type=int)
    args = parser.parse_args()
    convert(args.flat_dir, args.output_dir, args.part)