import matplotlib
matplotlib.use('Qt4Agg')
from pylab import *
from matplotlib.colors import LogNorm
from matplotlib.colors import LinearSegmentedColormap, LogNorm, Normalize
import numpy as np
from  PyQt4 import Qt, QtCore
from PyQt4.QtCore import SIGNAL, SLOT

referenceprefix  ="/data/scisofttmp/mirone/TEST_PYHST/RESULTS/2018bb_corrected/2gpu/tests/"

newversionprefix ="/data/scisofttmp/mirone/TEST_PYHST/RESULTS/2018bb_corrected/1gpu_fix/tests/"

# referenceprefix="/data/scisofttmp/mirone/TEST_PYHST/RESULTS/2018bb/2gpu_bis/tests/"




# /data/scisofttmp/mirone/TEST_PYHST/RESULTS/2015d/gpu2-0102/tests"

# newversionprefix ="/data/scisofttmp/mirone/TEST_PYHST/RESULTS/2017c_final/2gpu_bis/tests/MULTIPAGANIN/"
# referenceprefix="/data/scisofttmp/mirone/TEST_PYHST/RESULTS/2015d/gpu2-0102/tests/SIRT_LENA/"
# referenceprefix ="/data/scisofttmp/mirone/TEST_PYHST/RESULTS/2017c_final/1gpu_bis/tests/MULTIPAGANIN/"

import numpy
from mpi4py import MPI
myrank = MPI.COMM_WORLD.Get_rank()
nprocs = MPI.COMM_WORLD.Get_size()
procnm = MPI.Get_processor_name()
comm = MPI.COMM_WORLD
print "MPI LOADED , nprocs = ", nprocs

def compara_vols(ref, new  ):
    class parInfo:
        LOWBYTEFIRST=1
        s = open(ref+".info").read()
        s=s.replace("!","#")
        exec( s )
    print ref
    print new
    f_ref = open(ref,"r")
    f_new = open(new,"r")
    maxerr=-1
    overerr = 0
    miz=-1

    nblockfl = parInfo.NUM_Z/nprocs+1


    for iz in range(parInfo.NUM_Z):

        if (iz/nblockfl)%nprocs != myrank:
            continue
        sr = f_ref.read( parInfo.NUM_X*parInfo.NUM_Y * 4 )
        sn = f_new.read( parInfo.NUM_X*parInfo.NUM_Y * 4 )

        r  = numpy.reshape(numpy.fromstring(sr,"f"), [   parInfo.NUM_Y, parInfo.NUM_X ] )
        n  = numpy.reshape(numpy.fromstring(sn,"f"), [   parInfo.NUM_Y, parInfo.NUM_X ] )
        err = (r-n)
        err=(err*err).sum()
        rsum = (r*r).sum()
        nsum = (n*n).sum()

        if rsum==0 and nsum==0 :
            err = -1
            
        print iz, " err " , err
        if err> maxerr:
            maxerr=err
            dr = r
            dn = n
            overerr = (r*r).sum()
            miz = iz



    a    = numpy.array( [maxerr],"f" )
    a_r  = numpy.array( [0]    ,   "f"        )
    comm.Allreduce([a, MPI.FLOAT  ], [a_r, MPI.FLOAT  ], op=MPI.MAX)
    if a[0]==a_r[0]:
        a[0]=myrank
    else:
        a[0]=-1
    comm.Allreduce([a, MPI.FLOAT  ], [a_r, MPI.FLOAT  ], op=MPI.MAX)
    
    if myrank == a_r[0]:
        print "COMPARISON :   %s %s"%( ref, new  )
        print       " Relative Error : %e "% numpy.sqrt(maxerr/overerr)
        print " for slice ", miz
        mostra_slices( dr,dn, ref, new)
    comm.Barrier()

def compara_edfs(ref, new  ):
    import fabio
    
    maxerr=-1
    overerr = 0
    miz=-1
    done =0

    NN = len( zip (ref, new ))

    nblockfl = NN/nprocs+1



    for iz,(rname, nname) in enumerate(zip (ref, new )) :
        done=1

        if (iz/nblockfl)%nprocs != myrank:
            continue
     

        r  = fabio.open(rname).data
        n  = fabio.open(nname).data
        err = (r-n)
        err=(err*err).sum()
        if err> maxerr:
            maxerr=err
            dr = r
            dn = n
            overerr = (r*r).sum()
            mrname = rname
            mnname = nname 
    if done:



        a    = numpy.array( [maxerr],"f" )
        a_r  = numpy.array( [0]    ,   "f"        )
        comm.Allreduce([a, MPI.FLOAT  ], [a_r, MPI.FLOAT  ], op=MPI.MAX)
        if a[0]==a_r[0]:
            a[0]=myrank
        else:
            a[0]=-1
        comm.Allreduce([a, MPI.FLOAT  ], [a_r, MPI.FLOAT  ], op=MPI.MAX)
    
        if myrank == a_r[0]:
            print "COMPARISON :   %s %s"%(mrname , mnname  )
            print       " Relative Error : %e "% numpy.sqrt(maxerr/overerr)
            mostra_slices( dr,dn, mrname, mnname)
        comm.Barrier()


import PyMca5.PyMcaGui
from  PyMca5.PyMcaGui import MaskImageWidget as  pymcaMask



class MaskImageWidget(pymcaMask.MaskImageWidget):
    closealso=[]
    
    def __init__(self, *args, **argw):
        self.closealso.append(self)
        super(MaskImageWidget,self).__init__(*args, **argw)

        
    def closeEvent(self, event):
        print " in close ",  self.closealso
        
        if self in self.closealso:
            self.closealso.remove(self)
            
        for t in self.closealso:
            self.closealso.remove(t)
            t.close()
        event.accept() #

def     mostra_slices( dr,dn, namer, namen):
    app=Qt.QApplication([])

  
   
    maskW1 = MaskImageWidget(None , aspect=True,profileselection=True)
    maskW1.setImageData(dr  , xScale=(0.0, 1.0), yScale=(0., 1.))
    maskW2 = MaskImageWidget(None , aspect=True,profileselection=True)
    maskW2.setImageData(dn  , xScale=(0.0, 1.0), yScale=(0., 1.))
    maskW3 = MaskImageWidget(None , aspect=True,profileselection=True)
    maskW3.setImageData(abs(dn-dr)  , xScale=(0.0, 1.0), yScale=(0., 1.))
    maskW1.setDefaultColormap(0,0)
    maskW2.setDefaultColormap(0,0)
    maskW3.setDefaultColormap(2,1)
    maskW1.plotImage(True)
    maskW2.plotImage(True)
    maskW3.plotImage(True)

    maskW1.setWindowTitle("new")
    maskW2.setWindowTitle("old")
    maskW3.setWindowTitle("new - old")
    
    maskW1.show()
    maskW2.show()
    maskW3.show()
    app.connect(maskW1, SIGNAL("destroyed()"),
                app, SLOT("quit()"))

    
    app.exec_()
  
    # cdict = {'red': ((0.0, 0.0, 0.0),
    #                  (0.5, 0.0, 0.0),
    #                  (0.75, 1.0, 1.0),
    #                  (1.0, 1.0, 1.0)),
    #          'green': ((0.0, 0.0, 0.0),
    #                    (0.25, 1.0, 1.0),
    #                    (0.75, 1.0, 1.0),
    #                    (1.0, 0.0, 0.0)),
    #          'blue': ((0.0, 1.0, 1.0),
    #                   (0.25, 1.0, 1.0),
    #                   (0.5, 0.0, 0.0),
    #                   (1.0, 0.0, 0.0))}
    # temperature = LinearSegmentedColormap('temperature',cdict, 65536)
    # fig=plt.figure()
    # plt.subplot(1, 3, 1)
    # title(namer)
    # im = plt.imshow(dr,  cmap=temperature)
    # plt.subplot(1, 3, 2)
    # title(namen)
    # im = plt.imshow(dn,  cmap=temperature)
    # plt.subplot(1, 3, 3)
    # title('diff (abs)')
    # diff = abs( dr-dn  )
    # im = plt.imshow(diff,  cmap=temperature,
    #             norm=LogNorm(vmin=max(diff.min(), diff.max()*1.0e-8), vmax=diff.max()))
    # show()
    
def get_results(radice):
    N=len(radice)
    import os
    import glob
    result={}
    ws=os.walk(radice)
    leafs = [ t[0] for t in ws if len(t[1])==0 ]
    for l in leafs:
        result[l[N:]] = {}
        print l+"/abs.vol"
        result[l[N:]]["vol"] = glob.glob(l+"/vol*.vol")+glob.glob(l+"/reconstruction*.vol")
        lf = glob.glob(l+"/*.edf")
        result[l[N:]]["edf"] = [ t for t in lf if ("histo" not in t and "phase1" not in t and "phase2" not in t )]
    return result



ress_ref = get_results(referenceprefix)

print ress_ref
ress_new = get_results(newversionprefix)
print ress_new


leafs = list(ress_ref.keys())




for l in leafs:
    print "======================="
    print l
    print " continuo?"

    res=raw_input()
    if res=="1":

        print ress_ref[l]["vol"]
        print ress_new[l]["vol"]
        print ress_ref[l]["edf"]
        print ress_new[l]["edf"]
        
        if len(ress_ref[l]["vol"]):
            if(len(ress_new[l]["vol"])==0):
                print " ERRORE vol : new n'a pas produit de output pour ", l
            else:
                assert( len(ress_ref[l]["vol"])==1)
                ress_vol = ress_ref[l]["vol"][0]
                new_vol  = ress_new[l]["vol"][0]
                compara_vols(ress_vol,  new_vol   )
        else:

            if(len(ress_new[l]["edf"])==0):
                print " ERRORE edf : new n'a pas produit de output pour ", l
            else:
                # assert( len(ress_ref[l]["edf"] = len(ress_new[l]["edf"])==1) )

                
                ress_edf = ress_ref[l]["edf"]
                new_edf  = ress_new[l]["edf"]
                ress_edf.sort()
                new_edf .sort()
                compara_edfs(ress_edf,  new_edf   )
        
    
print ress_ref
print ress_new
print " END \n"*100
