# Nachfolgend sind zwei Verarbeitungsmöglichkeiten zur Shapefilerstellung gegeben: 
# 1: Erstellung des Shapelayers für eine bestimmte .txt Datei:
import os
import pyproj
from shapely.geometry import Point
import geopandas as gpd

# Klassen-Mapping
classes = {
    0: 'Großer Laubbaum',
    1: 'Kleiner Laubbaum'
}

# Bildparameter
pixel_size = 0.2  # 20 cm in Metern
img_width = 200
img_height = 200

# Pfad zur .txt Datei
file_name = '4362024_5320180_rgbi_20cm.txt'
file_path = 'Trained models/nano_500 Epochs_15_70_15/predictions/' + file_name

# Umwandlung der Koordinaten in WGS84
crs_31468 = pyproj.CRS('EPSG:31468')
crs_4326 = pyproj.CRS('EPSG:4326')
transformer = pyproj.Transformer.from_crs(crs_31468, crs_4326, always_xy=True)

# Lesen der .txt Datei
with open(file_path, 'r') as file:
    lines = file.readlines()

# Extrahieren der Koordinaten aus dem Dateinamen
filename = os.path.basename(file_path)
x_base, y_base = map(int, filename.split('_')[:2])

# Berechnung der Objektkoordinaten
objects = []
for i, line in enumerate(lines):
    class_id, x_center_rel, y_center_rel, width_rel, height_rel = map(float, line.split())
    x_center = x_base + (x_center_rel * img_width * pixel_size)
    y_center = y_base + ((1 - y_center_rel) * img_height * pixel_size)
    
    # Höhe abhängig von der Klasse als Text setzen
    if classes[int(class_id)] == 'Kleiner Laubbaum':
        height = "<= 12.19"  # Für Kleiner Laubbaum
    else:
        height = ">= 12.19"  # Für Großer Laubbaum

    objects.append({
        'ID': i + 1,
        'Klasse': classes[int(class_id)],
        'Höhe': height,
        'Längengrad': x_center,
        'Breitengrad': y_center
    })

# Konvertiere die Koordinaten in WGS84
for obj in objects:
    obj['Längengrad'], obj['Breitengrad'] = transformer.transform(obj['Längengrad'], obj['Breitengrad'])

# GeoDataFrame erstellen
gdf = gpd.GeoDataFrame(objects, geometry=[Point(xy) for xy in zip(
    [obj['Längengrad'] for obj in objects],
    [obj['Breitengrad'] for obj in objects]
)], crs="EPSG:4326")

# Shapefile speichern
shapefile_dir = 'shapefile_output'
if not os.path.exists(shapefile_dir):
    os.makedirs(shapefile_dir)

shapefile_path = os.path.join(shapefile_dir, 'detected_objects_in_' + file_name + '.shp')
gdf.to_file(shapefile_path)

# 2: Erstellung eines Shapelayers für alle Objekte aus 'predictions':
import os
import pyproj
import shapefile
import random
from shapely.geometry import Point
import geopandas as gpd

# Klassen-Mapping
classes = {
    0: 'Großer Laubbaum',
    1: 'Kleiner Laubbaum',
    2: 'Großer Nadelbaum',
    3: 'Kleiner Nadelbaum',
    4: 'Busch/Hecke (Laub/Hartlaub)',
    5: 'Busch/Hecke (Nadel)',
    6: 'Unbekannt'
}

# Bildparameter
pixel_size = 0.1875  # 20 cm in Metern
img_width = 128
img_height = 128

# Pfad zum predictions Ordner
predictions_dir = 'predictions'

# Liste aller .txt Dateien im predictions Ordner
txt_files = [f for f in os.listdir(predictions_dir) if f.endswith('.txt')]

# Umwandlung der Koordinaten in WGS84
crs_31468 = pyproj.CRS('EPSG:31468')
crs_4326 = pyproj.CRS('EPSG:4326')
transformer = pyproj.Transformer.from_crs(crs_31468, crs_4326, always_xy=True)

# Alle Objekte sammeln
all_objects = []

# Durch alle .txt Dateien iterieren
for txt_file in txt_files:
    file_path = os.path.join(predictions_dir, txt_file)
    
    # Lesen der .txt Datei
    with open(file_path, 'r') as file:
        lines = file.readlines()

    # Extrahieren der Koordinaten aus dem Dateinamen
    filename = os.path.basename(file_path)
    x_base, y_base = map(int, filename.split('_')[:2])

    # Berechnung der Objektkoordinaten
    for i, line in enumerate(lines):
        class_id, x_center_rel, y_center_rel, width_rel, height_rel = map(float, line.split())
        x_center = x_base + (x_center_rel * img_width * pixel_size)
        y_center = y_base + ((1-y_center_rel) * img_height * pixel_size)
        height = random.randint(5, 20)  # Zufällige Höhe

        all_objects.append({
            'ID': len(all_objects) + 1,
            'Klasse': classes[int(class_id)],
            'Höhe': height,
            'Längengrad': x_center,
            'Breitengrad': y_center
        })

# Konvertiere die Koordinaten in WGS84 und erstelle die Geometrien
for obj in all_objects:
    obj['Längengrad'], obj['Breitengrad'] = transformer.transform(obj['Längengrad'], obj['Breitengrad'])

# GeoDataFrame erstellen
gdf = gpd.GeoDataFrame(all_objects, geometry=[Point(xy) for xy in zip(
    [obj['Längengrad'] for obj in all_objects],
    [obj['Breitengrad'] for obj in all_objects]
)], crs="EPSG:4326")

# Shapefile speichern
shapefile_dir = 'shapefile_output'
if not os.path.exists(shapefile_dir):
    os.makedirs(shapefile_dir)

shapefile_path = os.path.join(shapefile_dir, 'detected_objects.shp')
gdf.to_file(shapefile_path)