import os
import csv
import json
import numpy as np
from osgeo import gdal, ogr
from tqdm import tqdm
from rasmas.converter import convert_dat_to_tiff  # Converts `.dat` to `.tif` if needed
import re
import rasterio

def get_wavelengths_from_hdr(hdr_file):
    """Extracts wavelength values from an ENVI metadata header, handling multi-line formats and extra metadata."""
    if not os.path.exists(hdr_file):
        print(f"[WARNING] Header file {hdr_file} not found. Assigning band numbers as headers.")
        return None  # Use default band numbers

    wavelengths = []
    inside_wavelength_section = False
    collected_lines = ""  # Stores accumulated wavelength lines

    with open(hdr_file, "r") as file:
        for line in file:
            line = line.strip()

            # Detect the start of the wavelength section
            if "wavelength" in line.lower():
                inside_wavelength_section = True
                line = line.split("=")[-1].strip()  # Get everything after "="

            if inside_wavelength_section:
                line = line.replace("{", "").replace("}", "").strip()
                collected_lines += " " + line  # Collect values from multiple lines

                if "}" in line:  # End of wavelength section
                    inside_wavelength_section = False

    # ✅ Remove any extra metadata that is not a float
    cleaned_values = []
    for item in collected_lines.split(","):
        item = item.strip()
        match = re.match(r"^\d+(\.\d+)?$", item)  # Match only valid numbers
        if match:
            cleaned_values.append(float(item))

    # ✅ Return cleaned wavelengths
    if cleaned_values:
        return cleaned_values
    else:
        print("[WARNING] No valid wavelengths extracted.")
        return None


def create_csv_with_headers(csv_path, num_bands, wavelengths=None):
    """Creates a CSV file with appropriate headers."""
    if not csv_path.lower().endswith('.csv'):
        csv_path += '.csv'

    # Use wavelengths if available; otherwise, assign band numbers
    if wavelengths and len(wavelengths) == num_bands:
        header = ['Plot_ID'] + [f"{w} nm" for w in wavelengths]
    else:
        header = ['Plot_ID'] + [f'Band_{i+1}' for i in range(num_bands)]

    # Create CSV file with headers
    with open(csv_path, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(header)
    print(f"[INFO] Created CSV: {csv_path}")


def append_data_to_csv(csv_path, plot_id, values):
    """Appends spectral data of a plot to the CSV file."""
    if not os.path.isfile(csv_path):
        print(f"[ERROR] CSV file not found: {csv_path}. Creating a new one.")

        # Create headers dynamically based on the number of bands
        num_bands = len(values)
        header = ['Plot_ID'] + [f'Band_{i+1}' for i in range(num_bands)]
        
        with open(csv_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(header)

    row = [plot_id] + list(values)

    with open(csv_path, 'a', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(row)


def extract_mean_spectra_without_clip(input_file, shp_file, output_csv, unique_field="Plotid"):
    """Processes hyperspectral images in memory using GDAL (without saving clipped files)."""
    
    # Convert `.dat` to `.tif` if needed
    if input_file.lower().endswith(".dat"):
        print(f"[INFO] Converting {input_file} to TIFF format...")
        tiff_file = convert_dat_to_tiff(input_file)
        hdr_file = input_file.replace(".dat", ".hdr")
        wavelengths = get_wavelengths_from_hdr(hdr_file)  
    else:
        tiff_file = input_file
        wavelengths = None  

    dataset = gdal.Open(tiff_file, gdal.GA_ReadOnly)
    if dataset is None:
        raise Exception(f"[ERROR] Failed to open image file: {tiff_file}")

    num_bands = dataset.RasterCount  

    shp_dataset = ogr.Open(shp_file)
    if shp_dataset is None:
        raise Exception(f"[ERROR] Failed to open shapefile: {shp_file}")

    layer = shp_dataset.GetLayer()
    print(f"[INFO] Opened shapefile '{shp_file}'. Found {layer.GetFeatureCount()} plots.")

    create_csv_with_headers(output_csv, num_bands, wavelengths)
    count = 0

    for feature in tqdm(layer, desc="Processing Plots"):
        plot_id = feature.GetField(unique_field)
        geometry = feature.GetGeometryRef()

        if plot_id is None:
            print(f"[WARNING] Skipping feature with missing `{unique_field}` attribute.")
            continue

        # **Clipping in Memory**
        options = gdal.WarpOptions(
            format="MEM",
            cutlineDSName=shp_file,
            cutlineLayer=layer.GetName(),
            cutlineWhere=f"{unique_field} = '{plot_id}'",
            cropToCutline=True,
            dstNodata=0  
        )
        clipped = gdal.Warp("", dataset, options=options)  

        if clipped is None:
            print(f"[WARNING] Skipping plot {plot_id} due to failed clipping.")
            continue

        band_data = []
        for band_idx in range(1, num_bands + 1):
            band = clipped.GetRasterBand(band_idx)
            band_array = band.ReadAsArray()
            band_data.append(np.nanmean(np.where(band_array > 0, band_array, np.nan)))

        if not np.all(np.isnan(band_data)):
            append_data_to_csv(output_csv, f"plot_id_{plot_id}", band_data)
            count += 1

    print(f"[INFO] Processing completed. Processed {count} plots. Data saved in {output_csv}.")
    return count

def extract_mean_spectra_with_clip(input_file, shp_file, output_csv, clip_folder, unique_field="Plotid"):
    """Processes hyperspectral images by clipping first (GDAL), then computing mean spectra using Rasterio (faster)."""
    
    if input_file.lower().endswith(".dat"):
        print(f"[INFO] Converting {input_file} to TIFF format...")
        tiff_file = convert_dat_to_tiff(input_file)
        hdr_file = input_file.replace(".dat", ".hdr")
        wavelengths = get_wavelengths_from_hdr(hdr_file)  
    else:
        tiff_file = input_file
        wavelengths = None  

    dataset = gdal.Open(tiff_file, gdal.GA_ReadOnly)
    if dataset is None:
        raise Exception(f"[ERROR] Failed to open image file: {tiff_file}")

    num_bands = dataset.RasterCount  

    shp_dataset = ogr.Open(shp_file)
    if shp_dataset is None:
        raise Exception(f"[ERROR] Failed to open shapefile: {shp_file}")

    layer = shp_dataset.GetLayer()
    print(f"[INFO] Opened shapefile '{shp_file}'. Found {layer.GetFeatureCount()} plots.")

    create_csv_with_headers(output_csv, num_bands, wavelengths)

    if not os.path.exists(clip_folder):
        os.makedirs(clip_folder)

    count = 0

    for feature in tqdm(layer, desc="Processing Plots"):
        plot_id = feature.GetField(unique_field)
        geometry = feature.GetGeometryRef()

        if plot_id is None:
            print(f"[WARNING] Skipping feature with missing `{unique_field}` attribute.")
            continue

        output_filename = os.path.join(clip_folder, f"Plot_{plot_id}.tif")

        # **Clipping and saving to disk**
        options = gdal.WarpOptions(
            format="GTiff",
            cutlineDSName=shp_file,
            cutlineLayer=layer.GetName(),
            cutlineWhere=f"{unique_field} = '{plot_id}'",
            cropToCutline=True,
            dstNodata=0  
        )
        clipped = gdal.Warp(output_filename, dataset, options=options)  

        if clipped is None:
            print(f"[WARNING] Skipping plot {plot_id} due to failed clipping.")
            continue

        # **Processing Clipped Image with Rasterio (Faster)**
        with rasterio.open(output_filename) as src:
            plot_numpy = src.read()
            plot_numpy = np.where(plot_numpy <= 0, np.nan, plot_numpy)  # Remove 0 and negative values
            plot_mean = np.nanmean(plot_numpy, axis=(1, 2))  

            if not np.all(np.isnan(plot_mean)):
                append_data_to_csv(output_csv, f"plot_id_{plot_id}", plot_mean)
                count += 1

    print(f"[INFO] Processing completed. Processed {count} plots. Data saved in {output_csv}.")
    return count