from datetime import datetime
from warnings import warn

import dask.array as da
import numpy as np
import pandas as pd
from shapely import MultiPolygon, Polygon
import contextlib
import os
import sys

def _get_expression_values(adata, X, key_type, key):
    # get expression values
    if key_type == "genes":
        try:
            gene_loc = adata.var_names.get_loc(key)
            color_value = X[:, gene_loc]
        except KeyError:
            if key in adata.obs.columns:
                color_value = adata.obs[key]
    elif key_type == "obs":
            color_value = adata.obs[key]
    elif key_type == "obsm":
        #TODO: Implement it for obsm
        obsm_key = key.split("#", maxsplit=1)[0]
        obsm_col = key.split("#", maxsplit=1)[1]
        data = adata.obsm[obsm_key]

        if isinstance(data, pd.DataFrame):
            color_value = data[obsm_col].values
        elif isinstance(data, np.ndarray):
            color_value = data[:, int(obsm_col)-1]
        else:
            warn("Data in `obsm` needs to be either pandas DataFrame or numpy array to be parsed.")
        pass
    else:
        print("Unknown key selected.", flush=True)

    return color_value

def _fill_multipolygon(mp):
    filled_polygons = [Polygon(p.exterior) for p in mp.geoms]
    filled_multipolygon = MultiPolygon(filled_polygons)
    return filled_multipolygon

def _format_multipolygon_coords(coords):
    # Convert to proper format
    formatted_coords = []
    for poly in coords:
        exterior = poly[0]
        holes = poly[1:] if len(poly) > 1 else []
        formatted_coords.append((exterior, holes))

    return formatted_coords

def _convert_to_float_coords(coords, mode):
    # Convert Decimal to float

    # if len(coords) == 1:
    if mode == "Polygon":
        float_coords = [[(float(x), float(y)) for x, y in ring] for ring in coords]
        poly = Polygon(float_coords[0])
    elif mode == "MultiPolygon":
        float_coords = [[[(float(x), float(y)) for x, y in ring] for ring in poly] for poly in coords]
        formatted_coords = _format_multipolygon_coords(float_coords)

        # Create MultiPolygon
        mp = MultiPolygon(formatted_coords)

        # fill holes in the polygons
        poly = _fill_multipolygon(mp)
        #poly = MultiPolygon(float_coords)
    else:
        raise ValueError(f"Unknown mode '{mode}'.")
    return poly

def _generate_mask(values, xmax, ymax, seg_mask_value):
    try:
        from rasterio.features import rasterize
    except ImportError:
        raise ImportError("This function requires the rasterio package, please install with `pip install rasterio`.")
    # rasterize polygons
    boundaries_mask = rasterize(
        list(zip(values, seg_mask_value)),
        out_shape=(ymax,xmax))
    boundaries_mask = da.from_array(boundaries_mask)

    return boundaries_mask


def sort_paths_by_datetime(paths):
    def extract_datetime(path):
        # Assumes ID format: "250805-115555000343-2c58ca86"
        parts = path.name.split("-")
        date_part = parts[0]  # "250805"
        time_part = parts[1]  # "115555000343"

        # Combine into full datetime string: "250805115555000343"
        full_dt_str = date_part + time_part

        # Parse as datetime: YYMMDDHHMMSSffffff
        return datetime.strptime(full_dt_str, "%y%m%d%H%M%S%f")

    return sorted(paths, key=extract_datetime, reverse=True)




@contextlib.contextmanager
def suppress_output():
    with open(os.devnull, 'w') as devnull:
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        try:
            sys.stdout = devnull
            sys.stderr = devnull
            yield
        finally:
            sys.stdout = old_stdout
            sys.stderr = old_stderr
