import numpy as np
import matplotlib.pyplot as plt
from numpy.fft import fft2, fftfreq, fftshift
import scipy.fftpack as sfft
from scipy import ndimage

fname = '_113403077_picture1_convert.png'

img = plt.imread(fname)
s0, s1 = img.shape[:2]
deg_per_pixel = 0.04 # This map of the CMB covers a section of the sky 50 times the Moon's width
pixel_per_deg = 1. / deg_per_pixel
print(img.shape)
r, g, b = np.moveaxis(img, 2, 0)
diff = r - b
f = fftshift(fft2(diff))
P = np.abs(f)**2
logP = np.log10(P)

FreqCompRows = np.fft.fftfreq(f.shape[0], d=2)
FreqCompCols = np.fft.fftfreq(f.shape[1], d=2)
FreqCompRows = np.fft.fftshift(FreqCompRows) # per pixel (vertical)
FreqCompCols = np.fft.fftshift(FreqCompCols) # per pixel (horizontal)
# https://dsp.stackexchange.com/a/50245/25659
extent = [0, s1 * deg_per_pixel, 0, s0 * deg_per_pixel]

extent_ft = [FreqCompCols.min() * pixel_per_deg, FreqCompCols.max() * pixel_per_deg,
             FreqCompRows.min() * pixel_per_deg, FreqCompRows.max() * pixel_per_deg]
if True:
    plt.figure()
    plt.imshow(logP[65:184, 163:462], cmap='gray', extent=extent_ft, vmin=3, vmax=7)
    plt.show()

if True:
    plt.figure()
    plt.subplot(2, 1, 1)
    # plt.imshow(diff, cmap='gray', extent=extent)
    plt.imshow(img,extent=extent)
    plt.xlabel('degrees')
    plt.ylabel('degrees')
    plt.subplot(2, 1, 2)
    plt.title('log10(abs(FT)^2)')
    plt.xlabel('inverse degrees')
    plt.ylabel('inverse degrees')
    plt.imshow(logP, cmap='gray', extent=extent_ft, vmin=3, vmax=7)
    plt.show()