import sys
import os
import SimpleITK
import numpy
import matplotlib.pyplot
import subprocess
sys.path.append(os.path.join(os.path.expanduser('~'),'software','src','nixSuite','wrapper'))
import nixWrapper
nixWrapper.loadLibrary('labkeyInterface')
import labkeyInterface
net=labkeyInterface.labkeyInterface()
fconfig=os.path.join(os.path.expanduser('~'),'.labkey','network.json')
net.init(fconfig)
net.getCSRF()
import labkeyFileBrowser
sys.path.append(os.getcwd())
import downloadPatient
fb=labkeyFileBrowser.labkeyFileBrowser(net)
project='dinamic_spect/Patients'
locDir=os.path.join(os.path.expanduser('~'),'temp','dynamicSPECT')
import labkeyDatabaseBrowser
db=labkeyDatabaseBrowser.labkeyDB(net)
ds=db.selectRows(project,'study','Imaging',[])
patients=[row['aliasID'] for row in ds['rows']]

#doAnalysis
import subprocess

patients=list(set(patients))
#this is if IVF is inferred from locations
runScript='doPixelAnalysis.sh'
#this is if IVF is taken from clusters
runScript='doPixelIVFAnalysis.sh'
scriptPath=os.path.join(os.path.expanduser('~'),\
                        'scripts','dynamicSPECT',runScript)
#patients1=['10KFMIR']
patients1=patients
for pId in patients1:
    if pId.find('8')==0:
        continue
    for sigma2 in [0.1,1.0,4.0]:
        print('{}/{}'.format(pId,sigma2))
        subprocess.run([scriptPath,pId,str(sigma2)], \
                         check=True, stdout=subprocess.PIPE)
    
   

def getPatientNIM(id):
    f=os.path.join(locDir,id,id+'Volume19.nrrd')
    im=SimpleITK.ReadImage(f)
    nim=SimpleITK.GetArrayFromImage(im)
    return nim

def getIWeights(locDir,p,nclass,realizationId,ic):
    fname='{}_{}_{}_center{}_centerWeigth.nrrd'.\
            format(p,nclass,realizationId+1,ic+1)
    uFile=os.path.join(locDir,p,fname)
            
    imU=SimpleITK.ReadImage(uFile)
    return SimpleITK.GetArrayFromImage(imU)



#gets weights by class for a particular realization and sigma2 averaging
def getWeights(project,locDir,p,nclass,realizationId,sigma2,na):
    #for w1, classes are in 0 to nclass-1 space
    idFilter={'variable':'aliasID','value':p,'oper':'eq'}
    ds=db.selectRows(project,'study','Segmentation',[idFilter])
    pts={r['regionId']:[float(x) for x in [r['x'],r['y'],r['z']]] for r in ds['rows']}
       
    w={region:numpy.zeros(nclass) for region in pts}
    na2=0.5*(na-1)
    
    for c in numpy.arange(nclass):
        fname='{}_{}_{}_center{}_centerWeigth.nrrd'.\
            format(p,nclass,realizationId+1,c+1)
        uFile=os.path.join(locDir,p,fname)
            
        imU=SimpleITK.ReadImage(uFile)
        nU=SimpleITK.GetArrayFromImage(imU)
        
        
        for region in w:
            #print(pts[region])
            idx=[int(x) for x in pts[region]]
            #print(nU[tuple(idx)])
            
            #gaussian weighted summation of surrounding pixels
            fidx=numpy.zeros(3)
            fs=0
            for i in numpy.arange(na):
                fidx[0]=idx[0]+(i-na2)
                for j in numpy.arange(na):
                    fidx[1]=idx[1]+(j-na2)
                    for k in numpy.arange(na):
                        fidx[2]=idx[2]+(k-na2)
                        fidx=[int(x) for x in fidx]
                        fw=numpy.array(fidx)-numpy.array(pts[region])
                        fw=numpy.exp(-0.5*sum(fw*fw)/sigma2);
                        fs+=fw
                        w[region][c]+=fw*nU[tuple(fidx)]
                        #print('{}/{}: {:.2f} {:.2g}'.format(idx,fidx,fw,nU[tuple(fidx)]))
            w[region][c]/=fs
            
            
    return w

#gets fitPar for a particular realization
def getFitPar(locDir,p,nclass,realizationId):
    #fitGoodnes A tau alpha delay [k1 BVF k2 delay]xNcenters
    fname='{}_{}_{}_fitParFinal.txt'.format(p,nclass,realizationId+1)
    uFile=os.path.join(locDir,p,fname)
    return numpy.genfromtxt(uFile,delimiter='\t')

def getFitParIVF(locDir,p,nclass,realizationId):
    #fitGoodnes A tau alpha delay [k1 BVF k2 delay]xNcenters
    fname='{}_{}_{}_IVF_fitParFinal.txt'.format(p,nclass,realizationId+1)
    uFile=os.path.join(locDir,p,fname)
    return numpy.genfromtxt(uFile,delimiter='\t')

def getFitParPixel(locDir,p,nclass,sigma2):
    sigmaCode='{:.2f}'.format(sigma2)
    sigmaCode=sigmaCode.replace(".","p")
    fname='{}_{}_{}_Pixel_fitParFinal.txt'.format(p,nclass,sigmaCode)
    uFile=os.path.join(locDir,p,fname)
    return numpy.genfromtxt(uFile,deliiter='\t')

def getFitParPixelIVF(locDir,p,nclass,sigma2):
    sigmaCode='{:.2f}'.format(sigma2)
    sigmaCode=sigmaCode.replace(".","p")
    fname='{}_{}_{}_PixelIVF_fitParFinal.txt'.format(p,nclass,sigmaCode)
    uFile=os.path.join(locDir,p,fname)
    return numpy.genfromtxt(uFile,delimiter='\t')

def getK1(fitPar,iclass):
    #fitGoodnes A tau alpha delay [k1 BVF k2 delay]xNcenters
    #0 to nclass-1 space
    return fitPar[4*iclass+5]

def calculateK1(w,fitPar):
    #calculateK1 for region weights
    #return the k1 that belongs to the 
    #maximum class in region (M) and 
    #a weighted average (W)
    k1={region:{'M':0,'W':0} for region in w}
    for region in w:
        cmax=numpy.argmax(w[region])
        k1[region]['M']=getK1(fitPar,cmax)
        fs=0
        for c in numpy.arange(len(w[region])):
            fs+=w[region][c]*getK1(fitPar,c)
        k1[region]['W']=fs
    return k1

def getPatientValuesByNclass(project,locDir,p,nclass,nrealizations,sigma2,na):
    #summary script
    #project/locDir and p identify patient
    #nclass and nrealizations select strategy
    #sigma2 is for combining output from adjacent pixels
    #na is neighborhood size where smoothing/combination is done
    k1={}
    for i in numpy.arange(nrealizations):
        w=getWeights(project,locDir,p,nclass,i,sigma2,na)
        fitPar=getFitParIVF(locDir,p,nclass,i)
        qk1=calculateK1(w,fitPar)
        for region in w:
            for type in qk1[region]:
                try:
                    k1[region][type].append(qk1[region][type])
                except KeyError:
                    k1={region:{type:[] for type in qk1[region]} for region in w}
                    print(type)
                    k1[region][type].append(qk1[region][type])
        print('[{}] {}/{}'.format(nclass,i+1,nrealizations))
    return k1   

def getSummaryPatientValuesByNclass(project,locDir,p,nclass,nrealizations,sigma2,na):
    #summary script, same arguments as above
    #also return deviation over realization
    k1=getPatientValuesByNclass(project,locDir,p,nclass,nrealizations,sigma2,na)
    avgType=['M','W']
    summaryK1={type:{region:{
        'mean':numpy.mean(k1[region][type]), 
        'std':numpy.std(k1[region][type]), 
        'median':numpy.median(k1[region][type])} for region in k1}
               for type in avgType}
    
    return summaryK1

def fullSummary(project,locDir,patientList,classes,nr,sigma2,na):
    return \
        {p:{c:getSummaryPatientValuesByNclass(project,locDir,p,c,nr,sigma2,na) for c in classes} for p in patientList}
          
def storeSummary(db,project,locDir,k1,sigma2,na):
    dsM=db.selectRows(project,'study','Imaging',[])
    for p in k1:
        idFilter={'variable':'aliasID','value':p,'oper':'eq'}
        for c in k1[p]:
            cFilter={'variable':'nclass','value':str(c),'oper':'eq'}
            for t in k1[p][c]:
                tFilter={'variable':'option','value':t,'oper':'eq'}
                for r in k1[p][c][t]:
                    rFilter={'variable':'regionId','value':str(r),'oper':'eq'}
                    ds=db.selectRows(project,'study','Summary',[idFilter,cFilter,tFilter,rFilter])
                    if len(ds['rows'])>0:
                        qrow=ds['rows'][0]
                        mode='update'
                    else:
                        frow=[r for r in dsM['rows'] if r['aliasID']==p]
                        qrow={'aliasID':p,'nclass':c,'option':t,'regionId':r}
                        qrow['SequenceNum']=frow[0]['SequenceNum']+c+0.001*r
                        if t=='M':
                            qrow['SequenceNum']+=0.0001
                        qrow['PatientId']=frow[0]['PatientId']
                        mode='insert'
                    for v in k1[p][c][t][r]:
                        qrow[v]=k1[p][c][t][r][v]
                    qrow['sigma2']=sigma2
                    qrow['na']=na
                    db.modifyRows(mode,project,'study','Summary',[qrow])
    
def summaryPixelIVF(project,locDir,patientList,sigma2List,na):
     return \
        {p:{sigma2:getPixelIVF(project,locDir,p,sigma2,na) for sigma2 in sigma2List} for p in patientList}
    
def storeIVF(db,project,locDir,k1):
    dsM=db.selectRows(project,'study','Imaging',[])
    for p in k1:
        idFilter={'variable':'aliasID','value':p,'oper':'eq'}
        for s2 in k1[p]:
            sigmaFilter={'variable':'sigma2','value':str(s2),'oper':'eq'}
            for r in k1[p][s2]:
                rFilter={'variable':'regionId','value':str(r),'oper':'eq'}
                typeFilter={'variable':'option','value':'D','oper':'eq'}
                ds=db.selectRows(project,'study','Summary',[idFilter,sigmaFilter,rFilter,typeFilter])
                if len(ds['rows'])>0:
                    qrow=ds['rows'][0]
                    mode='update'
                else:
                    frow=[r for r in dsM['rows'] if r['aliasID']==p]
                    qrow={'aliasID':p,'sigma2':s2,'regionId':r}
                    qrow['SequenceNum']=40+frow[0]['SequenceNum']+0.01*r+0.001*s2
                    qrow['PatientId']=frow[0]['PatientId']
                    mode='insert'
                qrow['mean']=k1[p][s2][r]
                qrow['na']=7
                qrow['nclass']=16
                qrow['option']='D'
                db.modifyRows(mode,project,'study','Summary',[qrow])

def getPixelIVF(project,locDir,p,sigma2,na):
    fitPar=getFitParPixelIVF(locDir,p,16,sigma2)
    idFilter={'variable':'aliasID','value':p,'oper':'eq'}
    ds=db.selectRows(project,'study','Segmentation',[idFilter])
    k1={r['regionId']:getK1(fitPar,r['regionId']) for r in ds['rows']}
    return k1

pat=[p for p in patients if p.find('8')!=0]
k1=summaryPixelIVF(project,locDir,pat,[0.1],7)
storeIVF(db,project,locDir,k1)

pat=[p for p in patients if p.find('8')!=0]
k1=fullSummary(project,locDir,pat,[10],20,0.1,3)
storeSummary(db,project,locDir,k1,0.1,3)

print(k1)

import matplotlib
p='7TMOBR'
sy=28
realizationId=5
nclass=30
nu=getIWeights(locDir,p,30,realizationId,4)
nim=getPatientNIM(p)
cut0=20
w0=20
cut1=0
fw=40
idx=[7,28,34]
print(nu[7,28,34])
print(nu[tuple(idx)])
w=getWeights(project,locDir,p,nclass,realizationId,0.1,7)
for r in w:
    print(numpy.max(w[r]))
matplotlib.pyplot.imshow(nu[:,sy,cut0:cut0+w0])
matplotlib.pyplot.show()
matplotlib.pyplot.imshow(nim[:,sy,cut0:cut0+w0])
matplotlib.pyplot.show()

k1=combinePatient(project,locDir,'7TMMIR',[10,20,30],20,4.0,7)

avgType=['M','W']
meanK1={type:{region:numpy.mean(k1[region][type]) for region in k1} for type in avgType}
stdK1={type:{region:numpy.std(k1[region][type]) for region in k1} for type in avgType}
print(meanK1['M'])
print(stdK1['M'])

#k1=combinePatient(project,locDir,'7TMMIR',[10],20,4.0,7)
realizationId=3
c=10
p='7TMOBR'
w=getWeights(project,locDir,p,c,realizationId,0.1,3)
wc={region:numpy.argmax(w[region]) for region in w}
fitPar=getFitParIVF(locDir,p,c,realizationId)
qk1=calculateK1(w,fitPar)
k1p=getPixelIVF(project,locDir,p,4.0,7)
pk=list(k1p.values())
print(pk)
ck=[qk1[region]['M'] for region in qk1]
matplotlib.pyplot.scatter(ck,pk)
for region in w:
    print('{} {} {}'.format(k1p[region],wc[region],qk1[region]['M']))

k1p=getPixelIVF(project,locDir,'7TMMIR',4.0,7)

for r in k1:
    print('{} {} {}/{}'.format(k1p[r], qk1[r]['M'] ,numpy.median(k1[r]['M'][10:20]),numpy.std(k1[r]['M'])))

numpy.zeros(3)