import sys, types, os, math
import numpy as np, torch
from PIL import Image, ImageDraw
L = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, f'{L}/lama')
# Stub the training-only imports so the generator module loads without kornia/lightning.
st = types.ModuleType('saicinpainting.training.modules.spatial_transform')
class LearnableSpatialTransformWrapper(torch.nn.Module): pass
st.LearnableSpatialTransformWrapper = LearnableSpatialTransformWrapper
sys.modules['saicinpainting.training.modules.spatial_transform'] = st
ut = types.ModuleType('saicinpainting.utils')
ut.get_shape = lambda t: tuple(t.shape)
sys.modules['saicinpainting.utils'] = ut
from saicinpainting.training.modules.ffc import FFCResNetGenerator

gen = FFCResNetGenerator(4, 3, ngf=64, n_downsampling=3, n_blocks=18, add_out_act='sigmoid',
                         init_conv_kwargs=dict(ratio_gin=0, ratio_gout=0, enable_lfu=False),
                         downsample_conv_kwargs=dict(ratio_gin=0, ratio_gout=0, enable_lfu=False),
                         resnet_conv_kwargs=dict(ratio_gin=0.75, ratio_gout=0.75, enable_lfu=False))
import pickle
class _Any(dict):
    def __init__(self, *a, **k): pass
    def __setstate__(self, state): pass
class _Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module.startswith('pytorch_lightning') or module.startswith('lightning'):
            return _Any
        return super().find_class(module, name)
_pm = types.ModuleType('lightning_free_pickle'); _pm.Unpickler = _Unpickler; _pm.load = pickle.load
ck = torch.load(f'{L}/big-lama/models/best.ckpt', map_location='cpu', weights_only=False, pickle_module=_pm)
sd = {k[len('generator.'):]: v for k, v in ck['state_dict'].items() if k.startswith('generator.')}
missing, unexpected = gen.load_state_dict(sd, strict=True), None
gen.eval()

import coremltools as ct
cm = ct.models.MLModel(f'{L}/tp/LaMa.mlpackage', compute_units=ct.ComputeUnit.CPU_ONLY)

N = 800
def scenes():
    rng = np.random.default_rng(7)
    # 1: smooth gradient with shapes
    a = np.zeros((N, N, 3), np.float32)
    yy, xx = np.mgrid[0:N, 0:N] / N
    a[..., 0] = xx; a[..., 1] = yy; a[..., 2] = 0.5 + 0.5 * np.sin(6 * xx)
    im = Image.fromarray((a * 255).astype(np.uint8)); d = ImageDraw.Draw(im)
    d.ellipse([200, 200, 420, 420], fill=(240, 40, 40)); d.rectangle([500, 450, 700, 700], fill=(30, 30, 200))
    yield 'gradient', im
    # 2: noisy texture
    yield 'texture', Image.fromarray((rng.random((N, N, 3)) * 80 + np.array([90, 140, 60])).astype(np.uint8))
    # 3: a stock photo (Apple's User Pictures)
    src = '/System/Library/Templates/Data/Library/User Pictures/Animals/Eagle.heic'
    os.system(f"sips -s format png -z {N} {N} '{src}' --out {L}/eagle.png >/dev/null 2>&1")
    yield 'eagle', Image.open(f'{L}/eagle.png').convert('RGB').resize((N, N))
def masks():
    m1 = Image.new('L', (N, N), 0); ImageDraw.Draw(m1).rectangle([250, 250, 480, 520], fill=255)
    m2 = Image.new('L', (N, N), 0); ImageDraw.Draw(m2).ellipse([380, 120, 640, 340], fill=255)
    return [('rect', m1), ('ellipse', m2)]

def psnr(a, b, sel=None):
    a = a.astype(np.float64); b = b.astype(np.float64)
    if sel is not None: a, b = a[sel], b[sel]
    mse = ((a - b) ** 2).mean()
    return 99.0 if mse == 0 else 10 * math.log10(255 ** 2 / mse)

for sname, img in scenes():
    for mname, m in masks():
        x = torch.from_numpy(np.asarray(img, np.float32) / 255).permute(2, 0, 1)[None]
        mk = (torch.from_numpy(np.asarray(m, np.float32) / 255) > 0.5).float()[None, None]
        with torch.no_grad():
            out = gen(torch.cat([x * (1 - mk), mk], 1))
        pt = (out[0].permute(1, 2, 0).numpy() * 255).clip(0, 255)
        cmo = np.asarray(cm.predict({'image': img, 'mask': m})['output'].convert('RGB'), np.float32)
        sel = np.asarray(m) > 127
        # Inside the mask is what the model invents; outside, either may pass the input through.
        print(f'{sname:9} {mname:8} PSNR inside mask {psnr(pt, cmo, sel):6.2f} dB | whole frame {psnr(pt, cmo):6.2f} dB')

# Outside the mask, the package should pass the input through unchanged.
for sname, img in scenes():
    m = masks()[0][1]
    cmo = np.asarray(cm.predict({'image': img, 'mask': m})['output'].convert('RGB'), np.float32)
    sel = np.asarray(m) <= 127
    print(f'{sname:9} outside mask vs input: PSNR {psnr(np.asarray(img, np.float32), cmo, sel):6.2f} dB, max diff {np.abs(np.asarray(img, np.float32)[sel] - cmo[sel]).max():.0f}')