#%%
import numpy as np
from pathlib import Path
from PIL import Image


w0 = np.array(
    [
        [
            [
                [-1.80224888e-02, 4.05534595e-01, 1.62569676e-02],
                [3.86995256e-01, -1.58828032e00, 6.36388600e-01],
                [-1.01250924e-01, 6.00708783e-01, -1.56698138e-01],
            ],
            [
                [-2.71480846e00, -1.42060733e00, -2.68766713e00],
                [-1.52638519e00, 1.73007298e01, -1.35628653e00],
                [-2.54135966e00, -1.04333639e00, -2.83616400e00],
            ],
            [
                [1.95330644e00, 2.82546377e00, 1.74052000e00],
                [2.87443542e00, -1.98198109e01, 3.01367259e00],
                [1.84999144e00, 2.56530333e00, 1.73608303e00],
            ],
        ],
        [
            [
                [-2.12553531e-01, -1.84073091e00, -3.59450318e-02],
                [-1.36595523e00, 6.61149931e00, -1.21200264e00],
                [-1.61612287e-01, -1.88850796e00, -4.96574976e-02],
            ],
            [
                [-1.69264674e00, 7.85137057e-01, -1.77036440e00],
                [3.79421353e-01, 5.04367065e00, 3.15313727e-01],
                [-1.68788207e00, 9.53136027e-01, -1.84203279e00],
            ],
            [
                [1.44280708e00, 2.15313840e00, 1.39907610e00],
                [2.01263666e00, -1.40174408e01, 1.93493605e00],
                [1.39259851e00, 2.11386871e00, 1.38634503e00],
            ],
        ],
    ],
    
)
b0 = np.array([-0.3287916, 0.75910544])
w1 = np.array([[[[53.911407]], [[-38.45282]]]])
b1 =np.array([15.667337])
threshold = 0.3323521614074707

def gelu(x):  # tanh approximation — ~1e-3 off from torch's exact erf-based GELU
    return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)))


def conv3x3_reflect(x : np.ndarray, w: np.ndarray, b: np.ndarray):
    """x:(Cin,H,W) w:(Cout,Cin,3,3) b:(Cout,) -> (Cout,H,W)"""
    Cin, H, W = x.shape
    Cout = w.shape[0]
    xp  = np.pad(x, ((0, 0), (1, 1), (1, 1)), mode="reflect")
    win : np.ndarray = np.lib.stride_tricks.sliding_window_view(xp, (3, 3), axis=(-2, -1))
    win = win.transpose(1, 2, 0, 3, 4).reshape(H * W, Cin * 9)
    out = win @ w.reshape(Cout, -1).T + b
    return out.reshape(H, W, Cout).transpose(2, 0, 1)


def conv1x1(x: np.ndarray, w: np.ndarray, b: np.ndarray):
    """x:(Cin,H,W) w:(Cout,Cin,1,1) b:(Cout,) -> (Cout,H,W)"""
    Cin, H, W = x.shape
    Cout = w.shape[0]
    return (w.reshape(Cout, Cin) @ x.reshape(Cin, -1) + b[:, None]).reshape(Cout, H, W)


def predict(input_image):
    """img_rgb_u8: (H, W, 3) uint8 -> (H, W) logits."""
    img_rgb_u8 = np.asarray(input_image.convert("RGB"))
    x = img_rgb_u8.astype(np.float32).transpose(2, 0, 1) / 255.0
    h = conv3x3_reflect(x, w0, b0)
    h = gelu(h)
    logits= conv1x1(h, w1, b1)[0]
    pred   = logits > threshold
    return pred


def fix(input_image):
    pred = predict(input_image)

    x = np.asarray(input_image.convert("RGB"))

    # Replace predicted-positive pixels with the average of their non-predicted
    # neighbours in a 3x3 window (zero-padded — matches the torch version).
    keep = (~pred).astype(np.float32)               # (H, W)
    xn   = x.astype(np.float32) / 255.0             # (H, W, 3) in [0, 1]

    x_t = (xn * keep[..., None]).transpose(2, 0, 1)            # (3, H, W)
    x_p = np.pad(x_t, ((0, 0), (1, 1), (1, 1)))                # zero pad
    k_p = np.pad(keep, ((1, 1), (1, 1)))

    wx :np.ndarray= np.lib.stride_tricks.sliding_window_view(x_p, (3, 3), axis=(-2, -1))  # (3,H,W,3,3)
    wk = np.lib.stride_tricks.sliding_window_view(k_p, (3, 3))                 # (H,W,3,3)

    sums   = wx.sum(axis=(-2, -1)).transpose(1, 2, 0)           # (H, W, 3)
    counts = wk.sum(axis=(-2, -1))                              # (H, W)
    avg    = sums / np.maximum(counts, 1)[..., None]

    out = np.where(pred[..., None], avg, xn)
    img = Image.fromarray(np.clip(out * 255, 0, 255).astype(np.uint8))
    return img

if __name__ == '__main__':
    path = Path("val4.png")
    input_image = Image.open(path)
    img = fix(input_image)    
    img.save('NP.png')

# %%