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

def visualize_daily_profile(filename, day_number):
    """
    Считывает 2D-профиль температуры для указанного дня из HDF5-файла 
    и визуализирует его.
    
    :param filename: Имя HDF5-файла с результатами моделирования.
    :param day_number: Номер дня (индекс профиля) для визуализации. 
                       Начинается с 0.
    """
    if not os.path.exists(filename):
        print(f"ОШИБКА: Файл '{filename}' не найден.")
        return

    try:
        with h5py.File(filename, 'r') as f:
            # 1. Чтение метаданных
            DZ = f.attrs['DZ']
            DX = f.attrs['DX']
            
            # 2. Чтение основных наборов данных
            temp_dset = f['temperature_profiles']
            mask_read = f['mask'][:]

            # Проверка границ дня
            if day_number < 0 or day_number >= temp_dset.shape[0]:
                print(f"ОШИБКА: День {day_number} находится вне диапазона.")
                print(f"Доступный диапазон дней: 0 до {temp_dset.shape[0] - 1}.")
                return

            # 3. Извлечение данных для конкретного дня
            T_profile = temp_dset[day_number, :, :]
            
            # 4. Расчет координат
            N_Z, N_X = T_profile.shape
            X_COORD = np.linspace(0, N_X * DX, N_X)
            Z_COORD = np.linspace(0, N_Z * DZ, N_Z) * -1 # Глубина отрицательна
            
            # 5. Определение диапазона цветов
            # Читаем все данные для правильного Vmin/Vmax (это может быть медленно для очень больших файлов)
            T_all = temp_dset[:].flatten()
            T_min = np.nanmin(T_all)
            T_max = np.nanmax(T_all)

            # 6. Настройка цветовой карты
            cmap_base = plt.colormaps['RdYlBu'] 
            cmap_custom = cmap_base.copy() 
            cmap_custom.set_bad('lightgray') # Области вне трапеции (NaN)
            
            # 7. Визуализация
            
            fig, ax = plt.subplots(figsize=(10, 5))
            fig.suptitle(f'Профиль температуры грунта - День {day_number}', fontsize=14)

            im = ax.imshow(T_profile, 
                           extent=[0, N_X * DX, -N_Z * DZ, 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'Температура на поверхности: {T_profile[0, int(N_X/2)]:.2f} °C')
            ax.set_xlabel('Длина, м')
            ax.set_ylabel('Глубина, м')

            # Добавление цветовой шкалы
            cbar = fig.colorbar(im, ax=ax, orientation='vertical', fraction=0.04, pad=0.04)
            cbar.set_label('Температура, °C')
            
            plt.tight_layout()
            plt.show()

    except Exception as e:
        print(f"Произошла ошибка при обработке данных HDF5: {e}")
        
# ----------------------------------------------------------------------
# 🌟 ПРИМЕР ИСПОЛЬЗОВАНИЯ ИНТЕРФЕЙСА
# ----------------------------------------------------------------------

FILENAME = 'temperature_simulation_results.hdf5'

# Пример 1: Профиль в разгар зимы (День 180)
print("Визуализация дня 180 (Пик холода)...")
visualize_daily_profile(FILENAME, 180)

# Пример 2: Профиль в конце моделирования (День 364)
print("\nВизуализация дня 364 (Конец года)...")
visualize_daily_profile(FILENAME, 364)

# Пример 3: Попытка визуализировать несуществующий день
# visualize_daily_profile(FILENAME, 400)