import ee
from .GetPalette import get_palette


def get_visualization_params(
    satellite, index, min_val=None, max_val=None, palette=None, scaled_bands=False
):
    """
    Function to define visualization paramaters for image visualization. Outputs an vis_params dictionary.

    Args:
        satellite (string): Denote satellite for visualization. Options: 'Landsat', 'landsat', 'Sentinel2', or 'sentinel2'
        index (string): Multispectral index to visualize. Options: 'TrueColor', 'NDVI', 'NDWI', 'halite', 'gypsum', 'LST', 'NDTI', 'KIVU', or '2BDA'. NOTE: LST and KIVU is not available for Sentinel2, and 2BDA is not available for Landsat.
        min_val (int or float): Optional override for minimum value to stretch raster
        max_val (int or float): Optional override for maximum value to stretch raster
        palette (string): Optional override for color palette used for image visualization. Must be from GetPalette.get_palette() function. Options: 'algae', 'dense', 'greens', 'haline', 'inferno', 'jet', 'matter', 'pubu', 'soft_blue_green_red', 'thermal', 'turbid', 'ylord'

    Returns:
        dictionary: vis_params dictionary to be used when visualizing an image on a map. Supplies min, max, band(s), and color palette (when appropriate).
    """
    # Define visualization parameters for various image types
    if satellite == "Landsat" or satellite == "landsat":
        if scaled_bands == False:
            params_dict = {
                "TrueColor": {"min": 0, "max": 30000, "bands": ["SR_B4", "SR_B3", "SR_B2"]},
                "NDVI": {
                    "min": 0,
                    "max": 0.5,
                    "bands": ["ndvi"],
                    "palette": get_palette("greens"),
                },
                "NDWI": {
                    "min": -0.2,
                    "max": 0.2,
                    "bands": ["ndwi"],
                    "palette": get_palette("inferno"),
                },
                "halite": {
                    "min": 0.1,
                    "max": 0.5,
                    "bands": ["halite"],
                    "palette": get_palette("haline"),
                },
                "gypsum": {
                    "min": 0.0,
                    "max": 0.5,
                    "bands": ["gypsum"],
                    "palette": get_palette("ylord"),
                },
                "LST": {
                    "min": 0,
                    "max": 40,
                    "bands": ["LST"],
                    "palette": get_palette("thermal"),
                },
                "NDTI": {
                    "min": -0.2,
                    "max": 0.2,
                    "bands": ["ndti"],
                    "palette": get_palette("turbid"),
                },
                "KIVU": {
                    "min": -0.5,
                    "max": 0.2,
                    "bands": ["kivu"],
                    "palette": get_palette("algae"),
                },
            }
        else:
            params_dict = {
                "TrueColor": {"min": 0, "max": 1, "bands": ["SR_B4", "SR_B3", "SR_B2"]},
                "NDVI": {
                    "min": 0,
                    "max": 0.5,
                    "bands": ["ndvi"],
                    "palette": get_palette("greens"),
                },
                "NDWI": {
                    "min": -0.2,
                    "max": 0.2,
                    "bands": ["ndwi"],
                    "palette": get_palette("inferno"),
                },
                "halite": {
                    "min": 0.1,
                    "max": 0.5,
                    "bands": ["halite"],
                    "palette": get_palette("haline"),
                },
                "gypsum": {
                    "min": 0.0,
                    "max": 0.5,
                    "bands": ["gypsum"],
                    "palette": get_palette("ylord"),
                },
                "LST": {
                    "min": 0,
                    "max": 40,
                    "bands": ["LST"],
                    "palette": get_palette("thermal"),
                },
                "NDTI": {
                    "min": -0.2,
                    "max": 0.2,
                    "bands": ["ndti"],
                    "palette": get_palette("turbid"),
                },
                "KIVU": {
                    "min": -0.5,
                    "max": 0.2,
                    "bands": ["kivu"],
                    "palette": get_palette("algae"),
                },
            }
    elif satellite == "Sentinel2" or satellite == "sentinel2":
        if scaled_bands == False:
            params_dict = {
                "TrueColor": {"min": 0, "max": 3500, "bands": ["B4", "B3", "B2"]},
                "NDVI": {
                    "min": 0.5,
                    "max": 0.9,
                    "bands": ["ndvi"],
                    "palette": get_palette("greens"),
                },
                "NDWI": {
                    "min": -0.2,
                    "max": 0.2,
                    "bands": ["ndwi"],
                    "palette": get_palette("inferno"),
                },
                "halite": {
                    "min": 0.1,
                    "max": 0.7,
                    "bands": ["halite"],
                    "palette": get_palette("haline"),
                },
                "gypsum": {
                    "min": 0.0,
                    "max": 0.7,
                    "bands": ["gypsum"],
                    "palette": get_palette("ylord"),
                },
                "NDTI": {
                    "min": -0.2,
                    "max": 0.5,
                    "bands": ["ndti"],
                    "palette": get_palette("turbid"),
                },
                "2BDA": {
                    "min": 0.5,
                    "max": 1.75,
                    "bands": ["2BDA"],
                    "palette": get_palette("algae"),
                },
            }
        else:
            params_dict = {
                "TrueColor": {"min": 0, "max": 1, "bands": ["B4", "B3", "B2"]},
                "NDVI": {
                    "min": 0.5,
                    "max": 0.9,
                    "bands": ["ndvi"],
                    "palette": get_palette("greens"),
                },
                "NDWI": {
                    "min": -0.2,
                    "max": 0.2,
                    "bands": ["ndwi"],
                    "palette": get_palette("inferno"),
                },
                "halite": {
                    "min": 0.1,
                    "max": 0.7,
                    "bands": ["halite"],
                    "palette": get_palette("haline"),
                },
                "gypsum": {
                    "min": 0.0,
                    "max": 0.7,
                    "bands": ["gypsum"],
                    "palette": get_palette("ylord"),
                },
                "NDTI": {
                    "min": -0.2,
                    "max": 0.5,
                    "bands": ["ndti"],
                    "palette": get_palette("turbid"),
                },
                "2BDA": {
                    "min": 0.5,
                    "max": 1.75,
                    "bands": ["2BDA"],
                    "palette": get_palette("algae"),
                },
            }
    else:
        raise ValueError(
            "Incorrect definition of satellite. Options: 'Landsat', 'landsat', 'Sentinel2', or 'sentinel2'"
        )

    # Get default visualization parameters for the given image name
    params = params_dict.get(index, None)
    if params is None:
        raise ValueError(
            "Incorrect definition of index. Options: 'TrueColor', 'NDVI', 'NDWI', 'halite', 'gypsum', 'LST', 'NDTI', or 'KIVU'. Note LST is not available for Sentinel2"
        )

    # Override default values if provided by the user
    if min_val is not None:
        params["min"] = min_val
    if max_val is not None:
        params["max"] = max_val
    if palette is not None:
        # Check if palette is a string and get the palette list if it is
        if isinstance(palette, str):
            palette = get_palette(palette)
            if palette is None:
                raise ValueError(f"Palette {palette} not found.")
        params["palette"] = palette

    return params
