from pathlib import Path
from pyPhasesRecordloader import Event
from pyPhasesRecordloader.recordLoaders.EDFRecordLoader import EDFRecordLoader

class RecordLoaderSleepEDF(EDFRecordLoader):    
    recordFileMap = {
        "SC4001": ("SC4001E0-PSG.edf", "SC4001EC-Hypnogram.edf"),
        "SC4002": ("SC4002E0-PSG.edf", "SC4002EC-Hypnogram.edf"),
        "SC4011": ("SC4011E0-PSG.edf", "SC4011EH-Hypnogram.edf"),
        "SC4012": ("SC4012E0-PSG.edf", "SC4012EC-Hypnogram.edf"),
        "SC4021": ("SC4021E0-PSG.edf", "SC4021EH-Hypnogram.edf"),
        "SC4022": ("SC4022E0-PSG.edf", "SC4022EJ-Hypnogram.edf"),
        "SC4031": ("SC4031E0-PSG.edf", "SC4031EC-Hypnogram.edf"),
        "SC4032": ("SC4032E0-PSG.edf", "SC4032EP-Hypnogram.edf"),
        "SC4041": ("SC4041E0-PSG.edf", "SC4041EC-Hypnogram.edf"),
        "SC4042": ("SC4042E0-PSG.edf", "SC4042EC-Hypnogram.edf"),
        "SC4051": ("SC4051E0-PSG.edf", "SC4051EC-Hypnogram.edf"),
        "SC4052": ("SC4052E0-PSG.edf", "SC4052EC-Hypnogram.edf"),
        "SC4061": ("SC4061E0-PSG.edf", "SC4061EC-Hypnogram.edf"),
        "SC4062": ("SC4062E0-PSG.edf", "SC4062EC-Hypnogram.edf"),
        "SC4071": ("SC4071E0-PSG.edf", "SC4071EC-Hypnogram.edf"),
        "SC4072": ("SC4072E0-PSG.edf", "SC4072EH-Hypnogram.edf"),
        "SC4081": ("SC4081E0-PSG.edf", "SC4081EC-Hypnogram.edf"),
        "SC4082": ("SC4082E0-PSG.edf", "SC4082EP-Hypnogram.edf"),
        "SC4091": ("SC4091E0-PSG.edf", "SC4091EC-Hypnogram.edf"),
        "SC4092": ("SC4092E0-PSG.edf", "SC4092EC-Hypnogram.edf"),
        "SC4101": ("SC4101E0-PSG.edf", "SC4101EC-Hypnogram.edf"),
        "SC4102": ("SC4102E0-PSG.edf", "SC4102EC-Hypnogram.edf"),
        "SC4111": ("SC4111E0-PSG.edf", "SC4111EC-Hypnogram.edf"),
        "SC4112": ("SC4112E0-PSG.edf", "SC4112EC-Hypnogram.edf"),
        "SC4121": ("SC4121E0-PSG.edf", "SC4121EC-Hypnogram.edf"),
        "SC4122": ("SC4122E0-PSG.edf", "SC4122EV-Hypnogram.edf"),
        "SC4131": ("SC4131E0-PSG.edf", "SC4131EC-Hypnogram.edf"),
        "SC4141": ("SC4141E0-PSG.edf", "SC4141EU-Hypnogram.edf"),
        "SC4142": ("SC4142E0-PSG.edf", "SC4142EU-Hypnogram.edf"),
        "SC4151": ("SC4151E0-PSG.edf", "SC4151EC-Hypnogram.edf"),
        "SC4152": ("SC4152E0-PSG.edf", "SC4152EC-Hypnogram.edf"),
        "SC4161": ("SC4161E0-PSG.edf", "SC4161EC-Hypnogram.edf"),
        "SC4162": ("SC4162E0-PSG.edf", "SC4162EC-Hypnogram.edf"),
        "SC4171": ("SC4171E0-PSG.edf", "SC4171EU-Hypnogram.edf"),
        "SC4172": ("SC4172E0-PSG.edf", "SC4172EC-Hypnogram.edf"),
        "SC4181": ("SC4181E0-PSG.edf", "SC4181EC-Hypnogram.edf"),
        "SC4182": ("SC4182E0-PSG.edf", "SC4182EC-Hypnogram.edf"),
        "SC4191": ("SC4191E0-PSG.edf", "SC4191EP-Hypnogram.edf"),
        "SC4192": ("SC4192E0-PSG.edf", "SC4192EV-Hypnogram.edf"),
        "SC4201": ("SC4201E0-PSG.edf", "SC4201EC-Hypnogram.edf"),
        "SC4202": ("SC4202E0-PSG.edf", "SC4202EC-Hypnogram.edf"),
        "SC4211": ("SC4211E0-PSG.edf", "SC4211EC-Hypnogram.edf"),
        "SC4212": ("SC4212E0-PSG.edf", "SC4212EC-Hypnogram.edf"),
        "SC4221": ("SC4221E0-PSG.edf", "SC4221EJ-Hypnogram.edf"),
        "SC4222": ("SC4222E0-PSG.edf", "SC4222EC-Hypnogram.edf"),
        "SC4231": ("SC4231E0-PSG.edf", "SC4231EJ-Hypnogram.edf"),
        "SC4232": ("SC4232E0-PSG.edf", "SC4232EV-Hypnogram.edf"),
        "SC4241": ("SC4241E0-PSG.edf", "SC4241EC-Hypnogram.edf"),
        "SC4242": ("SC4242E0-PSG.edf", "SC4242EA-Hypnogram.edf"),
        "SC4251": ("SC4251E0-PSG.edf", "SC4251EP-Hypnogram.edf"),
        "SC4252": ("SC4252E0-PSG.edf", "SC4252EU-Hypnogram.edf"),
        "SC4261": ("SC4261F0-PSG.edf", "SC4261FM-Hypnogram.edf"),
        "SC4262": ("SC4262F0-PSG.edf", "SC4262FC-Hypnogram.edf"),
        "SC4271": ("SC4271F0-PSG.edf", "SC4271FC-Hypnogram.edf"),
        "SC4272": ("SC4272F0-PSG.edf", "SC4272FM-Hypnogram.edf"),
        "SC4281": ("SC4281G0-PSG.edf", "SC4281GC-Hypnogram.edf"),
        "SC4282": ("SC4282G0-PSG.edf", "SC4282GC-Hypnogram.edf"),
        "SC4291": ("SC4291G0-PSG.edf", "SC4291GA-Hypnogram.edf"),
        "SC4292": ("SC4292G0-PSG.edf", "SC4292GC-Hypnogram.edf"),
        "SC4301": ("SC4301E0-PSG.edf", "SC4301EC-Hypnogram.edf"),
        "SC4302": ("SC4302E0-PSG.edf", "SC4302EV-Hypnogram.edf"),
        "SC4311": ("SC4311E0-PSG.edf", "SC4311EC-Hypnogram.edf"),
        "SC4312": ("SC4312E0-PSG.edf", "SC4312EM-Hypnogram.edf"),
        "SC4321": ("SC4321E0-PSG.edf", "SC4321EC-Hypnogram.edf"),
        "SC4322": ("SC4322E0-PSG.edf", "SC4322EC-Hypnogram.edf"),
        "SC4331": ("SC4331F0-PSG.edf", "SC4331FV-Hypnogram.edf"),
        "SC4332": ("SC4332F0-PSG.edf", "SC4332FC-Hypnogram.edf"),
        "SC4341": ("SC4341F0-PSG.edf", "SC4341FA-Hypnogram.edf"),
        "SC4342": ("SC4342F0-PSG.edf", "SC4342FA-Hypnogram.edf"),
        "SC4351": ("SC4351F0-PSG.edf", "SC4351FA-Hypnogram.edf"),
        "SC4352": ("SC4352F0-PSG.edf", "SC4352FV-Hypnogram.edf"),
        "SC4362": ("SC4362F0-PSG.edf", "SC4362FC-Hypnogram.edf"),
        "SC4371": ("SC4371F0-PSG.edf", "SC4371FA-Hypnogram.edf"),
        "SC4372": ("SC4372F0-PSG.edf", "SC4372FC-Hypnogram.edf"),
        "SC4381": ("SC4381F0-PSG.edf", "SC4381FC-Hypnogram.edf"),
        "SC4382": ("SC4382F0-PSG.edf", "SC4382FW-Hypnogram.edf"),
        "SC4401": ("SC4401E0-PSG.edf", "SC4401EC-Hypnogram.edf"),
        "SC4402": ("SC4402E0-PSG.edf", "SC4402EW-Hypnogram.edf"),
        "SC4411": ("SC4411E0-PSG.edf", "SC4411EJ-Hypnogram.edf"),
        "SC4412": ("SC4412E0-PSG.edf", "SC4412EM-Hypnogram.edf"),
        "SC4421": ("SC4421E0-PSG.edf", "SC4421EA-Hypnogram.edf"),
        "SC4422": ("SC4422E0-PSG.edf", "SC4422EA-Hypnogram.edf"),
        "SC4431": ("SC4431E0-PSG.edf", "SC4431EM-Hypnogram.edf"),
        "SC4432": ("SC4432E0-PSG.edf", "SC4432EM-Hypnogram.edf"),
        "SC4441": ("SC4441E0-PSG.edf", "SC4441EC-Hypnogram.edf"),
        "SC4442": ("SC4442E0-PSG.edf", "SC4442EV-Hypnogram.edf"),
        "SC4451": ("SC4451F0-PSG.edf", "SC4451FY-Hypnogram.edf"),
        "SC4452": ("SC4452F0-PSG.edf", "SC4452FW-Hypnogram.edf"),
        "SC4461": ("SC4461F0-PSG.edf", "SC4461FA-Hypnogram.edf"),
        "SC4462": ("SC4462F0-PSG.edf", "SC4462FJ-Hypnogram.edf"),
        "SC4471": ("SC4471F0-PSG.edf", "SC4471FA-Hypnogram.edf"),
        "SC4472": ("SC4472F0-PSG.edf", "SC4472FA-Hypnogram.edf"),
        "SC4481": ("SC4481F0-PSG.edf", "SC4481FV-Hypnogram.edf"),
        "SC4482": ("SC4482F0-PSG.edf", "SC4482FJ-Hypnogram.edf"),
        "SC4491": ("SC4491G0-PSG.edf", "SC4491GJ-Hypnogram.edf"),
        "SC4492": ("SC4492G0-PSG.edf", "SC4492GJ-Hypnogram.edf"),
        "SC4501": ("SC4501E0-PSG.edf", "SC4501EW-Hypnogram.edf"),
        "SC4502": ("SC4502E0-PSG.edf", "SC4502EM-Hypnogram.edf"),
        "SC4511": ("SC4511E0-PSG.edf", "SC4511EJ-Hypnogram.edf"),
        "SC4512": ("SC4512E0-PSG.edf", "SC4512EW-Hypnogram.edf"),
        "SC4522": ("SC4522E0-PSG.edf", "SC4522EM-Hypnogram.edf"),
        "SC4531": ("SC4531E0-PSG.edf", "SC4531EM-Hypnogram.edf"),
        "SC4532": ("SC4532E0-PSG.edf", "SC4532EV-Hypnogram.edf"),
        "SC4541": ("SC4541F0-PSG.edf", "SC4541FA-Hypnogram.edf"),
        "SC4542": ("SC4542F0-PSG.edf", "SC4542FW-Hypnogram.edf"),
        "SC4551": ("SC4551F0-PSG.edf", "SC4551FC-Hypnogram.edf"),
        "SC4552": ("SC4552F0-PSG.edf", "SC4552FW-Hypnogram.edf"),
        "SC4561": ("SC4561F0-PSG.edf", "SC4561FJ-Hypnogram.edf"),
        "SC4562": ("SC4562F0-PSG.edf", "SC4562FJ-Hypnogram.edf"),
        "SC4571": ("SC4571F0-PSG.edf", "SC4571FV-Hypnogram.edf"),
        "SC4572": ("SC4572F0-PSG.edf", "SC4572FC-Hypnogram.edf"),
        "SC4581": ("SC4581G0-PSG.edf", "SC4581GM-Hypnogram.edf"),
        "SC4582": ("SC4582G0-PSG.edf", "SC4582GP-Hypnogram.edf"),
        "SC4591": ("SC4591G0-PSG.edf", "SC4591GY-Hypnogram.edf"),
        "SC4592": ("SC4592G0-PSG.edf", "SC4592GY-Hypnogram.edf"),
        "SC4601": ("SC4601E0-PSG.edf", "SC4601EC-Hypnogram.edf"),
        "SC4602": ("SC4602E0-PSG.edf", "SC4602EJ-Hypnogram.edf"),
        "SC4611": ("SC4611E0-PSG.edf", "SC4611EG-Hypnogram.edf"),
        "SC4612": ("SC4612E0-PSG.edf", "SC4612EA-Hypnogram.edf"),
        "SC4621": ("SC4621E0-PSG.edf", "SC4621EV-Hypnogram.edf"),
        "SC4622": ("SC4622E0-PSG.edf", "SC4622EJ-Hypnogram.edf"),
        "SC4631": ("SC4631E0-PSG.edf", "SC4631EM-Hypnogram.edf"),
        "SC4632": ("SC4632E0-PSG.edf", "SC4632EA-Hypnogram.edf"),
        "SC4641": ("SC4641E0-PSG.edf", "SC4641EP-Hypnogram.edf"),
        "SC4642": ("SC4642E0-PSG.edf", "SC4642EP-Hypnogram.edf"),
        "SC4651": ("SC4651E0-PSG.edf", "SC4651EP-Hypnogram.edf"),
        "SC4652": ("SC4652E0-PSG.edf", "SC4652EG-Hypnogram.edf"),
        "SC4661": ("SC4661E0-PSG.edf", "SC4661EJ-Hypnogram.edf"),
        "SC4662": ("SC4662E0-PSG.edf", "SC4662EJ-Hypnogram.edf"),
        "SC4671": ("SC4671G0-PSG.edf", "SC4671GJ-Hypnogram.edf"),
        "SC4672": ("SC4672G0-PSG.edf", "SC4672GV-Hypnogram.edf"),
        "SC4701": ("SC4701E0-PSG.edf", "SC4701EC-Hypnogram.edf"),
        "SC4702": ("SC4702E0-PSG.edf", "SC4702EA-Hypnogram.edf"),
        "SC4711": ("SC4711E0-PSG.edf", "SC4711EC-Hypnogram.edf"),
        "SC4712": ("SC4712E0-PSG.edf", "SC4712EA-Hypnogram.edf"),
        "SC4721": ("SC4721E0-PSG.edf", "SC4721EC-Hypnogram.edf"),
        "SC4722": ("SC4722E0-PSG.edf", "SC4722EM-Hypnogram.edf"),
        "SC4731": ("SC4731E0-PSG.edf", "SC4731EM-Hypnogram.edf"),
        "SC4732": ("SC4732E0-PSG.edf", "SC4732EJ-Hypnogram.edf"),
        "SC4741": ("SC4741E0-PSG.edf", "SC4741EA-Hypnogram.edf"),
        "SC4742": ("SC4742E0-PSG.edf", "SC4742EC-Hypnogram.edf"),
        "SC4751": ("SC4751E0-PSG.edf", "SC4751EC-Hypnogram.edf"),
        "SC4752": ("SC4752E0-PSG.edf", "SC4752EM-Hypnogram.edf"),
        "SC4761": ("SC4761E0-PSG.edf", "SC4761EP-Hypnogram.edf"),
        "SC4762": ("SC4762E0-PSG.edf", "SC4762EG-Hypnogram.edf"),
        "SC4771": ("SC4771G0-PSG.edf", "SC4771GC-Hypnogram.edf"),
        "SC4772": ("SC4772G0-PSG.edf", "SC4772GC-Hypnogram.edf"),
        "SC4801": ("SC4801G0-PSG.edf", "SC4801GC-Hypnogram.edf"),
        "SC4802": ("SC4802G0-PSG.edf", "SC4802GV-Hypnogram.edf"),
        "SC4811": ("SC4811G0-PSG.edf", "SC4811GG-Hypnogram.edf"),
        "SC4812": ("SC4812G0-PSG.edf", "SC4812GV-Hypnogram.edf"),
        "SC4821": ("SC4821G0-PSG.edf", "SC4821GC-Hypnogram.edf"),
        "SC4822": ("SC4822G0-PSG.edf", "SC4822GC-Hypnogram.edf"),
        "ST7011": ("ST7011J0-PSG.edf", "ST7011JP-Hypnogram.edf"),
        "ST7012": ("ST7012J0-PSG.edf", "ST7012JP-Hypnogram.edf"),
        "ST7021": ("ST7021J0-PSG.edf", "ST7021JM-Hypnogram.edf"),
        "ST7022": ("ST7022J0-PSG.edf", "ST7022JM-Hypnogram.edf"),
        "ST7041": ("ST7041J0-PSG.edf", "ST7041JO-Hypnogram.edf"),
        "ST7042": ("ST7042J0-PSG.edf", "ST7042JO-Hypnogram.edf"),
        "ST7051": ("ST7051J0-PSG.edf", "ST7051JA-Hypnogram.edf"),
        "ST7052": ("ST7052J0-PSG.edf", "ST7052JA-Hypnogram.edf"),
        "ST7061": ("ST7061J0-PSG.edf", "ST7061JR-Hypnogram.edf"),
        "ST7062": ("ST7062J0-PSG.edf", "ST7062JR-Hypnogram.edf"),
        "ST7071": ("ST7071J0-PSG.edf", "ST7071JA-Hypnogram.edf"),
        "ST7072": ("ST7072J0-PSG.edf", "ST7072JA-Hypnogram.edf"),
        "ST7081": ("ST7081J0-PSG.edf", "ST7081JW-Hypnogram.edf"),
        "ST7082": ("ST7082J0-PSG.edf", "ST7082JW-Hypnogram.edf"),
        "ST7091": ("ST7091J0-PSG.edf", "ST7091JE-Hypnogram.edf"),
        "ST7092": ("ST7092J0-PSG.edf", "ST7092JE-Hypnogram.edf"),
        "ST7101": ("ST7101J0-PSG.edf", "ST7101JE-Hypnogram.edf"),
        "ST7102": ("ST7102J0-PSG.edf", "ST7102JE-Hypnogram.edf"),
        "ST7111": ("ST7111J0-PSG.edf", "ST7111JE-Hypnogram.edf"),
        "ST7112": ("ST7112J0-PSG.edf", "ST7112JE-Hypnogram.edf"),
        "ST7121": ("ST7121J0-PSG.edf", "ST7121JE-Hypnogram.edf"),
        "ST7122": ("ST7122J0-PSG.edf", "ST7122JE-Hypnogram.edf"),
        "ST7131": ("ST7131J0-PSG.edf", "ST7131JR-Hypnogram.edf"),
        "ST7132": ("ST7132J0-PSG.edf", "ST7132JR-Hypnogram.edf"),
        "ST7141": ("ST7141J0-PSG.edf", "ST7141JE-Hypnogram.edf"),
        "ST7142": ("ST7142J0-PSG.edf", "ST7142JE-Hypnogram.edf"),
        "ST7151": ("ST7151J0-PSG.edf", "ST7151JA-Hypnogram.edf"),
        "ST7152": ("ST7152J0-PSG.edf", "ST7152JA-Hypnogram.edf"),
        "ST7161": ("ST7161J0-PSG.edf", "ST7161JM-Hypnogram.edf"),
        "ST7162": ("ST7162J0-PSG.edf", "ST7162JM-Hypnogram.edf"),
        "ST7171": ("ST7171J0-PSG.edf", "ST7171JA-Hypnogram.edf"),
        "ST7172": ("ST7172J0-PSG.edf", "ST7172JA-Hypnogram.edf"),
        "ST7181": ("ST7181J0-PSG.edf", "ST7181JR-Hypnogram.edf"),
        "ST7182": ("ST7182J0-PSG.edf", "ST7182JR-Hypnogram.edf"),
        "ST7191": ("ST7191J0-PSG.edf", "ST7191JR-Hypnogram.edf"),
        "ST7192": ("ST7192J0-PSG.edf", "ST7192JR-Hypnogram.edf"),
        "ST7201": ("ST7201J0-PSG.edf", "ST7201JO-Hypnogram.edf"),
        "ST7202": ("ST7202J0-PSG.edf", "ST7202JO-Hypnogram.edf"),
        "ST7211": ("ST7211J0-PSG.edf", "ST7211JJ-Hypnogram.edf"),
        "ST7212": ("ST7212J0-PSG.edf", "ST7212JJ-Hypnogram.edf"),
        "ST7221": ("ST7221J0-PSG.edf", "ST7221JA-Hypnogram.edf"),
        "ST7222": ("ST7222J0-PSG.edf", "ST7222JA-Hypnogram.edf"),
        "ST7241": ("ST7241J0-PSG.edf", "ST7241JO-Hypnogram.edf"),
        "ST7242": ("ST7242J0-PSG.edf", "ST7242JO-Hypnogram.edf"),
    }

    eventMap = {
        "Sleep stage ?": "undefined",
        "Movement time": "undefined",
        "Sleep stage W": "W",
        "Sleep stage 1": "N1",
        "Sleep stage 2": "N2",
        "Sleep stage 3": "N3",
        "Sleep stage 4": "N3",
        "Sleep stage R": "R",
    }

    def getEventList(self, recordName, targetFrequency=1):
        filePath = self.getFilePathAnnotation(recordName)
        self.loadEdf(filePath, self.targetSignals)
        eventlist = [Event(self.eventMap[name], start=start, duration=duration) for start,duration,name in zip(*self.annotations)]
        return eventlist
    
    def existAnnotation(self, recordId):
        return Path(self.getFilePathAnnotation(recordId)).exists()
        
    def getFilePath(self, path):
        cohort = "sleep-cassette" if path[1] == "C" else "sleep-telemetry"
        return f"{self.filePath}/{cohort}/{path}"
    
    def getFilePathAnnotation(self, recordId):
        return self.getFilePath(self.recordFileMap[recordId][1])
    
    def getFilePathSignal(self, recordId):
        return self.getFilePath(self.recordFileMap[recordId][0])