from pathlib import Path
import shutil
from SleepHarmonizer.EDFWriter import EDFWriter
from SleepHarmonizer.PSGEventManager import PSGEventManager
from SleepHarmonizer.SignalPreprocessing import SignalPreprocessing
from pyPhases import Phase
from pyPhases.util import BatchProgress
from pyPhasesRecordloader import AnnotationNotFound, ChannelsNotPresent, RecordLoader, Event

class Export(Phase):
    skipExisting = False
    useMultiThreading = True
    useDigitalSignals = True
    prepareEventsFor = None

    def getAnnotationsFromDf(self, df):
        return df.T.to_dict().values()

    def getLightEvents(self, lightOff, lightOn):
        lightEvents = []
        if lightOn is not None:
            lightEvents.append(Event.fromdict({"start": lightOn, "duration": 0, "name": "Licht An"}))
        if lightOff > 0:
            lightEvents.append(Event.fromdict({"start": lightOff, "duration": 0, "name": "Licht Aus"}))

        return lightEvents

    def exportRecord(self, recordName):
        exportFilename = self.exportPath + "/" + recordName
        path = Path(exportFilename  + ".edf")
        if path.is_file() and self.skipExisting:
            self.log(f"Record {recordName} skipped because it already exists")
            return

        rl = self.recordLoader

        try:
            psgSignal, events = self.recordLoader.loadRecord(recordName)
            em = PSGEventManager()

            events += self.getLightEvents(rl.lightOff, rl.lightOn)
            events = em.getDataframeFromEvents(events)

            # tailor psg signal
            keepSignals = self.getConfig("export.channels", None)
            psgSignal.reduceSignals(keepSignals)

            allowEventsConfig = self.getConfig("export.annotations", None)
            allowEvents = []
            if allowEventsConfig is not None:
                for ev in allowEventsConfig:
                    if ev in em.eventGroups and ev :
                        allowEvents += em.eventGroups[ev]
                    else:
                        allowEvents.append(ev)
                        
                events = events.query("name in @allowEvents")

            annotations = self.getAnnotationsFromDf(events)

            # signal preprocessing / harmonizing
            targetFrequency = self.preProcessingConfig["targetFrequency"]
            signalProcessor = SignalPreprocessing(self.preProcessingConfig)
            psgSignal.targetFrequency = targetFrequency
            signalProcessor.preprocessingSignal(psgSignal)

            # write the harmonized signal into an edf
            rw = EDFWriter()
            tmpFilename = exportFilename + "_tmp"
            rw.writeSignals(tmpFilename, psgSignal.signals, annotations, signalIsDigital=self.useDigitalSignals)
            shutil.move(tmpFilename, exportFilename + ".edf")

        except AnnotationNotFound as e:
            self.logError(f"not all required annotation exist for {recordName}")
        except ChannelsNotPresent as e:
            self.logError(f"not all required channels exist for {recordName}: {e.channels}")

    def main(self):
        self.events = None
        self.annotations = []

        self.exportPath = self.getConfig("export-path")
        if not Path(self.exportPath).exists():
            Path(self.exportPath).mkdir(parents=True, exist_ok=True)


        recordId = self.getConfig("recordId", False)
        if not recordId:
            recordIds = self.project.getData("allDBRecordIds", list)
            # flatten grouped record ids
            recordIds = [recordId for groupedRecords in recordIds.values() for recordId in groupedRecords]
        else:
            self.useMultiThreading = False
            recordIds = [recordId]

        bp = BatchProgress()
        bp.asynchronous = True
        bp.useMultiThreading = self.useMultiThreading

        self.recordLoader = RecordLoader.get()
        self.recordLoader.useDigitalSignals = self.useDigitalSignals

        self.preProcessingConfig = self.getConfig("preprocessing")
        bp.start(self.exportRecord, batchList=recordIds)
