"""Written by @Jinerhal, adapted by @Xiaojieqiu.
"""

from typing import Tuple

import numpy as np
from anndata import AnnData

from ..configuration import SKM
from ..logging import logger_manager as lm
from .utils import *


@SKM.check_adata_is_type(SKM.ADATA_UMI_TYPE)
def digitize(
    adata: AnnData,
    ctrs: Tuple,
    ctr_idx: int,
    pnt_xy: Tuple[int, int],
    pnt_Xy: Tuple[int, int],
    pnt_xY: Tuple[int, int],
    pnt_XY: Tuple[int, int],
    spatial_key: str = "spatial",
    dgl_layer_key: str = "digital_layer",
    dgl_column_key: str = "digital_column",
    max_itr: int = 1e5,
    lh: float = 1,
    hh: float = 100,
) -> None:
    """Calculate the "heat" for a closed area of interests by solving a PDE, partial differential equation, the heat
        equation. Boundary conditions are defined upon four user provided coordinates that set the direction of heat
        diffusion. The value of "heat" will be used for define different spatial layers, domains and grids.

    Args:
        adata: The adata object to digitize.
        ctrs: Contours generated by `cv2.findContours`.
        ctr_idx: The index of the contour of interests.
        pnt_xy: Corner point to define an area of interest. pnt_xy corresponds to the point with minimal layer and
            minimal column value.
        pnt_Xy: Corner point corresponds to the point with maximal column value but minimal layer value.
        pnt_xY: Corner point corresponds to the point with minimal column value but maximal layer value.
        pnt_XY: Corner point corresponds to the point with maximal layer and maximal columns value.
        spatial_key: The key name in `adata.obsm` of the spatial coordinates. Default to "spatial".
        dgl_layer_key: The key name in `adata.obs` to store layer digital-heat (temperature). Default to
            "digital_layer".
        dgl_column_key: The key name to store column digital-heat (temperature).
        max_itr: Maximum number of iterations dedicated to solving the heat equation.
        lh: lowest digital-heat (temperature). Defaults to 1.
        hh: highest digital-heat (temperature). Defaults to 100.

    Returns:
        Nothing but update the `adata` object with the following keys in `.obs`:
        1. dgl_layer_key: The key in `adata.obs` points to the values of layer digital-heat (temperature).
        2. dgl_column_key: The key in `adata.obs` points to the values of column digital-heat (temperature).
    """

    lm.main_info("Initialize the field of the spatial domain of interests.")
    empty_field = np.zeros((int(max(adata.obsm[spatial_key][:, 0])) + 1, int(max(adata.obsm[spatial_key][:, 1])) + 1))

    field_border = np.zeros_like(empty_field)
    cv2.drawContours(field_border, ctrs, ctr_idx, ctr_idx + 1, 1)
    field_mask = np.zeros_like(empty_field)
    cv2.drawContours(field_mask, ctrs, ctr_idx, ctr_idx + 1, cv2.FILLED)

    lm.main_info("Prepare the isoline segments with either the highest/lower column or layer heat values.")
    min_line_l, max_line_l, min_line_c, max_line_c = field_contours(ctrs[ctr_idx], pnt_xy, pnt_Xy, pnt_xY, pnt_XY)

    lm.main_info("Solve the layer heat equation on spatial domain with the iso-layer-line conditions.")
    # of: the optimal field after solving the heat equation.
    of_layer = domain_heat_eqn_solver(
        empty_field,
        min_line_l,
        max_line_l,
        min_line_c,
        max_line_c,
        field_border,
        field_mask,
        lh=lh,
        hh=hh,
        max_itr=max_itr,
    )

    lm.main_info(f"Saving layer heat values to {dgl_layer_key}.")
    adata.obs[dgl_layer_key] = 0.0
    print(adata.obs[dgl_layer_key])
    for i in range(len(adata)):
        adata.obs[dgl_layer_key][i] = of_layer[int(adata.obsm[spatial_key][i, 0]), int(adata.obsm[spatial_key][i, 1])]

    lm.main_info("Solve the column heat equation on spatial domain with the iso-column-line conditions.")
    of_column = domain_heat_eqn_solver(
        empty_field,
        min_line_c,
        max_line_c,
        min_line_l,
        max_line_l,
        field_border,
        field_mask,
        lh=lh,
        hh=hh,
        max_itr=max_itr,
    )

    lm.main_info(f"Saving column heat values to {dgl_column_key}.")
    adata.obs[dgl_column_key] = 0.0
    for i in range(len(adata)):
        adata.obs[dgl_column_key][i] = of_column[int(adata.obsm[spatial_key][i, 0]), int(adata.obsm[spatial_key][i, 1])]


@SKM.check_adata_is_type(SKM.ADATA_UMI_TYPE)
def gridit(
    adata: AnnData,
    layer_num: int,
    column_num: int,
    lh: float = 1,
    hh: float = 100,
    dgl_layer_key: str = "digital_layer",
    dgl_column_key: str = "digital_column",
    layer_border_width: int = 2,
    column_border_width: int = 2,
    layer_label_key: str = "layer_label",
    column_label_key: str = "column_label",
    grid_label_key: str = "grid_label",
) -> None:
    """Segment the area of interests into specific layer/column number, according to precomputed digitization heat
    value.

    Args:
        adata: The adata object to do layer/column/grid segmentation.
        layer_num: Number of layers to segment.
        column_num: Number of columns to segment.
        lh: lowest digital-heat. Default to 1.
        hh: highest digi-heat. Default to 100.
        layer_border_width: Layer boundary width. Only affect grid_label.
        column_border_width: Column boundary width. Only affect grid_label.
        dgl_layer_key: The key name of layer digitization heat in `adata.obs`. Default to "digital_layer", precomputed.
        dgl_column_key: The key name of column digitization heat in `adata.obs`. Default to "digital_column",
            precomputed.
        layer_label_key: The key name to store layer labels in `adata.obs`. Default to "layer_label", will be added.
        column_label_key: The key name to store column labels in `adata.obs`. Default to "column_label", will be added.
        grid_label_key: The key name to store grid labels in `adata.obs`. Default to "grid_label", will be added.

    Returns:
        Nothing but update the adata object with the following keys in `.obs`:
         1. layer_label_key: this key points to layer labels.
         2. column_label_key: this key points to column labels.
         3. grid_label_key: this key points to grid labels.
    """

    lm.main_info(f"Prepare {layer_label_key}, {column_label_key} and {grid_label_key} keys.")
    adata.obs[layer_label_key] = 0
    adata.obs[column_label_key] = 0

    adata.obs[grid_label_key] = "NA"
    adata.obs[grid_label_key][adata.obs[dgl_layer_key] != 0] = "Grid Area"
    adata.obs[grid_label_key][adata.obs[dgl_column_key] != 0] = "Grid Area"

    region_mask = adata.obs[grid_label_key].copy()
    region_mask[region_mask == "Grid Area"] = "Region Boundary"

    #  Evenly dividing layer heat values into {layer_num} intervals and identify the index of the interval for each
    #  bucket
    lm.main_info(f"Set layer number for each bucket and identify buckets on the layer grids")
    value_list = np.linspace(lh, hh, layer_num + 1)
    for i in range(len(value_list) - 1):
        # fast method to identify within a list of range
        adata.obs[layer_label_key] = np.where(
            (adata.obs[dgl_layer_key] > value_list[i]) & (adata.obs[dgl_layer_key] <= value_list[i + 1]),
            i + 1,
            adata.obs[layer_label_key],
        )
        adata.obs[grid_label_key] = np.where(
            (adata.obs[dgl_layer_key] > (value_list[i] - layer_border_width / 2))
            & (adata.obs[dgl_layer_key] <= (value_list[i] + layer_border_width / 2)),
            region_mask,
            adata.obs[grid_label_key],
        )

    #  Evently dividing column heat values into {column_num} intervals and identify the index of the interval for each
    #  bucket
    lm.main_info(f"Set column number for each bucket and identify buckets on the column grids")
    value_list = np.linspace(lh, hh, column_num + 1)
    for i in range(len(value_list) - 1):
        adata.obs[column_label_key] = np.where(
            (adata.obs[dgl_column_key] > value_list[i]) & (adata.obs[dgl_column_key] <= value_list[i + 1]),
            i + 1,
            adata.obs[column_label_key],
        )
        adata.obs[grid_label_key] = np.where(
            (adata.obs[dgl_column_key] > (value_list[i] - column_border_width / 2))
            & (adata.obs[dgl_column_key] <= (value_list[i] + column_border_width / 2)),
            region_mask,
            adata.obs[grid_label_key],
        )
