import os
import pathlib
from osgeo import gdal
import numpy as np
import tqdm
import multiprocessing as mp
from functools import partial # Useful for passing fixed arguments to worker

import hazelbean as hb
from hazelbean.pog import *


def is_path_pog(path, check_tiled=True, full_check=False, raise_exceptions=False, verbose=False):
    
    is_pyramid = hb.is_path_global_pyramid(path, verbose=verbose)
    is_cog = hb.is_path_cog(path, check_tiled=check_tiled, full_check=full_check, raise_exceptions=raise_exceptions, verbose=verbose)

    if verbose:
        if is_pyramid:
            hb.log(f"Raster is a global pyramid: {path}")
        else:
            hb.log(f"Raster is not a global pyramid: {path}")
        if is_cog:
            hb.log(f"Raster is a COG: {path}")
        else:
            hb.log(f"Raster is not a COG: {path}")
            
    return is_pyramid and is_cog




# This worker function will be executed by each process in the pool.
# It needs to be defined at the top level of a module so it can be pickled.
def _worker_make_path_pog_starmap(input_file_path, specific_output_path, common_pog_options_dict):
    """
    Worker function for starmap.
    input_file_path: The specific input file for this task.
    specific_output_path: The specific output file path for this task (can be None).
    common_pog_options_dict: Dictionary of common options for hb.make_path_pog.
    """
    process_name = mp.current_process().name
    base_name = os.path.basename(input_file_path)

    try:
        hb.make_path_pog(
            input_file_path,
            output_raster_path=specific_output_path, # CRITICAL: Use the specific output path
            output_data_type=common_pog_options_dict.get('output_data_type', 'auto'),
            overview_resampling_method=common_pog_options_dict.get('overview_resampling_method'),
            ndv=common_pog_options_dict.get('ndv'),
            compression=common_pog_options_dict.get('compression', "ZSTD"),
            blocksize=common_pog_options_dict.get('blocksize', 512),
            verbose=common_pog_options_dict.get('verbose_hb_call', False)
        )
        output_msg_part = f" (output: {os.path.basename(specific_output_path)})" if specific_output_path else " (output: in-place/default)"
        return (input_file_path, True, f"Successfully processed by {process_name}{output_msg_part}")
    except Exception as e:
        return (input_file_path, False, f"Error in {process_name} for {base_name}: {type(e).__name__} - {e}")


def make_paths_pogs_in_parallel(
    input_folder,
    include_strings=None,
    include_extensions='.tif',
    exclude_strings=None,
    exclude_extensions=None,
    output_raster_path_suffix=None,
    output_data_type='auto',
    overview_resampling_method=None,
    ndv=None,
    compression="ZSTD",
    blocksize=512,
    verbose_hb_call=False,
    max_workers=None,
    dont_actually_do_it=False,
    force_reprocess_pog=False,
    seek_recursively=True, 
    verbose=0,
    verbose_pog_check=0, 
):
    """
    Lists raster files, checks POG status, and processes them in parallel
    using hazelbean's make_path_pog and starmap.
    Output filenames can be generated by appending a suffix.
    """
    abs_input_folder = os.path.abspath(input_folder)
    if not os.path.isdir(abs_input_folder):
        print(f"Error: Input folder not found: {abs_input_folder}")
        return []

    print(f"Scanning folder: {abs_input_folder}")
    if seek_recursively:
    
        initial_paths_found = hb.list_filtered_paths_recursively(
            abs_input_folder,
            include_strings=include_strings,
            include_extensions=include_extensions,
            exclude_strings=exclude_strings,
            exclude_extensions=exclude_extensions,
            return_only_filenames=False
        )
    else:
        initial_paths_found = hb.list_filtered_paths_nonrecursively(
            abs_input_folder,
            include_strings=include_strings,
            include_extensions=include_extensions,
            exclude_strings=exclude_strings,
            exclude_extensions=exclude_extensions,
            return_only_filenames=False
        )

    if not initial_paths_found:
        print("No files found matching the initial filter criteria.")
        return []

    print(f"\nFound {len(initial_paths_found)} candidate files. Checking POG status...")

    # Prepare lists for starmap: one for input paths, one for corresponding output paths
    input_paths_for_processing = []
    output_paths_for_processing = []

    for i, path_to_check in enumerate(initial_paths_found):
        base_name = os.path.basename(path_to_check)
        is_pog = hb.is_path_pog(path_to_check, verbose=verbose_pog_check)

        if is_pog and not force_reprocess_pog:
            print(f"  {i+1}. Skipping: {base_name} (is POG and not forcing reprocess)")
            continue
        
        action_reason = "(is POG but forcing reprocess)" if is_pog else "(not POG or POG status unknown)"
        print(f"  {i+1}. Queuing: {base_name} {action_reason}")
        
        input_paths_for_processing.append(path_to_check)

        # Determine the specific output path for this task
        if output_raster_path_suffix is not None:
            # Generate output path using the suffix
            specific_output_path = hb.suri(path_to_check, output_raster_path_suffix)
            output_paths_for_processing.append(specific_output_path)
        else:
            # No suffix means output_raster_path for hb.make_path_pog will be None
            # This typically implies in-place modification or hb default output naming.
            output_paths_for_processing.append(None)

    if not input_paths_for_processing:
        print("\nNo files to process after POG check.")
        return []

    print(f"\n{len(input_paths_for_processing)} files will be processed.")
    if dont_actually_do_it:
        print("  (Detailed list of tasks for processing if not in 'dont_actually_do_it' mode):")
        for i_path, o_path in zip(input_paths_for_processing, output_paths_for_processing):
            print(f"    Input: {os.path.basename(i_path)} -> Output: {os.path.basename(o_path) if o_path else 'None (in-place/default)'}")


    # Prepare common options for hb.make_path_pog
    # These are fixed for all tasks and will be passed via partial.
    common_pog_options = {
        'output_data_type': output_data_type,
        'overview_resampling_method': overview_resampling_method,
        'ndv': ndv,
        'compression': compression,
        'blocksize': blocksize,
        'verbose_hb_call': verbose_hb_call
    }

    # Prepare the arguments for starmap: a list of tuples,
    # where each tuple is (input_path, output_path)
    starmap_iterable = list(zip(input_paths_for_processing, output_paths_for_processing))

    # Use functools.partial to create a new worker function where common_pog_options_dict is pre-filled.
    # The resulting function will expect (input_file_path, specific_output_path) as arguments.
    worker_with_common_opts = partial(
        _worker_make_path_pog_starmap,
        common_pog_options_dict=common_pog_options
    )

    if max_workers is None:
        max_workers = os.cpu_count() - 2
    max_workers = min(max_workers, len(starmap_iterable)) # Don't use more processes than tasks

    results = []
    if max_workers <= 0:
        print("No tasks to process or max_workers is 0 or less. Skipping parallel execution.")
    else:
        print(f"\nStarting parallel processing with {max_workers} worker(s)...")
        if not dont_actually_do_it:
            with mp.Pool(processes=max_workers) as pool:
                # pool.starmap expects an iterable of argument tuples.
                # For each (in_path, out_path) in starmap_iterable, it effectively calls:
                #   _worker_make_path_pog_starmap(in_path, out_path, common_pog_options_dict=common_pog_options)
                results = pool.starmap(worker_with_common_opts, starmap_iterable)
        else:
            print(f"\nSKIPPING ACTUAL PROCESSING due to dont_actually_do_it=True.")
            # Simulate results if needed for testing downstream logic
            for in_path, out_path in starmap_iterable:
                results.append((in_path, True, f"Simulated success by DUMMY_PROCESS (output: {os.path.basename(out_path) if out_path else 'None'})"))


    print("\n--- Processing Summary ---")
    success_count = 0
    failure_count = 0
    for file_path, success, message in results: # Assuming worker returns this tuple
        base_name = os.path.basename(file_path)
        if success:
            print(f"SUCCESS: {base_name} - {message}")
            success_count += 1
        else:
            print(f"FAILURE: {base_name} - {message}")
            failure_count += 1
    
    print(f"\nFinished processing. {success_count} successful, {failure_count} failed.")
    return results
 
  
def make_path_pog(input_raster_path, output_raster_path=None, output_data_type=None, ndv=None, overview_resampling_method=None, compression="ZSTD", blocksize=512, force_rewrite=False, value_reclassification_dict=None, ndv_above=None, ndv_below=None, verbose=False):
    """ Create a Pog (pyramidal cog) from input_raster_path. Writes in-place if output_raster_path is not set. Chooses correct values for 
    everything else if not set."""

    # Check if input exists
    if not hb.path_exists(input_raster_path, verbose=verbose):
        raise FileNotFoundError(f"Input raster does not exist: {input_raster_path} at abs path {hb.path_abs(input_raster_path)}")
    
    # Do a fast check to see if it's pog.
    needs_censoring = False
    if not hb.is_path_global_pyramid(input_raster_path, verbose):
        if verbose:
            hb.log(f"Raster is not a global pyramid. {input_raster_path}")            
    else:        
        if ndv_above is not None or ndv_below is not None:
            stats_by_band = hb.get_stats_from_geotiff(input_raster_path)            
            if ndv_above < stats_by_band[1]['max']:
                if verbose:
                    hb.log(f"Raster is not a global pyramid because it has values above the ndv_above threshold. {input_raster_path}")
                needs_censoring = True
            if ndv_below > stats_by_band[1]['min']:
                if verbose:
                    hb.log(f"Raster is not a global pyramid because it has values below the ndv_below threshold. {input_raster_path}")
                needs_censoring = True  

        needs_reclassification = False
        if value_reclassification_dict is not None:
            needs_reclassification = True
                
        # Do a full check to see if the input is already a POG. If so, skip it.    
        if is_path_pog(input_raster_path, verbose) and not force_rewrite and not needs_censoring and not needs_reclassification:
            if verbose:
                hb.log(f"Raster is already a POG: {input_raster_path}")
            return
        
    # Make a local copy at a temp file to process on to avoid corrupting the original
    input_dir = os.path.dirname(input_raster_path)
    temp_copy_path = hb.temp('.tif', 'copy', True, folder=input_dir, tag_along_file_extensions=['.aux.xml'])
    temp_translate_path = hb.temp('.tif', 'translate', True, folder=input_dir, tag_along_file_extensions=['.aux.xml'])
    
    # Ensure output directory exists
    try:
        os.makedirs(os.path.dirname(temp_copy_path), exist_ok=True)       
    except:
        pass
    
    input_data_type = hb.get_datatype_from_uri(input_raster_path)
    
    ### WARNING: Forcing data type optimizations is a bad idea! Increment up the following number to represent the number of failed attempts to try it: attempts = 2
    # OUTDATED POG SPECIFIC HERE. Note the draconian forcing here of bytes into int64s
    # WAIT! Although gdal does support int64, the geotiff driver does not. It internally converts to float64 as a result.
    # One option would be to use ... dammit...
    # input_data_type = hb.get_datatype_from_uri(input_raster_path)
    # if output_data_type == 'auto':
    #     if input_data_type in [1, 2, 3, 4, 5, 12, 13]:
    #         output_data_type = 5
    #         output_data_type_object = gdal.GDT_Int32
    #     elif input_data_type in [6]:
    #         output_data_type = 6
    #         output_data_type_object = gdal.GDT_Float32
    #     elif input_data_type in [7, 8, 9, 10, 11]:
    #         output_data_type = 7
    #         output_data_type_object = gdal.GDT_Float64
    
    if output_data_type is None or output_data_type == 'auto':
        output_data_type = input_data_type 

    if output_data_type != input_data_type:
        if verbose:            
            hb.log(f"Changing data type from {input_data_type} to {output_data_type} for {input_raster_path}")
        output_data_type = hb.gdal_number_to_gdal_type[output_data_type]
        gdal.Translate(temp_translate_path, input_raster_path, outputType=output_data_type, options=['-co', 'COMPRESS=ZSTD'], callback=hb.make_gdal_callback(f"Translating to change datatype on {hb.path_filename(input_raster_path)} to {temp_translate_path} to be datatype {output_data_type}."))
        current_path = temp_translate_path
    else:
        if verbose:
            hb.log(f"Data type is already {output_data_type} for {input_raster_path}, so just copying to {temp_copy_path}.")
        hb.path_copy(input_raster_path, temp_copy_path) # Can just copy it direclty without accessing the raster.        
        current_path = temp_copy_path

    # Get the resolution from the src_ds
    degrees = hb.get_cell_size_from_path(current_path, force_to_pyramid=True)
    arcseconds = hb.get_cell_size_from_path_in_arcseconds(current_path, force_to_pyramid=True)
    
    original_output_raster_path = output_raster_path
    if output_raster_path is None:        
        output_raster_path = hb.temp('.tif', 'pog', remove_at_exit=False, folder=os.path.dirname(input_raster_path), tag_along_file_extensions=['.aux.xml'])
        
    ndv = hb.no_data_values_by_gdal_type[output_data_type][0] # NOTE AWKWARD INCLUSINO OF zero as second option to work with faster_zonal_stats
    
    if output_data_type <= 5 or output_data_type == 12 or output_data_type == 13:
        resample_method = 'nearest'
    else:
        resample_method = 'bilinear'
        
    # POG SPECIFIC DIFFERENCE HERE: Handles the case where the raster is not global.
    gt = hb.get_geotransform_path(input_raster_path)    
    gt_pyramid = hb.get_global_geotransform_from_resolution(degrees)
    
    user_dir = pathlib.Path.home()
    match_path = os.path.join(user_dir, 'Files', 'base_data', hb.ha_per_cell_ref_paths[arcseconds])
    if gt != gt_pyramid:
        resample_temp_path = hb.temp('.tif', 'resample', True, folder=os.path.dirname(input_raster_path), tag_along_file_extensions=['.aux.xml'])
        if verbose:
            hb.log(f"Resampling {current_path} to match {match_path}. Saving at {resample_temp_path}.")
        hb.resample_to_match(
            current_path,
            match_path,
            resample_temp_path,
            resample_method=resample_method,
            output_data_type=output_data_type,
            src_ndv=None,
            ndv=ndv,
            s_srs_wkt=None,
            compress=True,
            ensure_fits=False,
            gtiff_creation_options=hb.globals.PRECOG_GTIFF_CREATION_OPTIONS_LIST,
            calc_raster_stats=False,
            add_overviews=False,
            pixel_size_override=None,
            target_aligned_pixels=True,
            bb_override=None,
            verbose=False, 
        )
        current_path = resample_temp_path
    
    if (ndv_above is not None or ndv_below is not None) and needs_censoring:
        censor_temp_path = hb.temp('.tif', 'censor', True, folder=os.path.dirname(input_raster_path), tag_along_file_extensions=['.aux.xml'])
        if ndv_above and not ndv_below:
            def op(x):
                return np.where(x > ndv_above, ndv, x)
        elif ndv_below and not ndv_above:
            def op(x):
                return np.where(x < ndv_below, ndv, x)
        elif ndv_above and ndv_below:
            def op(x):
                return np.where((x > ndv_above) | (x < ndv_below), ndv, x)
            
        hb.log(f"Censoring {current_path} with ndv_above={ndv_above} and ndv_below={ndv_below}. Saving at {censor_temp_path}.")
        hb.raster_calculator_flex(current_path, op, censor_temp_path, output_data_type=output_data_type, ndv=ndv, compression=compression, verbose=verbose)
        current_path = censor_temp_path
    
    if value_reclassification_dict is not None:
        reclassify_temp_path = hb.temp('.tif', 'reclassify', True, '.', tag_along_file_extensions=['.aux.xml'])

        hb.log(f"Reclassifying {current_path} with {value_reclassification_dict}. Saving at")
        hb.reclassify_raster_hb(current_path, value_reclassification_dict, reclassify_temp_path)
     
    if not hb.raster_path_has_stats(current_path, approx_ok=False):
        if verbose:
            hb.log(f"Adding stats to {current_path}.")
        hb.add_stats_to_geotiff_with_gdalinfo(current_path, approx_ok=False, force_recompute=True, verbose=verbose)
    
    # Open the source raster in UPDATE MODE so it writes the overviews as internal
    if verbose:
        hb.log(f"Opening {current_path} for overview building.")
    src_ds = gdal.OpenEx(current_path, gdal.GA_Update, open_options=["IGNORE_COG_LAYOUT_BREAK=YES"])
    if not src_ds:
        raise ValueError(f"Unable to open raster: {current_path}")

    # Remove existing overviews (if any)
    src_ds.BuildOverviews(None, [])     
        
    if overview_resampling_method is None:
        overview_resampling_method = hb.pyramid_resampling_algorithms_by_data_type[output_data_type] 
    
    # Set the overview levels based on the pyramid arcseconds
    overview_levels = hb.pyramid_compatible_overview_levels[arcseconds]
    
    if verbose:
        hb.log(f"Building overviews for {current_path} with levels {overview_levels}...")
    src_ds.BuildOverviews(overview_resampling_method.upper(), overview_levels, hb.make_gdal_callback(f'Building overviews for {current_path}'))

    # Close the dataset to ensure overviews are saved
    del src_ds    
    
    # Reopen it to use it as a copy target
    if verbose:
        hb.log(f"Reopening {current_path} for COG creation...")
    src_ds = gdal.OpenEx(current_path, gdal.GA_ReadOnly)
    
    # Define creation options for COG
    creation_options = [
        f"COMPRESS={compression}",
        f"BLOCKSIZE={blocksize}",  
        f"BIGTIFF=YES", 
        f"OVERVIEW_COMPRESS={compression}",        
        f"RESAMPLING={overview_resampling_method}",
        # f"OVERVIEWS=IGNORE_EXISTING",
        f"OVERVIEW_RESAMPLING={overview_resampling_method}",
    ]

    cog_driver = gdal.GetDriverByName('COG')
    if cog_driver is None:
        raise RuntimeError("COG driver is not available in this GDAL build.")    
    
    if ndv is not None and src_ds is not None:
        for i in range(1, src_ds.RasterCount + 1):
            band = src_ds.GetRasterBand(i)
            band.SetNoDataValue(ndv)
            
    # Actually create the COG
    force_verbose = True
    if verbose or force_verbose:
        hb.log(f"Creating COG at {output_raster_path}. Abs path: {os.path.abspath(output_raster_path)}, Norm path: {os.path.normpath(output_raster_path)}")
    dst_ds = cog_driver.CreateCopy(
        output_raster_path,
        src_ds,
        strict=0,  # set to 1 to fail on any “creation option not recognized”
        options=creation_options
    )        
    dst_ds = None
    
    if not is_path_pog(output_raster_path, verbose=verbose) and verbose:
        hb.log(f"Failed to create COG: {output_raster_path} at abs path {hb.path_abs(output_raster_path)}")
    
    if original_output_raster_path is None:
        hb.displace_file(output_raster_path, input_raster_path)        


    
def write_pog_of_value_from_scratch(output_path, value, arcsecond_resolution, output_data_type, ndv=None, overview_resampling_method=None, compression='ZSTD', blocksize='512', verbose=False):
    # Define creation options for COG
    precog_gtiff_creation_options = [
        f"COMPRESS={compression}",
        f"BLOCKXSIZE={str(blocksize)}",  
        f"BLOCKYSIZE={str(blocksize)}",  
        f"BIGTIFF=YES", 
        "TILED=YES",
    ]

        
    if ndv is None:
        ndv = hb.no_data_values_by_gdal_type[output_data_type][0] 

    if overview_resampling_method is None:
            overview_resampling_method = hb.pyramid_resampling_algorithms_by_data_type[output_data_type] 

    geotransform = hb.pyramid_compatible_geotransforms[float(arcsecond_resolution)]
    projection = hb.wgs_84_wkt
    x_size = hb.pyramid_compatable_shapes[float(arcsecond_resolution)][0]
    y_size = hb.pyramid_compatable_shapes[float(arcsecond_resolution)][1]

    
    # Make a temp geotiff based on match
    temp_path = hb.temp('.tif', filename_start='temp_gtiff_b4_cog', remove_at_exit=True)
    driver = gdal.GetDriverByName('GTiff')
    tmp_ds = driver.Create(temp_path, x_size, y_size, 1, output_data_type, options=precog_gtiff_creation_options)
    tmp_ds.SetGeoTransform(geotransform)
    tmp_ds.SetProjection(projection)    
    
    # Memory safe (hopefully) way to write the value to the raster, row by row
    value_row = np.full((1, x_size), value, dtype=hb.gdal_number_to_numpy_type[output_data_type])

    # Initialize accumulators for statistics
    total_sum = 0.0
    total_sq_sum = 0.0
    total_count = 0
    global_min = np.inf
    global_max = -np.inf        
    
    band = tmp_ds.GetRasterBand(1)
    for row in tqdm.tqdm(range(y_size)):
        # hb.print_in_place('Writing row ' + str(row) + ' of ' + str(y_size))
        band.WriteArray(value_row, xoff=0, yoff=row)
      
        # Update statistics incrementally
        total_sum += value_row.sum()
        total_sq_sum += np.square(value_row).sum()
        total_count += value_row.size
        global_min = min(global_min, np.min(value_row))
        global_max = max(global_max, np.max(value_row))    

    mean = total_sum / total_count if total_count else 0
    variance = (total_sq_sum / total_count - mean ** 2) if total_count else 0
    variance = max(variance, 0)  # safeguard against negative variance
    stddev = np.sqrt(variance)
        
    # Build Overviews
    tmp_ds.GetRasterBand(1).SetNoDataValue(ndv)    
    tmp_ds.BuildOverviews(None, []) # Remove existing overviews (if any) 
    
    resampling_algorithm = hb.pyramid_resampling_algorithms_by_data_type[output_data_type]    
    if overview_resampling_method is None:
        overview_resampling_method = resampling_algorithm
    
    # Set the overview levels based on the pyramid arcseconds
    overview_levels = hb.pyramid_compatible_overview_levels[arcsecond_resolution]
    tmp_ds.BuildOverviews(overview_resampling_method.upper(), overview_levels, callback=hb.make_gdal_callback('Building overviews for ' + str(output_path)))
    
    tmp_ds.FlushCache()
    del tmp_ds  # Close temp dataset
    
    # Add statistics to the temp dataset
    hb.add_stats_to_geotiff_with_gdal(temp_path, approx_ok=False, force=True, verbose=verbose)

    # Step 2: Convert temporary GTiff to COG using CreateCopy
    cog_driver = gdal.GetDriverByName('COG')
    cog_creation_options = [
        f'COMPRESS={compression}',
        f'BLOCKSIZE={blocksize}',
        'BIGTIFF=YES',
        # 'OVERVIEWS=IGNORE_EXISTING'
    ]

    tmp_ds = gdal.Open(temp_path)
    cog_ds = cog_driver.CreateCopy(output_path, tmp_ds, options=cog_creation_options)

    if cog_ds is None:
        raise RuntimeError('Failed to create COG dataset.')

    # Cleanup
    del cog_ds
    del tmp_ds


def write_pog_of_value_from_match(output_path, match_path, value, output_data_type=None, ndv=None, overview_resampling_method=None, compression='ZSTD', blocksize='512', verbose=False):
    
    # Define creation options for COG
    precog_gtiff_creation_options = [
        f"COMPRESS={compression}",
        f"BLOCKXSIZE={str(blocksize)}",  
        f"BLOCKYSIZE={str(blocksize)}",  
        f"BIGTIFF=YES", 
        f"TILED=YES", 
    ]
    
    # Check if match exists
    if not hb.path_exists(match_path, verbose=verbose):
        raise FileNotFoundError(f"Input raster does not exist: {match_path} at abs path {hb.path_abs(match_path)}")
    
    if output_data_type is None:
        output_data_type = hb.get_datatype_from_uri(match_path)
        
    if ndv is None:
        ndv = hb.no_data_values_by_gdal_type[output_data_type][0] 

    if overview_resampling_method is None:
            overview_resampling_method = hb.pyramid_resampling_algorithms_by_data_type[output_data_type] 
            
    # Open existing raster to match
    src_ds = gdal.Open(match_path)
    geotransform = src_ds.GetGeoTransform()
    projection = src_ds.GetProjection()
    x_size = src_ds.RasterXSize
    y_size = src_ds.RasterYSize
    degrees = hb.get_cell_size_from_path(match_path)
    arcseconds = hb.get_cell_size_from_path_in_arcseconds(match_path)        
    
    # Make a temp geotiff based on match
    temp_path = hb.temp('.tif', filename_start='temp_gtiff_b4_cog', remove_at_exit=True)
    driver = gdal.GetDriverByName('GTiff')
    tmp_ds = driver.Create(temp_path, x_size, y_size, 1, output_data_type, options=precog_gtiff_creation_options)
    tmp_ds.SetGeoTransform(geotransform)
    tmp_ds.SetProjection(projection)    
    
    # Memory safe (hopefully) way to write the value to the raster, row by row
    value_row = np.full((1, x_size), value, dtype=hb.gdal_number_to_numpy_type[output_data_type])

    # Initialize accumulators for statistics
    total_sum = 0.0
    total_sq_sum = 0.0
    total_count = 0
    global_min = np.inf
    global_max = -np.inf        
    
    band = tmp_ds.GetRasterBand(1)
    for row in tqdm.tqdm(range(y_size)):
        # hb.print_in_place('Writing row ' + str(row) + ' of ' + str(y_size))
        band.WriteArray(value_row, xoff=0, yoff=row)
      
        # Update statistics incrementally
        total_sum += value_row.sum()
        total_sq_sum += np.square(value_row).sum()
        total_count += value_row.size
        global_min = min(global_min, np.min(value_row))
        global_max = max(global_max, np.max(value_row))    

    mean = total_sum / total_count if total_count else 0
    variance = (total_sq_sum / total_count - mean ** 2) if total_count else 0
    variance = max(variance, 0)  # safeguard against negative variance
    stddev = np.sqrt(variance)
        
    # Build Overviews
    tmp_ds.GetRasterBand(1).SetNoDataValue(ndv)    
    tmp_ds.BuildOverviews(None, []) # Remove existing overviews (if any) 
    
    resampling_algorithm = hb.pyramid_resampling_algorithms_by_data_type[output_data_type]    
    if overview_resampling_method is None:
        overview_resampling_method = resampling_algorithm
    
    # Set the overview levels based on the pyramid arcseconds
    overview_levels = hb.pyramid_compatible_overview_levels[arcseconds]
    tmp_ds.BuildOverviews(overview_resampling_method.upper(), overview_levels, callback=hb.make_gdal_callback('Building overviews for ' + str(output_path)))
    
    tmp_ds.FlushCache()
    del tmp_ds  # Close temp dataset

    # Set statistics, ensuring they are not approximate
    hb.add_stats_to_geotiff_with_gdal(temp_path, approx_ok=False, force=True, verbose=verbose)
    
    # Step 2: Convert temporary GTiff to COG using CreateCopy
    cog_driver = gdal.GetDriverByName('COG')
    cog_creation_options = [
        f'COMPRESS={compression}',
        f'BLOCKSIZE={blocksize}',
        'BIGTIFF=YES',
        'NUM_THREADS=ALL_CPUS',
        # 'OVERVIEWS=IGNORE_EXISTING'
    ]

    tmp_ds = gdal.Open(temp_path)
    cog_ds = cog_driver.CreateCopy(output_path, tmp_ds, options=cog_creation_options)

    if cog_ds is None:
        raise RuntimeError('Failed to create COG dataset.')

    # Cleanup
    del cog_ds
    del tmp_ds
    del src_ds

