import argparse
import datetime
import logging
import sys
import shlex
import cmd
from pathlib import Path
import torch
import json
import traceback

from PIL import Image
from slugify import slugify
from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    PNDMScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    DDIMScheduler,
    LMSDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    KDPM2DiscreteScheduler,
    KDPM2AncestralDiscreteScheduler
)
from transformers import CLIPTextModel

from data.keywords import prompt_to_keywords, keywords_to_prompt
from models.clip.embeddings import patch_managed_embeddings
from models.clip.tokenizer import MultiCLIPTokenizer
from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
from util import load_config, load_embeddings_from_dir


torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True


default_args = {
    "model": "stabilityai/stable-diffusion-2-1",
    "precision": "fp32",
    "ti_embeddings_dir": "embeddings",
    "output_dir": "output/inference",
    "config": None,
}


default_cmds = {
    "project": "",
    "scheduler": "dpmsm",
    "template": "{}",
    "prompt": None,
    "negative_prompt": None,
    "shuffle": False,
    "image": None,
    "image_noise": .7,
    "width": 768,
    "height": 768,
    "batch_size": 1,
    "batch_num": 1,
    "steps": 30,
    "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(
        "--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(
        "--project",
        type=str,
        default=None,
        help="The name of the current project.",
    )
    parser.add_argument(
        "--scheduler",
        type=str,
        choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a"],
    )
    parser.add_argument(
        "--template",
        type=str,
    )
    parser.add_argument(
        "--prompt",
        type=str,
        nargs="+",
    )
    parser.add_argument(
        "--negative_prompt",
        type=str,
        nargs="*",
    )
    parser.add_argument(
        "--shuffle",
        type=bool,
    )
    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:
        conf_args = load_config(args.config)
        conf_args = parser.parse_known_args(namespace=argparse.Namespace(**conf_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(pipeline, embeddings_dir):
    added_tokens, added_ids = load_embeddings_from_dir(
        pipeline.tokenizer,
        pipeline.text_encoder.text_model.embeddings,
        Path(embeddings_dir)
    )
    print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")


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

    tokenizer = MultiCLIPTokenizer.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)
    scheduler = DDIMScheduler.from_pretrained(model, subfolder='scheduler', torch_dtype=dtype)

    patch_managed_embeddings(text_encoder)

    pipeline = VlpnStableDiffusion(
        text_encoder=text_encoder,
        vae=vae,
        unet=unet,
        tokenizer=tokenizer,
        scheduler=scheduler,
    )
    pipeline.enable_xformers_memory_efficient_attention()
    pipeline.enable_vae_slicing()
    pipeline.to("cuda")

    print("Pipeline loaded.")

    return pipeline


def shuffle_prompts(prompts: list[str]) -> list[str]:
    return [keywords_to_prompt(prompt_to_keywords(prompt), shuffle=True) for prompt in prompts]


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

    args.prompt = [args.template.format(prompt) for prompt in args.prompt]

    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
    image_dir = []

    if len(args.prompt) != 1:
        if len(args.project) != 0:
            output_dir = output_dir.joinpath(f"{now}_{slugify(args.project)}")
        else:
            output_dir = output_dir.joinpath(now)

        for prompt in args.prompt:
            dir = output_dir.joinpath(slugify(prompt)[:100])
            dir.mkdir(parents=True, exist_ok=True)
            image_dir.append(dir)
    else:
        output_dir = output_dir.joinpath(f"{now}_{slugify(args.prompt[0])[:100]}")
        output_dir.mkdir(parents=True, exist_ok=True)
        image_dir.append(output_dir)

    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

    if args.scheduler == "plms":
        pipeline.scheduler = PNDMScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "klms":
        pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "ddim":
        pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "dpmsm":
        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "dpmss":
        pipeline.scheduler = DPMSolverSinglestepScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "euler_a":
        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "kdpm2":
        pipeline.scheduler = KDPM2DiscreteScheduler.from_config(pipeline.scheduler.config)
    elif args.scheduler == "kdpm2_a":
        pipeline.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipeline.scheduler.config)

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

        seed = args.seed + i
        prompt = shuffle_prompts(args.prompt) if args.shuffle else args.prompt
        generator = torch.Generator(device="cuda").manual_seed(seed)
        images = pipeline(
            prompt=prompt,
            negative_prompt=args.negative_prompt,
            height=args.height,
            width=args.width,
            num_images_per_prompt=args.batch_size,
            num_inference_steps=args.steps,
            guidance_scale=args.guidance_scale,
            generator=generator,
            image=init_image,
            strength=args.image_noise,
        ).images

        for j, image in enumerate(images):
            basename = f"{seed}_{j // len(args.prompt)}"
            dir = image_dir[j % len(args.prompt)]

            image.save(dir.joinpath(f"{basename}.png"))
            image.save(dir.joinpath(f"{basename}.jpg"), quality=85)
            with open(dir.joinpath(f"{basename}.txt"), 'w') as f:
                f.write(prompt[j % len(args.prompt)])

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


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

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

        self.output_dir = output_dir
        self.ti_embeddings_dir = ti_embeddings_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))
            return

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

        if elements[0] == 'reload_embeddings':
            load_embeddings(self.pipeline, self.ti_embeddings_dir)
            return

        try:
            args = run_parser(self.parser, default_cmds, elements)

            if len(args.prompt) == 0:
                print('Try again with a prompt!')
                return
        except SystemExit:
            traceback.print_exc()
            self.parser.print_help()
            return
        except Exception as e:
            traceback.print_exc()
            return

        try:
            generate(self.output_dir, self.pipeline, args)
        except KeyboardInterrupt:
            print('Generation cancelled.')
        except Exception as e:
            traceback.print_exc()
            return

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


if __name__ == "__main__":
    main()