# ---------------------------------------------------
# TP Analyse d'un MNT
# ---------------------------------------------------
from docutils.nodes import header
from osgeo import gdal
import numpy as np
import scipy
import matplotlib.pyplot as plt
import sys
import os

# -------------------
# -------------------
# -------------------
# PART 1: Anlayse MNT
# -------------------
# -------------------
# -------------------

# -------------------
# Ouverture de l’image
# -------------------

demfile = "MNT_LiDAR_Ribaute_crop.tif"
filepath = "data/" + demfile
if not os.path.isfile("data/" + demfile):
    print(f"Error: File '{demfile}' not found.")
    sys.exit(1)

ds = gdal.Open(filepath, gdal.GA_ReadOnly)
band = ds.GetRasterBand(1)
topo = band.ReadAsArray()

# -------------------
# Étape 1 - Métadonnées
# -------------------

ncols, nlines = ds.RasterXSize, ds.RasterYSize
print("Taille (X,Y) :", ncols, nlines)
gt = ds.GetGeoTransform()
print("Georeferencement :", gt)
proj = ds.GetProjection()
print("Système de coordonnées :", proj)
print("Type de données :", gdal.GetDataTypeName(band.DataType))

# -------------------
# Étape 2 - Extraires coordoonnées
# -------------------

# Extraire latitudes
x = gt[0] + (np.arange(ncols) * gt[1])
y = gt[3] + (np.arange(nlines) * gt[5])
# une liste en comprehension fonctionne aussi
# x = [gt[0] + col*gt[1] for col in ncols]
print("x, y :", x, y)
print(np.min(x), np.max(x))

# Extraire les limites du raster (xmin, xmax, ymin, ymax)
xmin = x.min()
xmax = x.max()
ymin = y.min()
ymax = y.max()
print("xmin, xmax, ymin, ymax :", xmin, xmax, ymin, ymax)

# -------------------
# Étape 3 - Filtrage, calcul de gradient et calcule de la pente et de l'orientation
# -------------------

# Filter to smooth data and remove noise
filter_size = 2
print(f"Filtering DEM with a {filter_size} window size.")
filtered_topo = scipy.ndimage.gaussian_filter(topo, sigma=(filter_size, filter_size))

# Determine resolution based on coordinate system
res_x, res_y = gt[1], gt[5]
print(f"Resolution: dx={res_x:.2f}, dy={res_y:.2f}")

# Compute gradient (first axis = Y, second axis = X)
# dsm_dy is positive towards north (because pixelHeight < 0)
# dsm_dx is positive towards east (because pixelWidth > 0)
dsm_dy, dsm_dx = np.gradient(filtered_topo, -res_y, -res_x)

# Calculate slope and aspect
slope = np.sqrt(dsm_dx**2 + dsm_dy**2)
slope_deg = np.rad2deg(slope) # conversion en degree
aspect = (np.pi/2) - np.arctan2(dsm_dy, dsm_dx)  # Clockwise from north, arctan2 renvoie dans [-180, 180].
aspect_deg = (np.rad2deg(aspect) + 360) % 360 # forcer dans [0, 360]

# -------------------
# Étape 4 - Calcule de l'ombrage
# -------------------

# calcul de l'ombrage
# Paramètres du soleil (azimut et hauteur en degrés)
azimuth = 315      # soleil à l'ouest-nord-ouest
altitude = 45      # soleil à mi-hauteur

# Conversion en radians
azimuth_rad = np.deg2rad(azimuth)
altitude_rad = np.deg2rad(altitude)

# Calcul du hillshade
hillshade = 255 * (
    np.cos(altitude_rad) * np.cos(slope) +
    np.sin(altitude_rad) * np.sin(slope) *
    np.cos(azimuth_rad - aspect)
)
# print(np.min(hillshade), np.max(hillshade))
# hillshade = np.clip(hillshade, 0, 255)
# print(np.min(hillshade), np.max(hillshade))
# sys.exit(0)
# -------------------
# Étape 5 - Sauvegarde des résultats
# -------------------

# Save outputs to GeoTIFF
def save_raster(filename, array):
    drv = gdal.GetDriverByName('GTiff')
    dst = drv.Create(filename, ncols, nlines, 1, gdal.GDT_Float32)
    dst.SetGeoTransform(gt)
    dst.SetProjection(proj)
    dst.GetRasterBand(1).WriteArray(array)
    dst.FlushCache()
    print(f"Saved: {filename}")

 # extract basename
filename = os.path.splitext(os.path.basename(demfile))[0]
save_raster('data/' + filename + '_slope.tif', slope_deg)
save_raster('data/' + filename + '_aspect.tif', aspect)
save_raster('data/' + filename + '_hillshade.tif', hillshade)

plot = True
if plot:

    # -------------------
    # Étape 6 - Affichage graphique
    # ------
    fig, axes = plt.subplots(2, 2, figsize=(11, 7))
    titles = ['Slope', 'Gradient X', 'Gradient Y', 'Aspect']
    #titles = ['Slope', 'Gradient X', 'Gradient Y', 'Hillshade']
    datasets = [
        slope_deg[:, :],
        dsm_dx[:, :],
        dsm_dy[:, :],
        aspect[:, :]
        #hillshade[:, :]
    ]
    cmaps = ['Greys_r', 'coolwarm', 'coolwarm', 'Greys_r']

    for ax, title, data, cmap in zip(axes.flatten(), titles, datasets, cmaps):
        im = ax.imshow(data, cmap=cmap, vmax=np.nanpercentile(data, 98), vmin=np.nanpercentile(data, 2))
        ax.set_title(title)
        fig.colorbar(im, ax=ax, orientation="vertical")

    fig.savefig('data/' + 'dem_slope_aspect.png', dpi=300)

# plt.show()
# sys.exit()
# -------------------
# -------------------
# -------------------
# PART 2: Geologie scructurale
# -------------------
# -------------------
# -------------------

# Charger les 3 CSV
x1, y1, z1 = np.loadtxt('data/besson_banc1.csv', delimiter=',', skiprows=4, unpack=True)
x2, y2, z2 = np.loadtxt('data/besson_banc2.csv', delimiter=',', skiprows=4, unpack=True)

def fit_plane(x, y, z):
    """
    Ajuste un plan z = ax + by + c à partir d'un nuage de points (x, y, z).
    x, y, x: coordonnées des points
    Retourne (a, b, c).
    """
    A = np.ones((len(x), 3))
    A[:, 0] = x; A[:, 1] = y
    coefficients = np.linalg.lstsq(A, z, rcond=None)[0]
    a, b, c = coefficients
    return a, b, c

# Ajuster un plan au banc 1 et 2
a1_fit, b1_fit, c1_fit = fit_plane(x1, y1, z1)
print(f"Les coefficients du plan ajusté pour le banc 1 sont : {a1_fit:.2f}x + {b1_fit:.2f}y + {c1_fit:.2f} = z")
a2_fit, b2_fit, c2_fit = fit_plane(x2, y2, z2)
print(f"Les coefficients du plan ajusté pour le banc 2 sont : {a2_fit:.2f}x + {b2_fit:.2f}y + {c2_fit:.2f} = z")

def plane2strike_dip(a, b, c):
    """
    Convertit un plan z = ax + by + c en (strike, dip, dip direction).
    Angles en degrés.
    """
    # Pendage
    norm_n = np.sqrt(a**2 + b**2 + c**2)
    cos_theta = c / norm_n  # Utiliser abs(c) pour s'assurer que le pendage est positif
    theta_rad = np.arccos(cos_theta)
    dip = np.degrees(theta_rad)
    # Strike
    strike = np.degrees(np.arctan2(b, a))
    if strike < 0:
        strike += 360
    if dip > 90:
        dip = 180 - dip
    return strike, dip

strike1, dip1 = plane2strike_dip(a1_fit, b1_fit, -1)
print(f"Banc 1 : Strike={strike1:.1f}°, Dip={dip1:.1f}°")

strike2, dip2 = plane2strike_dip(a2_fit, b2_fit, -1)
print(f"Banc 2 : Strike={strike2:.1f}°, Dip={dip2:.1f}°")

# Créer une grille régulière
X, Y = np.meshgrid( np.linspace(xmin, xmax, ncols), np.linspace(ymin, ymax, nlines) )
Z1 = a1_fit*X + b1_fit*Y + c1_fit
Z2 = a2_fit*X + b2_fit*Y + c2_fit

# --- Visualisation ---
fig3, ax = plt.subplots(figsize=(8, 6))
# Image en coordonnées géoréférencées
im = ax.imshow(hillshade, cmap="gray", origin="upper", extent=[xmin, xmax, ymin, ymax])
# Courbes de niveau (même extent)
contours = ax.contour(
    topo, levels=5, colors='black', linewidths=0.5,
    origin='upper', extent=[xmin, xmax, ymin, ymax])
ax.clabel(contours, inline=True, fontsize=8, fmt="%d")

# Points de mesures (déjà en coordonnées réelles x,y)
ax.scatter(x1, y1, c="dodgerblue", s=10, alpha=0.9, label="Banc 1")
ax.scatter(x2, y2, c="mediumseagreen", s=10, alpha=0.9,label="Banc 2")
ax.set_title("MNT ombragé")
ax.legend()
fig3.savefig('data/dem_besson.png', dpi=300)

# Affichage en 3D
fig = plt.figure(figsize=(12,8))
ax = fig.add_subplot(111, projection='3d')

# Affichage des plans ajustés
ax.plot_surface(X, Y, Z1, color="dodgerblue", alpha=0.4, label="Plan 1")
ax.plot_surface(X, Y, Z2, color="mediumseagreen", alpha=0.4, label="Plan 2")

# Affichage des points de mesure
ax.scatter(x1, y1, z1, c="dodgerblue", s=20)
ax.scatter(x2, y2, z2, c="mediumseagreen", s=20)

ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Altitude (m)")
ax.legend()
plt.show()