import os

import numpy
import matplotlib.pyplot as plt

ratios = [1, 4, 10, 25, 30, 40, 50]
epoch_oris = [25, 170, 150, 185, 195, 190, 195]
epoch_sams = [15, 20, 25, 35, 30, 35, 35]

sam = ['_SAM', ""]
ldir = "result/Set11/"
rd = "ISTA%s_NET_ratio_%d_epoch_%d"

flist = sorted(os.listdir(ldir))
for (ratio, epocho, epochs) in zip(ratios, epoch_oris, epoch_sams):
    res = {
        "ISTA_NET": {},
        "SAISTA_NET": {}
    }
    os.mkdir(f'ista_net_pdfs/{ratio}')
    for fn in flist:
        if rd % ("", ratio, epocho) in fn:
            res["ISTA_NET"][fn.split('.tif')[0]] = {}
            res["ISTA_NET"][fn.split('.tif')[0]]["PSNR"] = fn.split('PSNR_')[1].split('_')[0]
            res["ISTA_NET"][fn.split('.tif')[0]]["SSIM"] = fn.split('SSIM_')[1].split('.')[0]
            res["ISTA_NET"][fn.split('.tif')[0]]["image"] = plt.imread(ldir + fn)
        if rd % ("_SAM", ratio, epochs) in fn:
            res["SAISTA_NET"][fn.split('.tif')[0]] = {}
            res["SAISTA_NET"][fn.split('.tif')[0]]["PSNR"] = fn.split('PSNR_')[1].split('_')[0]
            res["SAISTA_NET"][fn.split('.tif')[0]]["SSIM"] = fn.split('SSIM_')[1].split('.')[0]
            res["SAISTA_NET"][fn.split('.tif')[0]]["image"] = plt.imread(ldir + fn)

    for k in res["ISTA_NET"].keys():
        net = res["ISTA_NET"][k]
        sam = res["SAISTA_NET"][k]
        plt.imsave(f"ista_net_pdfs/{ratio}/net_{k}_{net['PSNR']}_{net['SSIM']}.pdf", net["image"])
        plt.imsave(f"ista_net_pdfs/{ratio}/sanet_{k}_{sam['PSNR']}_{sam['SSIM']}.pdf", sam["image"])
