from numpy import array as _np_array
from pandas import (DataFrame as _pd_DataFrame, cut as _pd_cut, concat as _pd_concat) 
from aabpl.utils.misc import DataFrameRelation, arr_to_tpls, find_column_name
from aabpl.utils.cell_geometry import classify_disk_cells
from .region_tree import gen_weak_order_rel_to_convex_set
from ..point_assignment import assign_points_to_cells, aggregate_point_data_to_cells, aggregate_point_data_to_cells_adaptive_nd
from ..point_region_assignment import assign_points_to_mirco_regions
from .disk_aggregation_chunk_adaptive_nd import search_and_aggregate as _search_and_aggregate_adaptive

from aabpl.testing.test_performance import time_func_perf
from aabpl import config as _cfg

def search_and_aggregate(*args, **kwargs):
    # FIXED_NEST_DEPTH routes through the adaptive file: nd_choice.py's
    # best_nd_tag short-circuits to the fixed tag when set, so the adaptive
    # file's planning (col-splitting, chunking, etc.) degenerates to "always
    # this one nd" instead of choosing per-chunk. Keeps a single
    # column-split/chunking implementation to maintain instead of parallel ones.
    return _search_and_aggregate_adaptive(*args, **kwargs)


################ DiskSearchSource ######################################################################################
class DiskSearchBase(object):

    def assign_pts_to_cells(
        self,
        silent:bool = False
    ):
        return assign_points_to_cells(
            grid=self.grid,
            pts=self.pts,
            y=self.y,
            x=self.x,
            row_name=self.row_name,
            col_name=self.col_name,
            silent=silent,
        )
    #
#

################ DiskSearchSource ######################################################################################
class DiskSearchSource(DiskSearchBase):
    # @time_func_perf
    def __init__(
        self,
        grid,
        pts:_pd_DataFrame,
        c:list=[],
        y:str='lat',
        x:str='lon',
        row_name:str='id_y',
        col_name:str='id_x',
        suffix:str='_750m',
    ):
        self.grid = grid 
        self.pts = pts
        self.y = y
        self.x = x
        self.row_name = row_name
        self.col_name = col_name
        self.cell_region_name = find_column_name('cell_reg', existing_columns=pts.columns)
        self.off_x = find_column_name('offset_x', existing_columns=pts.columns)
        self.off_y = find_column_name('offset_y', existing_columns=pts.columns)
        
        self.suffix = suffix
        self.aggregate_columns = [str(column)+suffix for column in c]

    #next(('helper_col'+i for i in (['']+list(range(len(pts_target.columns))))))

    def assign_pts_to_cell_regions(
            self,
            plot_cell_reg_assign:dict=None,
            plot_offset_checks:dict=None,
            plot_offset_regions:dict=None,
            plot_offset_raster:dict=None,
            silent:bool=False,
    ):
        # TODO 
        return assign_points_to_mirco_regions(
        # return assign_points_to_cell_regions(
            grid=self.grid,
            pts=self.pts,
            r=self.grid._search_class.r,
            nest_depth=self.grid._search_internals.nest_depth,
            include_boundary=self.grid._search_class.include_boundary,
            y=self.y,
            x=self.x,
            off_x=self.off_x,
            off_y=self.off_y,
            row_name=self.row_name,
            col_name=self.col_name,
            cell_region_name=self.cell_region_name,
            plot_cell_reg_assign=plot_cell_reg_assign,
            plot_offset_checks=plot_offset_checks,
            plot_offset_regions=plot_offset_regions,
            plot_offset_raster=plot_offset_raster,
            silent=silent,
        )
    #
#

################ DiskSearchTarget ######################################################################################
class DiskSearchTarget(DiskSearchBase):
    @time_func_perf
    def __init__(
        self,
        grid,
        pts:_pd_DataFrame,
        c:list=[],
        y:str='lat',
        x:str='lon',
        row_name:str='id_y',
        col_name:str='id_x',
    ):
        self.grid = grid
        self.pts = pts
        self.c = c
        self.x = x
        self.y = y
        self.row_name = row_name
        self.col_name = col_name
    #
    
    def aggregate_pt_data_to_cells(
            self,
            silent,
            depth_by_block=None,
            needs_level0=None,
    ):
        if getattr(_cfg, 'USE_ADAPTIVE_ND_AGGREGATION', False):
            _r = self.grid._search_class.r
            _spacing = self.grid._search_internals.spacing
            return aggregate_point_data_to_cells_adaptive_nd(
                grid=self.grid,
                pts=self.pts,
                y=self.y,
                x=self.x,
                c=self.c,
                row_name=self.row_name,
                col_name=self.col_name,
                nest_depth=self.grid._search_internals.nest_depth,
                sr=_r / _spacing,
                silent=silent,
                depth_by_block=depth_by_block,
                needs_level0=needs_level0,
            )
        return aggregate_point_data_to_cells(
            grid=self.grid,
            pts=self.pts,
            y=self.y,
            x=self.x,
            c=self.c,
            row_name=self.row_name,
            col_name=self.col_name,
            nest_depth=self.grid._search_internals.nest_depth,
            silent=silent,
        )
    #
#
   

################ DiskSearch ######################################################################################
class DiskSearch(object):
    # @time_func_perf
    def __init__(
        self,
        grid,
        r:float=0.0075,
        nest_depth:int=None,
        exclude_self:bool=True,
        area_weight:str=None,
        weight_valid_area:str=None,  # deprecated alias
        include_boundary:bool=False,
    ):

        """

        """
        if weight_valid_area is not None and area_weight is None:
            from aabpl.utils.progress import progress_print
            progress_print("Deprecation: 'weight_valid_area' has been renamed to 'area_weight'.")
            area_weight = weight_valid_area
        # link to grid
        grid._search_class = self
        self.grid = grid
        self.exclude_self = exclude_self
        self.area_weight = area_weight
        self.include_boundary = include_boundary

        self.update_search_params(
            grid=grid,
            exclude_self=exclude_self,
            area_weight=area_weight,
            r=r,
            nest_depth=nest_depth,
            include_boundary=include_boundary,
        )
        #
    #

    @time_func_perf
    def update_search_params(
        self,
        grid,
        exclude_self:bool=None,
        area_weight:str=None,
        r:float=None,
        nest_depth:int=None,
        include_boundary:bool=None,
        relation_tgt_to_src:str=None,
    ):
        if exclude_self is not None:
            self.exclude_self = exclude_self
        if area_weight is not None:
            self.area_weight = area_weight
        if relation_tgt_to_src is not None:
            self.relation_tgt_to_src = relation_tgt_to_src
        
        # early return if r and include_boundary have not changed
        if hasattr(self, 'r') and hasattr(self, 'include_boundary'):
            if self.r == r and self.include_boundary == include_boundary:
                return
        
        # store params
        self.r = r
        self.nest_depth = nest_depth if not nest_depth is None else grid._search_internals.nest_depth
        self.include_boundary = include_boundary
        self.overlap_checks = []
        self.contain_checks = []
        
        # get relative position of cells that are always included within r for current gridsize    
        (
        self.cells_cntd_in_all_disks, 
        self.cells_cntd_in_all_trgl_disks, 
        self.cells_maybe_overlapping_a_disk, 
        self.cells_maybe_overlapping_a_trgl_disk
        ) = classify_disk_cells(
                grid_spacing=grid._search_internals.spacing,
                r=r,
                include_boundary=include_boundary,
        )
        # hierarchically order all cells with respect to any point in triangle. 
        # Some cells are at least as far away as others 
        # e.g. (2,2) is weakly closer than (-2,-2) as any pt in triangle is P(x>=0,y>=0)
        # e.g. (2,3) is weakly closer than (?,?) as any pt in triangle is P(x,y<=0.5*grid._search_internals.spacing)
        triangle_1_vertices = _np_array([[0,0],[0.5,0],[0.5,0.5]])
        vertices_is_inside_triangle_1 = _np_array([True,True,False],dtype=bool)
        # TODO r,grid_spacing, include_boundary could be removed from weak_order_tree generation
        # TODO THIS HAS BECOME OBSOLETE
        self.weak_order_tree = gen_weak_order_rel_to_convex_set(
                cells=self.cells_maybe_overlapping_a_trgl_disk,
                convex_set_vertices = triangle_1_vertices,
                vertex_is_inside_convex_set = vertices_is_inside_triangle_1,
                r=r,
                grid_spacing=grid._search_internals.spacing,
                include_boundary=include_boundary,
        )
        
        # TODO lvl is hard coded to 0. IF MULTIPLE LEVELS ARE CONSIDERED THIS MIGHT NEED TO BE ADJUSTED in classify_disk_cells
        self.cells_cntd_in_all_disks = [(0,(row,col)) for (row,col) in arr_to_tpls(self.cells_cntd_in_all_disks,int)]
    #
    
    # @time_func_perf
    def check_if_tgt_df_contains_src_df(
        self,
        silent:bool=False,
    )->bool:
        if not hasattr(self, 'target'): return False
        if not hasattr(self, 'source'): return False
        if not hasattr(self.target, 'pts'): return False
        if not hasattr(self.source, 'pts'): return False
        return DataFrameRelation.check_if_df_is_contained(self.source.pts, self.target.pts,silent=silent)
    
    # @time_func_perf
    def check_if_search_obj_already_exist(
        self,
        pts:_pd_DataFrame,
        obj:str=['source','target'],
        silent:bool=False,
        **kwargs
    ):
        """
        check if search sortarget already created
        kwarg at pos 0 is pandas.DataFrame and will thus be checked for
        equlity by .equals insted of ==  
        Args:
        self : DiskSearch
          Checks on disk searchs object
        Returns:
        
        """
        
        # check if this attribute is already set as source
        alr_added_pts_to_grid = (
            hasattr(self, 'target') and 
            all([hasattr(self.target, k) and 
                v == getattr(self.target, k)
                for k,v in kwargs.items()]) and
                hasattr(self.target, 'pts') and 
                DataFrameRelation.check_if_df_is_contained(pts, self.target.pts,silent=silent)
        )
        alr_assg_to_cell_regions = obj=='source' and (
            hasattr(self, 'source') and all([
                hasattr(self.source, k) and 
                v == getattr(self.source, k)
                for k,v in kwargs.items()]) and
                hasattr(self.source, 'pts') and 
                DataFrameRelation.check_if_df_is_contained(pts, self.source.pts,silent=silent)
        )
        alr_assg_to_cells = (alr_added_pts_to_grid or alr_assg_to_cell_regions)
        
        return (alr_assg_to_cells, alr_assg_to_cell_regions, alr_added_pts_to_grid)
    #

    @time_func_perf
    def set_source(
        self,
        pts:_pd_DataFrame,
        c:list=[],
        y:str='lat',
        x:str='lon',
        row_name:str='id_y',
        col_name:str='id_x',
        suffix:str='_750m',
        plot_cell_reg_assign:dict=None,
        plot_offset_checks:dict=None,
        plot_offset_regions:dict=None,
        plot_offset_raster:dict=None,
        silent:bool=False,
    ):
        """
        TODO also make shortcut if  grid.search.tgt_df_contains_src_df
        """
        (alr_assg_to_cells,
        alr_assg_to_cell_regions,
        alr_added_pts_to_grid) = self.check_if_search_obj_already_exist(
            **dict(list(locals().items())[1:8])
        )
        
        self.source = DiskSearchSource(
            grid=self.grid,
            pts=pts,
            c=c,
            y=y,
            x=x,
            row_name=row_name,
            col_name=col_name,
            suffix=suffix,
        )

        self.tgt_df_contains_src_df = self.check_if_tgt_df_contains_src_df(silent=silent)

        if not alr_assg_to_cells:
            self.source.assign_pts_to_cells(silent=silent,)
            # self.source.pts.sort_values(
            #     [self.source.row_name, self.source.col_name],
            #     inplace=True
            # )
        #
        if not alr_assg_to_cell_regions:
            self.source.assign_pts_to_cell_regions(
                plot_cell_reg_assign=plot_cell_reg_assign,
                plot_offset_checks=plot_offset_checks,
                plot_offset_regions=plot_offset_regions,
                plot_offset_raster=plot_offset_raster,
                silent=silent,
            )
            # self.source.pts.sort_values(
            #     [self.source.row_name, self.source.col_name, self.source.cell_region_name],
            #     inplace=True
            # )
        #
    #

    @time_func_perf
    def set_target(
        self,

        pts:_pd_DataFrame,
        c:list=['employment'],
        y:str='lat',
        x:str='lon',
        row_name:str='id_y',
        col_name:str='id_x',

        silent:bool=False,
    ):
        
        (alr_assg_to_cells,
        alr_assg_to_cell_regions,
        alr_added_pts_to_grid) = self.check_if_search_obj_already_exist(
            **dict(list(locals().items())[1:8])
        )

        self.target = DiskSearchTarget(
            grid=self.grid,
            pts=pts,
            c=c,
            y=y,
            x=x,
            row_name=row_name,
            col_name=col_name,
        )

        self.tgt_df_contains_src_df = self.check_if_tgt_df_contains_src_df(silent=silent)

        if not alr_assg_to_cells:
            self.target.assign_pts_to_cells(silent=silent)

            # also sort according to grid cell region!
            # self.target.pts.sort_values(
            #     [self.target.row_name, self.target.col_name],
            #     inplace=True
            # )
        #

        # Aggregation (Layer 1, aggregate_pt_data_to_cells) is deliberately NOT
        # called here anymore -- it used to run as a side effect of set_target,
        # before any chunk planning happened, forcing it to independently
        # re-estimate what nd/depth each region needs (a second, sometimes-
        # disagreeing copy of the same decision the search's own chunk
        # dispatch makes later). It's deferred to perform_search instead,
        # which is now the single place aggregation gets triggered from.
        self._pending_aggregate = not alr_added_pts_to_grid
        #


    #

    @time_func_perf
    def perform_search(
            self,
            plot_pt_disk:dict=None,
            silent:bool=False,
    ):
        _kwargs = dict(
            grid=self.grid,
            pts_source=self.source.pts,
            pts_target=self.target.pts,
            r=self.r,
            c=self.target.c,
            y=self.source.y,
            x=self.source.x,
            off_x=self.source.off_x,
            off_y=self.source.off_y,
            row_name=self.source.row_name,
            col_name=self.source.col_name,
            cell_region_name=self.source.cell_region_name,
            suffix=self.source.suffix,
            exclude_self=self.exclude_self,
            area_weight=self.area_weight,
            plot_pt_disk=plot_pt_disk,
            silent=silent,
        )
        _pending = getattr(self, '_pending_aggregate', True)
        if getattr(_cfg, 'USE_ADAPTIVE_ND_AGGREGATION', False):
            # Single call: search_and_aggregate plans, then (if run_aggregation)
            # runs Layer 1 aggregation inline using that same plan's real
            # per-chunk nd decision, then executes using the same plan -- no
            # second, wasted re-planning pass, and nothing before the planning
            # point (e.g. _parse_area_weight) runs twice.
            #
            # An earlier attempt at this caused a confirmed correctness
            # regression in self-search (pts_target is pts_source): Layer 1
            # aggregation (aggregate_point_data_to_cells_adaptive_nd) used to
            # call pts.sort_values(inplace=True), physically reordering the
            # shared pts_source object mid-function, after this function's own
            # source-side arrays (rows/cols/group boundaries) had already been
            # built from pts_source's original order -- final results then got
            # written to the wrong row positions (same total, permuted values;
            # confirmed via matching aggregate sums with per-point mismatches).
            # Fixed at the root in point_assignment.py: aggregation now sorts a
            # local index array instead of the DataFrame itself, so it never
            # mutates a caller-visible object's row order.
            result = search_and_aggregate(run_aggregation=_pending, **_kwargs)
        else:
            # Non-adaptive aggregation has no dependency on planning at all, so
            # it's still triggered externally, before search_and_aggregate.
            if _pending:
                self.target.aggregate_pt_data_to_cells(silent=silent)
            result = search_and_aggregate(**_kwargs)
        self._pending_aggregate = False
        return result
    #
#
    
    
