"""
dheeren.pipeline — Diffusion class.
Uses ONNX Runtime / OpenVINO for fast CPU inference.
Falls back to pure PyTorch automatically if neither is installed.
"""

import os, shutil, torch
import numpy as np
from pathlib import Path
from PIL import Image
from tqdm import tqdm

from .models import SimpleTokenizer, VAE, TextEncoder, UNet
from .config  import CFG
from .inference_backend import create_backend


class Diffusion:
    """
    Text-to-image pipeline using Flow Matching diffusion.

    Automatically uses the fastest available backend:
      OpenVINO  >  ONNX Runtime  >  PyTorch

    Usage
    -----
    >>> from dheeren import Diffusion
    >>> pipe  = Diffusion.from_pretrained("dheeren/flow-weights")
    >>> image = pipe.generate("a sunset over mountains")
    >>> image.save("output.png")
    """

    def __init__(self, backend, tokenizer):
        self.backend   = backend
        self.tokenizer = tokenizer

    # ── Loading ───────────────────────────────────────────────────────

    @classmethod
    def from_pretrained(cls, repo_id_or_path: str):
        """
        Load from HuggingFace Hub  OR  a local folder.

        HF repo should contain:
            tokenizer.pt
            text_encoder.onnx  unet.onnx  vae_decoder.onnx   (ONNX)
            text_encoder.xml   unet.xml   vae_decoder.xml     (OpenVINO, optional)

        Fallback: if only .pt files are present, uses PyTorch backend.
        """
        path = Path(repo_id_or_path)

        if not path.exists():
            try:
                from huggingface_hub import snapshot_download
            except ImportError:
                raise ImportError("Run:  pip install huggingface_hub")

            print(f"Downloading weights: {repo_id_or_path}")
            local = snapshot_download(
                repo_id   = repo_id_or_path,
                local_dir = Path.home() / ".cache" / "dheeren" /
                            repo_id_or_path.replace("/", "_"),
            )
            path = Path(local)

        print("Loading pipeline...")

        # Load tokenizer (always from .pt — it's just a dict)
        tok_data  = torch.load(str(path / "tokenizer.pt"),
                               map_location="cpu", weights_only=False)
        tokenizer = SimpleTokenizer(tok_data["word2idx"], tok_data["idx2word"])

        # Try ONNX/OpenVINO first, fall back to PyTorch if needed
        te, unet, vae = None, None, None
        if not (path / "unet.onnx").exists() and not (path / "unet.xml").exists():
            # Only .pt files present — load PyTorch models for fallback
            print("  No ONNX files found — loading PyTorch models...")
            te   = cls._load_text_encoder(path)
            unet = cls._load_unet(path)
            vae  = cls._load_vae(path)

        backend = create_backend(path, te=te, unet=unet, vae=vae)
        print("✓ Ready\n")
        return cls(backend, tokenizer)

    # ── PyTorch model loaders (used only as fallback) ─────────────────

    @staticmethod
    def _load_state(path):
        ckpt = torch.load(str(path), map_location="cpu", weights_only=False)
        return ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt

    @classmethod
    def _load_text_encoder(cls, path):
        te    = TextEncoder().eval()
        state = cls._load_state(path / "clip.pt")
        te_s  = {k.replace("text_encoder.", ""): v
                 for k, v in state.items() if k.startswith("text_encoder.")}
        te.load_state_dict(te_s)
        return te

    @classmethod
    def _load_unet(cls, path):
        unet = UNet().eval()
        unet.load_state_dict(cls._load_state(path / "unet.pt"))
        return unet

    @classmethod
    def _load_vae(cls, path):
        vae = VAE().eval()
        vae.load_state_dict(cls._load_state(path / "vae.pt"), strict=False)
        return vae

    # ── Inference ─────────────────────────────────────────────────────

    @torch.no_grad()
    def generate(
        self,
        prompt:    str,
        steps:     int   = 20,
        guidance:  float = 7.5,
        seed:      int   = 42,
        save_path: str   = None,
    ) -> Image.Image:
        """
        Generate an image from a text prompt.

        Parameters
        ----------
        prompt    : What to generate.
        steps     : Euler integration steps. Default 20.
        guidance  : CFG scale — how closely to follow prompt. Default 7.5.
        seed      : Random seed. Default 42.
        save_path : Auto-save PNG path (optional).

        Returns
        -------
        PIL.Image.Image
        """
        np.random.seed(seed)

        # ── 1. Encode prompts ─────────────────────────────────────────
        def _encode(text):
            ids = self.tokenizer.encode(text).numpy().astype(np.int64)
            ids = ids[np.newaxis, :]                          # (1, 77)
            emb = self.backend.encode_text(ids)               # (1, 512)
            return emb[:, np.newaxis, :]                      # (1, 1, 512)

        text_emb = _encode(prompt)
        uncond   = _encode("")
        context  = np.concatenate([uncond, text_emb], axis=0).astype(np.float32)  # (2, 1, 512)

        # ── 2. Initial noise ──────────────────────────────────────────
        latents = np.random.randn(
            1, CFG.unet_in_channels,
            CFG.unet_image_size, CFG.unet_image_size
        ).astype(np.float32)

        # ── 3. Flow Matching denoising loop ───────────────────────────
        dt        = 1.0 / steps
        timesteps = np.linspace(0, CFG.num_diffusion_steps - 1, steps).astype(np.int64)

        for t in tqdm(timesteps,
                      desc=f'Generating "{prompt[:35]}{"…" if len(prompt)>35 else ""}"',
                      unit="step", ncols=70,
                      bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} steps"):

            # Batch latents for CFG (uncond + cond in one forward pass)
            latent_in  = np.concatenate([latents, latents], axis=0)  # (2, 4, 64, 64)
            timestep_in = np.array([t, t], dtype=np.int64)

            velocity = self.backend.run_unet(latent_in, timestep_in, context)

            # Classifier-Free Guidance
            v_uncond, v_cond = velocity[0:1], velocity[1:2]
            guided = v_uncond + guidance * (v_cond - v_uncond)

            # Euler step
            latents = latents + guided * dt

        # ── 4. Decode ─────────────────────────────────────────────────
        decoded = self.backend.decode_vae(latents)            # (1, 3, H, W)
        decoded = np.clip((decoded + 1) / 2, 0, 1)
        decoded = (decoded[0].transpose(1, 2, 0) * 255).astype(np.uint8)
        image   = Image.fromarray(decoded)

        if save_path:
            image.save(save_path)
            print(f"Saved → {save_path}")

        return image

    # ── show_source_code ──────────────────────────────────────────────

    @classmethod
    def show_source_code(cls, dest: str = "."):
        """
        Copy all library source files into `dest` folder.
        Call this in your exam to show the examiner the full implementation.

        Files saved
        -----------
        pipeline.py          — Diffusion class + inference loop
        inference_backend.py — OpenVINO / ONNX Runtime / PyTorch backends
        models.py            — VAE, CLIP TextEncoder, UNet, Tokenizer
        config.py            — all model hyperparameters
        """
        dest_path = Path(dest)
        dest_path.mkdir(parents=True, exist_ok=True)

        pkg_dir  = Path(__file__).parent
        files    = ["pipeline.py", "inference_backend.py",
                    "models.py", "config.py", "__init__.py"]

        print("Source files:")
        for fname in files:
            src = pkg_dir / fname
            if src.exists():
                shutil.copy2(src, dest_path / fname)
                print(f"  ✓ {fname}")

        print(f"\nAll files saved to: {dest_path.resolve()}")
        print("These implement the full diffusion pipeline from scratch.")
        print(f"\nBackend active: {cls._last_backend_name}")

    _last_backend_name = "unknown"