# -*- coding: utf-8 -*-

import numpy as np
from tkinter import filedialog
from tkinter import *
from PIL import Image
from scipy import signal
import open3d
from sklearn.decomposition import PCA


class Point3D:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z
        self.adjacent = []

    def add_adjacent(self, point):
        present = 0
        for adj in self.adjacent:
            if point == adj:
                present = 1
                break
        if present == 0:
            self.adjacent.append(point)

    def translate(self, dx, dy, dz):
        self.x += dx
        self.y += dy
        self.z += dz


class Harris3D:
    def __init__(self, filename):
        self.pointmap = []
        print('Nacitam')
        print(filename)
        self.mesh = open3d.read_triangle_mesh(filename)
        self.points = np.asarray(self.mesh.vertices)
        self.triangles = np.asarray(self.mesh.triangles)
        self.create_map()
        self.response = []

    def draw(self, maxima):
        pcd = open3d.PointCloud()
        points = []
        colors = []

        red = [1, 0, 0]

        for arg in maxima:
            x = self.pointmap[arg].x
            y = self.pointmap[arg].y
            z = self.pointmap[arg].z
            point = [x, y, z]
            points.append(point)
            colors.append(red)

        points = np.asarray(points)
        colors = np.asarray(colors)
        pcd.points = open3d.Vector3dVector(points)
        pcd.colors = open3d.Vector3dVector(colors)
        self.mesh.compute_vertex_normals()

        open3d.draw_geometries([self.mesh, pcd])

    def create_map(self):
        for point in self.points:
            self.pointmap.append(Point3D(point[0], point[1], point[2]))
        for triangle in self.triangles:
            self.pointmap[triangle[0]].add_adjacent(triangle[1])
            self.pointmap[triangle[1]].add_adjacent(triangle[2])
            self.pointmap[triangle[2]].add_adjacent(triangle[0])
            self.pointmap[triangle[1]].add_adjacent(triangle[0])
            self.pointmap[triangle[2]].add_adjacent(triangle[1])
            self.pointmap[triangle[0]].add_adjacent(triangle[2])

    def find_points_around_point(self, point_i, k):
        neighborhood = []
        if k == 0:
            present = 0
            for i in neighborhood:
                if i == point_i:
                    present = 1
                    break
            if present == 0:
                neighborhood.append(point_i)
        else:
            test_neighborhood = self.pointmap[point_i].adjacent
            for i in test_neighborhood:
                present = 0
                for j in neighborhood:
                    if i == j:
                        present = 1
                        break
                if present == 0:
                    neighborhood.append(i)
                next_level = self.find_points_around_point(i, k-1)
                for l in next_level:
                    present = 0
                    for m in neighborhood:
                        if l == m:
                            present = 1
                            break
                    if present == 0:
                        neighborhood.append(l)

        return neighborhood

    def proceed(self, k, s, n):
        i = 0
        for point in self.pointmap:
            neighborhood = self.find_points_around_point(i, k)
            self.response.append(self.compute_Harris(neighborhood, point, s))
            i += 1

        maxima = []
        min = np.min(self.response)
        for i in range(0, n):
            argmax = np.argmax(self.response)
            argmax = int(argmax)
            maxima.append(argmax)
            self.response[argmax] = min

        self.draw(maxima)

    def compute_Harris(self, points, v, s):
        local = []
        sumx = 0
        sumy = 0
        sumz = 0
        for i in points:
            x = self.pointmap[i].x
            sumx += x
            y = self.pointmap[i].y
            sumy += y
            z = self.pointmap[i].z
            sumz += z
            point = Point3D(x, y, z)
            local.append(point)
        n = len(local)
        if n == 0:
            return 0
        tx = sumx/n
        ty = sumy/n
        tz = sumz/n

        X = []

        for point in local:
            point.translate(-tx, -ty, -tz)
            x = point.x
            y = point.y
            z = point.z
            X.append([x, y, z])

        Y = self.pca_transform(X)

        vx = v.x
        vy = v.y
        vz = v.z

        local = []

        for coordinates in Y:
            point = Point3D(coordinates[0], coordinates[1], coordinates[2])
            point.translate(-vx, -vy, -vz)
            local.append(point)

        D = self.approximation(Y)

        a = D[3][0]*D[3][0]+D[0][0]*D[0][0]+2*D[1][0]*D[1][0]
        b = D[4][0]*D[4][0]+2*D[1][0]*D[1][0]+D[2][0]*D[2][0]
        c = D[3][0]*D[4][0]+D[0][0]*D[1][0]+D[1][0]*D[2][0]

        det = a*b - c*c
        tr = a + b

        R = det - s*tr*tr

        return R

    def pca_transform(self, X):
        pca = PCA()
        pca.fit(X)
        Y = pca.transform(X)
        return Y

    def approximation(self, Y):
        A = []
        B = []
        for y in Y:
            A.append([y[0]*y[0], y[0]*y[1], y[1]*y[1], y[0], y[1], 1])
            B.append([y[2]])
        A = np.array(A)
        B = np.array(B)
        A = np.linalg.pinv(A)
        C = np.dot(A, B)
        return C



class Program:
    def __init__(self):
        self.filename = ''

    def main(self):
        window = Tk()
        window.minsize(400, 230)
        window.title('Harrisův detektor rohů')
        header = Label(window, text='HARRISŮV DETEKTOR ROHŮ PRO 2D')
        header2 = Label(window, text='HARRISŮV DETEKTOR ROHŮ PRO 3D')
        constant_label = Label(window, text='Konstanta k')
        numofcorners_label = Label(window, text='Počet rohů')
        constant_label2 = Label(window, text='Konstanta k')
        numofcorners_label2 = Label(window, text='Počet rohů')
        width_label = Label(window, text='Počet kruhů v okolí')
        file_button = Button(window, text='Zvolit obrázek', command=self.select_file_button_pressed)
        file3d_button = Button(window, text='Zvolit model', command=self.select_file3d_button_pressed)
        self.constant_entry = Entry(window)
        self.constant_entry.insert(0, '0.05')
        self.numofcorners_entry = Entry(window)
        self.numofcorners_entry.insert(0, '30')
        self.constant_entry2 = Entry(window)
        self.constant_entry2.insert(0, '0.05')
        self.numofcorners_entry2 = Entry(window)
        self.numofcorners_entry2.insert(0, '30')
        self.width_entry = Entry(window)
        self.width_entry.insert(0, '2')
        confirm_button = Button(window, text='Vyhledat rohy', command=self.confirm_button_pressed)
        confirm3d_button = Button(window, text='Vyhledat rohy ve 3D', command=self.confirm3d_button_pressed)

        header.grid(row=0, column=0)
        file_button.grid(row=1, column=0)
        constant_label.grid(row=2, column=0)
        self.constant_entry.grid(row=3, column=0, pady=2)
        numofcorners_label.grid(row=4, column=0)
        self.numofcorners_entry.grid(row=5, column=0, pady=2)
        confirm_button.grid(row=8, column=0)

        header2.grid(row=0, column=1)
        file3d_button.grid(row=1, column=1)
        constant_label2.grid(row=2, column=1)
        self.constant_entry2.grid(row=3, column=1)
        numofcorners_label2.grid(row=4, column=1)
        self.numofcorners_entry2.grid(row=5, column=1)
        width_label.grid(row=6, column=1)
        self.width_entry.grid(row=7, column=1, pady=2)
        confirm3d_button.grid(row=8, column=1, pady=2)
        window.mainloop()

    def select_file_button_pressed(self):
        self.filename = filedialog.askopenfilename(
            initialdir="/", title="Zvolte obrázek",
            filetypes=(("jpeg files", "*.jpg"), ("png files", "*.png"), ("all files", "*.*")))

    def select_file3d_button_pressed(self):
        self.filename = filedialog.askopenfilename(
            initialdir="/", title="Zvolte model",
            filetypes=(("ply files", "*.ply"), ("all files", "*.*")))
        print(self.filename)

    def confirm_button_pressed(self):
        print(self.filename)
        if self.filename != '':
            arr = self.image_load()
        else:
            return
        arrIx, arrIy = self.diff_image(arr)
        Sx2, Sy2, Sxy = self.diffsums(arrIx, arrIy, self.weightfun())
        response_matrix = self.responsematrix(Sx2,Sy2,Sxy, float(self.constant_entry.get()))
        img = self.new_image(arr, response_matrix)
        img.save("output.png")
        print('HOTOVO!')

    def confirm3d_button_pressed(self):
        print(self.filename)
        if self.filename != '':
            an3d = Harris3D(self.filename)
            an3d.proceed(int(self.width_entry.get()),
                         float(self.constant_entry2.get()),
                         int(self.numofcorners_entry2.get()))
            print('HOTOVO!')
        else:
            return

    def image_load(self):
        img = Image.open(self.filename)
        arr = np.array(img)
        print(arr)
        return arr/255

    def weightfun(self):
        weight = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
        return weight

    def diff_image(self, arr):
        sx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
        sy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
        arrIx = signal.convolve2d(arr, sx, boundary='symm', mode='same')
        arrIy = signal.convolve2d(arr, sy, boundary='symm', mode='same')
        return arrIx, arrIy

    def diffsums(self, arrIx, arrIy, weight):
        Ix2 = arrIx * arrIx
        Iy2 = arrIy * arrIy
        Ixy = arrIx * arrIy
        Sx2 = signal.convolve2d(Ix2, weight, boundary='fill', mode='same')
        Sy2 = signal.convolve2d(Iy2, weight, boundary='fill', mode='same')
        Sxy = signal.convolve2d(Ixy, weight, boundary='fill', mode='same')
        return Sx2, Sy2, Sxy

    def responsematrix(self, Sx2, Sy2, Sxy, k):
        m = len(Sx2)
        n = len(Sx2[0])
        print(m)
        print(n)
        matrix = np.zeros((m, n))
        for i in range(0, m):
            for j in range(0, n):
                det, trace = self.harrisprop(Sx2[i,j], Sy2[i,j], Sxy[i,j])
                response = self.response(det, trace, k)
                matrix[i, j] = response
        # print(matrix)
        # print(len(matrix))
        # print(len(matrix[0]))

        return matrix

    def harrisprop(self, Sx, Sy, Sxy):
        det = Sx*Sy - Sxy*Sxy
        trace = Sx + Sy
        # print(det)
        # print(trace)
        return det, trace

    def response(self, det, trace, k):
        return det-k*trace*trace

    def red_placer(self, image, response):
        red = np.max(response)
        blue = np.min(response)
        for i in range(0, len(response)):
            for j in range(0, len(response[0])):
                if response[i, j] == red:
                    response[i, j] = blue
                    image[i, j] = (255, 0, 0)
                    print(i, j)
                    return image

    def new_image(self, arr, response):
        m = len(arr)
        n = len(arr[0])
        arr = np.int8(arr * 255)
        new_img = np.zeros((m, n), dtype=(np.int8, 3))
        for i in range(0, m):
            for j in range(0, n):
                new_img[i, j] = (arr[i, j], arr[i, j], arr[i, j])
        for i in range(0, int(self.numofcorners_entry.get())):
            print(i)
            new_img = self.red_placer(new_img, response)
        print(new_img)
        img = Image.fromarray(new_img, 'RGB')
        return img

    def test(self, arr):
        m = len(arr)
        n = len(arr[0])
        new_img = np.zeros((m, n), dtype=(np.int8, 3))
        for i in range(0, m):
            for j in range(0, n):
                new_img[i, j] = (arr[i, j], arr[i, j], arr[i, j])
        print(new_img)
        img = Image.fromarray(new_img, 'RGB')
        img.save('output.png')

p = Program()
p.main()





