import argparse
import datetime
import logging
import sys
import shlex
import cmd
from pathlib import Path
from torch import autocast
import torch
import json
from PIL import Image
from diffusers import StableDiffusionPipeline, AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
from slugify import slugify
from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
from schedulers.scheduling_euler_a import EulerAScheduler


default_args = {
    "model": None,
    "scheduler": "euler_a",
    "precision": "bf16",
    "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(
        "--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,
    )
    parser.add_argument(
        "--negative_prompt",
        type=str,
    )
    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 create_pipeline(model, scheduler, dtype):
    print("Loading Stable Diffusion pipeline...")

    tokenizer = CLIPTokenizer.from_pretrained(model + '/tokenizer', torch_dtype=dtype)
    text_encoder = CLIPTextModel.from_pretrained(model + '/text_encoder', torch_dtype=dtype)
    vae = AutoencoderKL.from_pretrained(model + '/vae', torch_dtype=dtype)
    unet = UNet2DConditionModel.from_pretrained(model + '/unet', torch_dtype=dtype)
    feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=dtype)

    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 = EulerAScheduler(
            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,
        feature_extractor=feature_extractor
    )
    # pipeline.enable_attention_slicing()
    pipeline.to("cuda")

    print("Pipeline loaded.")

    return pipeline


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

    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 autocast("cuda"):
        for i in range(args.batch_num):
            pipeline.set_progress_bar_config(desc=f"Batch {i + 1} of {args.batch_num}")

            generator = torch.Generator(device="cuda").manual_seed(seed + i)
            images = pipeline(
                prompt=[args.prompt] * args.batch_size,
                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"{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()

        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, dtype)
    cmd_parser = create_cmd_parser()
    cmd_prompt = CmdParse(output_dir, pipeline, cmd_parser)
    cmd_prompt.cmdloop()


if __name__ == "__main__":
    main()