"""
Description:
The RasterConverter module provides a set of functions for efficient data conversion and manipulation of raster datasets.
"""

import os
import numpy as np
import rasterio
from tqdm import tqdm

def raster_list(file_path):
    """
    Read raster data from a file, convert NoData values to None, and convert the data into a Python list.

    Args:
        file_path (str): Path to the raster file.

    Returns:
        list: Raster data as a Python list with NoData values converted to None.
    """
    # Open the raster file
    dataset = rasterio.open(file_path)

    # Check if the dataset was opened successfully
    if dataset is None:
        raise Exception("Failed to open the raster file.")

    # Get the raster band
    band = dataset.read(1)

    # Get the NoData value
    nodata_value = dataset.nodata

    # Get the number of rows and columns of raster data
    rows, cols = band.shape

    # Convert raster data to matrix
    raster_matrix = band.reshape(rows, cols)

    # Converts data types to object types to support storing None
    raster_matrix = raster_matrix.astype(object)

    # Convert NoData values to None
    raster_matrix[raster_matrix == nodata_value] = None

    # Convert NumPy array to a Python list
    raster_list = [[raster_matrix[row][col] for col in range(cols)] for row in range(rows)]

    return raster_list


def multiRaster_listDict(folder_path, index_list):
    """
    Reads multiple raster data (.tif) from a folder and stores it in a dictionary based on specific index values.

    Args:
        folder_path (str): The path to the folder containing raster data (.tif).
        index_list (list): A list of indices corresponding to the raster files.

    Returns:
        dict: A dictionary with a specific index value as the key and a list of stored raster pixel values as the value.
    """
    # Create a list of .tif files in the specified folder
    tif_list = list()
    for file in os.listdir(folder_path):
        if file[-4:] == ".tif":
            tif_list.append(file)

    # Check that the raster file list is the same length as the index list
    if len(tif_list) != len(index_list):
        raise Exception("The length of the raster file list does not match the index list..")

    print("Data reading...")
    # Create a data dictionary with keys as index values for the corresponding positions in the index list and values as the data list generated by the raster file conversion.
    data_dict = dict()
    for i in tqdm(range(len(tif_list))):
        index, tif = index_list[i], tif_list[i]
        tif_path = os.path.join(folder_path, tif)
        data_dict[index] = raster_list(tif_path)
    print("Data reading is complete!")

    return data_dict


def reorganize_multiRaster_listDict(data_dict, index_list):
    """
    Reorganizes data from a list dictionary representing multiple raster data on the basis of the multiRaster_listDict function.

    Args:
        data_dict (dict): Resulting dictionary file generated by multiRaster_listDict function.
        index_list (list): A list of indexes corresponding to the list of keys of the dictionary file in a defined order,
                           the same as the index_list parameter of the multiRaster_listDict function.

    Returns:
        tuple: The first parameter is a two-dimensional indexed list with the same number of rows and columns as the original
               raster data, and the list element values are tuples of horizontal and vertical coordinates. The second parameter
               is a data dictionary with the key being the coordinate tuple (index list element value) and the value being a list
               of all raster image element values corresponding to the coordinates (in index_list order).
    """
    key1 = list(data_dict.keys())[0]
    list1 = data_dict[key1]
    rows, cols = len(list1), len(list1[0])

    data_index_list, data_value_dict = list(), dict()
    print("Data conversion...")
    for i in tqdm(range(rows)):
        inner_list = list()
        for j in range(cols):
            row, col = i + 1, j + 1
            inner_list.append((row, col))

            value_list = list()
            for index in index_list:
                data_list = data_dict[index]
                value_list.append(data_list[i][j])
            data_value_dict[(row, col)] = value_list

        data_index_list.append(inner_list)
    print("Data conversion is complete!")

    return data_index_list, data_value_dict


def process_raster_data(input_raster_path, output_raster_path, new_data, nodata_value):
    """
    Create a single-band raster template based on the input raster, process the raster data by replacing pixel values with the new data.

    Args:
        input_raster_path (str): Path to the original raster file.
        output_raster_path (str): Path to save the output raster file.
        new_data (list): New data in the form of a 2D list.
        nodata_value (num): NoData value in new_data.

    Returns:
        None
    """
    # Open the original raster dataset
    input_raster = rasterio.open(input_raster_path)

    # Set the NoData parameter of the output raster
    nodata_raster = input_raster.nodata
    # If nodata_value is not None, set the nodata_raster to the value of nodata_value
    if nodata_value is not None:
        nodata_raster = nodata_value

    # Create a new raster from the original raster dataset
    output_raster = rasterio.open(output_raster_path, "w", driver="GTiff",
                                  height=input_raster.height, width=input_raster.width, count=1,
                                  dtype=input_raster.dtypes[0], crs=input_raster.crs, transform=input_raster.transform,
                                  nodata=nodata_raster)

    # Replace pixel values with new data，and set the NoData value
    masked_array = new_data == nodata_value
    output_raster.write(np.array(new_data), 1, masked=masked_array)

    print(f"Processed raster data saved to {output_raster_path}.")

    # Close datasets
    input_raster.close()
    output_raster.close()
