# -*- coding: utf-8 -*-
"""
Created on Nov 3 21:18:26 2020

@author: l.z.y
@e-mail: li.zhenye@qq.com
"""
import logging
import sys
import numpy as np
import cv2
from sklearn.cluster import KMeans
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from scipy.stats import binom
import matplotlib.pyplot as plt
import time
import pickle
import os
sys.path.append(os.getcwd())
from root_dir import ROOT_DIR
import utils

FEATURE_INDEX = [0, 1, 2, 6, 7]


class WoodClass(object):
    def __init__(self, load_from=None, w=2048, h=12450, n=5000, p1=0.3, pur=0.99999, single_pick_mode=False,
                 debug_mode=False):
        """
        初始化.

        :param w: 图像的尺寸w
        :param h: 图像的尺寸h
        :param p1: 木板色彩在图像中的比例p1
        :param n: 采集用于识别的样本点的个数
        :param single_pick_mode: 是否使用单点提取方案
        """
        if load_from is None:
            if w is None or h is None:
                print("It will damage your performance if you don't set w and h and use single_pick_mode!")
                raise ValueError("w or h is None")
            self.pur, self.p1, self.k = pur, p1, 1
            self.w, self.h, self.n = w, h, n
            self.ww, self.hh = None, None
            self.width = None
            self._single_pick = single_pick_mode
            self.set_purity(self.pur)
            self.change_pick_mode(single_pick_mode)
            self.model = LogisticRegression(C=1e5)
        else:
            self.load(load_from)
        self.isCorrect = False
        self.correct_color = None
        self.log = utils.Logger(is_to_file=debug_mode)
        self.debug_mode = debug_mode
        self.image_num = 0

    def change_pick_mode(self, single_pick_mode):
        """
        更改图像提取方法，

        :param single_pick_mode:若True， 则为单点随机抽取模式
        :return: None
        """
        w, h, n = self.w, self.h, self.n

        if single_pick_mode:
            self._single_pick = True
            width = int(np.floor(np.sqrt(w * h / n)))
            w0, h0 = np.arange(0, w-width, width), np.arange(0, h-width, width)
            self.ww, self.hh = np.meshgrid(w0, h0)
            self.width = width
        else:
            self._single_pick = False
            ratio = np.sqrt(n / (w * h))
            self.ww, self.hh = int(ratio * w), int(ratio * h)

    def get_rand_sample(self, x):
        """
        在图像中进行随机抽取，如果single_pick_mode为True，则为真随机抽取，反之为假的抽取.

        :param x:
        :return:
        """
        if self._single_pick:
            offset_w, offset_h = np.random.randint(0, self.width), np.random.randint(0, self.width)
            sample = x[self.hh+offset_h, self.ww+offset_w,  ...]
        else:
            sample = cv2.resize(x, (self.ww, self.hh))
        return sample

    def fit_pictures(self, data_path=ROOT_DIR):
        """
        根据给出的data_path 进行 fit.如果没有给出data目录，那么将会使用当前文件夹
        :param data_path:
        :return:
        """
        # 训练数据文件位置
        result = self.get_train_data(data_path)
        if result is False:
            return 0
        x, y = result
        score = self.fit(x, y)
        self.save()
        return score

    def fit(self, x,  y, test_size=0.1):
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=test_size, random_state=0)
        self.model.fit(x_train, y_train)
        y_pred = self.model.predict(x_test)
        pre_score = accuracy_score(y_test, y_pred)
        self.log.log("Test accuracy is:"+str(pre_score * 100) + "%.")
        y_pred = self.model.predict(x_train)
        pre_score = accuracy_score(y_train, y_pred)
        self.log.log("Train accuracy is:"+str(pre_score * 100) + "%.")
        y_pred = self.model.predict(x)
        pre_score = accuracy_score(y, y_pred)
        self.log.log("Total accuracy is:"+str(pre_score * 100) + "%.")
        return int(pre_score*100)

    def calculate_p1(self, x, remove_background=False):
        """

        :param x:
        :param remove_background:
        :return:
        """
        if remove_background:
            x = self.remove_background(x)
        kmeans = KMeans(n_clusters=2, init='k-means++')
        kmeans.fit(x)
        result = kmeans.predict(x)  # 聚类结果

    def predict(self, img):
        """
        :param img: 输入图像
        :return: 分类值
        """
        if self.debug_mode:
            cv2.imwrite(str(self.image_num) + ".bmp", img)
            self.image_num += 1
        feature = self.extract_feature(img, remove_background=False, debug_mode=False)
        feature = feature.reshape(1, -1)[:, FEATURE_INDEX]
        if self.isCorrect:
            feature = feature / (self.correct_color+1e-4)
        pred_color = self.model.predict(feature)
        if self.debug_mode:
            self.log.log(feature)
        return int(pred_color[0])

    def correct(self, img=None, img_path=None):
        """
        记录校准值，将校准值记录到类别内
        :param img: 校准板图片
        :param img_path: 用于校准的图片路径
        :return: 0 if correct success, 1 if failed
        """
        if img is None:
            path = os.path.join(ROOT_DIR, "data", "correct")
            utils.mkdir_if_not_exist(path)
            file_list = os.listdir(path)
            if len(file_list) == 0:
                return 1
            if img_path is None:
                file_name = os.path.join(ROOT_DIR, "data", "correct", file_list[-1])
            else:
                file_name = img_path
            img = cv2.imread(file_name)

        feature = self.extract_feature(img)[FEATURE_INDEX]
        self.correct_color = feature
        self.isCorrect = True
        self.log.log("Correct Successfully!")
        return 0

    def set_purity(self, purity):
        self.pur = purity
        vs_pur = 1 - self.pur
        for i in range(self.n):
            vs_pur_i = binom.cdf(k=i, p=self.p1, n=self.n)
            if vs_pur_i > vs_pur:
                self.k = i
                return i

    def remove_background(self, x):
        # TODO: 利用色度？饱和度或者明度？亮度？去除背景
        # 去背景的方法效果不好，太慢了，所以没弄了。。。
        # x = x[2000:16000, 300:1600, :]
        x = x[2000:10000, 300:1600, :]
        return x

    def save(self, file_name=None):
        """
        保存当前文件下的classify.model文件模型
        save_parameters 为要保存的参数
        :return: None
        """
        if file_name is None:
            file_name = "model_" + time.strftime("%Y-%m-%d_%H-%M") + ".p"
            file_name = os.path.join(ROOT_DIR, "models", file_name)
        model_dic = {"n": self.n, "k": self.k, "p1": self.p1, "pur": self.pur, "model": self.model,
                     "ww": self.ww, "hh": self.hh, "width": self.width, "w": self.w, "h": self.h,
                     "mode": self._single_pick, "isCorrect": self.isCorrect}
        with open(file_name, "wb") as f:
            pickle.dump(model_dic, f)
        self.log.log("Save file to '" + str(file_name) + "'")

    def load(self, path=None):
        if path is None:
            path = os.path.join(ROOT_DIR, "models")
            utils.mkdir_if_not_exist(path)
            model_files = os.listdir(path)
            if len(model_files) == 0:
                self.log.log("No model found!")
                return 1
            self.log.log("./ Models Found:")
            _ = [self.log.log("├--"+str(model_file)) for model_file in model_files]
            file_times = [model_file[6:-2] for model_file in model_files]
            latest_model = model_files[int(np.argmax(file_times))]
            self.log.log("└--Using the latest model: "+str(latest_model))
            path = os.path.join(ROOT_DIR, "models", str(latest_model))
        if not os.path.isabs(path):
            logging.warning('给的是相对路径')
            return -1
        if not os.path.exists(path):
            logging.warning('文件不存在')
            return -1
        with open(path, "rb") as f:
            model_dic = pickle.load(f)
        self.n, self.k, self.p1, self.pur = model_dic["n"], model_dic["k"], model_dic["p1"], model_dic["pur"]
        self.ww, self.hh, self.width = model_dic["ww"], model_dic["hh"], model_dic["width"]
        self.w, self.h, self.model = model_dic["w"], model_dic["h"], model_dic["model"]
        self.isCorrect = model_dic["isCorrect"]
        self._single_pick = model_dic["mode"]
        self.set_purity(self.pur)
        self.change_pick_mode(self._single_pick)
        return 0

    def extract_feature(self, x, correct_color=False, remove_background=False, debug_mode=False):
        """
        获取图片的特征,色彩值的mean和var【l, a, b, s_l, s_a, s_b】.
        :param x: 图片
        :param correct_color: 是否进行颜色校准
        :param remove_background:是否需要移除背景
        :param debug_mode: 是否使用debug模式
        :return:
        """
        if remove_background:
            x = self.remove_background(x)
        x = self.get_rand_sample(x)
        if correct_color is True:
            x = x / self.correct_color
        if debug_mode:
            plt.figure()
            plt.subplot(211)
            plt.imshow(cv2.cvtColor(x, cv2.COLOR_BGR2RGB))
        x_hsv = cv2.cvtColor(x, cv2.COLOR_BGR2HSV)
        x = cv2.cvtColor(x, cv2.COLOR_BGR2LAB)
        x = np.concatenate((x, x_hsv), axis=2)
        x = np.reshape(x, (x.shape[0]*x.shape[1], x.shape[2]))
        x = x[np.argsort(x[:, 0])]
        x = x[-self.k:, :]
        if debug_mode:
            # self.log.log(x)
            self.log.log(x.shape)
            # self.log.log(self.k)
            # self.log.log(x)
            self.log.log(x.shape)
        mean_value = np.mean(x, axis=0)
        if debug_mode:
            self.log.log("mean color:"+str(mean_value))
            plt.subplot(212)
            color_img = np.asarray(np.ones((100, 100, 3), dtype=np.uint8) * mean_value[:3], dtype=np.uint8)
            color_img = cv2.cvtColor(color_img, cv2.COLOR_LAB2RGB)
            plt.imshow(color_img)
            plt.show()
        var_value = np.var(x, axis=0)
        feature = np.hstack((mean_value, var_value))
        if debug_mode:
            self.log.log("var: "+str(var_value))
        return feature

    def get_image_data(self, img_dir="./data/dark"):
        """
        :param img_dir: 图像文件的路径
        :return: 图像数据
        """
        img_data = []
        utils.mkdir_if_not_exist(img_dir)
        files = os.listdir(img_dir)
        if len(files) == 0:
            return False
        for file in files:
            path = os.path.join(img_dir, file)
            if self.debug_mode:
                self.log.log(path)
            train_img = cv2.imread(path)
            data = self.extract_feature(train_img)
            img_data.append(data)
        img_data = np.array(img_data)
        return img_data

    def get_train_data(self, data_dir=None, plot_2d=False, plot_data_3d=False, save_data=False):
        """
        获取图像数据
        :return: x_data, y_data
        """
        data_dir = ROOT_DIR if data_dir is None else data_dir
        dark_data = self.get_image_data(img_dir=os.path.join(data_dir, "data", "dark"))
        middle_data = self.get_image_data(img_dir=os.path.join(data_dir, "data", "middle"))
        light_data = self.get_image_data(img_dir=os.path.join(data_dir, "data", "light"))
        if (dark_data is False) or (middle_data is False) or (light_data is False):
            return False
        x_data = np.vstack((dark_data, middle_data, light_data))
        dark_label = np.zeros(len(dark_data)).T
        middle_label = np.ones(len(middle_data)).T
        light_label = 2 * np.ones(len(light_data)).T
        y_data = np.hstack((dark_label, middle_label, light_label))
        x_data = x_data[:, FEATURE_INDEX]
        # 进行色彩数据校正
        if self.isCorrect:
            x_data = x_data / (self.correct_color+1e-4)
        if plot_data_3d:
            fig = plt.figure()
            ax = fig.add_subplot(1, 1, 1, projection="3d")
            ax.scatter(x_data[:, 1], x_data[:, 2], x_data[:, 0], c=y_data, edgecolors="k")
            ax.set_xlabel("a*")
            ax.set_ylabel("b*")
            ax.set_zlabel("l")
            plt.show()
        if plot_2d:
            plt.figure()
            plt.scatter(x_data[:, 0], x_data[:, 1], c=y_data)
            plt.show()
            # 尝试最合适的特征组合，保存提取出的特征的方法
            # 0: l, 1: a, 2: b, 3: var(l), 4: var(a), 5: var(s), 6: h,  7: s, 8: v, 9: var(h) 10: var(s): 11: var(v)
            # 全部：0.941
            # [:, [0, 1, 2, 3, 4, 5, 6, 7]] : 0.88
            # [:, [0, 1, 2]] : 0.911
            # [:, [0, 1, 2, 6, 7, 8]] : 0.941
            # [:, [1, 2, 6, 7, 8]] : 0.9705
            # [:, [1, 2, 6, 7]] : 0.9705
            # [:, [1, 2, 4, 5, 6, 7]] : 0.941
            # [:, [0, 1, 2, 6, 7]] : 0.8529
        if save_data:
            with open(os.path.join("data", "data.p"), "rb") as f:
                pass
        return x_data, y_data


if __name__ == '__main__':
    # 初始化wood
    wood = WoodClass(w=4096, h=1200, n=3000, debug_mode=False)
    print("色彩纯度控制量{}/{}".format(wood.k, wood.n))
    wood.correct()
    # wood.load()
    # fit 相应的文件夹
    wood.fit_pictures(data_path=r"C:\Users\FEIJINTI\PycharmProjects\wood_color")

    # 测试单张图片的预测，predict_mode=True表示导入本地的model, False为现场训练的
    pic = cv2.imread(r"./data/dark/rgb60.png")
    start_time = time.time()
    for i in range(100):
        wood_color = wood.predict(pic)
    end_time = time.time()
    print("time consume:"+str((end_time - start_time)/100))
    print("wood_color:"+str(wood_color))