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)
  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):
    img = resize(img,VALUE)
    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 = []
        for np in [color_points[j] for j in knn_p]:
            t = img.getpixel((np[0],np[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):
    c_p, n_p = get_points(img)
    return restore(img,k,c_p,n_p)