import argparse
import json
import os
from pprint import pprint

import torch
import torch.distributed as dist
from guided_diffusion.script_util import (classifier_defaults,
                                          create_gaussian_diffusion,
                                          diffusion_defaults)

from tada.unet import UNetModel

NUM_CLASSES = 1000


def augment_defaults():
    """
    Defaults for data augmentation.
    """
    return dict(
        adaptive_augment=False, p=0.5, max_apply=2, augmentations="tada",
        translate_int_max=0.125,
        scale_std=0.2, rotate_frac_max=1.0, aniso_std=0.2, aniso_rotate_prob=0.5, translate_frac_std=0.125,
        brightness_std=0.2, contrast_std=0.5, hue_max=1.0, saturation_std=1.0,
        imgfilter_bands=[1, 1, 1, 1], imgfilter_std=1.0,
        cutout_size=0.5,
    )


def model_and_diffusion_defaults():
    """
    Defaults for image training.
    """
    res = dict(
        image_size=64,
        num_channels=128,
        num_res_blocks=2,
        num_heads=4,
        num_heads_upsample=-1,
        num_head_channels=-1,
        attention_resolutions="16,8",
        channel_mult="",
        dropout=0.0,
        augment_cond=False,
        class_cond=False,
        use_checkpoint=False,
        use_scale_shift_norm=True,
        resblock_updown=False,
        use_fp16=False,
        use_new_attention_order=False,
    )
    res.update(diffusion_defaults())
    return res


def create_model_and_diffusion(
    image_size,
    augment_cond,
    class_cond,
    learn_sigma,
    num_channels,
    num_res_blocks,
    channel_mult,
    num_heads,
    num_head_channels,
    num_heads_upsample,
    attention_resolutions,
    dropout,
    diffusion_steps,
    noise_schedule,
    timestep_respacing,
    use_kl,
    predict_xstart,
    rescale_timesteps,
    rescale_learned_sigmas,
    use_checkpoint,
    use_scale_shift_norm,
    resblock_updown,
    use_fp16,
    use_new_attention_order,
    augment_class=None,
):
    model = create_model(
        image_size,
        num_channels,
        num_res_blocks,
        channel_mult=channel_mult,
        learn_sigma=learn_sigma,
        augment_cond=augment_cond,
        class_cond=class_cond,
        use_checkpoint=use_checkpoint,
        attention_resolutions=attention_resolutions,
        num_heads=num_heads,
        num_head_channels=num_head_channels,
        num_heads_upsample=num_heads_upsample,
        use_scale_shift_norm=use_scale_shift_norm,
        dropout=dropout,
        resblock_updown=resblock_updown,
        use_fp16=use_fp16,
        use_new_attention_order=use_new_attention_order,
        augment_class=augment_class,
    )
    diffusion = create_gaussian_diffusion(
        steps=diffusion_steps,
        learn_sigma=learn_sigma,
        noise_schedule=noise_schedule,
        use_kl=use_kl,
        predict_xstart=predict_xstart,
        rescale_timesteps=rescale_timesteps,
        rescale_learned_sigmas=rescale_learned_sigmas,
        timestep_respacing=timestep_respacing,
    )
    return model, diffusion


def create_model(
    image_size,
    num_channels,
    num_res_blocks,
    channel_mult="",
    learn_sigma=False,
    augment_cond=False,
    class_cond=False,
    use_checkpoint=False,
    attention_resolutions="16",
    num_heads=1,
    num_head_channels=-1,
    num_heads_upsample=-1,
    use_scale_shift_norm=False,
    dropout=0,
    resblock_updown=False,
    use_fp16=False,
    use_new_attention_order=False,
    augment_class=None,
):
    if channel_mult == "":
        if image_size == 512:
            channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
        elif image_size == 256:
            channel_mult = (1, 1, 2, 2, 4, 4)
        elif image_size == 128:
            channel_mult = (1, 1, 2, 3, 4)
        elif image_size == 64:
            channel_mult = (1, 2, 3, 4)
        elif image_size == 32:
            channel_mult = (2, 2, 2)
        else:
            raise ValueError(f"unsupported image size: {image_size}")
    else:
        channel_mult = tuple(float(ch_mult) for ch_mult in channel_mult.split(","))

    attention_ds = []
    for res in attention_resolutions.split(","):
        try:
            attention_ds.append(image_size // int(res))
        except ValueError:
            continue

    if augment_class in ("none", "tada-v2", None):
        dim_augment_labels = None
    elif augment_class == "edm":
        dim_augment_labels = 9
    elif augment_class in ("augreg", "aug-reg", "aug_reg"):
        dim_augment_labels = 19
    else:
        raise ValueError(f"Unknown augment_class: {augment_class}")
    return UNetModel(
        image_size=image_size,
        in_channels=3,
        model_channels=num_channels,
        out_channels=(3 if not learn_sigma else 6),
        num_res_blocks=num_res_blocks,
        attention_resolutions=tuple(attention_ds),
        dropout=dropout,
        channel_mult=channel_mult,
        dim_augment_labels=(dim_augment_labels if augment_cond else None),
        num_classes=(NUM_CLASSES if class_cond else None),
        use_checkpoint=use_checkpoint,
        use_fp16=use_fp16,
        num_heads=num_heads,
        num_head_channels=num_head_channels,
        num_heads_upsample=num_heads_upsample,
        use_scale_shift_norm=use_scale_shift_norm,
        resblock_updown=resblock_updown,
        use_new_attention_order=use_new_attention_order,
    )


def set_seed(seed, max_seed=2 ** 32 - 1):
    """
    Set random seed for torch, numpy, and python.
    """
    if seed is None:
        seed = torch.randint(0, max_seed, (1,), device="cuda")
    elif isinstance(seed, (int, str)):
        seed = torch.as_tensor([int(seed)], device="cuda")
    dist.broadcast(seed, src=0)
    seed = int(seed)

    if dist.is_initialized():
        shift = dist.get_rank()
    else:
        shift = 0
    torch.manual_seed(seed+shift)  # For gaussian noise

    import numpy as np
    np.random.seed(seed)

    import random
    random.seed(seed)

    os.environ['PYTHONHASHSEED'] = str(seed)
    return seed


def find_config(model_path):
    dirname = os.path.dirname(model_path)
    return os.path.join(dirname, "config.json")


def save_config(args, config_path, verbose=True):
    if dist.get_rank() == 0:
        model_diffusion_keys = model_and_diffusion_defaults().keys()
        diffusion_keys = diffusion_defaults().keys()
        classifier_keys = classifier_defaults().keys()
        augment_keys = augment_defaults().keys()

        config = {
            "model": {},
            "diffusion": {},
            "classifier": {},
            "augment": {},
        }

        for k, v in vars(args).items():
            if k in model_diffusion_keys:
                if k in diffusion_keys:
                    config["diffusion"][k] = v
                else:
                    config["model"][k] = v
            elif k in classifier_keys:
                config["classifier"][k] = v
            elif k in augment_keys:
                config["augment"][k] = v
            else:
                config[k] = v
        if verbose:
            pprint(config, sort_dicts=False)
        with open(config_path, "w") as f:
            json.dump(config, f, indent=4, sort_keys=False)
    dist.barrier()


def update_config(config, args, unknown_args):
    def _update(cfg, key, val):
        for k, v in cfg.items():
            if isinstance(v, dict):
                _update(v, key, val)
            elif k == key:
                cfg[k] = val
                break

    # Update config with args
    for k, v in vars(args).items():
        _update(config, k, v)

    # Update config with unknown args
    while len(unknown_args) > 0:
        key = unknown_args.pop(0)
        if key.startswith("--"):
            key = key[2:]
            val = unknown_args.pop(0)
        elif "=" in key:
            key, val = key.split("=")
        else:
            raise argparse.ArgumentError(f"Unknown argument: {key}")
        _update(config, key, val)
    return config


def save_env_vars(env_path, verbose=True):
    env_vars = {
        "CALIBRATE_SNR": "1",
        "SNR_BOUNDARY": "-3,0",
        "ADAPT_METHOD": "quad",
        "KAPPA": "",
        "AUG_COND_TYPE": "edm",  # "edm", "vdm"
    }
    for env in env_vars:
        if env not in os.environ:
            os.environ[env] = env_vars[env]
        env_vars[env] = os.environ.get(env, env_vars[env])

    if dist.get_rank() == 0:
        if verbose:
            pprint(env_vars, sort_dicts=False)
        with open(env_path, "w") as f:
            json.dump(env_vars, f, indent=4, sort_keys=False)
    dist.barrier()


def find_env_vars(model_path):
    dirname = os.path.dirname(model_path)
    return os.path.join(dirname, "env.json")


def int2str(v, unit=""):
    assert isinstance(v, int)
    assert unit in ("", "k", "m", "b")

    if unit == "k":
        return f"{v // 1000}k"
    elif unit == "m":
        return f"{v // 1000000}m"
    elif unit == "b":
        return f"{v // 1000000000}b"
    else:
        return str(v)
