"""Stable Diffusion 1.5 engine with the same interface and rules as the FLUX engine.

The loop runs Euler in sigma space, so the update reads like the FLUX one:
x = x + (sigma_next - sigma) * eps. The network is told the time as a timestep in [0, 1000).
"""

from __future__ import annotations

import threading
from collections import OrderedDict
from contextlib import contextmanager

import torch
from PIL import Image

import engine as E
from engine import Clock, Settings

DECODE_PIXELS = E.DECODE_PIXELS


class SD15Engine(E.Engine):
    latent_factor = 8               # VAE scale, no patching
    standard_layers = (12,)         # CLIP's last layer, as in the pipeline

    def __init__(self, pipe, device: str | torch.device = "cuda"):
        from diffusers import EulerDiscreteScheduler

        self.pipe = pipe
        self.device = torch.device(device)
        self.clock = Clock()
        self._emb_cache: OrderedDict = OrderedDict()
        self._lock = threading.Lock()
        self.post_step: list = []       # functions latent -> latent after every sampler step (Tab 4)
        self.encode_tag = ""            # part of the encoding cache key while a tool changes CLIP
        self.run_notes: list = []
        self.dtype = pipe.unet.dtype
        # Euler in sigma space, configured from the checkpoint's scheduler (same betas, same offset).
        self._scheduler_config = EulerDiscreteScheduler.from_config(pipe.scheduler.config).config
        self._Scheduler = EulerDiscreteScheduler
        for m in (pipe.unet, pipe.text_encoder, pipe.vae):
            m.requires_grad_(False)
            m.eval()

    @classmethod
    def from_pretrained(cls, model_id: str = E.MODEL_ID, device: str = "cuda") -> "SD15Engine":
        from diffusers import StableDiffusionPipeline

        # bfloat16 rather than float16: bent activations overflow less.
        dtype = torch.bfloat16 if (device != "cuda" or torch.cuda.is_bf16_supported()) else torch.float16
        pipe = StableDiffusionPipeline.from_pretrained(model_id, dtype=dtype, variant="fp16", safety_checker=None,
                                                       feature_extractor=None, requires_safety_checker=False)
        pipe = pipe.to(device)
        pipe.vae.to(torch.float32)
        # Typed sizes go up to 1536 px, where the fp32 decoder would run out of memory. SD's VAE would
        # tile from 512 px on, so tiles start only above 1024 px and smaller images decode as before.
        pipe.vae.enable_tiling()
        pipe.vae.tile_sample_min_size = 1024
        pipe.vae.tile_latent_min_size = 1024 // 8
        return cls(pipe, device)

    # ---------------------------------------------------------------- text

    def encode(self, prompt: str, layers: tuple = (12,), max_len: int = 77,
               shuffle_seed: int | None = None) -> torch.Tensor:
        """CLIP hidden states (1, 77, 768) after the final layer norm, read at one layer. Cached.

        shuffle_seed permutes the prompt's tokens before the encoder reads them (R4).
        """
        if len(layers) != 1:
            raise ValueError("SD 1.5 reads the text encoder at exactly one layer")
        key = (prompt, tuple(layers), max_len, shuffle_seed, self.encode_tag)
        with self._lock:
            if key in self._emb_cache:
                self._emb_cache.move_to_end(key)
                return self._emb_cache[key]
        te, tok = self.pipe.text_encoder, self.pipe.tokenizer
        ids = tok(prompt, padding="max_length", max_length=tok.model_max_length, truncation=True,
                  return_tensors="pt").input_ids
        pos = self.content_positions(prompt)
        if shuffle_seed is not None and len(pos) > 1:
            perm = torch.randperm(len(pos), generator=torch.Generator("cpu").manual_seed(shuffle_seed))
            idx = torch.tensor(pos)
            ids[0, idx] = ids[0, idx[perm]]
        ids = ids.to(self.device)
        with torch.inference_mode():
            out = te(ids, output_hidden_states=True)
            # Layer 12 gives exactly the pipeline's prompt_embeds (last_hidden_state).
            # transformers 5 keeps the norm on the model itself, older versions on .text_model.
            norm = te.final_layer_norm if hasattr(te, "final_layer_norm") else te.text_model.final_layer_norm
            emb = norm(out.hidden_states[layers[0]]).to(self.dtype)
        with self._lock:
            self._emb_cache[key] = emb
            while len(self._emb_cache) > E.EMB_CACHE_SIZE:
                self._emb_cache.popitem(last=False)
        return emb

    def prompt_embeddings(self, s: Settings, max_len: int) -> tuple[torch.Tensor, torch.Tensor]:
        # The UNet's cross-attention has no positions on the text side, so permuting the encoded
        # vectors would change nothing. R4 shuffles the word order before CLIP reads it instead.
        seed = s.seed + E.SEED_SHUFFLE if s.emb_shuffle else None
        return (self.encode(s.prompt_a, s.te_layers, max_len, seed).float(),
                self.encode(s.prompt_b, s.te_layers, max_len, seed).float())

    def shuffle_cond(self, s: Settings, cond: torch.Tensor, max_len: int) -> torch.Tensor:
        return cond                                        # already shuffled before the encoder

    def content_positions(self, prompt: str, max_len: int = 77) -> list[int]:
        """Positions of the prompt's tokens: after the start token, before the end token and the padding."""
        n = len(self.pipe.tokenizer(prompt, add_special_tokens=False).input_ids)
        return list(range(1, 1 + min(n, self.pipe.tokenizer.model_max_length - 2)))

    # ---------------------------------------------------------------- blocks

    def block(self, block_id: str):
        unet = self.pipe.unet
        if block_id == "mid":
            return unet.mid_block
        kind, k = block_id.split(".")
        return {"down": unet.down_blocks, "up": unet.up_blocks}[kind][int(k)]

    def block_ids(self) -> list[str]:
        unet = self.pipe.unet
        return ([f"down.{k}" for k in range(len(unet.down_blocks))] + ["mid"]
                + [f"up.{k}" for k in range(len(unet.up_blocks))])

    def reads_text(self, block_id: str) -> bool:
        return bool(getattr(self.block(block_id), "has_cross_attention", False))

    def text_bend_note(self, s: Settings) -> str | None:
        if s.bend_stream == "text" and not self.reads_text(s.bend_block):
            return ("Dieser Block liest den Text nicht, er hat keine Cross-Attention. "
                    "Der Eingriff ändert deshalb nichts. Wählt einen anderen Block.")
        return None

    @contextmanager
    def interventions(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
        """R8 and R9 as in the FLUX engine, plus the tools of Tab 4."""
        import tools_sd15

        with super().interventions(s, lat_h, lat_w, txt_len), tools_sd15.apply(self, s) as notes:
            self.run_notes += notes
            yield

    @contextmanager
    def bending(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
        """R8. Image stream: the block's output feature map. Text stream: the text the block attends to."""
        if not s.bend_block or s.bend_op == "none" or s.bend_strength == 0:
            yield
            return
        op, st, image = s.bend_op, float(s.bend_strength), s.bend_stream == "image"
        if op not in E.BEND_OPS:
            raise ValueError(f"unknown bending op {op!r}")
        if not image and op in E.SPATIAL_OPS:
            raise ValueError(f"{op} works on the image stream only")
        if not image and not self.reads_text(s.bend_block):
            yield                                          # nothing to bend, render() adds a note
            return
        block = self.block(s.bend_block)
        masks: dict[int, torch.Tensor] = {}
        gen = torch.Generator(self.device).manual_seed(s.seed + E.SEED_BEND_NOISE) if op == "noise" else None
        clock = self.clock

        def mask(dim: int, device, dtype) -> torch.Tensor | None:
            if op != "ablate":
                return None
            if dim not in masks:                           # fixed per call, one per channel count
                off = torch.randperm(dim, generator=torch.Generator("cpu").manual_seed(s.seed + E.SEED_ABLATE))
                m = torch.ones(dim)
                m[off[:round(min(max(st, 0.0), 1.0) * dim)]] = 0.0   # typed values may leave 0..1
                masks[dim] = m
            return masks[dim].to(device, dtype)

        def bend(h: torch.Tensor) -> torch.Tensor:
            if h.dim() == 4:                               # feature map (B, C, H, W) as tokens (B, H*W, C)
                b, c, hh, ww = h.shape
                tokens = h.permute(0, 2, 3, 1).reshape(b, hh * ww, c)
                out = E.bend_tensor(tokens, op, st, (hh, ww), mask(c, h.device, h.dtype), gen)
                return out.reshape(b, hh, ww, c).permute(0, 3, 1, 2).contiguous()
            return E.bend_tensor(h, op, st, None, mask(h.shape[-1], h.device, h.dtype), gen)

        def active() -> bool:
            return s.bend_from <= clock.progress <= s.bend_to

        if image:
            def hook(module, args, output):
                if not active():
                    return None
                if isinstance(output, tuple):              # down blocks: (hidden states, skip connections)
                    return (bend(output[0]),) + tuple(output[1:])
                return bend(output)

            handle = block.register_forward_hook(hook)
        else:
            def pre_hook(module, args, kwargs):
                if not active() or kwargs.get("encoder_hidden_states") is None:
                    return None
                return args, {**kwargs, "encoder_hidden_states": bend(kwargs["encoder_hidden_states"])}

            handle = block.register_forward_pre_hook(pre_hook, with_kwargs=True)
        try:
            yield
        finally:
            handle.remove()

    # ---------------------------------------------------------------- loop

    def denoise(self, s: Settings, cond: torch.Tensor, uncond: torch.Tensor, seeds: list[int],
                guidances: list[float], on_step=None) -> torch.Tensor:
        """Euler in sigma space for one chunk. Returns latents (n, 4, h, w) in float32."""
        dev, N = self.device, s.steps
        n = len(seeds)
        lat_h, lat_w = s.height // 8, s.width // 8
        noise = torch.stack([torch.randn((4, lat_h, lat_w), generator=torch.Generator("cpu").manual_seed(sd))
                             for sd in seeds]).to(dev)        # float32, device independent
        sched = self._Scheduler.from_config(self._scheduler_config)   # fresh: the scheduler is stateful
        sched.set_timesteps(N, device=dev)
        sig = sched.sigmas.float().cpu().tolist()             # N + 1 values ending in 0
        lat = noise * sched.init_noise_sigma
        n_run = max(1, round(N * s.stop_at))

        enc = torch.cat([cond, uncond.expand(n, -1, -1)]).to(self.dtype)
        g = torch.tensor(guidances, device=dev, dtype=torch.float32).view(-1, 1, 1, 1)
        x0_pred = None
        for i, t in enumerate(sched.timesteps[:n_run]):
            self.clock.progress = i / N
            sigma = sig[i]
            t_model = float(t) * s.sigma_mult + s.sigma_offset * 1000   # R6: the model is told a different time
            x = sched.scale_model_input(lat, t)                         # x / sqrt(sigma^2 + 1)
            x = torch.cat([x, x]).to(self.dtype)
            eps = self.pipe.unet(x, torch.full((2 * n,), t_model, device=dev), encoder_hidden_states=enc,
                                 return_dict=False)[0].float()
            e_c, e_u = eps.chunk(2)
            eps = e_u + g * (e_c - e_u)                                 # R3: per-sample guidance
            x0_pred = lat - sigma * eps
            lat = sched.step(eps, t, lat, return_dict=False)[0]         # x + (sigma_next - sigma) * eps
            for fn in self.post_step:                                   # Tab 4, e.g. the mirror
                lat = fn(lat)
            if on_step is not None:
                on_step()
        return x0_pred if (s.show_x0 and n_run < N) else lat

    def decode(self, final: torch.Tensor, lat_h: int, lat_w: int) -> tuple[list[Image.Image], list[bool]]:
        vae = self.pipe.vae
        n = final.shape[0]
        finite = torch.isfinite(final).flatten(1).all(1).tolist()
        lat = torch.nan_to_num(final.float(), 0.0, 0.0, 0.0) / vae.config.scaling_factor
        per_call = max(1, DECODE_PIXELS // (lat_h * lat_w * 64))
        images = []
        for k in range(0, n, per_call):
            img = vae.decode(lat[k:k + per_call].to(vae.dtype), return_dict=False)[0]
            img = torch.nan_to_num(img.float(), 0.0, 0.0, 0.0)
            images += self.pipe.image_processor.postprocess(img, output_type="pil",
                                                           do_denormalize=[True] * img.shape[0])
        return images, finite