#!/bin/env python
import taichi as ti
import numpy as np
import cv2
from copy import deepcopy
ti.init(arch=ti.cpu)

DEBUG=True
MAX=255
np.set_printoptions(threshold=np.inf, linewidth=180)  # numpy打印选项
img_from='/home/n/photo/Portal_Companion_Cube.jpg'
# img_from='/home/n/photo/neco godness.jpg'
# img_from='/media/n/data/download/firefox/updated/browser/chrome/icons/default/default32.png'

def show_image(img):
    if isinstance(img, str):
        img = cv2.imread(img)
    elif img.dtype != 'uint8':
        img = np.clip(img, 0, MAX)  # 将图片像素值限制在 0~255 之间
        img = img.astype(np.uint8)
    
    # 显示图片
    cv2.imshow('Image', img)
    while True:
        key = cv2.waitKey(1) & 0xFF
        if key == 27:  # 27 是 ESC 键的 ASCII 码
            break
    cv2.destroyAllWindows()


PALETTE=[
0x00,0xFF
]
# img2d=ti.types.ndarray(dtype=ti.math.vec3, ndim=2)
type_img2d = ti.types.ndarray(element_dim=0,ndim=2)
type_bayerM = ti.types.ndarray(element_dim=0,ndim=2)

@ti.func
def clamp(x,min=-MAX,max=MAX):
    """控制出血阈值，建议min in [-256,0]"""
    return ti.math.clamp(x,min,max)
    # return np.clip(x,min,max)

@ti.kernel
def dither_basic(img:type_img2d):
    h,w = img.shape
    for i in range(h):
      for j in range(w):
        min_distance = MAX
        oldpixel = img[i, j]
        newpixel = 0
        for c in ti.static(PALETTE):
            distance = abs(img[i, j] - c)
            if distance < min_distance:
                min_distance = distance
                newpixel = c
        img[i, j] = newpixel
        if j + 1 < w:
            img[i, j+1] += oldpixel - newpixel

@ti.kernel
def dither_floyd(img:type_img2d):
    h,w = img.shape
    for i in range(h):
      for j in range(w):
    # for i,j in ti.ndrange(h,w):  # taichi的ti.ndrange有bug，与下面的结果不同！
        oldpixel = img[i, j]
        newpixel = 0 if oldpixel < 128 else MAX
        img[i, j] = newpixel
        quant_error = oldpixel - newpixel
        if j + 1 < img.shape[1]:
            img[i, j + 1] += quant_error * 7 >> 4
        if i + 1 < img.shape[0]:
            if j - 1 >= 0:
                img[i + 1, j - 1] += quant_error * 3 >> 4
            img[i + 1, j] += quant_error * 5 >> 4

def bit_reverse(x, n):
    return int(bin(x)[2:].zfill(n)[::-1], 2)

def bit_interleave(x, y, n):
    x = bin(x)[2:].zfill(n)
    y = bin(y)[2:].zfill(n)
    return int(''.join(''.join(i) for i in zip(x, y)), 2)

def bayer_entry(x, y, n):
    return bit_reverse(bit_interleave(x ^ y, y, n), 2*n)

def bayer_matrix(n):
    """https://gamedev.stackexchange.com/questions/130696/how-to-generate-bayer-matrix-of-arbitrary-size"""
    r = range(2**n)
    return [[bayer_entry(x, y, n) for x in r] for y in r]

@ti.kernel
def dither_bayer(img:type_img2d, bayerM:type_bayerM):
    h,w = img.shape
    n = bayerM.shape[0]
    for i in range(h):
      for j in range(w):
        threshold = bayerM[i % n, j % n] * MAX // (n**2)
        if img[i, j] > threshold:
            img[i, j] = MAX
        else:
            img[i, j] = 0

@ti.kernel
def dither_atkinson(img:type_img2d):
    h,w = img.shape
    for i in range(h):
      for j in range(w):
        oldpixel = img[i, j]
        newpixel = 0 if oldpixel < 128 else MAX
        img[i, j] = newpixel
        quant_err = oldpixel - newpixel

        if j + 1 < img.shape[1]:
            img[i, j + 1] += quant_err >> 3
        if j + 2 < img.shape[1]:
            img[i, j + 2] += quant_err >> 3
        if i + 1 < img.shape[0]:
            if j - 1 >= 0:
                img[i + 1, j - 1] += quant_err >> 3
            img[i + 1, j] += quant_err >> 3
            if j + 1 < img.shape[1]:
                img[i + 1, j + 1] += quant_err >> 3
            if j + 2 < img.shape[1]:
                img[i + 1, j + 2] += quant_err >> 3
        if i + 2 < img.shape[0]:
            if j - 1 >= 0:
                img[i + 2, j - 1] += quant_err >> 3
            img[i + 2, j] += quant_err >> 3
            if j + 1 < img.shape[1]:
                img[i + 2, j + 1] += quant_err >> 3
            if j + 2 < img.shape[1]:
                img[i + 2, j + 2] += quant_err >> 3

def diff(img1,img2):
    return np.sum(np.abs(img1-img2))

img = cv2.imread(img_from)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_u8 = deepcopy(img) if DEBUG else None
img = img.astype(np.int32)  # 转换为 int32 类型

if DEBUG:
    cmd = 'diff(img,img_u8)'
    print('i32<-u8',cmd,'=', eval(cmd))

    # print('Raw:')
    # print(img, '\n')

n=4
bayerM = ti.ndarray(shape=(2**n, 2**n), dtype=ti.u8)
bayerM.from_numpy(np.matrix(bayer_matrix(n), dtype=np.uint8))
dither_bayer(img, bayerM)
if DEBUG:
    img_i32 = deepcopy(img)

    # print('dithered:',diff(img,img_u8),np.max(img), np.min(img), img.shape)
    # print(img, '\n')

if DEBUG:
    img = np.clip(img, 0, MAX)
    img = img.astype(np.uint8)
    cmd = 'diff(img,img_i32)'
    print('u8<-i32',cmd,'=', eval(cmd),np.max(img), np.min(img))

    # print(img, '\n')
show_image(img)



# neighbors = [(-1, -1), (-1, ), (-1, +1),
#              (0, -1),             (0, +1),
#              (+1, -1), (+1, ), (+1, +1)] # 简化写法

# values = [1, 2, 3, 4, 5, 6, 7, 8]  # 示例值
# # a = np.arange(100).reshape(10, 10)
# a = np.zeros((10, 10))
# print(a,'\n')
# for i in range(10-3):
#   for j in range(10-3):
#     for (x, y), value in zip(neighbors, values):
#         a[i+x, j+y] = value
# print(a,'\n')