---

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import RegularGridInterpolator

# --- Параметры и исходные данные (как в вашем коде) ---
hbar = 1.0545718e-34
me = 9.1093837e-31
kB = 1.380649e-23

T = (hbar**2 / (2 * me)) * 4 / kB

mx = -4.0 * me
my = -me / 4.0

def energy(kx, ky):
    return (hbar**2 * kx**2) / (2 * mx) + (hbar**2 * ky**2) / (2 * my)

def distribution(kx, ky, T):
    return np.exp(energy(kx, ky) / (kB * T))

K = 5
dx, dy = 1.0, 1.0
kx_coarse = np.arange(-K, K + dx, dx)
ky_coarse = np.arange(-K, K + dy, dy)
KX_coarse, KY_coarse = np.meshgrid(kx_coarse, ky_coarse, indexing='ij')
n_coarse = distribution(KX_coarse, KY_coarse, T)

interp = RegularGridInterpolator((kx_coarse, ky_coarse), n_coarse, method='linear', bounds_error=False, fill_value=None)

dk_fine = 0.1
kx_fine = np.arange(-K, K + dk_fine, dk_fine)
ky_fine = np.arange(-K, K + dk_fine, dk_fine)
KX_fine, KY_fine = np.meshgrid(kx_fine, ky_fine, indexing='ij')

points_fine = np.array([KX_fine.ravel(), KY_fine.ravel()]).T
n_interpolated = interp(points_fine).reshape(KX_fine.shape)
n_fine_real = distribution(KX_fine, KY_fine, T)
difference = n_fine_real - n_interpolated

# --- 3D визуализация погрешности ---
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# Построение поверхности
surf = ax.plot_surface(KX_fine, KY_fine, difference, cmap='RdBu_r', 
                       linewidth=0, antialiased=True, alpha=0.9)

# Настройка осей
ax.set_xlabel('kx')
ax.set_ylabel('ky')
ax.set_zlabel('Погрешность')
ax.set_title('3D визуализация погрешности интерполяции\n(чередование знаков)')

# Добавление цветовой шкалы
fig.colorbar(surf, ax=ax, shrink=0.5, aspect=10, label='Разница (истина - интерполяция)')

plt.show()

# --- Дополнительно: 3D самого распределения n(k) для сравнения ---
fig2 = plt.figure(figsize=(10, 8))
ax2 = fig2.add_subplot(111, projection='3d')
surf2 = ax2.plot_surface(KX_fine, KY_fine, n_fine_real, cmap='hot', linewidth=0, antialiased=True)
ax2.set_xlabel('kx')
ax2.set_ylabel('ky')
ax2.set_zlabel('n(k)')
ax2.set_title('Истинное распределение (экспоненциальный колокол)')
fig2.colorbar(surf2, ax=ax2, shrink=0.5, aspect=10)
plt.show()