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

def initialize_trapezoid_simulation(depth_max, length_bottom, length_top, T_initial, dz_dx=0.1):
    """
    Создает заготовку для 2D-моделирования теплопереноса в трапециевидном грунтовом массиве.
    
    :param depth_max: Максимальная глубина (Z_MAX), м.
    :param length_bottom: Ширина трапеции на дне (при Z_MAX), м.
    :param length_top: Ширина трапеции на поверхности (при Z=0), м.
    :param T_initial: Начальная равномерная температура грунта, °C.
    :param dz_dx: Размер ячейки сетки (DX = DZ), м.
    :return: Словарь с инициализированными данными (Temperature, Mask, Coordinates).
    """
    
    # 1. Определение параметров сетки
    DX = DZ = dz_dx
    
    N_Z = int(depth_max / DZ) + 1
    N_X = int(length_bottom / DX) + 1

    # 2. Инициализация и геометрическая маска
    
    # Создаем 2D-массив температуры, заполненный начальной температурой T_initial
    Temperature = np.full((N_Z, N_X), T_initial, dtype=float)
    
    # Создаем булеву маску для определения области трапеции
    mask = np.full((N_Z, N_X), False)

    for i in range(N_Z):
        depth_i = i * DZ
        
        # Линейное изменение ширины по глубине Z
        # W(Z) = L_bottom - (L_bottom - L_top) * (1 - Z/Z_MAX)
        W_z = length_bottom - (length_bottom - length_top) * (1 - depth_i / depth_max)
        
        # Центрирование: отступ слева
        offset_x = (length_bottom - W_z) / 2
        
        # Индексы X, которые попадают в трапецию
        start_j = int(offset_x / DX)
        end_j = int((length_bottom - 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. Визуализация Дня 0
    
    X_COORD = np.linspace(0, length_bottom, N_X)
    Z_COORD = np.linspace(-depth_max, 0, N_Z)

    fig, ax = plt.subplots(figsize=(10, 5))
    fig.suptitle(f'Начальное состояние моделирования (День 0)', fontsize=14)

    # Настройка цветовой карты для NaN
    cmap_base = plt.colormaps['RdYlBu'] 
    cmap_custom = cmap_base.copy() 
    cmap_custom.set_bad('lightgray')

    im = ax.imshow(Temperature, 
                   extent=[0, length_bottom, -depth_max, 0], 
                   origin='upper', 
                   cmap=cmap_custom, 
                   aspect='auto', 
                   vmin=T_initial - 1,  # Устанавливаем Vmin/Vmax для старта
                   vmax=T_initial + 1)
    
    # Наложение контура трапеции
    ax.contour(X_COORD, Z_COORD, mask.astype(int), levels=[0.5], colors='k', linewidths=1.5, linestyles='--')
    
    ax.set_title(f'Начальная температура: {T_initial:.1f} °C')
    ax.set_xlabel('Длина (X), м')
    ax.set_ylabel('Глубина (Z), м')

    cbar = fig.colorbar(im, ax=ax, orientation='vertical', fraction=0.04, pad=0.04)
    cbar.set_label('Температура, °C')
    
    plt.tight_layout()
    plt.show()
    
    # 4. Возврат данных для дальнейших расчетов
    return {
        'Temperature_Initial': Temperature,
        'Mask': mask,
        'DX': DX,
        'DZ': DZ,
        'N_X': N_X,
        'N_Z': N_Z
    }

# ----------------------------------------------------------------------
## 🚀 Пример использования
# ----------------------------------------------------------------------

# ВХОДНЫЕ ПАРАМЕТРЫ:
MAX_DEPTH = 2.0         # Глубина 2 м
WIDTH_BOTTOM = 22.0     # Ширина на дне 22 м
WIDTH_TOP = 10.0        # Ширина на поверхности 10 м
INITIAL_TEMP = 5.0      # Начальная температура грунта 5 °C
CELL_SIZE = 0.1         # Размер ячейки 10 см

# Выполнение инициализации и визуализация
initial_data = initialize_trapezoid_simulation(
    depth_max=MAX_DEPTH, 
    length_bottom=WIDTH_BOTTOM, 
    length_top=WIDTH_TOP, 
    T_initial=INITIAL_TEMP,
    dz_dx=CELL_SIZE
)

print("\n--- Сводка по инициализации ---")
print(f"Размеры сетки (Z x X): {initial_data['N_Z']} x {initial_data['N_X']}")
print(f"Размер ячейки: {initial_data['DX']} м")

# Теперь initial_data можно передать в цикл моделирования 
# (который мы создавали ранее), чтобы начать расчеты.