import tkinter as tk
from tkinter import filedialog
from tkinter import messagebox

import os

from PIL import Image
from PIL.Image import EXTENT
import colorsys
from haishoku.haishoku import Haishoku
import math

import matplotlib.pyplot as plt

# ---------------------- 对话框输入参数 ----------------------

# 设置主窗口
window = tk.Tk()
window.title('彩虹色带生成器')
window.geometry('400x630')

l = tk.Label(window, width=50, text='')
l.pack()

# 输入存放图片的文件夹地址
def directory():
    global directory_var
    directory_var.set("")
    #获得选择好的文件夹
    directory_name = filedialog.askdirectory()
    directory_var.set(directory_name)

# 绘制按钮
directory_button = tk.Button(window, 
    text='选择文件夹',
    width=15,
    command=directory,
    bg='mistyrose')
directory_button.pack()

# 文字变量储存器
directory_var = tk.StringVar()
directory_var.set("")
l = tk.Label(window, 
    textvariable=directory_var,
    font=('Arial', 10), width=50)
l.pack()

# 选择辨色模式
ALL_COLOR = tk.BooleanVar()
ALL_COLOR.set(True)
l = tk.Label(window, width=50, text='选择辨色模式', bg='peachpuff')
l.pack()
r1 = tk.Radiobutton(window, text='所有颜色加权', variable=ALL_COLOR, value=True)
r1.pack()
r2 = tk.Radiobutton(window, text='单一主色调', variable=ALL_COLOR, value=False)
r2.pack()

# 选择是否去背选择按钮
REMOVE = tk.BooleanVar()
REMOVE.set(True)
l = tk.Label(window, width=50, text='是否去除底色（仅识别白色）', bg='lightgoldenrodyellow')
l.pack()
r1 = tk.Radiobutton(window, text='是', variable=REMOVE, value=True)
r1.pack()
r2 = tk.Radiobutton(window, text='否', variable=REMOVE, value=False)
r2.pack()

# 选择深浅判断模式
LINE = tk.IntVar()
LINE.set(1)
l = tk.Label(window, width=50, text='颜色排序方式', bg='honeydew')
l.pack()
r1 = tk.Radiobutton(window, text='色相优先（适用于彩图偏多的情况，如绘画）', variable=LINE, value=1)
r1.pack()
r2 = tk.Radiobutton(window, text='饱和度优先（适用于颜色偏灰的情况，如照片）', variable=LINE, value=2)
r2.pack()
r3 = tk.Radiobutton(window, text='亮度优先（适用于黑白偏多的情况，如灰阶图）', variable=LINE, value=3)
r3.pack()

# 选择出图类型
TYPE = tk.IntVar()
TYPE.set(1)
l = tk.Label(window, width=50, text='出图类型', bg='lightcyan')
l.pack()
r1 = tk.Radiobutton(window, text='正方形图（适用于图片很多的情况）', variable=TYPE, value=1)
r1.pack()
r2 = tk.Radiobutton(window, text='圈形图（色相展示效果好，但生成图较大）', variable=TYPE, value=2)
r2.pack()
r3 = tk.Radiobutton(window, text='散点图（适用于色彩均匀且图多的情况）', variable=TYPE, value=3)
r3.pack()

# 每一张图的缩图尺寸（横像素）
SIZE = tk.IntVar()
SIZE.set(100)
l = tk.Label(window, width=50, text='成品图上每一张单图的横向像素', bg='lightskyblue')
l.pack()
e2 = tk.Entry(window, width=15, textvariable=SIZE)
e2.pack()
SIZE = e2.get()

# ---------------------- 提取所有图片 ----------------------

def get_all_images(directory_name):
    # 储存图片地址的数组
    array_of_img = []
    # 遍历提供的文件夹
    for parent, dirnames, filenames in os.walk(directory_name):
        # 获取所有的文件名
        for filename in filenames:
           # 提取出图片格式的文件
            if filename.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
                # 储存至图片地址数组
                array_of_img.append(os.path.join(parent, filename))
    notice_var.set("已完成图片获取……")
    window.update()
    # 返回值
    return array_of_img

# ---------------------- 图片按照色调排序 ----------------------

max_v = 0

def get_all_colors(array_of_img, ALL_COLOR, REMOVE, LINE, get_type):
    global notice_var, window, max_v
    # 储存排序后图像的数组
    new_images = {}
    palettes1 = 0
    palettes2 = 0
    palettes3 = 0

    # 遍历提取出来的图像
    m = 1
    for image in array_of_img:
        # 判断主要颜色加权
        if ALL_COLOR == True:
            # 获取调色板
            getcolor = Haishoku.getPalette(image)
            # 获取主色调
            n = 1.0
            # 遍历调色板
            for color in getcolor:
                # 权重太低则跳过
                if color[0] < 0.01:
                    continue
                # 去掉白色
                if REMOVE == True:
                    if color[1][0] > 250 & color[1][1] > 250 & color[1][2] > 250:
                        continue
                # 计算 颜色权重 * 颜色数据 的总和
                palettes1 += color[0] * color[1][0]
                palettes2 += color[0] * color[1][1]
                palettes3 += color[0] * color[1][2]
                # 计算个数
                n = n + 1.0
            # 对新调色板进行平均计算
            palettes1 = palettes1 / n
            palettes2 = palettes2 / n
            palettes3 = palettes3 / n
            # 记录新的色调为主色调
            getcolor = (palettes1, palettes2, palettes3)
        # 判断单一主色调
        else:
            # 获取主色调
            getcolor = Haishoku.getDominant(image)
        # 排序色相、明暗和饱和度
        getcolor = colorsys.rgb_to_hsv(getcolor[0]/255.0, getcolor[1]/255.0, getcolor[2]/255.0)
        # 非散点图
        if get_type != 3:
            if LINE == 1:
                # 色相优先排序
                getcolor = (
                    format(getcolor[0], '.2f'),
                    format(1 - getcolor[2], '.1f'),
                    format(getcolor[1], '.1f')
                    )
            elif LINE == 2:
                # 饱和度优先排序
                getcolor = (
                    format(getcolor[0] * 3, '.1f'),
                    format(getcolor[1], '.2f'),
                    format(1 - getcolor[2], '.1f')
                    )
            elif LINE == 3:
                # 明度优先排序
                getcolor = (
                    format((1 - getcolor[2])*getcolor[0], '.3f'),
                    format(getcolor[0], '.1f'),
                    )
        else:
            # 散点图遍历最大的饱和度值
            if getcolor[1] > max_v:
                max_v = getcolor[1]
    	# 重新列表新的数组
        new_images[image] = getcolor
        # 提示文本
        num = str(m) + "/" + str(len(array_of_img))
        notice_var.set("正在处理图像色卡……" + num)
        window.update()
        m = m + 1

    # 按照字典值进行排序
    new_images = sorted(new_images.items(), key = lambda kv:kv[1])

    notice_var.set("已完成颜色提取……")
    window.update()

    # 输出内容
    return new_images

# ---------------------- 合并并导出图像 ----------------------

def get_result(new_images, size, get_type, directory_name, line):
    global notice_var, window, max_v

    r = size * 5
    # 正方形图
    if get_type == 1:
        # 计算一排放多少张图片
        x = math.ceil(len(new_images) ** 0.5)
        r = size * x / 2
    # 圈形图
    elif get_type == 2:
        r = size * len(new_images) / math.pi * 0.5
    # 圆形图
    elif get_type == 3:
        r = size * max_v * len(new_images) * 0.2
    # 设置最大尺寸
    if r > 3000:
        if get_type == 1:
            size = size * 3000 / r
        r = 3000
    # 设置最小尺寸
    if r < size:
        r = size * 5
    # 生成画布（正方形）
    result = Image.new("RGB", (int(r * 2), int(r * 2)), (255,255,255))

    # 初始化图片的坐标
    m = 1
    i = 0
    j = 0
    # 遍历已排序后的图片数组
    for item in new_images:
        # 获取图像
        getimage = Image.open(item[0])
        # 裁剪成正方形
        if getimage.size[0] > getimage.size[1]:
            offset=int(getimage.size[0]-getimage.size[1])/2
            getimage=getimage.transform((getimage.size[1],getimage.size[1]),EXTENT,(offset,0,int(getimage.size[0]-offset),getimage.size[1]))
        elif getimage.size[0] < getimage.size[1]:
            offset=int(getimage.size[1]-getimage.size[0])/2
            getimage=getimage.transform((getimage.size[0],getimage.size[0]),EXTENT,(0,offset,getimage.size[0],(getimage.size[1]-offset)))
        # 缩小图像
        getimage.thumbnail((size, size))
        # 正方形图
        if get_type == 1:
            # 若横向排版超过画布尺寸则提新行
            if i >= x:
               i = 0
               j = j + 1
            # 在画布上放置图像
            result.paste(getimage, box=(size * i, size * j))
        # 圈形图
        elif get_type == 2:
            # 计算角度
            s = i / len(new_images) * 2 * math.pi
            # 计算图片坐标
            x = r + r * 0.85 * math.sin(s) - size * 0.5
            y = r - r * 0.85 * math.cos(s) - size * 0.5
            # 在画布上放置图像
            result.paste(getimage, box=(int(x), int(y)))
        # 圆形色板图
        elif get_type == 3:
            # 获取色相
            rate1 = float(item[1][0])
            # 获取饱和度
            rate2 = float(item[1][1]) / max_v
            # 计算角度
            s = rate1 * 2 * math.pi
            # 计算图片坐标
            x = r + r * rate2 * math.sin(s) - size * 0.5
            y = r - r * rate2 * math.cos(s) - size * 0.5
            # 在画布上放置图像
            result.paste(getimage, box=(int(x), int(y)))
        # 设置下一张图像的排版位置
        i = i + 1
        # 提示文本
        num = str(m) + "/" + str(len(new_images))
        notice_var.set("正在绘制色带图……" + num)
        window.update()
        m = m + 1

    # 储存生成的图片
    if get_type == 1:
        name = "_方形图"
    elif get_type == 2:
        name = "_圈形图"
    elif get_type == 3:
        name = "_散点图"
    result.save("./色带图_" + directory_name.split("/")[-1] + name + ".jpg")
    plt.imshow(result)

    notice_var.set("已生成色带图！")
    window.update()

    # 在Python编辑器中显示生成的图片
    plt.show()

# ---------------------- 开始计算按钮 ----------------------

# 逐步调用计算方法
def begin():
    global notice_var, directory_var
    notice_var.set("")
    # 参数错误则返回
    if directory_var.get() == "":
        notice_var.set("未选择文件夹")
        return
    # 获取文件夹内所有图片
    array_of_img = get_all_images(directory_var.get())
    # 按照颜色排序图片
    new_images = get_all_colors(array_of_img, ALL_COLOR.get(), REMOVE.get(), LINE.get(), TYPE.get())
    # 最后输出结果图
    get_result(new_images, int(SIZE), TYPE.get(), directory_var.get(), LINE.get())

# 绘制按钮
l = tk.Label(window, width=50, text='')
l.pack()
begin_button = tk.Button(window, 
    text='开始计算',
    width=15, height=2, 
    command=begin, 
    bg='plum')
begin_button.pack()

# 显示提示文字
notice_var = tk.StringVar()
notice_var.set("")
notice = tk.Label(window, 
    textvariable=notice_var,
    font=('Arial', 12), width=100, height=2)
notice.pack()

l = tk.Label(window, width=50, text='By 野生的羽狼', anchor="e")
l.pack()

# 呼出窗口
window.mainloop()