#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GEPARD - Gepard-Enabled PARticle Detection
Copyright (C) 2018  Lars Bittrich and Josef Brandt, Leibniz-Institut für 
Polymerforschung Dresden e. V. <bittrich-lars@ipfdd.de>    

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program, see COPYING.  
If not, see <https://www.gnu.org/licenses/>.
"""
from PyQt5 import QtWidgets, QtCore, QtGui
import numpy as np
import cv2

class ParticlePainter(QtWidgets.QGraphicsItem):
    def __init__(self, editorParent, img, topLeft):
        super(ParticlePainter, self).__init__()
        self.editorParent = editorParent
        self.viewparent = self.editorParent.viewparent
        self.setZValue(5)
        self.topLeft = topLeft
        self.img = img
        self.pixmap = None
        
        self.mousePos = None
        self.painting = False
        self.erasing = False
        
        self.minRadius = 5
        self.maxRadius = 500
        self.radius = 30
        self.brect = QtCore.QRectF(0,0,1,1)
        
        self.setPixmap()
        self.setBrect()

    def setPixmap(self):
        img = self.img.repeat(3).reshape(self.img.shape[0], self.img.shape[1], 3)
        height, width, channel = img.shape
        bytesPerLine = 3 * width
        self.pixmap = QtGui.QPixmap()
        self.pixmap.convertFromImage(QtGui.QImage(img.data, width, height, bytesPerLine, QtGui.QImage.Format_RGB888))

    def setBrect(self):
        x0 = self.topLeft[1]
        y0 = self.topLeft[0]
        x1 = self.topLeft[1] + self.img.shape[1]
        y1 = self.topLeft[0] + self.img.shape[0]
        self.brect.setCoords(x0,y0,x1,y1)
    
    def boundingRect(self):
        return self.brect
    
    def mousePressEvent(self, event):
        self.mousePos = self.viewparent.mapToScene(event.pos())
        self.erasing = self.painting = False
        if event.modifiers()==QtCore.Qt.ControlModifier:
            self.painting = True
        elif event.modifiers()==QtCore.Qt.ShiftModifier:
            self.erasing = True
            
        if self.painting or self.erasing:
            drawPos = self.viewparent.mapToScene(event.pos()) - self.brect.topLeft()
            self.drawParticle(drawPos)
            
    def mouseMoveEvent(self, event):
        self.mousePos = self.viewparent.mapToScene(event.pos())
        self.update()
        if self.painting or self.erasing:
            drawPos = self.viewparent.mapToScene(event.pos()) - self.brect.topLeft()
            self.drawParticle(drawPos)
        
    def wheelEvent(self, event):
        if event.angleDelta().y() < 0:
            self.radius = int(np.clip(self.radius+self.radius*0.1, self.minRadius, self.maxRadius))
        else:
            self.radius = int(np.clip(self.radius-self.radius*0.1, self.minRadius, self.maxRadius))
            
        self.update()
        
    def mouseReleaseEvent(self, event):
        self.erasing = self.painting = False
            
    def keyPressEvent(self, event):
        if event.key() == QtCore.Qt.Key_Escape:
            self.editorParent.destroyParticlePainter()
        elif event.key() in [QtCore.Qt.Key_Return, QtCore.Qt.Key_Enter]:
            self.editorParent.acceptPaintedResult()
    
    def drawParticle(self, pixelPos):
        x_min, x_max = round(pixelPos.x()-self.radius), round(pixelPos.x()+self.radius)
        y_min, y_max = round(pixelPos.y()-self.radius), round(pixelPos.y()+self.radius)
        
        x_shift = y_shift = int(0)
        
        if x_min < 0:
            x_shift = int(x_min)
            self.topLeft[1] -= abs(x_min)
        elif x_max >= self.img.shape[1]:
            x_shift = int(x_max + 1 - self.img.shape[1])
        
        if y_min < 0:
            self.topLeft[0] -= abs(y_min)
            y_shift = int(y_min)
        elif y_max >= self.img.shape[0]:
            y_shift = int(y_max + 1 - self.img.shape[0])
        
        if x_shift != 0 or y_shift != 0:
            x_range = int(self.img.shape[1] + abs(x_shift))
            y_range = int(self.img.shape[0] + abs(y_shift))
            newImg = np.zeros((y_range, x_range))
            if x_shift < 0:
                if y_shift < 0:
                    newImg[abs(y_shift):, abs(x_shift):] = self.img
                elif y_shift == 0:
                    newImg[:, abs(x_shift):] = self.img
                elif y_shift > 0:
                    newImg[:self.img.shape[0], abs(x_shift):] = self.img
        
            elif x_shift == 0:
                if y_shift < 0:
                    newImg[abs(y_shift):, :] = self.img
                elif y_shift == 0:
                    newImg[:, :] = self.img
                elif y_shift > 0:
                    newImg[:self.img.shape[0], :] = self.img
                
            elif x_shift > 0:
                if y_shift < 0:
                    newImg[abs(y_shift):, :self.img.shape[1]] = self.img
                elif y_shift == 0:
                    newImg[:, :self.img.shape[1]] = self.img
                elif y_shift > 0:
                    newImg[:self.img.shape[0], :self.img.shape[1]] = self.img
            
            self.img = np.uint8(newImg)
            self.setBrect()
        
        center = (int(pixelPos.x()), int(pixelPos.y()))
        
        if self.painting:
            cv2.circle(self.img, center, self.radius, 255, -1)
        elif self.erasing:
            cv2.circle(self.img, center, self.radius, 0, -1)
            
        self.setPixmap()
        self.update()

    def paint(self, painter, option, widget):   
        painter.setPen(QtCore.Qt.white)
        painter.drawRect(self.brect)
        
        if self.mousePos is not None:
            p = [self.mousePos.x(), self.mousePos.y(), self.radius]
            painter.drawEllipse(p[0]-p[2], p[1]-p[2], 2*p[2], 2*p[2])            
        
        if self.pixmap is not None:
            painter.setOpacity(0.4)
            painter.drawPixmap(self.topLeft[1], self.topLeft[0], self.pixmap)