import numpy as np
import scipy
from PIL import Image

VALUE = 512

def resize(value,img):
  img = Image.open(img)
  #img = img.resize((value,value), Image.Resampling.LANCZOS)
  img.thumbnail((VALUE,VALUE), Image.Resampling.LANCZOS)
  return img

def get_mask(img,p):
    w,h=img.size
    return np.random.choice(a=[False, True], size=(w, h), p=[p, 1-p])

def generate_points(mask):
    (w,h) = mask.shape
    noise_points = []
    color_points = []
    for x in range(w):
        for y in range(h):
            if mask[x,y]:
                color_points.append(np.array([x,y]))
            else:
                noise_points.append(np.array([x,y]))
    return color_points, noise_points

def encoder_cp(img,color_points):
    w,h=img.size
    img2=Image.new('RGB',(w,h))
    for p in color_points:
        t = img.getpixel((p[0],p[1]))
        img2.putpixel((p[0],p[1]),(t[0],t[1],t[2]))
    return img2

def encoder(img,p=0.95):
    img = resize(VALUE,img)
    mask = get_mask(img,p)
    c_p, n_p = generate_points(mask)
    return encoder_cp(img, c_p)


def get_points(img):
    w,h=img.size
    noise_points = []
    color_points = []
    for x in range(w):
        for y in range(h):
            t = img.getpixel((x,y))
            if np.sum(t[:3]) > 0 :
                color_points.append(np.array([x,y]))
            else:
                noise_points.append(np.array([x,y]))
    return color_points, noise_points

def restore(img, k, color_points, noise_points):
    kdtree = scipy.spatial.KDTree(color_points)
    for p in noise_points:
        _, knn_p = kdtree.query(p, k)
        r_m = []
        v_m = []
        b_m = []
        if k == 1:
            c_p = color_points[knn_p]
            t = img.getpixel((c_p[0],c_p[1]))
            img.putpixel((p[0],p[1]),(t[0],t[1],t[2]))
        else:
            for c_p in [color_points[j] for j in list(knn_p)]:
                t = img.getpixel((c_p[0],c_p[1]))
                r_m.append(t[0])
                v_m.append(t[1])
                b_m.append(t[2])
            r_m = int(sum(r_m)/k)
            v_m = int(sum(v_m)/k)
            b_m = int(sum(b_m)/k)
            img.putpixel((p[0],p[1]),(r_m,v_m,b_m))
    return img

def decoder(img,k=1):
    img = resize(VALUE,img)
    c_p, n_p = get_points(img)
    return restore(img,int(k),c_p,n_p)

def decoder_noise(img,k=1):
    img = Image.fromarray(img)
    #img = resize(VALUE,img)
    c_p, n_p = get_points(img)
    return restore(img,int(k),c_p,n_p)