import argparse
import datetime
import logging
import sys
import shlex
import cmd
from pathlib import Path
import torch
import json
from PIL import Image
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer
from slugify import slugify

from schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion


torch.backends.cuda.matmul.allow_tf32 = True


default_args = {
    "model": None,
    "scheduler": "euler_a",
    "precision": "fp32",
    "ti_embeddings_dir": "embeddings_ti",
    "output_dir": "output/inference",
    "config": None,
}


default_cmds = {
    "prompt": None,
    "negative_prompt": None,
    "image": None,
    "image_noise": .7,
    "width": 512,
    "height": 512,
    "batch_size": 1,
    "batch_num": 1,
    "steps": 50,
    "guidance_scale": 7.0,
    "seed": None,
    "config": None,
}


def merge_dicts(d1, *args):
    d1 = d1.copy()

    for d in args:
        d1.update({k: v for (k, v) in d.items() if v is not None})

    return d1


def create_args_parser():
    parser = argparse.ArgumentParser(
        description="Simple example of a training script."
    )
    parser.add_argument(
        "--model",
        type=str,
    )
    parser.add_argument(
        "--scheduler",
        type=str,
        choices=["plms", "ddim", "klms", "euler_a"],
    )
    parser.add_argument(
        "--precision",
        type=str,
        choices=["fp32", "fp16", "bf16"],
    )
    parser.add_argument(
        "--ti_embeddings_dir",
        type=str,
    )
    parser.add_argument(
        "--output_dir",
        type=str,
    )
    parser.add_argument(
        "--config",
        type=str,
    )

    return parser


def create_cmd_parser():
    parser = argparse.ArgumentParser(
        description="Simple example of a training script."
    )
    parser.add_argument(
        "--prompt",
        type=str,
        nargs="+",
    )
    parser.add_argument(
        "--negative_prompt",
        type=str,
        nargs="*",
    )
    parser.add_argument(
        "--image",
        type=str,
    )
    parser.add_argument(
        "--image_noise",
        type=float,
    )
    parser.add_argument(
        "--width",
        type=int,
    )
    parser.add_argument(
        "--height",
        type=int,
    )
    parser.add_argument(
        "--batch_size",
        type=int,
    )
    parser.add_argument(
        "--batch_num",
        type=int,
    )
    parser.add_argument(
        "--steps",
        type=int,
    )
    parser.add_argument(
        "--guidance_scale",
        type=float,
    )
    parser.add_argument(
        "--seed",
        type=int,
    )
    parser.add_argument(
        "--config",
        type=str,
    )

    return parser


def run_parser(parser, defaults, input=None):
    args = parser.parse_known_args(input)[0]
    conf_args = argparse.Namespace()

    if args.config is not None:
        with open(args.config, 'rt') as f:
            conf_args = parser.parse_known_args(
                namespace=argparse.Namespace(**json.load(f)["args"]))[0]

    res = defaults.copy()
    for dict in [vars(conf_args), vars(args)]:
        res.update({k: v for (k, v) in dict.items() if v is not None})

    return argparse.Namespace(**res)


def save_args(basepath, args, extra={}):
    info = {"args": vars(args)}
    info["args"].update(extra)
    with open(f"{basepath}/args.json", "w") as f:
        json.dump(info, f, indent=4)


def load_embeddings_ti(tokenizer, text_encoder, embeddings_dir):
    print(f"Loading Textual Inversion embeddings")

    embeddings_dir = Path(embeddings_dir)
    embeddings_dir.mkdir(parents=True, exist_ok=True)

    placeholder_tokens = [file.stem for file in embeddings_dir.iterdir() if file.is_file()]
    tokenizer.add_tokens(placeholder_tokens)

    text_encoder.resize_token_embeddings(len(tokenizer))

    token_embeds = text_encoder.get_input_embeddings().weight.data

    for file in embeddings_dir.iterdir():
        if file.is_file():
            placeholder_token = file.stem
            placeholder_token_id = tokenizer.convert_tokens_to_ids(placeholder_token)

            data = torch.load(file, map_location="cpu")

            assert len(data.keys()) == 1, 'embedding file has multiple terms in it'

            emb = next(iter(data.values()))
            if len(emb.shape) == 1:
                emb = emb.unsqueeze(0)

            token_embeds[placeholder_token_id] = emb

            print(f"Loaded {placeholder_token}")


def create_pipeline(model, scheduler, ti_embeddings_dir, dtype):
    print("Loading Stable Diffusion pipeline...")

    tokenizer = CLIPTokenizer.from_pretrained(model, subfolder='tokenizer', torch_dtype=dtype)
    text_encoder = CLIPTextModel.from_pretrained(model, subfolder='text_encoder', torch_dtype=dtype)
    vae = AutoencoderKL.from_pretrained(model, subfolder='vae', torch_dtype=dtype)
    unet = UNet2DConditionModel.from_pretrained(model, subfolder='unet', torch_dtype=dtype)

    load_embeddings_ti(tokenizer, text_encoder, ti_embeddings_dir)

    if scheduler == "plms":
        scheduler = PNDMScheduler(
            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True
        )
    elif scheduler == "klms":
        scheduler = LMSDiscreteScheduler(
            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
        )
    elif scheduler == "ddim":
        scheduler = DDIMScheduler(
            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False
        )
    else:
        scheduler = EulerAncestralDiscreteScheduler(
            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
        )

    pipeline = VlpnStableDiffusion(
        text_encoder=text_encoder,
        vae=vae,
        unet=unet,
        tokenizer=tokenizer,
        scheduler=scheduler,
    )
    pipeline.aesthetic_gradient_iters = 20
    pipeline.to("cuda")

    print("Pipeline loaded.")

    return pipeline


def generate(output_dir, pipeline, args):
    if isinstance(args.prompt, str):
        args.prompt = [args.prompt]

    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
    output_dir = output_dir.joinpath(f"{now}_{slugify(args.prompt[0])[:100]}")
    output_dir.mkdir(parents=True, exist_ok=True)

    args.seed = args.seed or torch.random.seed()

    save_args(output_dir, args)

    if args.image:
        init_image = Image.open(args.image)
        if not init_image.mode == "RGB":
            init_image = init_image.convert("RGB")
    else:
        init_image = None

    with torch.autocast("cuda"), torch.inference_mode():
        for i in range(args.batch_num):
            pipeline.set_progress_bar_config(
                desc=f"Batch {i + 1} of {args.batch_num}",
                dynamic_ncols=True
            )

            generator = torch.Generator(device="cuda").manual_seed(args.seed + i)
            images = pipeline(
                prompt=args.prompt * (args.batch_size // len(args.prompt)),
                height=args.height,
                width=args.width,
                negative_prompt=args.negative_prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance_scale,
                generator=generator,
                latents=init_image,
                strength=args.image_noise,
            ).images

            for j, image in enumerate(images):
                image.save(output_dir.joinpath(f"{args.seed + i}_{j}.jpg"))

    if torch.cuda.is_available():
        torch.cuda.empty_cache()


class CmdParse(cmd.Cmd):
    prompt = 'dream> '
    commands = []

    def __init__(self, output_dir, pipeline, parser):
        super().__init__()

        self.output_dir = output_dir
        self.pipeline = pipeline
        self.parser = parser

    def default(self, line):
        line = line.replace("'", "\\'")

        try:
            elements = shlex.split(line)
        except ValueError as e:
            print(str(e))

        if elements[0] == 'q':
            return True

        try:
            args = run_parser(self.parser, default_cmds, elements)
        except SystemExit:
            self.parser.print_help()
        except Exception as e:
            print(e)
            return

        if len(args.prompt) == 0:
            print('Try again with a prompt!')

        try:
            generate(self.output_dir, self.pipeline, args)
        except KeyboardInterrupt:
            print('Generation cancelled.')

    def do_exit(self, line):
        return True


def main():
    logging.basicConfig(stream=sys.stdout, level=logging.WARN)

    args_parser = create_args_parser()
    args = run_parser(args_parser, default_args)

    output_dir = Path(args.output_dir)
    dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.precision]

    pipeline = create_pipeline(args.model, args.scheduler, args.ti_embeddings_dir, dtype)
    cmd_parser = create_cmd_parser()
    cmd_prompt = CmdParse(output_dir, pipeline, cmd_parser)
    cmd_prompt.cmdloop()


if __name__ == "__main__":
    main()