import streamlit as st
import rasterio
import numpy as np
import pyvista as pv
from pyvistaqt import BackgroundPlotter
import tempfile

# Fonction pour charger un raster
def load_raster(file_path):
    with rasterio.open(file_path) as src:
        raster = src.read(1)  # Lecture de la première bande
        transform = src.transform  # Transformation affine du raster
    return raster, transform

# Fonction pour créer un modèle 3D
def create_3d_model(raster, transform, vertical_exaggeration=1.0):
    rows, cols = raster.shape
    x = np.linspace(transform[2], transform[2] + transform[0] * cols, cols)
    y = np.linspace(transform[5], transform[5] + transform[4] * rows, rows)
    z = raster * vertical_exaggeration

    # Création des coordonnées
    xx, yy = np.meshgrid(x, y)
    xx = xx.flatten()
    yy = yy.flatten()
    zz = z.flatten()

    # Création de la grille structurée
    points = np.c_[xx, yy, zz]
    grid = pv.StructuredGrid()
    grid.points = points
    grid.dimensions = (cols, rows, 1)  # Dimensions de la grille
    return grid

# Streamlit App
st.title("Maquette 3D à partir d'un Raster")
uploaded_file = st.file_uploader("Téléchargez un fichier raster (MNT ou MES)", type=["tif"])

if uploaded_file:
    raster, transform = load_raster(uploaded_file)
    st.success("Raster chargé avec succès!")

    # Configuration utilisateur
    vertical_exaggeration = st.slider("Échelle verticale (exagération)", 0.1, 10.0, 1.0)

    # Création du modèle 3D
    st.subheader("Modèle 3D")
    grid = create_3d_model(raster, transform, vertical_exaggeration)

    # Visualisation avec PyVista
    plotter = pv.Plotter(off_screen=True)
    plotter.add_mesh(grid, scalars=raster.flatten(), cmap="terrain", show_edges=False)
    plotter.add_scalar_bar("Élévation (m)")
    plotter.view_isometric()
    screenshot_path = tempfile.NamedTemporaryFile(suffix=".png").name
    plotter.show(screenshot=screenshot_path)

    # Affichage dans Streamlit
    st.image(screenshot_path, caption="Modèle 3D", use_column_width=True)

    # Exportation du modèle
    st.subheader("Exportation du Modèle")
    export_format = st.selectbox("Format d'exportation", ["STL", "OBJ"])
    if st.button("Exporter"):
        export_file = tempfile.NamedTemporaryFile(suffix=f".{export_format.lower()}", delete=False).name
        grid.save(export_file)
        st.success(f"Modèle exporté sous forme de fichier {export_format}.")
        st.download_button("Télécharger le fichier exporté", data=open(export_file, "rb").read(), file_name=f"model.{export_format.lower()}")