import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.colors as mcolors

# Определяем новые параметры трапеции
height = 200000  # Высота трапеции увеличена до 200000 см
top_width = 8  # Ширина верхней стороны
bottom_width = 12  # Ширина нижней стороны
depth = 8  # Глубина трапеции

# Генерация температуры на каждом уровне z (реверсированная)
def calculate_temperature(z):
    # Температура изменяется от 0°C у основания до 100°C на вершине
    return 100 * (z / height)

# Сетка для разбиения трапеции на кубы
# Чтобы избежать слишком больших затрат памяти, увеличим шаг кубов до, например, 10000 см.
x_step, y_step, z_step = 1000, 1000, 10000
x_range = np.arange(0, bottom_width, x_step)
y_range = np.arange(0, depth, y_step)
z_range = np.arange(0, height, z_step)

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# Нормализация и палитра цветов
cmap = plt.get_cmap("coolwarm")
norm = mcolors.Normalize(vmin=0, vmax=100)

# Функция для добавления куба к визуализации
def add_cube(ax, x, y, z, temperature):
    vertices = [
        [x, y, z], [x+x_step, y, z], [x+x_step, y+y_step, z], [x, y+y_step, z],  # нижняя грань
        [x, y, z+z_step], [x+x_step, y, z+z_step], [x+x_step, y+y_step, z+z_step], [x, y+y_step, z+z_step]  # верхняя грань
    ]
    edges = [
        [vertices[j] for j in [0, 1, 2, 3]],  # нижняя грань
        [vertices[j] for j in [4, 5, 6, 7]],  # верхняя грань
        [vertices[j] for j in [0, 1, 5, 4]],  # боковая грань
        [vertices[j] for j in [2, 3, 7, 6]],  # боковая грань
        [vertices[j] for j in [0, 3, 7, 4]],  # передняя грань
        [vertices[j] for j in [1, 2, 6, 5]]   # задняя грань
    ]
    
    # Определяем цвет в зависимости от температуры
    color = cmap(norm(temperature))
    
    ax.add_collection3d(Poly3DCollection(edges, color=color, edgecolor='black', linewidths=0.1, alpha=0.7))

# Заполняем трапецию кубами
for z in z_range:
    # Линейная интерполяция для ширины на каждом уровне высоты
    width = top_width + (bottom_width - top_width) * (height - z) / height
    x_min = (bottom_width - width) / 2
    x_max = x_min + width
    for x in np.arange(x_min, x_max, x_step):
        for y in y_range:
            temperature = calculate_temperature(z)
            add_cube(ax, x, y, z, temperature)

# Настройки графика
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.set_xlim(0, bottom_width)
ax.set_ylim(0, depth)
ax.set_zlim(0, height)

# Создаём цветовую шкалу
mappable = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
mappable.set_array([])
plt.colorbar(mappable, ax=ax, label="Температура (°C)")

plt.show()