import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import h5py 

# --- 1. Параметры и сетка (2D) ---
DX = 0.1  
DZ = 0.1

Z_MAX = 2.0  
X_MAX = 10.0 

N_Z = int(Z_MAX / DZ) + 1
N_X = int(X_MAX / DX) + 1

# Время и теплофизические параметры
N_DAYS = 365 
DT = 1800  # 30 минут (сек).
ALPHA = 7.5e-7 
R = ALPHA * DT / (DX**2) 
N_STEPS = int(N_DAYS * 24 * 3600 / DT) 

# Параметры сохранения данных
HDF5_FILENAME = 'temperature_simulation_results.hdf5' 
RECORD_INTERVAL_STEPS = int(24 * 3600 / DT) 

print(f"2D-сетка: {N_Z}x{N_X} узлов.")
print(f"Устойчивость соблюдена: R = {R:.4f} <= 0.25.")

# --- 2. Инициализация и геометрическая маска (Трапеция) ---

T_INITIAL = 5.0  # °C
Temperature = np.full((N_Z, N_X), T_INITIAL, dtype=float) 

# Создание маски для области ТРАПЕЦИИ:
mask = np.full((N_Z, N_X), False)

# Параметры трапеции (ширина на дне 10м, на поверхности 5м)
W_bottom = X_MAX
W_top = 5.0

for i in range(N_Z):
    depth_i = i * DZ
    W_z = W_bottom - (W_bottom - W_top) * (1 - depth_i / Z_MAX)
    offset_x = (X_MAX - W_z) / 2
    
    start_j = int(offset_x / DX)
    end_j = int((X_MAX - offset_x) / DX)
    
    start_j = max(0, start_j)
    end_j = min(N_X, end_j)
    
    mask[i, start_j:end_j] = True
    
# Обнуляем температуру вне маски (заполняем NaN)
Temperature[~mask] = np.nan

# --- 3. Граничные условия (Упрощенные) ---

def surface_temp(step):
    """Температура поверхности (Z=0) - синусоида 1 год."""
    time_in_seconds = step * DT
    time_in_years = time_in_seconds / (N_DAYS * 24.0 * 3600.0)
    T_avg = 0.0
    T_amp = 20.0
    return T_avg + T_amp * np.sin(2 * np.pi * time_in_years + np.pi/2)

T_BOTTOM = T_INITIAL 

# --- 4. Цикл моделирования и сохранение в HDF5 (Без изменений) ---

print(f"Начало моделирования ({N_STEPS} шагов). Результаты будут сохранены в {HDF5_FILENAME}...")

with h5py.File(HDF5_FILENAME, 'w') as f:
    
    f.attrs['DZ'] = DZ
    f.attrs['DX'] = DX
    f.attrs['N_DAYS'] = N_DAYS
    f.attrs['DT_seconds'] = DT
    f.create_dataset('mask', data=mask, dtype='bool')
    
    temp_dset = f.create_dataset('temperature_profiles', 
                                 shape=(0, N_Z, N_X), 
                                 maxshape=(N_DAYS + 1, N_Z, N_X),
                                 dtype='f4', 
                                 chunks=(1, N_Z, N_X),
                                 compression="gzip")

    record_counter = 0

    for n in range(N_STEPS):
        T_old = Temperature.copy()

        T_surface = surface_temp(n)
        Temperature[0, mask[0, :]] = T_surface
        Temperature[N_Z - 1, mask[N_Z - 1, :]] = T_BOTTOM

        for i in range(1, N_Z - 1):
            for j in range(1, N_X - 1):
                if mask[i, j]:
                    T_i_minus_1 = T_old[i-1, j] if mask[i-1, j] else T_old[i, j]
                    T_i_plus_1 = T_old[i+1, j] if mask[i+1, j] else T_old[i, j]
                    T_j_minus_1 = T_old[i, j-1] if mask[i, j-1] else T_old[i, j]
                    T_j_plus_1 = T_old[i, j+1] if mask[i, j+1] else T_old[i, j]

                    T_new = T_old[i, j] * (1 - 4 * R) + R * (T_i_minus_1 + T_i_plus_1 + T_j_minus_1 + T_j_plus_1)

                    Temperature[i, j] = T_new
                    
        if n % RECORD_INTERVAL_STEPS == 0:
            temp_dset.resize(record_counter + 1, axis=0)
            temp_dset[record_counter, :, :] = Temperature
            record_counter += 1
            
            if record_counter % 50 == 0:
                 print(f"Записан профиль за день {record_counter}.")


print(f"Моделирование завершено. Сохранено {record_counter} профилей температуры в файле {HDF5_FILENAME}.")

# --- 5. Пример визуализации (Чтение из HDF5 с исправлениями) ---

try:
    with h5py.File(HDF5_FILENAME, 'r') as f:
        temp_dset_read = f['temperature_profiles']
        mask_read = f['mask'][:]
        
        days_to_plot = [0, 90, 180, 270, N_DAYS - 1] # N_DAYS - 1, так как 365-й день это индекс 364
        indices_to_plot = [day for day in days_to_plot if day < temp_dset_read.shape[0]]

        if not indices_to_plot:
            print("Нет сохраненных данных для визуализации.")
        else:
            X_COORD = np.linspace(0, X_MAX, N_X)
            Z_COORD = np.linspace(-Z_MAX, 0, N_Z)

            fig, axes = plt.subplots(1, len(indices_to_plot), figsize=(18, 5))
            fig.suptitle(f'2D Теплоперенос (Чтение из {HDF5_FILENAME})', fontsize=14)

            T_all = temp_dset_read[:].flatten()
            T_min = np.nanmin(T_all)
            T_max = np.nanmax(T_all)

            # ИСПРАВЛЕНИЕ 1: Использование plt.colormaps
            cmap_base = plt.colormaps['RdYlBu'] 
            # Создаем копию для изменения "плохих" значений
            cmap_custom = cmap_base.copy() 
            cmap_custom.set_bad('lightgray')

            for idx, day_index in enumerate(indices_to_plot):
                T_profile = temp_dset_read[day_index, :, :]
                # Проверяем, что оси - это массив, если только одна ось
                ax = axes[idx] if len(indices_to_plot) > 1 else axes 

                im = ax.imshow(T_profile, 
                               extent=[0, X_MAX, -Z_MAX, 0], 
                               origin='upper', 
                               cmap=cmap_custom, 
                               aspect='auto', 
                               vmin=T_min, 
                               vmax=T_max)
                
                ax.contour(X_COORD, Z_COORD, mask_read.astype(int), levels=[0.5], colors='k', linewidths=1.5, linestyles='--')
                ax.set_title(f'День: {day_index}')
                ax.set_xlabel('Длина, м')
                ax.set_ylabel('Глубина, м')

            # ИСПРАВЛЕНИЕ 2: Удаление rect=[0, 0, 0.98, 1] и let axes handle spacing
            plt.tight_layout() 
            
            # Добавление цветовой шкалы
            # Убедимся, что cbar создается корректно, используя массив осей
            cbar = fig.colorbar(im, ax=axes.ravel().tolist() if len(indices_to_plot) > 1 else axes, orientation='vertical', fraction=0.02, pad=0.04)
            cbar.set_label('Температура, °C')
            
            plt.show()

except Exception as e:
    print(f"Ошибка при чтении или визуализации HDF5: {e}")