import numpy as np
import matplotlib.pyplot as plt

# --- 1. Определение параметров моделирования ---

# Геометрические параметры (1D модель: только глубина Z)
DZ = 0.1  # Размер ячейки по глубине (Z) в метрах (10 см)
Z_MAX = 2.0  # Максимальная глубина моделирования в метрах
L_MAX = 10.0  # Максимальная длина (не используется в 1D модели)
W_MAX = 22.0  # Максимальная ширина (не используется в 1D модели)

# Количество узлов (ячеек) по глубине
N_NODES = int(Z_MAX / DZ) + 1
DEPTHS = np.linspace(0, Z_MAX, N_NODES)

# Теплофизические параметры грунта (приближенные значения для талого грунта)
LAMBDA = 1.5  # Теплопроводность грунта (λ), Вт/(м·К)
RHO_C = 2.0e6  # Объемная теплоемкость (ρ·C), Дж/(м³·К)
ALPHA = LAMBDA / RHO_C  # Температуропроводность (α), м²/с

# Параметры времени
DT_HOUR = 3600  # Шаг по времени (Δt) в секундах (1 час)
N_DAYS = 365  # Продолжительность моделирования в днях (1 год)
N_STEPS = N_DAYS * 24  # Общее количество временных шагов

# Критерий устойчивости (Число Фурье, r)
R = (ALPHA * DT_HOUR) / (DZ**2)

# Проверка устойчивости явной схемы
if R > 0.5:
    print(f"ВНИМАНИЕ: Нарушен критерий устойчивости (R = {R:.4f}).")
    print("Уменьшите шаг по времени (DT_HOUR) или шаг по пространству (DZ).")
    R = 0.5  # Принудительное ограничение для демонстрации
else:
    print(f"Критерий устойчивости соблюден (R = {R:.4f} <= 0.5).")

# --- 2. Инициализация температуры ---

# Начальная температура грунта (равномерное распределение)
T_INITIAL = 5.0  # °C
Temperature = np.full(N_NODES, T_INITIAL)

# Массив для записи результатов (для построения графика)
temp_history = [Temperature.copy()]

# --- 3. Определение граничных условий (ГКУ) ---

def surface_temperature(step):
    """
    Граничное условие (ГКУ) на поверхности (z=0).
    Используется синусоидальная функция для имитации сезонного цикла (1 год).
    """
    time_in_hours = step * DT_HOUR
    time_in_years = time_in_hours / (365 * 24 * 3600)

    # Параметры: T_avg = 0°C, Amplitude = 20°C, сдвиг фазы (начало зимы)
    T_avg = 0.0
    T_amp = 20.0
    # Фаза: +pi/2, чтобы начать симуляцию с T(0) = T_avg
    return T_avg + T_amp * np.sin(2 * np.pi * time_in_years + np.pi/2)

# ГКУ на нижней границе (z=Z_MAX) - постоянная температура
T_BOTTOM = T_INITIAL

# --- 4. Цикл моделирования ---

print(f"Начало моделирования: {N_DAYS} дней ({N_STEPS} шагов по 1 часу)...")

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

    # 1. Применение граничных условий
    # Верхняя граница (i=0) - Температура воздуха (Dirichlet BC)
    T_surface = surface_temperature(n)
    Temperature[0] = T_surface

    # Нижняя граница (i=N_NODES-1) - Постоянная температура (Dirichlet BC)
    Temperature[-1] = T_BOTTOM

    # 2. Расчет температуры для внутренних узлов (i=1 до N_NODES-2)
    # Формула явной схемы FDM: T_i^{n+1} = T_i^n (1 - 2r) + r (T_{i-1}^n + T_{i+1}^n)
    for i in range(1, N_NODES - 1):
        # --- Здесь должно быть место для учета фазовых переходов (мерзлота/талый грунт) ---
        # В реальной модели R и T_i^{n+1} должны рассчитываться с учетом
        # изменения LAMBDA, RHO_C и скрытой теплоты L (L_v, формула 3.13)
        # в зависимости от T_i^n и температуры фазового перехода (T_bf, формула 3.7).
        # В этой упрощенной модели LAMBDA, RHO_C, R - константы.
        # ---------------------------------------------------------------------------------

        T_new = T_old[i] * (1 - 2 * R) + R * (T_old[i-1] + T_old[i+1])
        Temperature[i] = T_new

    # Запись истории температуры каждые 90 дней для графика
    if (n % (90 * 24) == 0) and (n > 0):
        temp_history.append(Temperature.copy())

temp_history.append(Temperature.copy())
print("Моделирование завершено.")

# --- 5. Визуализация результатов ---

plt.figure(figsize=(10, 6))
time_labels = [f"День {int(i / 24)}" for i in np.linspace(0, N_STEPS, len(temp_history))]

# Добавление начального состояния
plt.plot(temp_history[0], -DEPTHS, label=f't=0 (Начальное)', linestyle='--')

# Построение графиков для разных моментов времени
for i, T_profile in enumerate(temp_history[1:]):
    plt.plot(T_profile, -DEPTHS, label=f'{time_labels[i+1]}', linewidth=2)

plt.xlabel('Температура, °C')
plt.ylabel('Глубина, м')
plt.title('Одномерное моделирование теплопереноса в грунте (1 год)')
plt.grid(True, linestyle=':', alpha=0.6)
plt.legend(title='Время')
plt.axvline(x=0, color='grey', linestyle='-')
plt.show()