import os
import numpy as np 

datafolder="/home/jliu447/lossycompression/MIRANDA/extracted"

ebs=[i*1e-4 for i in range(1,10)]+[i*1e-3 for i in range(1,10)]+[i*1e-2 for i in range(1,11)]

cr=np.zeros((29,51),dtype=np.float32)
psnr=np.zeros((29,51),dtype=np.float32)
filelist=os.listdir(datafolder)
filelist=[x for x in filelist if x.split('.')[-1]=='f32']
filelist=sorted(filelist)
for i,eb in enumerate(ebs):
    cr[i+1][0]=eb
    psnr[i+1][0]=eb
    for j,filename in enumerate(filelist):
        idx=int(filename.split('-')[1])
        cr[0][j+1]=idx
        psnr[0][j+1]=idx
        filepath=os.path.join(datafolder,filename)
        comm="sz -z -f -i %s -M REL -R %f -3 384 384 256" % (filepath,eb)
        os.system(comm)
        szpath=filepath+".sz"
        comm="sz -x -f -i %s -s %s -3 384 384 256 -a>tempmirandasz.txt" % (filepath,szpath)
        os.system(comm)
        with open("tempmirandasz.txt","r") as f:
            lines=f.read().splitlines()
            p=eval(lines[4].split(',')[0].split('=')[1])
            
            r=eval(lines[7].split('=')[1])
            
            cr[i+1][j+1]=r
            psnr[i+1][j+1]=p

        comm="rm -f %s" % szpath
        os.system(comm)
        comm="rm -f %s" % szpath+".out"
        os.system(comm)
        os.system("rm -f tempmirandasz.txt")


np.savetxt("sz_miranda_cr.txt",cr,delimiter='\t')
np.savetxt("sz_miranda_psnr.txt",psnr,delimiter='\t')