import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms


def resize_pad(img_rgb: Image.Image, image_size: int) -> Image.Image:
    img_rgb = img_rgb.convert("RGB")
    w, h = img_rgb.size
    size = int(image_size)
    if w == 0 or h == 0:
        return Image.new("RGB", (size, size), (0, 0, 0))
    scale = max(size / float(w), size / float(h))
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    img_r = img_rgb.resize((new_w, new_h), resample=Image.BICUBIC)
    left = (new_w - size) // 2
    top = (new_h - size) // 2
    return img_r.crop((left, top, left + size, top + size))


def cond_tensor_from_pil(img_rgb: Image.Image, device: torch.device) -> torch.Tensor:
    t = transforms.ToTensor()(img_rgb).unsqueeze(0).to(device)
    t = t * 2.0 - 1.0
    return t


def full_inference(model, img_rgb: Image.Image, image_size: int, device: torch.device, num_passes: int = 5, noise_std: float = 0.01):
    img_rgb = resize_pad(img_rgb, int(image_size))
    x = cond_tensor_from_pil(img_rgb, device)

    map_names = ['basecolor', 'normal', 'rmd']
    stacks = {k: [] for k in map_names}

    for p in range(num_passes):
        noise = torch.randn_like(x) * noise_std if noise_std > 0 else 0
        preds = model(x + noise)

        for k in map_names:
            stacks[k].append(preds[k])

    merged = {}
    for k in map_names:
        merged[k] = torch.median(torch.stack(stacks[k]), dim=0).values

    inv_input = torch.cat([merged['basecolor'], merged['normal'], merged['rmd']], dim=1)
    with torch.no_grad():
        merged['rgb'] = model(inv_input, mode=1)['rgb']

    def to_pil(tensor):
        out = (tensor + 1.0) / 2.0
        out = out.clamp(0, 1)
        out_np = out[0].detach().cpu().permute(1, 2, 0).numpy()
        return Image.fromarray((out_np * 255.0 + 0.5).astype("uint8"))

    outputs = {k: to_pil(v) for k, v in merged.items()}
    return img_rgb, outputs


def tiled_inference(model, img_rgb: Image.Image, tile: int, device: torch.device, overlap: int = 16):
    img_rgb = img_rgb.convert("RGB")

    overlap = int(overlap)
    if overlap < 0:
        overlap = 0
    if overlap >= tile:
        overlap = max(0, tile - 1)
    stride = max(1, tile - overlap)

    w, h = img_rgb.size
    pad_w = (tile - (w % stride)) % stride
    pad_h = (tile - (h % stride)) % stride

    if pad_w or pad_h:
        new_img = Image.new("RGB", (w + pad_w, h + pad_h), (0, 0, 0))
        new_img.paste(img_rgb, (0, 0))
        src_padded = new_img
    else:
        src_padded = img_rgb

    pw, ph = src_padded.size

    map_names = ['basecolor', 'normal', 'rmd']
    acc = {k: np.zeros((ph, pw, 3), dtype=np.float32) for k in map_names}
    wsum = np.zeros((ph, pw, 1), dtype=np.float32)

    xs = list(range(0, max(1, pw - tile + 1), stride))
    ys = list(range(0, max(1, ph - tile + 1), stride))
    if xs[-1] != pw - tile:
        xs.append(pw - tile)
    if ys[-1] != ph - tile:
        ys.append(ph - tile)

    for top in ys:
        for left in xs:
            patch_img = src_padded.crop((left, top, left + tile, top + tile))
            cond = cond_tensor_from_pil(patch_img, device)

            with torch.no_grad():
                preds = model(cond)

            def tensor_to_np(t):
                t = (t + 1.0) / 2.0
                t = t.clamp(0, 1)
                return t[0].detach().cpu().permute(1, 2, 0).numpy()

            ramp_x = np.ones((tile,), dtype=np.float32)
            ramp_y = np.ones((tile,), dtype=np.float32)
            if overlap > 0:
                if left > 0:
                    ramp_x[:overlap] = np.linspace(0.0, 1.0, overlap, endpoint=False, dtype=np.float32)
                if left + tile < pw:
                    ramp_x[-overlap:] = np.linspace(1.0, 0.0, overlap, endpoint=False, dtype=np.float32)
                if top > 0:
                    ramp_y[:overlap] = np.linspace(0.0, 1.0, overlap, endpoint=False, dtype=np.float32)
                if top + tile < ph:
                    ramp_y[-overlap:] = np.linspace(1.0, 0.0, overlap, endpoint=False, dtype=np.float32)

            weight = (ramp_y[:, None] * ramp_x[None, :])[:, :, None]

            for k in map_names:
                np_pred = tensor_to_np(preds[k])
                acc[k][top : top + tile, left : left + tile, :] += np_pred * weight
            wsum[top : top + tile, left : left + tile, :] += weight

    def acc_to_pil(out_np):
        out_np = out_np / np.maximum(wsum, 1e-8)
        out_np = np.clip(out_np, 0.0, 1.0)
        return Image.fromarray((out_np * 255.0 + 0.5).astype("uint8"))

    outputs = {k: acc_to_pil(acc[k]) for k in map_names}

    if pad_w or pad_h:
        for k in map_names:
            outputs[k] = outputs[k].crop((0, 0, w, h))

    return img_rgb, outputs


def _draw_label(img: Image.Image, label: str, bar_color=(0, 0, 0)) -> Image.Image:
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("arial.ttf", 18)
    except OSError:
        font = ImageFont.load_default()
    draw.rectangle((0, 0, img.width, 24), fill=bar_color)
    draw.text((4, 2), label, fill=(255, 255, 255), font=font)
    return img


def _draw_arrow(img: Image.Image, color=(180, 180, 180)) -> Image.Image:
    draw = ImageDraw.Draw(img)
    cx, cy = img.width // 2, img.height // 2
    r = 8
    draw.line((0, cy, img.width - r, cy), fill=color, width=3)
    draw.polygon([(img.width - r, cy - r), (img.width - r, cy + r), (img.width, cy)], fill=color)
    return img


def _draw_label(img: Image.Image, label: str, bar_color=(0, 0, 0)) -> Image.Image:
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("arial.ttf", 18)
    except OSError:
        font = ImageFont.load_default()
    draw.rectangle((0, 0, img.width, 24), fill=bar_color)
    draw.text((4, 2), label, fill=(255, 255, 255), font=font)
    return img


def _draw_arrow(img: Image.Image, color=(180, 180, 180)) -> Image.Image:
    draw = ImageDraw.Draw(img)
    cx, cy = img.width // 2, img.height // 2
    r = 8
    draw.line((0, cy, img.width - r, cy), fill=color, width=3)
    draw.polygon([(img.width - r, cy - r), (img.width - r, cy + r), (img.width, cy)], fill=color)
    return img


def make_side_by_side(inp_img: Image.Image, outputs: dict) -> Image.Image:
    inp_img = inp_img.convert("RGB")

    outputs = {k: v.convert("RGB") for k, v in outputs.items()}

    r, g, b = outputs['rmd'].split()

    cell_size = 200
    arrow_w = 48
    gap = 8

    stage1_w = cell_size
    stage2_w = cell_size * 3 + gap * 2
    stage3_w = cell_size

    total_w = stage1_w + arrow_w + stage2_w + arrow_w + stage3_w
    total_h = cell_size * 2 + gap

    canvas = Image.new("RGB", (total_w, total_h), (35, 35, 35))

    def place(cvs, img, x, y, sz=cell_size):
        img = img.resize((sz, sz), Image.BICUBIC)
        cvs.paste(img, (x, y))

    cx = 0

    stage1 = _draw_label(inp_img.resize((cell_size, cell_size), Image.BICUBIC), "INPUT", (50, 100, 200))
    place(canvas, stage1, cx, (total_h - cell_size) // 2)
    cx += cell_size

    arrow1 = Image.new("RGB", (arrow_w, total_h), (35, 35, 35))
    arrow1 = _draw_arrow(arrow1)
    canvas.paste(arrow1, (cx, 0))
    cx += arrow_w

    ordered = [
        ("BASECOLOR", outputs['basecolor'], (50, 160, 80)),
        ("NORMAL", outputs['normal'], (50, 160, 80)),
        ("DEPTH", b.convert("RGB"), (50, 160, 80)),
        ("ROUGHNESS", r.convert("RGB"), (50, 160, 80)),
        ("METALLIC", g.convert("RGB"), (50, 160, 80)),
    ]

    for i, (label, img, color) in enumerate(ordered):
        col = i % 3
        row = i // 3
        px = cx + col * (cell_size + gap)
        py = row * (cell_size + gap)
        panel = _draw_label(img.resize((cell_size, cell_size), Image.BICUBIC), label, color)
        canvas.paste(panel, (px, py))

    cx += cell_size * 3 + gap * 2

    arrow2 = Image.new("RGB", (arrow_w, total_h), (35, 35, 35))
    arrow2 = _draw_arrow(arrow2)
    canvas.paste(arrow2, (cx, 0))
    cx += arrow_w

    stage3 = _draw_label(outputs['rgb'].resize((cell_size, cell_size), Image.BICUBIC), "RECON RGB", (200, 120, 50))
    place(canvas, stage3, cx, (total_h - cell_size) // 2)

    return canvas