"""
Author: Supath Dhital
Date Updated: January 2026
"""
import os
import re
import numpy as np
from pathlib import Path
import geopandas as gpd
import rasterio
import shutil
import subprocess
import platform
import pandas as pd
from rasterio.warp import reproject, Resampling
from rasterio.io import MemoryFile
from rasterio import features
from shapely.geometry import shape
from rasterio.mask import mask

os.environ["CHECK_DISK_FREE_SPACE"] = "NO"

import warnings

warnings.filterwarnings("ignore", category=rasterio.errors.ShapeSkipWarning)

from .methods import AOI, convex_hull, smallest_extent, get_smallest_raster_path
from .metrics import evaluationmetrics
from .water_bodies import ExtractPWB
from ..utilis import MakeFIMsUniform, benchmark_name, find_best_boundary
from ..setup_benchFIM import ensure_benchmark


# giving the permission to the folder
def is_writable(path):
    """Check if the directory and its contents are writable."""
    path = Path(path)
    return os.access(path, os.W_OK)


def fix_permissions(path):
    path = Path(path).resolve()

    if is_writable(path):
        return

    uname = platform.system()

    try:
        if uname in ["Darwin", "Linux"]:
            subprocess.run(
                ["chmod", "-R", "u+rwX", str(path)],
                check=True,
                capture_output=True,
                text=True,
            )
            print(f"Permissions granted for user (u+rwX): {path}")

        elif "MINGW" in uname or "MSYS" in uname or "CYGWIN" in uname:
            subprocess.run(
                ["icacls", str(path), "/grant", "Everyone:F", "/T"],
                check=True,
                capture_output=True,
                text=True,
            )
            print(f"Permissions granted for working folder: {path}")

        else:
            print(f"Unsupported OS: {uname}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to fix permissions for {path}:\n{e.stderr}")


# Function for the evalution of the model
def evaluateFIM(
    benchmark_path, candidate_paths, PWB_Dir, folder, method, output_dir, shapefile=None
):
    # Lists to store evaluation metrics
    csi_values = []
    TN_values = []
    FP_values = []
    FN_values = []
    TP_values = []
    TPR_values = []
    FNR_values = []
    Acc_values = []
    Prec_values = []
    sen_values = []
    F1_values = []
    POD_values = []
    FPR_values = []
    Merged = []
    Unique = []
    FAR_values = []

    # Dynamically call the specified method
    method = globals().get(method)
    if method is None:
        raise ValueError(f"Method '{method}' is not defined.")

    # Save the smallest extent boundary and cliped FIMS
    save_dir = os.path.join(output_dir, os.path.basename(folder), f"{method.__name__}")
    os.makedirs(save_dir, exist_ok=True)

    # Get the smallest matched raster extent and make a boundary shapefile
    smallest_raster_path = get_smallest_raster_path(benchmark_path, *candidate_paths)

    # If method is AOI, and direct shapefile directory is not provided, then it will search for the shapefile in the folder
    if method.__name__ == "AOI":
        # Ubest-matching boundary file, prefer .gpkg from benchFIM downloads
        if shapefile is None:
            shapefile_path = find_best_boundary(Path(folder), Path(benchmark_path))
            if shapefile_path is None:
                raise FileNotFoundError(
                    f"No boundary file (.gpkg, .shp, .geojson, .kml) found in {folder}. "
                    "Either provide a shapefile path or place a boundary file in the folder."
                )
            shapefile = str(shapefile_path)
        else:
            shapefile = str(shapefile)

        # Run AOI with the found or provided shapefile
        bounding_geom = AOI(benchmark_path, shapefile, save_dir)

    else:
        print(f"**{method.__name__} is processing**")
        bounding_geom = method(smallest_raster_path, save_dir=save_dir)

    # Read and process benchmark raster
    with rasterio.open(benchmark_path) as src1:
        out_image1, out_transform1 = mask(
            src1, bounding_geom, crop=True, all_touched=True
        )
        benchmark_nodata = src1.nodata
        benchmark_crs = src1.crs
        b_profile = src1.profile

        #Getting the correct geometry shape and crs to extract PWB
        boundary_shape = shape(bounding_geom[0])
        boundary_gdf = gpd.GeoDataFrame(geometry=[boundary_shape], crs=benchmark_crs)

        #Proceed the masking
        out_image1[out_image1 == benchmark_nodata] = 0
        out_image1 = np.where(out_image1 > 0, 2, 0).astype(np.float32)

        #If PWB_Dir is provided, use the local PWB shapefile, else download from ArcGIS API
        if PWB_Dir is not None:
            gdf = gpd.read_file(PWB_Dir)
        else:
            #Get the permanent water bodies from ArcGIS REST API
            pwb_obj = ExtractPWB(boundary = boundary_gdf, save = False)
            gdf = pwb_obj.gdf

        gdf = gdf.to_crs(benchmark_crs)
        shapes1 = [
            geom for geom in gdf.geometry if geom is not None and not geom.is_empty
        ]
        mask1 = features.geometry_mask(
            shapes1,
            transform=out_transform1,
            invert=True,
            out_shape=out_image1.shape[1:],
        )
        extract_b = np.where(mask1, out_image1, 0)
        extract_b = np.where(extract_b > 0, 1, 0)
        idx_pwb = np.where(extract_b == 1)
        out_image1[idx_pwb] = 0

        benchmark_basename = os.path.basename(benchmark_path).split(".")[0]
        clipped_dir = os.path.join(save_dir, "MaskedFIMwithBoundary")
        if not os.path.exists(clipped_dir):
            os.makedirs(clipped_dir)

        clipped_benchmark = os.path.join(
            clipped_dir, f"{benchmark_basename}_clipped.tif"
        )
        b_profile.update(
            {
                "height": out_image1.shape[1],
                "width": out_image1.shape[2],
                "transform": out_transform1,
            }
        )

        with rasterio.open(clipped_benchmark, "w", **b_profile) as dst:
            dst.write(np.squeeze(out_image1), 1)

    def resize_image(
        source_image,
        source_transform,
        source_crs,
        target_crs,
        target_shape,
        target_transform,
    ):
        target_image = np.empty(target_shape, dtype=source_image.dtype)
        reproject(
            source=source_image,
            destination=target_image,
            src_transform=source_transform,
            dst_transform=target_transform,
            src_crs=source_crs,
            dst_crs=target_crs,
            resampling=Resampling.nearest,
        )
        return target_image

    # Process each candidate file
    for idx, candidate_path in enumerate(candidate_paths):
        base_name = os.path.splitext(os.path.basename(candidate_path))[0]
        with rasterio.open(candidate_path) as src2:
            candidate = src2.read(1)
            candidate_nodata = src2.nodata
            candidate_transform = src2.transform
            candidate_meta = src2.meta.copy()
            candidate_crs = src2.crs
            c_profile = src2.profile
            candidate[candidate == src2.nodata] = 0
            candidate = np.where(candidate > 0, 2, 1).astype(np.float32)
            with MemoryFile() as memfile:
                with memfile.open(**candidate_meta) as mem2:
                    mem2.write(candidate, 1)
                    dst_transform, width, height = (
                        rasterio.warp.calculate_default_transform(
                            mem2.crs,
                            benchmark_crs,
                            mem2.width,
                            mem2.height,
                            *mem2.bounds,
                        )
                    )
                    dst_meta = mem2.meta.copy()
                    dst_meta.update(
                        {
                            "crs": benchmark_crs,
                            "transform": dst_transform,
                            "width": width,
                            "height": height,
                        }
                    )

                    with MemoryFile() as memfile_reprojected:
                        with memfile_reprojected.open(**dst_meta) as mem2_reprojected:
                            for i in range(1, mem2.count + 1):
                                reproject(
                                    source=rasterio.band(mem2, i),
                                    destination=rasterio.band(mem2_reprojected, i),
                                    src_transform=mem2.transform,
                                    src_crs=mem2.crs,
                                    dst_transform=dst_transform,
                                    dst_crs=benchmark_crs,
                                    resampling=Resampling.nearest,
                                )
                            out_image2, out_transform2 = mask(
                                mem2_reprojected,
                                bounding_geom,
                                crop=True,
                                all_touched=True,
                            )
                            out_image2 = np.where(
                                out_image2 == candidate_nodata, 0, out_image2
                            )

                            # Save the clipped candidate raster
                            candidate_basename = os.path.basename(candidate_path).split(
                                "."
                            )[0]
                            clipped_candidate = os.path.join(
                                clipped_dir, f"{candidate_basename}_clipped.tif"
                            )
                            b_profile.update(
                                {
                                    "height": out_image1.shape[1],
                                    "width": out_image1.shape[2],
                                    "transform": out_transform1,
                                }
                            )
                            with rasterio.open(
                                clipped_candidate, "w", **b_profile
                            ) as dst:
                                dst.write(np.squeeze(out_image2), 1)

                            mask2 = features.geometry_mask(
                                shapes1,
                                transform=out_transform2,
                                invert=True,
                                out_shape=(out_image2.shape[1], out_image2.shape[2]),
                            )
                            extract_c = np.where(mask2, out_image2, 0)
                            extract_c = np.where(extract_c > 0, 1, 0)
                            idx_pwc = np.where(extract_c == 1)
                            out_image2[idx_pwc] = 5
                            out_image2_resized = resize_image(
                                out_image2,
                                out_transform2,
                                mem2_reprojected.crs,
                                benchmark_crs,
                                out_image1.shape,
                                out_transform1,
                            )
                            merged = out_image1 + out_image2_resized
                            merged[merged == 7] = 5

            # Get Evaluation Metrics
            (
                unique_values,
                TN,
                FP,
                FN,
                TP,
                TPR,
                FNR,
                Acc,
                Prec,
                sen,
                CSI,
                F1_score,
                POD,
                FPR,
                merged,
                FAR,
            ) = evaluationmetrics(out_image1, out_image2_resized)

            # Append values to the lists
            csi_values.append(CSI)
            TN_values.append(TN)
            FP_values.append(FP)
            FN_values.append(FN)
            TP_values.append(TP)
            TPR_values.append(TPR)
            FNR_values.append(FNR)
            Acc_values.append(Acc)
            Prec_values.append(Prec)
            sen_values.append(sen)
            F1_values.append(F1_score)
            POD_values.append(POD)
            FPR_values.append(FPR)
            Merged.append(merged)
            Unique.append(unique_values)
            FAR_values.append(FAR)

    results = {
        "CSI_values": csi_values,
        "TN_values": TN_values,
        "FP_values": FP_values,
        "FN_values": FN_values,
        "TP_values": TP_values,
        "TPR_values": TPR_values,
        "FNR_values": FNR_values,
        "Acc_values": Acc_values,
        "Prec_values": Prec_values,
        "sen_values": sen_values,
        "F1_values": F1_values,
        "POD_values": POD_values,
        "FPR_values": FPR_values,
        # 'Merged': Merged,
        #  'Unique': Unique
        "FAR_values": FAR_values,
    }
    for candidate_idx, candidate_path in enumerate(candidate_paths):
        candidate_BASENAME = os.path.splitext(os.path.basename(candidate_path))[0]
        merged_raster = Merged[candidate_idx]
        if merged_raster.ndim == 3:
            band = merged_raster.squeeze()
        elif merged_raster.ndim == 2:
            band = merged_raster
        else:
            raise ValueError(
                f"Unexpected number of dimensions in Merged[{candidate_idx}]."
            )

        # Construct the contingency file name dynamically
        contigency_dir = os.path.join(save_dir, "ContingencyMaps")
        os.makedirs(contigency_dir, exist_ok=True)
        output_filename = os.path.join(
            contigency_dir, f"ContingencyMAP_{candidate_BASENAME}.tif"
        )
        with rasterio.open(output_filename, "w", **b_profile) as dst:
            dst.write(band, 1)
            dst.transform = out_transform1
            dst.crs = benchmark_crs

    # Saving it into dataframe
    candidate_names = [
        os.path.splitext(os.path.basename(path))[0] for path in candidate_paths
    ]
    df = pd.DataFrame.from_dict(results, orient="index")
    df.columns = candidate_names
    df.reset_index(inplace=True)
    df.rename(columns={"index": "Metrics"}, inplace=True)

    # Save the DataFrame
    evaluationMetrics_DIR = os.path.join(save_dir, "EvaluationMetrics")
    os.makedirs(evaluationMetrics_DIR, exist_ok=True)

    csv_file = os.path.join(evaluationMetrics_DIR, "EvaluationMetrics.csv")
    df.to_csv(csv_file, index=False)
    print(f"Evaluation metrics saved to {csv_file}")
    return results


# Safely deleting the folder
def safe_delete_folder(folder_path):
    fix_permissions(folder_path)
    try:
        shutil.rmtree(folder_path)
    except PermissionError:
        print(f"Permission denied: Could not delete {folder_path}")
    except FileNotFoundError:
        print(f"Folder not found: {folder_path}")
    except Exception as e:
        print(f"Error deleting {folder_path}: {e}")


def EvaluateFIM(
    main_dir,
    method_name=None,
    output_dir=None,
    PWB_dir=None,
    shapefile_dir=None,
    target_crs=None,
    target_resolution=None,
    benchmark_dict=None,
):
    if output_dir is None:
        output_dir = os.path.join(os.getcwd(), "Evaluation_Results")

    main_dir = Path(main_dir)

    # Grant the permission to the main directory
    fix_permissions(main_dir)

    # runt the process
    def process_TIFF(tif_files, folder_dir):
        benchmark_path = None
        candidate_path = []

        for tif_file in tif_files:
            if benchmark_name(tif_file):
                benchmark_path = tif_file
            else:
                candidate_path.append(tif_file)

        if benchmark_path and candidate_path:
            if method_name is None:
                local_method = "AOI"

                # For single case, if user have explicitly send boundary, use that, else use the boundary from the benchmark FIM evaluation
                if shapefile_dir is not None:
                    local_shapefile = shapefile_dir
                else:
                    boundary = find_best_boundary(folder_dir, benchmark_path)
                    if boundary is None:
                        print(
                            f"Skipping {folder_dir.name}: no boundary file found "
                            f"and method_name is None (auto-AOI)."
                        )
                        return
                    local_shapefile = str(boundary)
            else:
                local_method = method_name
                local_shapefile = shapefile_dir

            print(f"**Flood Inundation Evaluation of {folder_dir.name}**")
            try:
                Metrics = evaluateFIM(
                    benchmark_path,
                    candidate_path,
                    PWB_dir,
                    folder_dir,
                    local_method,
                    output_dir,
                    shapefile=local_shapefile,
                )
                
                # Print results in structured table format with 3 decimal points
                candidate_names = [os.path.splitext(os.path.basename(path))[0] for path in candidate_path]
                df_display = pd.DataFrame.from_dict(Metrics, orient='index')
                df_display.columns = candidate_names
                df_display.reset_index(inplace=True)
                df_display.rename(columns={'index': 'Metrics'}, inplace=True)
                print("\n")
                print(df_display.to_string(index=False, float_format='%.3f'))
                print("\n")
            except Exception as e:
                print(f"Error evaluating {folder_dir.name}: {e}")
        else:
            print(
                f"Skipping {folder_dir.name} as it doesn't have a valid benchmark and candidate configuration."
            )

    # Check if main_dir directly contains tif files
    TIFFfiles_main_dir = list(main_dir.glob("*.tif"))
    if TIFFfiles_main_dir:

        # Ensure benchmark is present if needed
        TIFFfiles_main_dir = ensure_benchmark(
            main_dir, TIFFfiles_main_dir, benchmark_dict
        )

        processing_folder = main_dir / "processing"
        try:
            MakeFIMsUniform(
                main_dir, target_crs=target_crs, target_resolution=target_resolution
            )

            # processing folder
            TIFFfiles = list(processing_folder.glob("*.tif"))

            process_TIFF(TIFFfiles, main_dir)
        except Exception as e:
            print(f"Error processing {main_dir}: {e}")
        finally:
            safe_delete_folder(processing_folder)
    else:
        for folder in main_dir.iterdir():
            if folder.is_dir():
                tif_files = list(folder.glob("*.tif"))

                if tif_files:
                    processing_folder = folder / "processing"
                    try:
                        # Ensure benchmark is present if needed
                        tif_files = ensure_benchmark(folder, tif_files, benchmark_dict)

                        MakeFIMsUniform(
                            folder,
                            target_crs=target_crs,
                            target_resolution=target_resolution,
                        )

                        TIFFfiles = list(processing_folder.glob("*.tif"))

                        process_TIFF(TIFFfiles, folder)
                    except Exception as e:
                        print(f"Error processing folder {folder.name}: {e}")
                    finally:
                        safe_delete_folder(processing_folder)
                else:
                    print(
                        f"Skipping {folder.name} as it doesn't contain any tif files."
                    )