"""Faithful PyTorch port of Mordvintsev's original DeepDream (dream.ipynb).

Same algorithm: normalized-L2 gradient ascent on a chosen GoogLeNet layer, across octaves,
with random jitter — plus the notebook's guided-dream objective. Uses the original BVLC caffe
weights (BGR, mean [104,116,122], 0-255) so the aesthetic matches the 2015 original.
"""
import numpy as np
import torch
import torch.nn.functional as F

MEAN = torch.tensor([104.0, 116.0, 122.0]).view(1, 3, 1, 1)


class _Stop(Exception):
    pass


def preprocess(img_rgb, device):
    x = torch.from_numpy(np.asarray(img_rgb, dtype=np.float32)).permute(2, 0, 1)[None].to(device)
    return x[:, [2, 1, 0]] - MEAN.to(device)  # RGB->BGR, subtract mean


def deprocess(x):
    x = (x + MEAN.to(x.device))[:, [2, 1, 0]]
    return x[0].permute(1, 2, 0).clamp(0, 255).byte().cpu().numpy()


def _capture(net, layer_name, x):
    """Forward only up to layer_name (short-circuit) and return its activation."""
    box = {}

    def hook(m, i, o):
        box["a"] = o
        raise _Stop()

    h = dict(net.named_modules())[layer_name].register_forward_hook(hook)
    try:
        net(x)
    except _Stop:
        pass
    finally:
        h.remove()
    return box["a"]


def _make_step(net, layer, x, step, jitter, guide_feat):
    ox, oy = [int(v) for v in np.random.randint(-jitter, jitter + 1, 2)]
    x = torch.roll(x, shifts=(oy, ox), dims=(2, 3)).requires_grad_(True)
    a = _capture(net, layer, x)
    if guide_feat is None:
        loss = (a ** 2).sum() * 0.5          # maximize L2 (objective_L2)
    else:
        ch = a.shape[1]
        xf = a[0].reshape(ch, -1)
        yf = guide_feat[0].reshape(ch, -1)
        idx = (xf.t() @ yf).argmax(1)        # best-matching guide features
        loss = (xf * yf[:, idx]).sum()       # objective_guide
    g, = torch.autograd.grad(loss, x)
    x = x.detach() + step / g.abs().mean().clamp(min=1e-8) * g
    x = torch.roll(x, shifts=(-oy, -ox), dims=(2, 3))
    mean = MEAN.to(x.device)
    return torch.clamp(x, -mean, 255 - mean)


def deepdream(net, base_rgb, layer="inception_4c_output", iter_n=10, octave_n=4,
              octave_scale=1.4, step=1.5, jitter=32, guide_rgb=None, device="cuda"):
    """Generator: yields (octave_done, octave_total, rgb_image) after each octave.
    The final yield holds the finished dream."""
    guide_feat = None
    if guide_rgb is not None:
        guide_feat = _capture(net, layer, preprocess(guide_rgb, device)).detach()

    octaves = [preprocess(base_rgb, device)]
    for _ in range(octave_n - 1):
        octaves.append(F.interpolate(octaves[-1], scale_factor=1.0 / octave_scale,
                                     mode="bilinear", align_corners=False, recompute_scale_factor=True))
    detail = torch.zeros_like(octaves[-1])
    for o, octave_base in enumerate(octaves[::-1]):
        h, w = octave_base.shape[-2:]
        if o > 0:
            detail = F.interpolate(detail, size=(h, w), mode="bilinear", align_corners=False)
        x = octave_base + detail
        for _ in range(iter_n):
            x = _make_step(net, layer, x, step, jitter, guide_feat)
        detail = x - octave_base
        yield o + 1, octave_n, deprocess(x)


def dream_once(net, base_rgb, **kw):
    """Run to completion, return the final image (for the zoom loop)."""
    img = base_rgb
    for _, _, img in deepdream(net, base_rgb, **kw):
        pass
    return img