#! /usr/bin/env python3
# -*- coding:utf8 -*-

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
from math import *

def shepp_logan(x):
    return 2. / (pi * pi * (1 - 4. * x * x))

N = 16
xlim = -6, 6
xd = np.arange(N)
kernel = shepp_logan(np.minimum(xd, N - xd))

fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(7, 3.5))

ax0.plot(xd - N // 2, np.roll(kernel, N // 2), "o-", label="Shepp-\nLogan")

xa = np.linspace(*xlim, 1201)
ya = np.divide(2. / (pi * pi), (1 - 4. * xa * xa), out=None, where=(xa!=0.25))
ya[xa==0] = float(nan)
ax0.plot(xa, ya, "--", color="#bbbbbb", zorder=1.9,
    label=r"$\frac{2}{\pi^2 (1 - 4x^2)}$")

annotate_args = {"va":"center", "ha":"left", "xytext":(12, 0), "zorder":1.8,
    "textcoords":"offset points", "fontsize":14,
    "bbox":{"fc":"white", "ec":"none"} }
ax0.annotate(r"$\frac{2}{\pi^2}$", (0, kernel[0]), **annotate_args)
ax0.annotate(r"$\frac{-2}{3\pi^2}$", (1, kernel[1]), **annotate_args)

ax0.grid(True)
ax0.set_xlim(*xlim)
ax0.set_ylim(-0.10, 0.24)
ax0.set_xlabel("x")
ax0.legend(loc="upper left", framealpha=1, borderpad=0.6)

kernel_spectrum = np.fft.fft(kernel, n=N).real
xa2 = np.arange(-N//2, N//2+1)
ax1.plot(xa2 / N, np.roll(np.concatenate((kernel_spectrum[:N//2+1], kernel_spectrum[N//2:])), N//2), "o-", label="FT(Shepp-Logan)")
ax1.plot(xa2 / N, 1./pi * np.sin(pi / N * np.abs(xa2)), "--", color="#888888", zorder=1.9, label=r"$\sin(\pi\cdot|k|\,/\,N)\,/\,\pi$")
ax1.grid(True)
ax1.set_xlim(-0.5, 0.5)
ax1.set_ylim(-0.03, 0.4)
ax1.xaxis.set_major_locator(ticker.MultipleLocator(0.25))
ax1.yaxis.set_major_locator(ticker.MultipleLocator(0.1))
ax1.set_xlabel("$k/N$")
ax1.set_ylabel(r"$a_k$")
ax1.legend(loc="upper center", framealpha=1, borderpad=0.6)

plt.tight_layout()
plt.savefig("Shepp-Logan kernel + fourier.svg")