# -*- 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/>.
"""
import os
from PyQt5 import QtCore, QtWidgets
from .zeissxml import ZeissHandler, make_parser
from .opticalscan import PointCoordinates
from .helperfunctions import cv2imread_fix, cv2imwrite_fix
from .ramancom.ramancontrol import defaultPath
from .dataset import DataSet
from scipy.optimize import least_squares
from itertools import permutations
import cv2
import numpy as np

class ZeissImporter(QtWidgets.QDialog):
    def __init__(self, fname, ramanctrl, parent=None):
        super().__init__(parent)
        
        if not ramanctrl.connect() or not self.readImportData(fname):
            msg = QtWidgets.QMessageBox()
            msg.setText('Connection failed! Please enable remote control.')
            msg.exec()
            self.validimport = False
            return
        else:
            self.validimport = True
            
        self.ramanctrl = ramanctrl
        
        vbox = QtWidgets.QVBoxLayout()
        pointgroup = QtWidgets.QGroupBox('Marker coordinates at Raman spot [µm]', self)
        self.points = PointCoordinates(len(self.markers), self.ramanctrl, self,
                                       names = [m.name for m in self.markers])
        self.points.pimageOnly.setVisible(False)
        pointgroup.setLayout(self.points)
        self.points.readPoint.connect(self.takePoint)
        
        self.pconvert = QtWidgets.QPushButton('Convert', self)
        self.pexit = QtWidgets.QPushButton('Cancel', self)
        self.pconvert.released.connect(self.convert)
        self.pexit.released.connect(self.reject)
        self.pconvert.setEnabled(False)
        
        self.xinvert = QtWidgets.QCheckBox('Invert x-axis')
        self.yinvert = QtWidgets.QCheckBox('Invert y-axis')
        self.zinvert = QtWidgets.QCheckBox('Invert z-axis')
        
        btnLayout = QtWidgets.QHBoxLayout()
        btnLayout.addStretch()
        btnLayout.addWidget(self.pconvert)
        btnLayout.addWidget(self.pexit)
        
        label = QtWidgets.QLabel("Z-Image blur radius", self)
        self.blurspinbox = QtWidgets.QSpinBox(self)
        self.blurspinbox.setMinimum(3)
        self.blurspinbox.setMaximum(99)
        self.blurspinbox.setSingleStep(2)
        self.blurspinbox.setValue(5)
        blurlayout = QtWidgets.QHBoxLayout()
        blurlayout.addWidget(label)
        blurlayout.addWidget(self.blurspinbox)
        blurlayout.addStretch()
        
        vbox.addWidget(pointgroup)
        vbox.addLayout(blurlayout)
        vbox.addWidget(self.xinvert)
        vbox.addWidget(self.yinvert)
        vbox.addWidget(self.zinvert)
        vbox.addLayout(btnLayout)
        self.setLayout(vbox)
        
    def readImportData(self, fname):
        path = os.path.split(fname)[0]
        self.edfimgname, self.zmapimgname, xmlname = '', '', ''
        for name in os.listdir(path):
            if name.lower().endswith('_meta.xml'):
                xmlname = os.path.join(path, name)
            elif name.lower().endswith('_c1.tif'):
                self.edfimgname = os.path.join(path, name)
            elif name.lower().endswith('_c2.tif'):
                self.zmapimgname = os.path.join(path, name)
                
        errmsges = []
        if not os.path.exists(self.zmapimgname):
            errmsges.append('Depth map image not found: NAME_c2.tif')
        if not os.path.exists(self.edfimgname):
            errmsges.append('EDF image not found: NAME_c1.tif')
        if not os.path.exists(xmlname):
            errmsges.append('XML metadata not found: NAME_meta.xml')
        else:
            parser = make_parser()
            z = ZeissHandler()
            parser.setContentHandler(z)
            parser.parse(xmlname)
            
            if len(z.markers)<3:
                errmsges.append('Fewer than 3 markers found to adjust coordinates!')
            if None in [z.region.centerx, z.region.centery, 
                        z.region.width, z.region.height]:
                errmsges.append('Image dimensions incomplete or missing!')
            if None in [z.zrange.z0, z.zrange.zn, z.zrange.dz]:
                errmsges.append('ZStack information missing or incomplete!')
    
        if len(errmsges)>0:
            QtWidgets.QMessageBox.critical(self, 'Error!',
                                '\n'.join(errmsges),
                                QtWidgets.QMessageBox.Ok, 
                                QtWidgets.QMessageBox.Ok)
            return False
        
        self.region = z.region
        self.zrange = z.zrange
        self.markers = z.markers
        print(self.region)
        print(self.zrange)
        print(self.markers, flush=True)
        return True
        
    @QtCore.pyqtSlot(float, float, float) 
    def takePoint(self, x, y, z):
        points = self.points.getPoints()
        if len(points)>=3:
            self.pconvert.setEnabled(True)
     
    @QtCore.pyqtSlot()
    def convert(self):
        T, pc, zpc, accept = self.getTransform()
        if accept:
            fname = QtWidgets.QFileDialog.getSaveFileName(self,
                               'Create New GEPARD Project', defaultPath, '*.pkl')[0]
            if fname=='':
                return
            dataset = DataSet(fname, newProject=True)
            self.convertZimg(dataset, T, pc, zpc)
            self.convertImage(dataset)
            dataset.save()
            self.gepardname = dataset.fname
            self.accept()
        
    def convertImage(self, dataset):
        img = cv2imread_fix(self.edfimgname)
        cv2imwrite_fix(dataset.getImageName(), img)
    
    def convertZimg(self, dataset, T, pc, zpc):
        N = int(round(abs(self.zrange.zn-self.zrange.z0)/self.zrange.dz))
        z0, zn = self.zrange.z0, self.zrange.zn
        if zn<z0: 
            zn, z0 = z0, zn
        # zeiss z axis has large values at the surface and smaller for large particles
        # that is why we need to invert the zpositions: zimg black corresponds to zn
        # zimg white corresponds to z0
        dataset.zpositions = np.linspace(z0, zn, N)[::-1]-zpc[2]
        dataset.heightmap = np.zeros(3)
        dataset.signy = 1.
        zimg = cv2imread_fix(self.zmapimgname, cv2.IMREAD_GRAYSCALE)
        print("zimg shape:", zimg.shape, flush=True)
        radius = self.blurspinbox.value()
        blur = cv2.GaussianBlur(zimg, (radius, radius), 0)

        pshift = self.ramanctrl.getRamanPositionShift()
        dataset.pshift = pshift
        pixelscale = self.region.scalex*1.e6
        
        # use input image as single image aquired in one shot
        dataset.imagedim_df = (zimg.shape[1]*pixelscale, 
                               zimg.shape[0]*pixelscale, 0.0)
        dataset.pixelscale_df = pixelscale
        
        dataset.imagedim_bf = (zimg.shape[1]*pixelscale, 
                               zimg.shape[0]*pixelscale, 0.0)
        dataset.pixelscale_bf = pixelscale
        print('scale:', pixelscale)
        print('dim df:', dataset.imagedim_df)
        print('dim bf:', dataset.imagedim_bf)
        
        dataset.coordinatetransform = T, pc
        
        # set image center as reference point (assume just one tile) 
        # in data set (use Zeiss coordinates)
        p0center = np.array([self.region.centerx, self.region.centery]) - zpc[:2]
        dataset.readin = False
        dataset.lastpos = p0center
        dataset.maxdim = p0center
        
        zmin, zmax = dataset.zpositions.min(), dataset.zpositions.max()
        blur = (blur)*(255.)
        blur[blur>255.] = 255.
        blur[np.isnan(blur)] = 0.
        blur = np.uint8(blur)
        
        zimgname = dataset.getZvalImageName()
        cv2imwrite_fix(zimgname, blur)
        dataset.zvalimg = "saved"
    
    def getTransform(self):
        points = self.points.getPoints()
        pshift = self.ramanctrl.getRamanPositionShift()
        points[:,0] -= pshift[0]
        points[:,1] -= pshift[1]
        Parity = np.mat(np.diag([-1. if self.xinvert.isChecked() else 1.,
                          -1. if self.yinvert.isChecked() else 1.,
                          -1. if self.zinvert.isChecked() else 1.]))
        zpoints = np.array([m.getPos() for m in self.markers], dtype=np.double)
        pc = points.mean(axis=0)
        zpc = zpoints.mean(axis=0)
        
        points -= pc[np.newaxis,:]
        zpoints -= zpc[np.newaxis,:]
        
        def getRotMat(angles):
            c1, s1 = np.cos(angles[0]), np.sin(angles[0])
            c2, s2 = np.cos(angles[1]), np.sin(angles[1])
            c3, s3 = np.cos(angles[2]), np.sin(angles[2])
            return np.mat([[c1*c3-s1*c2*s3, -c1*s3-s1*c2*c3, s1*s2],
                           [s1*c3+c1*c2*s3, -s1*s3+c1*c2*c3, -c1*s2],
                           [s1*s3, s2*c3, c2]])
        
        # find the transformation matrix with best fit for small angles in 
        # [-45°,45°] for all permutation of markers
        permbest = None
        pointsbest = None
        ppoints = points[:,:].copy()
        
        def err(angles_shift):
            T = (getRotMat(angles_shift[:3]).T*Parity).A
            return (np.dot(zpoints, T) - angles_shift[np.newaxis,3:] \
                    - ppoints).ravel()

        angle = np.zeros(3)
        opt = least_squares(err, np.concatenate((angle, np.zeros(3))),
                            bounds=(np.array([-np.pi/4]*3+[-np.inf]*3),
                                    np.array([np.pi/4]*3+[np.inf]*3)),
                            method='dogbox')
        permbest = opt  
        pointsbest = ppoints
        
        optangles = permbest.x[:3]
        shift = permbest.x[3:]
        T = (getRotMat(optangles).T*Parity).A
        e = (np.dot(zpoints, T)-shift[np.newaxis,:]-pointsbest)
        print("Transformation angles:", optangles, flush=True)
        print("Transformation shift:", shift, flush=True)
        print("Transformation err:", e, flush=True)
        d = np.linalg.norm(e, axis=1)
        accept = True
        if np.any(d>1.):
            ret = QtWidgets.QMessageBox.warning(self, 'Warning!',
                                f'Transformation residuals are large:{d}',
                                QtWidgets.QMessageBox.Ok|QtWidgets.QMessageBox.Cancel, 
                                QtWidgets.QMessageBox.Ok)
            if ret==QtWidgets.QMessageBox.Cancel:
                accept = False
        return T, pc-shift, zpc, accept