#!/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/>.
"""

import numpy as np
import os
from PyQt5 import QtWidgets, QtCore

class LoadTrueMatchResults(QtWidgets.QWidget):
    def __init__(self, particleContainer, analysisParent):
        super(LoadTrueMatchResults, self).__init__()
        self.setGeometry(400, 400, 200, 300)
        self.setWindowTitle('Get Truematch Results')
        self.layout = QtWidgets.QGridLayout()
        self.setLayout(self.layout)

        self.particleContainer = particleContainer
        self.analysisParent = analysisParent
        self.analysisParent.setDisabled(True)
        self.trueMatchResults = None
        self.polymertypes = None
        self.additives = None
        self.hqis = None
        self.addhqis = None
        
        self.loadBtn = QtWidgets.QPushButton('LoadTrueMatchResults')
        self.loadBtn.resize(self.loadBtn.sizeHint()) 
        self.loadBtn.clicked.connect(self.getImportFiles)
        
        self.reviewGroup = QtWidgets.QGroupBox('Review Changes')
        self.reviewGroup.setDisabled(True)
        reviewLayout = QtWidgets.QVBoxLayout()
        
        self.otherBtn = QtWidgets.QPushButton('manual overwrite')
        self.otherBtn.clicked.connect(self.show3FlagsReview)
        reviewLayout.addWidget(self.otherBtn)
        reviewLayout.addStretch()
        self.reviewGroup.setLayout(reviewLayout)

        self.layout.addWidget(self.loadBtn, 0, 0)
        self.layout.addWidget(self.reviewGroup, 2, 0)
        
        self.manualPolymers = {}
        self.manualAdditives = {}
        self.editEntryWindow = None
        self.sortWindow = None

        self.numFlagForTakeSpectrum = 1
        self.numFlagForSetUnknown = 2
        self.numFlagForPrompt = 3
        self.show()
    
    def show3FlagsReview(self):
        self.editEntryWindow = ModifyManualEdits(self, self.manualPolymers, self.manualAdditives)        
        self.setDisabled(True)
        self.editEntryWindow.show()

    def getImportFiles(self):
        dsetpath = self.analysisParent.dataset.path
        self.fnames =QtWidgets.QFileDialog.getOpenFileNames(self, 'Select TrueMatch result file', dsetpath, 'text file (*.txt)')[0]
        self.trueMatchResults = []
        if len(self.fnames) > 1:
            self.sortWindow = SortImportFiles(self, self.fnames)
        else:
            self.runCalculations()
            
    def runCalculations(self):
        allResults = []
        numSpectra = 0
        for fileindex, fname in enumerate(self.fnames):
            with open(fname) as file:
                currentResults = []
                for lineIndex, line in enumerate(file):
                    if lineIndex == 0:
                        line = line.strip().split(';')[:-1]         #disregard the last entry of each line, as each line ends with ; <- that produces an empty entry...
                        assert line[0] == 'Search Spectrum Name', 'Assertion Error in loadTrueMatchResults'
                        #detect, whether one- or multicomponent-search was done
                        if line[-1] == 'IsMarked':
                            numhits = np.int(line[-2].split(' ')[-1])
                            numcomps = 1
                        else:
                            numhits = np.int(line[-1].split(' ')[-1])
                            numcomps = np.int(line[-1].split(' ')[-3])
                    else:
                        currentResults.append(line)
                        numSpectra += 1
                        
                allResults.append(currentResults)

        if numSpectra > 0:
            print('{} components, {} hits per sample, {} spectra'.format(numcomps, numhits, numSpectra))        
        
        self.polymertypes = [None]*numSpectra
        self.hqis = [None]*numSpectra
        if numcomps > 1:
            self.additives = [None]*numSpectra
            self.addhqis = [None]*numSpectra
            
        addedIndices = []
        numSpectraInPreviousBatches = 0
        
        for resBatchIndex, resBatch in enumerate(allResults):
            print('spectra in batch:', len(resBatch))
            for result in resBatch:
                specIndex, polymertype, additive, hqi, addhqi = self.interpretEntry(resBatchIndex, numSpectraInPreviousBatches, result, numhits, numcomps)
                self.polymertypes[specIndex] = polymertype
                self.hqis[specIndex] = hqi
                if numcomps > 1:
                    self.additives[specIndex] = additive
                    self.addhqis[specIndex] = addhqi
                    
                addedIndices.append(specIndex)               
        
            numSpectraInPreviousBatches = len(addedIndices)
            print(numSpectraInPreviousBatches)
    
        assert sorted(addedIndices) == list(range(numSpectra)), f'mismatch in sorted indices in loadResults\bAdded {len(addedIndices)} assignments to {numSpectra} spectra'
        assert not None in self.polymertypes, 'wrong assignments in loadResults'
        assert not None in self.hqis, 'wrong assignments in loadResults'

        self.particleContainer.applyAssignmentListToParticleMeasurements(self.polymertypes)
        self.particleContainer.applyHQIListToParticleMeasurements(self.hqis)
        
        if len(self.manualPolymers) > 0:
            self.reviewGroup.setDisabled(False)
        
        self.analysisParent.applyHQIThresholdToResults()
        self.loadBtn.setText('Data loaded')
        QtWidgets.QMessageBox.about(self, 'Success', 'Results were imported successfully.\nYou can close this window or modify manual edits.')
    
    def interpretEntry(self, resBatchIndex, numSpectraInPreviousBatches, result, numhits, numcomps):
        result = result.split(';')
        specIndex, polymertype, additive, hqi, addhqi = None, None, None, None, 0   #assign defaults
        
        #specNameFormat is: SampleName_000_Spec.Data 1 (SpecIndex)
        specName = result[0]
        curSpecIndex = self.getSpecIndexFromName(specName)
        
        specIndex = numSpectraInPreviousBatches + curSpecIndex
        
        #find yes-flags
        flags = np.where(np.array(result) == 'yes')[0]
        if len(flags) == 0:
            #take highest HQI entry
            if numcomps == 1:
                polymertype = result[2]
                hqi = result[1]
            else:
                polymertype = result[5]
                hqi = result[1]
                additive = result[7]
                addhqi = result[6]

        elif len(flags) == 1:
            #exactly one flag was placed, take this entry
            if numcomps == 1:
                polymertype = result[int(flags[0])-1]
                hqi = 100
            else:
                polymertype = result[int(flags[0])+2]
                additive = result[int(flags[0])+4]
                hqi = 100
                addhqi = 100
            
        elif len(flags) == 2:
            polymertype = 'unknown'
            hqi = 0
            if numcomps > 1:
                additive = 'none'
                addhqi = 0
                
        elif len(flags) == 3:
            hqi = 100
            if (curSpecIndex, resBatchIndex) not in self.manualPolymers:
                polymertype, ok = QtWidgets.QInputDialog.getText(self, 'Name of main component', 'Spectrum at index {} of result file {} is:'.format(curSpecIndex, resBatchIndex+1))
                self.manualPolymers[(curSpecIndex, resBatchIndex)] = polymertype
            else:
                polymertype = self.manualPolymers[(curSpecIndex, resBatchIndex)]
            if numcomps > 1:
                addhqi = 100
                if (curSpecIndex, resBatchIndex) not in self.manualAdditives:
                    additive, ok = QtWidgets.QInputDialog.getText(self, 'Name of additive', 'Additive at index {} of result file {} is:'.format(curSpecIndex, resBatchIndex+1))
                    self.manualAdditives[(curSpecIndex, resBatchIndex)] = additive
                else:
                    additive = self.manualAdditives[(curSpecIndex, resBatchIndex)]
        else: 
            QtWidgets.QMessageBox.about(self, 'Error!', 'No rule for {} flags, found at spectrum index {}'.format(len(flags), specIndex))
            
        return specIndex, polymertype, additive, float(hqi), float(addhqi)
    
    def getSpecIndexFromName(self, specName):
        def specNameHasBrackets():
            if specName.find('(') != -1 and specName.find(')') != -1:
                return True
            else:
                return False
        
        def isNumber(string):
            try:
                int(string)
                return True
            except:
                return False
        
        index = ''
        if specNameHasBrackets():
            for i in reversed(range(len(specName))):
                if specName[i] == ')':
                    continue
                if specName[i] != '(':
                    index = specName[i] + index
                else:
                    break
            index = int(index)
        else:
            for i in reversed(range(len(specName))):
                if isNumber(specName[i]):
                    index = specName[i] + index
                else:
                    break
            index = int(index) - 1
        return index
    
    def closeEvent(self, event):
        for win in [self.editEntryWindow, self.sortWindow]:
            if win is not None:
                win.close()
                
        self.analysisParent.setEnabled(True)
        event.accept()


class SortImportFiles(QtWidgets.QWidget):
    def __init__(self, loadResultsParent, importFiles):
        super(SortImportFiles, self).__init__()
        self.setWindowTitle('Sort Import Files')
        self.loadResultsParent = loadResultsParent
        self.loadResultsParent.setDisabled(True)
        self.importFiles = importFiles
        self.spinBoxes = []
        layout = QtWidgets.QGridLayout()
        self.setLayout(layout)
        
        label1 = QtWidgets.QLabel('File name')
        label1.setMinimumWidth(150)
        label2 = QtWidgets.QLabel('Order')
        label2.setMinimumWidth(50)
        
        layout.addWidget(label1, 0, 0, QtCore.Qt.AlignLeft)
        layout.addWidget(label2, 0, 1, QtCore.Qt.AlignRight)
        
        for fileIndex, importFile in enumerate(self.importFiles):
            layout.addWidget(QtWidgets.QLabel(os.path.basename(importFile)), 1+fileIndex, 0, QtCore.Qt.AlignLeft)
            newSpinBox = QtWidgets.QSpinBox()
            newSpinBox.setMaximumWidth(50)
            newSpinBox.setMinimum(1)
            newSpinBox.setMaximum(len(self.importFiles))
            newSpinBox.setValue(fileIndex+1)
            layout.addWidget(newSpinBox, 1+fileIndex, 1, QtCore.Qt.AlignRight)
            self.spinBoxes.append(newSpinBox)
        
        self.acceptBtn = QtWidgets.QPushButton('Accept and Close')
        self.acceptBtn.clicked.connect(self.sortResultFiles)
        layout.addWidget(self.acceptBtn, fileIndex+2, 0, QtCore.Qt.AlignLeft)
        
        self.show()
    
    def sortResultFiles(self):
        order = [box.value()-1 for box in self.spinBoxes]
        if list(np.unique(order)) != list(range(len(self.spinBoxes))):
            QtWidgets.QMessageBox.critical(self, 'Error', 'Invalid order indicated')
            return
        else:
            newFiles = []
            for index in order:
                newFiles.append(self.importFiles[index])
            
            self.loadResultsParent.fnames = newFiles
            self.loadResultsParent.runCalculations()
            self.loadResultsParent.setDisabled(False)
            self.close()
            
            
class ModifyManualEdits(QtWidgets.QWidget):
    def __init__(self, parentWindow, polymerEdits, additiveEdits = None):
        super(ModifyManualEdits, self).__init__()
        self.setWindowTitle('Edit Manual Changes')
        self.setGeometry(200, 200, 800, 500)
        layout = QtWidgets.QVBoxLayout()
        self.setLayout(layout)
        
        self.parent = parentWindow
        self.polymerEdits = polymerEdits
        self.additiveEdits = additiveEdits
        self.labels = []
        self.polymerLineEdits = []
        self.additiveLineEdits = []
        
        groupBox = QtWidgets.QGroupBox('Manually entered edits')
        groupLayout = QtWidgets.QGridLayout()

        #create Labels and LineEdits        
        for curSpecIndex, resBatchIndex in self.polymerEdits:
            self.labels.append(QtWidgets.QLabel(f'Index of Spectrum: {curSpecIndex}, in file number: {resBatchIndex+1}'))
            self.polymerLineEdits.append(QtWidgets.QLineEdit())
            self.polymerLineEdits[-1].setText(self.polymerEdits[(curSpecIndex, resBatchIndex)])
            if len(self.additiveEdits) > 0:
               self.additiveLineEdits.append(QtWidgets.QLineEdit()) 
               self.additiveLineEdits[-1].setText(self.additiveEdits[(curSpecIndex, resBatchIndex)])
        
        for i in range(len(self.labels)):
            groupLayout.addWidget(self.labels[i], i, 0)
            groupLayout.addWidget(self.polymerLineEdits[i], i, 1)
            if len(self.additiveEdits) > 0:
                groupLayout.addWidget(self.additiveLineEdits[i], i, 2)
        
        groupBox.setLayout(groupLayout)
        
        scrollarea = QtWidgets.QScrollArea()
        scrollarea.setWidget(groupBox)
        scrollarea.setMaximumHeight(600)
        layout.addWidget(scrollarea)
        
        saveAndCloseBtn = QtWidgets.QPushButton('Save and Exit')
        saveAndCloseBtn.clicked.connect(self.saveAndClose)
        layout.addWidget(saveAndCloseBtn)

    def saveAndClose(self):
        #read out LineEdits:
        for index, key in enumerate(self.polymerEdits):
            self.polymerEdits[key] = self.polymerLineEdits[index].text()
            if len(self.additiveEdits) > 0:
                self.additiveEdits[key] = self.additiveLineEdits[index].text()
        
        #copy values to parent Window
        self.parent.manualPolymers, self.parent.manualAdditives = self.polymerEdits, self.additiveEdits
        self.parent.runCalculations()
        self.close()