import numpy as np
import scipy.stats
import matplotlib.pyplot as plt
from pathlib import Path
import sys
import math

VISUALIZE_2D = True
REF_MODE = 'ts'

scene_path = Path(sys.argv[1])
assert scene_path.is_dir()
scene = scene_path.name

# Compure radial profile from 2D slice
def radial(data):
  x = np.linspace(-1, 1, data.shape[0])
  y = np.linspace(-1, 1, data.shape[1])
  xx, yy = np.meshgrid(x, y)
  r = np.sqrt(xx**2. + yy**2.)
  a, b, c = scipy.stats.binned_statistic(r.flatten(), data.flatten(), bins=50)
  return (b[1:] + b[:-1])/2, a

# Compute mean of all samples
mean = []
for folder in scene_path.iterdir():
  for file in folder.iterdir():
    data = plt.imread(str(file))
    mean.append(np.mean(data))
mean = sum(mean) / len(mean)

# Process the samples of one sampling mode
def process(folder, ref):
  mode = folder.name
  fftmean = None
  num = 0
  mse = []
  for file in folder.iterdir():
    # Load PNG
    data = plt.imread(str(file))
    # Normalize
    data = (data / mean) - 1

    # Accumulate MSE
    mse.append(np.mean(data**2))

    # Compute spectrum
    data = np.fft.fft2(data)
    # Mask DC term
    data[0, 0] = math.nan
    data = np.fft.fftshift(data)
    data = abs(data)

    # Accumulate spectrum
    fftmean = data if fftmean is None else fftmean + data
    num += 1

  # Combine MSE
  mse = sum(mse) / len(mse)
  psnr = -10.0 * math.log10(mse)

  # Combine spectra, compute radial profile
  fftmean /= num
  x, y = radial(fftmean)
  filter = ~np.isnan(y)
  x = x[filter]
  y = y[filter]

  # Normalize to reference mode
  if ref is None:
    ref = np.mean(y), psnr
  y /= ref[0]
  psnr -= ref[1]

  # Visualize
  if VISUALIZE_2D:
    fig = plt.figure(mode, frameon=False)
    fig.set_size_inches(fftmean.shape[0], fftmean.shape[1])
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    ax.imshow(fftmean, aspect='auto')
    fig.savefig(f'fft_{scene}_{mode}.png', dpi=1)
    fig.show()
  else:
    plt.plot(x, y, label=mode)

  # Output
  print(f"{scene} {mode} {num} {psnr}")
  return ref

# Process REF_MODE first, then others in sorted order (to get consistent legend)
ref = None
modes = {folder.name for folder in scene_path.iterdir()}
modes.discard(REF_MODE)
for mode in [REF_MODE] + list(sorted(modes)):
  ref = process(scene_path / mode, ref)

# Show plot
if not VISUALIZE_2D:
  plt.legend()
  plt.ylim(bottom=0)
  plt.savefig(f'fft_{scene}_radial.png')
plt.show()