import numpy as np
import matplotlib.pylab as plt
import ot
import ot.plot
from PIL import Image

from sklearn.cluster import KMeans
import gradio as gr





def kclus(data,n_clusters,resize1, resize2):
    # img = np.resize(data,(resize2, resize1,3))
    # data = np.array(img,dtype="float64").reshape(-1, 3)
    data = Image.fromarray(data)
    # print(np.array(data.resize((resize2, resize1)), dtype="float64"))
    data = np.array(data.resize((resize2, resize1)), dtype="float64").reshape(-1, 3)
    kmeans = KMeans(n_clusters, random_state=0)
    cluster_pred = kmeans.fit_predict(data)
    # print(cluster_pred.shape)
    ans = [0 for i in range(n_clusters)]
    for i in range(n_clusters):
        x = data[cluster_pred==i]
        color_ave = np.sum(x, axis=0)/len(x)
        ans[i] = color_ave
    # print(ans.shape)
    return ans,cluster_pred

def saiteki(source,target,gaso,cluster_pred,resize1,resize2):
    xs = np.array(source, dtype="float64")
    xt = np.array(target, dtype="float64")
    # print(xt.shape)
    a, b = np.ones((gaso,)) / gaso, np.ones((gaso,)) / gaso
    # print(a.shape,b.shape)
    M = ot.dist(xs, xt)
    M /= M.max()
    G0 = ot.emd(a, b, M)
    P = G0 * (gaso)
    # print(P.shape)
    kansei = np.einsum('id,ji -> jd', xt, P)
    # cluster_pred = cluster_pred.reshape(100,100)
    kansei = kansei.reshape(gaso,3)
    # print(kansei.shape)
    new_data = [0 for i in range(len(cluster_pred))]
    for i in range(len(cluster_pred)):
        new_data[i] = kansei[cluster_pred[i]]
    # print(new_data.shape)
    new_data = np.array(new_data, dtype="uint8").reshape(resize1,resize2,3)
    # plt.figure(1)
    # plt.imshow(new_data)
    # plt.show()
    return new_data


# resize1 = 200
# resize2 = 350
#
# im1 = Image.open('data/gohho.jpeg')
# im2 = Image.open('data/yuya_sea.jpeg')

# source = np.array(im1.resize((resize2,resize1)),dtype="float64").reshape(-1,3)
# target = np.array(im2.resize((resize2,resize1)),dtype="float64").reshape(-1,3)
# sourc = np.array(source, dtype="uint8").reshape(resize1,resize2,3)
# plt.figure(1)
# plt.imshow(sourc)
# plt.show()
# gaso = 200
# so,cluster_pred1 = kclus(source,gaso)
# re,cluster_pred2 = kclus(target,gaso)
# saiteki(so, re,gaso,cluster_pred1,resize1,resize2)

def sepia(source,target):
    resize1 = 200
    resize2 = 300
    gaso = 150
    # a = np.array(source.resize((resize2, resize1)), dtype="float64").reshape(-1, 3)
    # b = np.array(target.resize((resize2, resize1)), dtype="float64").reshape(-1, 3)
    so, cluster_pred1 = kclus(source, gaso,resize1, resize2)
    re, cluster_pred2 = kclus(target, gaso,resize1, resize2)
    sepia_img = saiteki(so, re, gaso, cluster_pred1, resize1, resize2)
    # plt.figure(1)
    # plt.imshow(sepia_img)
    # plt.show()

    return sepia_img

demo = gr.Interface(fn = sepia, inputs = ["image","image"], outputs = ["image"])

demo.launch()