---

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import RectBivariateSpline

# ========== 1. Параметры задачи ==========
# Единицы: ħ = 1, m_e = 1
# Условие ħ^2/(2 m_e T) = 1/4 => 1/(2T) = 1/4 => T = 2
T = 2.0

# Эффективные массы дырок (берём по модулю, т.к. энергия должна быть положительной)
m_x = 4.0   # |m_x| = 4 m_e
m_y = 0.25  # |m_y| = m_e/4

# Функция точной энергии (ħ=1): ε = (kx^2/(2 m_x) + ky^2/(2 m_y))
def energy(kx, ky):
    return 0.5 * (kx**2 / m_x + ky**2 / m_y)

# Функция распределения Максвелла: n = exp(-ε / T)
def exact_n(kx, ky):
    return np.exp(-energy(kx, ky) / T)

# ========== 2. Грубая сетка (шаг 1) ==========
K = 5                     # полуширина области
step_coarse = 1.0
kx_coarse = np.arange(-K, K + step_coarse, step_coarse)
ky_coarse = np.arange(-K, K + step_coarse, step_coarse)
KX_coarse, KY_coarse = np.meshgrid(kx_coarse, ky_coarse, indexing='ij')
n_coarse = exact_n(KX_coarse, KY_coarse)

# Построение тепловой карты для грубой сетки (задание а)
plt.figure(figsize=(6,5))
plt.imshow(n_coarse.T, origin='lower', extent=[-K, K, -K, K],
           cmap='hot', interpolation='nearest')
plt.colorbar(label='$n(k_x, k_y)$')
plt.title('(a) Распределение на грубой сетке (шаг 1)')
plt.xlabel('$k_x$')
plt.ylabel('$k_y$')
plt.tight_layout()
plt.savefig('task_a_coarse.png', dpi=150)
plt.show()

# ========== 3. Мелкая сетка для интерполяции ==========
step_fine = 0.2
kx_fine = np.arange(-K, K + step_fine, step_fine)
ky_fine = np.arange(-K, K + step_fine, step_fine)
KX_fine, KY_fine = np.meshgrid(kx_fine, ky_fine, indexing='ij')

# Интерполяция с помощью RectBivariateSpline (сплайны)
# Обучаем на грубой сетке (координаты должны быть возрастающими)
spline = RectBivariateSpline(kx_coarse, ky_coarse, n_coarse, kx=3, ky=3)
n_interp = spline(kx_fine, ky_fine)

# Построение тепловой карты интерполированных значений (задание б)
plt.figure(figsize=(6,5))
plt.imshow(n_interp.T, origin='lower', extent=[-K, K, -K, K],
           cmap='hot', interpolation='bilinear')
plt.colorbar(label='$n_{\\text{interp}}(k_x, k_y)$')
plt.title('(b) Интерполяция на мелкую сетку (шаг 0.2)')
plt.xlabel('$k_x$')
plt.ylabel('$k_y$')
plt.tight_layout()
plt.savefig('task_b_interp.png', dpi=150)
plt.show()

# ========== 4. Точные значения на мелкой сетке и разница ==========
n_exact_fine = exact_n(KX_fine, KY_fine)
difference = n_interp - n_exact_fine   # разница между интерполяцией и точным значением

# Построение карты разницы
plt.figure(figsize=(6,5))
im = plt.imshow(difference.T, origin='lower', extent=[-K, K, -K, K],
                cmap='RdBu', interpolation='bilinear', vmin=-0.05, vmax=0.05)
plt.colorbar(im, label='$n_{\\text{interp}} - n_{\\text{exact}}$')
plt.title('(c) Разница между интерполяцией и точным значением')
plt.xlabel('$k_x$')
plt.ylabel('$k_y$')
plt.tight_layout()
plt.savefig('task_c_diff.png', dpi=150)
plt.show()

# ========== 5. Оценка погрешности интерполяции (RMSE) ==========
rmse = np.sqrt(np.mean((n_interp - n_exact_fine)**2))
print(f'Среднеквадратичное отклонение (RMSE) на мелкой сетке: {rmse:.6f}')

# Дополнительно: гистограмма ошибок
plt.figure(figsize=(5,4))
plt.hist(difference.flatten(), bins=30, edgecolor='black')
plt.xlabel('Ошибка интерполяции')
plt.ylabel('Частота')
plt.title('Гистограмма ошибок')
plt.tight_layout()
plt.savefig('task_c_hist.png', dpi=150)
plt.show()