import tomopy
import matplotlib.pyplot as plt

rec = dict()  # Store all of the different reconstructions here

import dxchange as dx

prj, flat, dark, ang = dx.read_aps_32id('data/data-simulated.h5')

rec = dict()
rec['before'] = tomopy.recon(prj, ang, algorithm='gridrec', num_gridx=128, num_gridy=128)

z = 1  # slice that has the poisson noise
prj_corrected = tomopy.median_filter(prj, axis=0)

z = 1
plt.figure(dpi=150)
plt.subplot(1, 2, 1)
plt.title('before')
plt.imshow(prj[:, z, :])
plt.subplot(1, 2, 2)
plt.title('after')
plt.imshow(prj_corrected[:, z, :])
plt.show()

rec['after'] = tomopy.recon(prj_corrected, ang, algorithm='gridrec', num_gridx=128, num_gridy=128)

plt.figure(dpi=150)
plt.subplot(1, 2, 1)
plt.title('before')
plt.imshow(rec['before'][z], vmin=0, vmax=255)
plt.subplot(1, 2, 2)
plt.title('after')
plt.imshow(rec['after'][z], vmin=0, vmax=255)
plt.show()

z = 2  # slice that has the zingers

plt.figure(dpi=150)
plt.subplot(1, 2, 1)
plt.title('before')
plt.imshow(prj[:, z, :])
plt.subplot(1, 2, 2)
plt.title('after')
plt.imshow(prj_corrected[:, z, :])
plt.show()

plt.figure(dpi=150)
plt.subplot(1, 2, 1)
plt.title('before')
plt.imshow(rec['before'][z], vmin=0, vmax=255)
plt.subplot(1, 2, 2)
plt.title('after')
plt.imshow(rec['after'][z], vmin=0, vmax=255)
plt.show()

z = 3  # slice that has the rings
prj_corrected = tomopy.remove_stripe_based_filtering(prj)

plt.figure(dpi=150)
plt.subplot(1, 2, 1)
plt.title('before')
plt.imshow(prj[:, z, :])
plt.subplot(1, 2, 2)
plt.title('after')
plt.imshow(prj_corrected[:, z, :])
plt.show()

rec['after'] = tomopy.recon(prj_corrected, ang, algorithm='gridrec', num_gridx=128, num_gridy=128)

plt.figure(dpi=150)
plt.subplot(1, 2, 1)
plt.title('before')
plt.imshow(rec['before'][z], vmin=0, vmax=255)
plt.subplot(1, 2, 2)
plt.title('after')
plt.imshow(rec['after'][z], vmin=0, vmax=255)
plt.show()