from pathlib import Path

import torch
import yaml


class NoAliasSafeDumper(yaml.SafeDumper):
    # Disable anchor/alias in yaml because looks ugly
    def ignore_aliases(self, data):
        return True


def yaml_no_alias_safe_dump(data, stream=None, **kwargs):
    """Safe-dump in yaml with no anchor/alias"""
    return yaml.dump(
        data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs
    )


def gen_conf(file, out_dir):
    conf = torch.load(file)["config"]
    conf["oss_bucket"] = "null"
    print(conf)
    output_dir = Path(out_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    with (output_dir / "config.yaml").open("w", encoding="utf-8") as f:
        yaml_no_alias_safe_dump(conf, f, indent=4, sort_keys=False)


if __name__ == "__main__":
    import sys

    in_f = sys.argv[1]
    out_f = sys.argv[2]
    gen_conf(in_f, out_f)