import numpy as np
import matplotlib.pyplot as plt

# =====================================================
# PARAMETRI FISICI
# =====================================================

lunghezza_d_onda = 0.05
k = 2*np.pi / lunghezza_d_onda

distanza_fenditure = 2.0
larghezza_fenditura = 0.1

# coordinate fenditure
x_fenditura = 0.0

y1 = -distanza_fenditure/2
y2 = +distanza_fenditure/2

# =====================================================
# SPAZIO DI SIMULAZIONE
# =====================================================

x = np.linspace(0.05, 10, 1600)
y = np.linspace(-5, 5, 1600)

X, Y = np.meshgrid(x, y)

# =====================================================
# FENDITURA CONTINUA (HUYGENS-FRESNEL)
# =====================================================

def fenditura(y_centro,
              larghezza,
              numero_sorgenti=120):

    ys = np.linspace(
        y_centro - larghezza/2,
        y_centro + larghezza/2,
        numero_sorgenti
    )

    psi = np.zeros_like(X, dtype=complex)

    for ysorg in ys:

        r = np.sqrt(
            (X - x_fenditura)**2 +
            (Y - ysorg)**2
        )

        theta = np.arctan2(
            Y - ysorg,
            X - x_fenditura
        )

        obliquita = np.maximum(
            np.cos(theta),
            0
        )

        diffrazione = np.sinc(
            larghezza_fenditura *
            np.sin(theta) /
            lunghezza_d_onda
        )

        psi += (
            obliquita
            * diffrazione
            * np.exp(1j * k * r)
            / np.sqrt(r)
        )

    psi /= numero_sorgenti

    return psi

# =====================================================
# AMPIEZZE DELLE DUE FENDITURE
# =====================================================

psi1 = fenditura(
    y1,
    larghezza_fenditura
)

psi2 = fenditura(
    y2,
    larghezza_fenditura
)

psi_tot = psi1 + psi2

# =====================================================
# DENSITA' DI PROBABILITA'
# =====================================================

rho = np.abs(psi_tot)**2

rho = rho**0.25

rho /= rho.max()

# =====================================================
# CORRENTE DI PROBABILITA'
# j = Im( psi* grad psi )
# =====================================================

dx = x[1] - x[0]
dy = y[1] - y[0]

dpsi_dx = np.gradient(
    psi_tot,
    dx,
    axis=1
)

dpsi_dy = np.gradient(
    psi_tot,
    dy,
    axis=0
)

jx = np.imag(
    np.conj(psi_tot) * dpsi_dx
)

jy = np.imag(
    np.conj(psi_tot) * dpsi_dy
)

# normalizzazione per il plot
mod_j = np.sqrt(jx**2 + jy**2)

jx_plot = jx / (mod_j + 1e-12)
jy_plot = jy / (mod_j + 1e-12)

# =====================================================
# GRAFICO
# =====================================================

plt.figure(figsize=(12,8))

plt.imshow(
    rho,
    extent=[
        x.min(),
        x.max(),
        y.min(),
        y.max()
    ],
    origin="lower",
    cmap="inferno",
    aspect="auto"
)

plt.colorbar(
    label="|ψ|²"
)

# fenditure

plt.plot(
    [0,0],
    [
        y1-larghezza_fenditura/2,
        y1+larghezza_fenditura/2
    ],
    color="cyan",
    linewidth=5
)

plt.plot(
    [0,0],
    [
        y2-larghezza_fenditura/2,
        y2+larghezza_fenditura/2
    ],
    color="cyan",
    linewidth=5
)

# corrente di probabilità

step = 24

plt.streamplot(
    X[::step,::step],
    Y[::step,::step],
    jx_plot[::step,::step],
    jy_plot[::step,::step],
    color="white",
    density=2.0,
    linewidth=0.8
)

plt.xlabel("x")
plt.ylabel("y")

plt.title(
    "Doppia fenditura: densita' di probabilita' e corrente di probabilita'"
)

plt.tight_layout()

plt.savefig(
    "due_fenditure_corrente.png",
    dpi=300
)

print()
print("Salvato: due_fenditure_corrente.png")
print("Mostro il grafico...")
print("Chiudi la finestra per terminare il programma.")

plt.show()

print("Grafico chiuso.")