import statUtils
import os
import radiomics
import SimpleITK
import sys
import json
import numpy

def main(parFile='../templates/statistics.json'):


    featureExtractor=radiomics.featureextractor.RadiomicsFeatureExtractor
 
    setup=statUtils.loadSetup(parFile)
    rFile='radiomics.json'

    forceLiver=setup.get('forceLiver',False)
    forceSUVMax=setup.get('forceSUVMax',False)
    forceLiver15=setup.get('forceLiver15',False)
    forceGlobal=setup.get('forceGlobal',False)

    doUpload=setup.get('doUpload',True)

    #update threshold values if needed
    with open(rFile,'w') as f:
        f.write(json.dumps(setup['radiomics']))
    setup['db'],setup['fb']=statUtils.connectDB('onko-nix')
    users=statUtils.getUsers(setup['db'],setup['project'])
    p=setup['project']
    qFilter=[]
    try:
        vList=';'.join(setup['participants'])    
        qFilter.append({'variable':'ParticipantId','value':vList,'oper':'in'})
    except KeyError:
        pass

    try:
        vList=';'.join(setup['visits'])
        qFilter.append({'variable':'visitCode','value':vList,'oper':'in'})
    except KeyError:
        pass

    ds=setup['db'].selectRows(p,'study',setup['imagingDataset'],qFilter)
    if not os.path.isdir(setup['localDir']):
        os.mkdir(setup['localDir'])
    #select just the first row; debugging
    rows=ds['rows']
    setup['values']=['COM','MTV','TLG','SUVmean','SUVmax','voxelCount','SUVSD']
    setup['featureExtractor']=featureExtractor(rFile)

    #make sure we calculate Variance
    firstOrder=setup['radiomics']['featureClass']['firstorder']
    if 'Variance' not in firstOrder:
        firstOrder.append('Variance')

    n=setup.get('n',-1)
    if n>0:
        rows=rows[:n]

    for r in rows:

        setup['SUVdataset']='SUVanalysis'
        globalDone=not forceGlobal and checkData(setup,r)

        #check if we have to do calculation
        setup['SUVdataset']='SUVanalysis_liver'
        liverDone=not forceLiver and checkData(setup,r)

        setup['SUVdataset']='SUVanalysis_SUVmax'
        suvMaxDone=not forceSUVMax and checkData(setup,r) 

        setup['SUVdataset']='SUVanalysis_liver1p5'
        liver1p5Done=not forceLiver15 and checkData(setup,r)

        doneCode=f'({globalDone}/{liverDone}/{liver1p5Done}/{suvMaxDone})'
        print(f'Done: (global/liver/liver1p5/suvMax): {doneCode}')

        if globalDone and liverDone and suvMaxDone and liver1p5Done:
            print('Skipping {} {}'.format(r['ParticipantId'],r['visitCode']))
            continue

        #PET
        for q in ['petResampled']:
            localPath=statUtils.getImage(setup,r,q)
        if localPath=="NONE":
            continue
        pet=SimpleITK.ReadImage(localPath)

        #Seg
        segPaths=statUtils.getSegmentations(setup,r)
        if "NONE" in segPaths.values():
            os.remove(localPath)
            continue
        segKeys=list(segPaths.keys())
        for x in segPaths:
            print('Loaded {}/{}'.format(users[x],segPaths[x]))

        seg={x:SimpleITK.ReadImage(segPaths[x]) for x in segPaths}


        globalThreshold={x:0 for x in seg}
        if not globalDone:
            #get value for maximum in organs or liver mean and std
            setup['SUVdataset']='SUVanalysis'
            globalOutputs=thresholdAnalysis(setup,r,pet,seg,globalThreshold)
            if doUpload:
                uploadData(setup,r,globalOutputs)

        default={'SUVmax':0,'SUVmean':0,'SUVSD':0}
        setup['SUVdataset']='SUVanalysis'
        outputs=loadGlobals(setup,r,seg)

        #liver threshold
        liverId=1
        liverThreshold={x:outputs[x].get(liverId,default)['SUVmean']
                +2*outputs[x].get(liverId,default)['SUVSD'] 
                for x in outputs}
        liver1p5Threshold={x:1.5*outputs[x].get(liverId,default)['SUVmean']
                +2*outputs[x].get(liverId,default)['SUVSD'] 
                for x in outputs}
        lesionId=4
        bmId=3
        suvMax={x:numpy.max([outputs[x].get(lesionId,default)['SUVmax'],
                    outputs[x].get(bmId,default)['SUVmax']])
                for x in outputs}
        suvMaxThreshold={x:0.41*suvMax[x] for x in suvMax}
    
        ftStr='thr[liver]={} thr[liver/1.5]={} thr(suvmax)={}'
        print(ftStr.format(liverThreshold,liver1p5Threshold,suvMaxThreshold))
       

        if not liverDone:
            setup['SUVdataset']='SUVanalysis_liver'
            liverOutputs=thresholdAnalysis(setup,r,pet,seg,liverThreshold)
            if doUpload:
                uploadData(setup,r,liverOutputs)

        #also for threshold=1.5
        if not liver1p5Done:
            setup['SUVdataset']='SUVanalysis_liver1p5'
            liver1p5Outputs=thresholdAnalysis(setup,r,pet,seg,liver1p5Threshold)
            if doUpload:
                uploadData(setup,r,liver1p5Outputs)

        if not suvMaxDone:
            setup['SUVdataset']='SUVanalysis_SUVmax'
            suvMaxOutputs=thresholdAnalysis(setup,r,pet,seg,suvMaxThreshold)
            if doUpload:
                uploadData(setup,r,suvMaxOutputs)

        #skip threshold of 4
        doThreshold4=False
        if doThreshold4:
            #threshold of 4
            setup['radiomics']['setting']['resegmentRange']=[4]
            setup['radiomics']['setting']['resegmentShape']=True
            outputs4=getValues(setup,r,pet,seg)
            setup['SUVdataset']='SUVanalysis_SUV4'
            uploadData(setup,r,outputs4)

        
        #cleanup
        os.remove(localPath)

        for x in segPaths:
            os.remove(segPaths[x])


def loadGlobals(setup,r,seg):
    matchingVisits={
            'VISIT_1':['VISIT_1','VISIT_11'],
            'VISIT_11':['VISIT_1','VISIT_11'],
            'VISIT_2':['VISIT_2','VISIT_12'],
            'VISIT_12':['VISIT_2','VISIT_12'],
            'VISIT_3':['VISIT_3','VISIT_13'],
            'VISIT_13':['VISIT_3','VISIT_13'],
            'VISIT_4':['VISIT_4','VISIT_14'],
            'VISIT_14':['VISIT_4','VISIT_14']}
    ids=[1,2,3,4,5,6]
    p=setup['project']
    d=setup['SUVdataset']
    idVar='ParticipantId'
    vCodes=';'.join(matchingVisits[r['visitCode']])
    vars=['MTV','TLG','SUVmean','SUVmax','SUVSD']
    outputs={}
    for x in seg:
        outputs[x]={}
        for i in ids:
            qFilter=[]
            qFilter.append({'variable':idVar,'value':r[idVar],'oper':'eq'})
            qFilter.append({'variable':'visitCode','value':vCodes,'oper':'in'})
            qFilter.append({'variable':'User','value':f'{x}','oper':'eq'})
            qFilter.append({'variable':'segment','value':f'{i}','oper':'eq'})
            ds=setup['db'].selectRows(p,'study',d,qFilter)
            outputs[x][i]={v:combineData(ds['rows'],v) for v in vars}
    return outputs

def combineData(rows,x):
    dt=numpy.array([float(r[x] or 0) for r in rows])
    #print('Combine data[{}] {}'.format(x,dt))
    r=0
    if x=='MTV' or x=='TLG':
        r=numpy.sum(dt)
    if x=='SUVmean':
        m=combineData(rows,'MTV')
        if m==0:
            r=0
        else:
            r=combineData(rows,'TLG')/m
    if x=='SUVmax':
        try:
            r=numpy.max(dt)
        except ValueError:
            #empty array
            r=0
    if x=='SUVSD':
        #this should be probably be done right, 
        #here I only want it to work also if one of the components is 0
        r=numpy.sqrt(numpy.sum(dt*dt))
    #print(f'Result: {r}')
    return r
    
def thresholdAnalysis(setup,r,pet,seg,thrs):
    #thresholds thrs are by segmentation author
    outputs={}
    for s in thrs:
        t=thrs[s]
        print(f'Thr[{s}]: {t}')
        #update radiomics setting
        setup['radiomics']['setting']['resegmentRange']=[thrs[s]]
        setup['radiomics']['setting']['resegmentShape']=True
        outputs[s]=getValuesForSegmentation(setup,r,pet,seg[s])
            
        _=[outputs[s][y].update({'threshold':thrs[s]}) for y in outputs[s]]
    print(outputs)
    return outputs


def getValues(setup,row,pet,seg):
    #seg is a list of segments

    rFile='radiomics.json'
    #short function names
    featureExtractor=radiomics.featureextractor.RadiomicsFeatureExtractor
    getStats=statUtils.getRadiomicsComponentStats

    with open(rFile,'w') as f:
        f.write(json.dumps(setup['radiomics']))

    setup['featureExtractor']=featureExtractor(rFile)


    #find labels associated with each (non-overlaping) segmentation
    ids=statUtils.getSegments(list(seg.values())[0])
    outputs={x:{} for x in seg}
    for x in seg:
        
        for id in ids:
            print('{} {}'.format(id,ids[id]))
            try:
                output=getStats(setup,pet,seg[x],ids[id])
            except ValueError:
                continue
            outputs[x][ids[id]]=output
        
    os.remove(rFile)
    return outputs

def getValuesForSegmentation(setup,row,pet,seg):

    #short function names
    featureExtractor=radiomics.featureextractor.RadiomicsFeatureExtractor
    getStats=statUtils.getRadiomicsComponentStats

    rFile='radiomics.json'

    with open(rFile,'w') as f:
        f.write(json.dumps(setup['radiomics']))

    setup['featureExtractor']=featureExtractor(rFile)


    #find labels associated with each (non-overlaping) segmentation
    ids=statUtils.getSegments(seg)
    outputs={}

    for id in ids:
        print('{} {}'.format(id,ids[id]))
        try:
            output=getStats(setup,pet,seg,ids[id])
        except ValueError:
            continue
        outputs[ids[id]]=output
        
    os.remove(rFile)
    return outputs


def uploadData(setup,r,outputs):
    baseVar=['ParticipantId','SequenceNum','patientCode','visitCode']
    p=setup['project']
    d=setup['SUVdataset']
    for x in outputs:
        for s in outputs[x]:
            output=outputs[x][s]
            output.update({x:r[x] for x in baseVar})
            output['User']=x
            output['segment']=s
            statUtils.updateDatasetRows(setup['db'],p,d,[output])
    
def checkData(setup,r):
    p=setup['project']
    d=setup['SUVdataset']
    qFilter=[]
    idVar='ParticipantId'
    qFilter.append({'variable':idVar,'value':r[idVar],'oper':'eq'})
    qFilter.append({'variable':'visitCode','value':r['visitCode'],'oper':'eq'})
    ds=setup['db'].selectRows(p,'study',d,qFilter)
    n=len(ds['rows'])
    print('[{}:{}/{}] got {} rows.'.format(d,r[idVar],r['visitCode'],n))
    return n>0



if __name__=='__main__':
    main(sys.argv[1])