from pandas import DataFrame as _pd_DataFrame
from numpy import (
    array as _np_array,
    zeros as _np_zeros,
    linspace as _np_linspace,
    searchsorted as _np_searchsorted,
    sort as _np_sort,
)
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable as _make_axes_locatable
from matplotlib.pyplot import close as _plt_close
from matplotlib.colors import LogNorm as _plt_LogNorm, Normalize as _plt_Normalize, LinearSegmentedColormap as _plt_LinearSegmentedColormap, ListedColormap as _plt_ListedColormap
from matplotlib.pyplot import (subplots as _plt_subplots, colorbar as _plt_colorbar, get_cmap as _plt_get_cmap)
from matplotlib.patches import Patch as _plt_Patch
from aabpl.illustrations.plot_utils import add_color_bar_ax, set_map_frame, truncate_colormap, plot_polygon
from shapely.geometry import Polygon as _shapely_Polygon

def create_distribution_plot(
        pts:_pd_DataFrame,
        rndm_pts:_pd_DataFrame,
        cluster_threshold_values:list,
        k_th_percentile:list,
        radius_sum_columns:list,
        grid:dict=None,
        r:float=None,
        x:str='lon',
        y:str='lat',
        filename:str='',
        plot_kwargs:dict={},
        show:bool=True,
        display_dpi:int=100,
):
    """
    Four-panel plot comparing the random-point null distribution to the
    observed radius-sum distribution. Called via ``grid.plot.rand_dist()``.

    Parameters
    ----------
    pts : DataFrame
        Source points with radius-sum columns already appended.
    rndm_pts : DataFrame
        Random points with radius-sum columns appended by ``compute_null_distribution``.
    cluster_threshold_values : list[float]
        One threshold per column — the k-th percentile cutoff value.
    k_th_percentile : list[float]
        Percentile values used to derive ``cluster_threshold_values``.
    radius_sum_columns : list[str]
        Column names of the radius sums to plot.
    grid
        Grid object (used to draw sample-area overlay).
    r : float
        Search radius in metres (used for the plot title).
    x, y : str
        Coordinate column names in ``pts`` and ``rndm_pts``.
    filename : str
        Save path. Empty string skips saving (default ``''``).
    plot_kwargs : dict
        Supported keys:
        - ``figsize`` : tuple (default ``(10, 10)``)
        - ``s`` : float — scatter marker size (default ``0.8``)
        - ``color`` : str — marker colour (default ``'#eaa'``)
        - ``hlines`` : dict — kwargs for ``ax.hlines`` threshold line
        - ``vlines`` : dict — kwargs for ``ax.vlines`` percentile line
        - ``fig``, ``axs`` — existing Figure/Axes to draw into
        - ``sample_area_linewidth`` : float — border linewidth of the valid-area polygon (default ``0.2``)
        - ``color_scale`` : str — ``'log'``, ``'linear'``, or ``'auto'`` (log when vmin>0, default ``'auto'``)
    show : bool
        Display the figure (default ``True``).
    display_dpi : int
        Resolution for inline/screen display (default ``100``).
        Use ``save_kwargs={'dpi': 300}`` to control saved-file resolution.

    Returns
    -------
    matplotlib.figure.Figure
    """
    x_coord_name, y_coord_name = x, y
    for _df, _col, _label in [
        (pts, x_coord_name, 'pts'), (pts, y_coord_name, 'pts'),
        (rndm_pts, x_coord_name, 'rndm_pts'), (rndm_pts, y_coord_name, 'rndm_pts'),
    ]:
        if _col not in _df.columns:
            raise KeyError(
                f"Coordinate column '{_col}' not found in {_label}. "
                f"The column must be the local-projection coordinate (in metres). "
                f"Pass the correct column name via x= / y= to the plot call, "
                f"or check that keep_cols did not drop it. "
                f"Available columns: {list(_df.columns)}"
            )
    _pts_x = pts[x_coord_name]
    _pts_y = pts[y_coord_name]
    _rnd_x = rndm_pts[x_coord_name]
    _rnd_y = rndm_pts[y_coord_name]
    disk_sums_for_random_points = rndm_pts[radius_sum_columns]
    (n_random_points, ncols) = disk_sums_for_random_points.shape
    # specify default plot kwargs and add defaults
    default_kwargs = {
        's':0.8,
        'color':'#eaa',

        'figsize': (10,10),
        'fig':None,
        'axs':None,

        'title': None,           # override the auto suptitle; use {auto} to embed it: "My prefix: {auto}"
        'suptitle': {},          # extra kwargs forwarded to fig.suptitle(), e.g. fontsize, fontweight, y
        'ax_titles': {},         # dict keyed by panel name to override individual panel titles
                                 # keys: 'rand_dist', 'pts_dist', 'rand_map', 'pts_map'
        'hlines':{'color':'red', 'linewidth':1},
        'vlines':{'color':'red', 'linewidth':1},
        'sample_area_linewidth': 0.2,
        'color_scale': 'auto',   # 'log', 'linear', or 'auto' (log when vmin>0, else linear)
    }
    # Callers (grid.plot.rand_dist) forward **plot_kwargs; extract control args
    # that must not leak into ax.scatter().
    show = plot_kwargs.pop('show', show)
    display_dpi = plot_kwargs.pop('display_dpi', display_dpi)
    plot_kwargs.pop('filename', None)
    kwargs = {}
    for k in list(plot_kwargs.keys()):
        if k in [k for k,v in default_kwargs.items() if type(v)==dict]:
            kwargs[k] = {**default_kwargs.pop(k), **plot_kwargs.pop(k)}
    kwargs.update(default_kwargs)
    kwargs.update(plot_kwargs)
    figsize = kwargs.pop('figsize')
    fig = kwargs.pop('fig')
    axs = kwargs.pop('axs')
    title_override = kwargs.pop('title')
    suptitle_kwargs = kwargs.pop('suptitle')
    ax_titles = kwargs.pop('ax_titles')
    sample_area_lw = kwargs.pop('sample_area_linewidth')
    color_scale = kwargs.pop('color_scale', 'auto')
    for k in ['fig', 'axs', 'figsize', 'title', 'suptitle', 'ax_titles', 'sample_area_linewidth', 'color_scale']:
        plot_kwargs.pop(k,None)

    if fig is None or axs is None:
        fig = plt.figure(figsize=figsize, dpi=display_dpi, constrained_layout=True)
        outer = gridspec.GridSpec(ncols, 1, figure=fig, hspace=0.08)

    _auto_title = "Aggregate for indicator" + ("" if ncols==1 else "s") + " within " + str(r) + " meters"
    _suptitle = title_override.replace('{auto}', _auto_title) if title_override else _auto_title
    fig.suptitle(_suptitle, **suptitle_kwargs)
    non_valid_area = _shapely_Polygon([
        (grid.sample_grid_bounds[0], grid.sample_grid_bounds[1]),
        (grid.sample_grid_bounds[2], grid.sample_grid_bounds[1]),
        (grid.sample_grid_bounds[2], grid.sample_grid_bounds[3]),
        (grid.sample_grid_bounds[0], grid.sample_grid_bounds[3])
        ]).difference(grid.sample_area)

    pct_xmin, pct_xmax = 0, 100
    xs_random_pts = _np_linspace(pct_xmin, pct_xmax, n_random_points)
    xs_pts = _np_linspace(pct_xmin, pct_xmax, len(pts))
    random_vals = disk_sums_for_random_points.values
    pts_vals = pts[radius_sum_columns].values

    for (i, colname, cluster_threshold_value, k) in zip(
        range(ncols), radius_sum_columns, cluster_threshold_values, k_th_percentile):
        # Layout: top row = merged distribution axis; bottom row = two maps + shared colorbar
        # Map orientation: side-by-side when portrait (H > 1.2*W), stacked when landscape/flat
        map_w = max((_pts_x.max() if hasattr(_pts_x, 'max') else max(_pts_x)) - (_pts_x.min() if hasattr(_pts_x, 'min') else min(_pts_x)), 1e-9)
        map_h = max((_pts_y.max() if hasattr(_pts_y, 'max') else max(_pts_y)) - (_pts_y.min() if hasattr(_pts_y, 'min') else min(_pts_y)), 1e-9)
        maps_side_by_side = map_h > 1.05 * map_w
        panel = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=outer[i], hspace=0.12, height_ratios=[1, 4])
        if maps_side_by_side:
            bottom_gs = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec=panel[1], wspace=0.08)
        else:
            bottom_gs = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=panel[1], hspace=0.08)

        # CUMULATIVE DISTRIBUTIONS — both series on a single merged axis
        ys_rnd = _np_sort(random_vals[:,i])
        ys_pts = _np_sort(pts_vals[:,i])
        ymin = min(ys_rnd.min(), ys_pts.min())
        ymax = max(ys_rnd.max(), ys_pts.max())

        idx = _np_searchsorted(ys_rnd, cluster_threshold_value)
        next_smaller_val = ys_rnd[max(0, idx - 1)]
        next_larger_val  = ys_rnd[idx] if idx < len(ys_rnd) else ys_rnd[-1]
        sufficient_digits = next((
            d for d in range(100) if (
                (
                    (next_smaller_val == next_larger_val or cluster_threshold_value == next_smaller_val) and
                    round(next_smaller_val, d) == next_smaller_val
                ) or (
                    round(next_larger_val, d) != round(cluster_threshold_value, d) and
                    round(next_smaller_val, d) != round(cluster_threshold_value, d)
                )
            )
        ), 100)

        ax_dist = fig.add_subplot(panel[0])
        # ax_title = ax_titles.get('rand_dist', "{{col}}: {{k}}th-percentile threshold = {thresh:.6g}")
        ax_title = ax_titles.get('rand_dist', "{col}: {k}th-percentile threshold = {thresh}")
        ax_title = ax_title.format(col=colname, k=k, thresh=round(cluster_threshold_value, sufficient_digits), r=r, n=len(rndm_pts))
        ax_dist.set_title(ax_title, fontdict={'fontsize': 6})
        ax_dist.set_ylabel("Radius sum", fontsize=7)
        ax_dist.set_xlabel("Percentile rank", fontsize=7)

        # ticks — x shared, y spanning both series
        xtick_steps, ytick_steps = 5, 5
        xticks = _np_array(sorted(
            [x for x in _np_linspace(pct_xmin, pct_xmax, xtick_steps) if abs(x - k) > (pct_xmax - pct_xmin) / (xtick_steps * 2)] + [k]
        ))
        ax_dist.set_xticks(xticks, labels=xticks)
        yticks = _np_array(sorted(
            [y for y in _np_linspace(ymin, ymax, ytick_steps) if abs(cluster_threshold_value - y) > (ymax - ymin) / (ytick_steps * 10)] +
            [cluster_threshold_value]
        ))
        ax_dist.set_yticks(yticks, labels=[round(t, sufficient_digits) for t in yticks])

        # threshold lines
        ax_dist.hlines(y=cluster_threshold_value, xmin=pct_xmin, xmax=pct_xmax, **kwargs['hlines'])
        ax_dist.vlines(x=k, ymin=ymin, ymax=ymax, **kwargs['vlines'])

        # both distribution series
        _s = plot_kwargs.get('s', 0.8)
        rnd_color = plot_kwargs.get('color', '#eeaaaa')
        ax_dist.scatter(x=xs_random_pts, y=ys_rnd, s=_s, marker='.', label=f"Random pts (n={len(rndm_pts):,})", **{k: v for k, v in plot_kwargs.items() if k != 's'})
        _pts_color = '#66aabb'
        ax_dist.scatter(x=xs_pts, y=ys_pts, color=_pts_color, s=_s, marker='.', label=f"Dataset pts (n={len(pts):,})")
        ax_dist.set_xlim([pct_xmin, pct_xmax])
        if ymin != ymax:
            ax_dist.set_ylim([ymin, ymax])
        ax_dist.legend(fontsize=6, markerscale=2)

        # MAP COLORMAP SETUP
        map_xmin = min(_pts_x.min(), _rnd_x.min())
        map_xmax = max(_pts_x.max(), _rnd_x.max())
        map_ymin = min(_pts_y.min(), _rnd_y.min())
        map_ymax = max(_pts_y.max(), _rnd_y.max())

        _pts_nonzero = pts_vals[:,i][pts_vals[:,i] != 0]
        _rnd_nonzero = random_vals[:,i][random_vals[:,i] != 0]
        _has_range = len(_pts_nonzero) > 0 and len(_rnd_nonzero) > 0
        if _has_range:
            vmin = min(_pts_nonzero.min(), _rnd_nonzero.min())
            vmax = max(pts_vals[:,i].max(), random_vals[:,i].max())
            _has_range = vmin < vmax and vmin > 0
        if not _has_range:
            vmin = 1.0
            vmax = max(float(cluster_threshold_value), 2.0)

        non_valid_area_color = "#bedbe6"
        sample_area_color = "#ffffff"
        color_cluster = "#2a07ee"
        cmap_scatter = _plt_get_cmap('Reds')
        minval = 0.2
        color_under = cmap_scatter(minval / 2)
        cmap_scatter = truncate_colormap(cmap=cmap_scatter, minval=minval, maxval=1.0, n=100)
        cmap_scatter.set_under(color_under)
        cmap_scatter.set_bad(color_under)
        cmap_scatter.set_over(color_cluster)
        _norm_vmax = float(cluster_threshold_value) if cluster_threshold_value > vmin else vmax
        _use_log = (color_scale == 'log') or (color_scale == 'auto' and vmin > 0)
        norm = _plt_LogNorm(vmin=vmin, vmax=_norm_vmax, clip=False) if _use_log else _plt_Normalize(vmin=vmin, vmax=_norm_vmax, clip=False)
        s = 0.2 * figsize[0] / 10

        # build grid overlay once (reused by both map axes)
        X = None
        extent = None
        cmap_binary = None
        if grid is not None:
            cells_rndm_sample = grid.cells_rndm_sample
            col_min = int(round((grid.sample_grid_bounds[0] - grid.total_bounds.xmin) / grid._search_spacing, 0))
            row_min = int(round((grid.sample_grid_bounds[1] - grid.total_bounds.ymin) / grid._search_spacing, 0))
            col_max = int(round((grid.sample_grid_bounds[2] - grid.total_bounds.xmin) / grid._search_spacing - 1, 0))
            row_max = int(round((grid.sample_grid_bounds[3] - grid.total_bounds.ymin) / grid._search_spacing - 1, 0))
            n_rows_x = row_max - row_min + 1
            n_cols_x = col_max - col_min + 1
            if type(cells_rndm_sample) == bool:
                X = _np_zeros((n_rows_x, n_cols_x), dtype=bool)
                if not cells_rndm_sample:
                    X[:] = True
            else:
                X = _np_zeros((n_rows_x, n_cols_x), dtype=bool)
                X[:] = True
                for (row, col) in cells_rndm_sample:
                    ri = row_max - row
                    ci = col - col_min
                    if 0 <= ri < n_rows_x and 0 <= ci < n_cols_x:
                        X[ri, ci] = False
            cmap_binary = _plt_ListedColormap([sample_area_color, non_valid_area_color])
            extent = [grid.sample_grid_bounds[0], grid.sample_grid_bounds[2],
                      grid.sample_grid_bounds[1], grid.sample_grid_bounds[3]]

        # MAP — random points (first map panel)
        ax_rnd = fig.add_subplot(bottom_gs[0])
        ax_rnd.set_facecolor(sample_area_color)
        ax_rnd.set_title(ax_titles.get('rand_map', "Random points (n={n})").format(col=colname, k=k, thresh=round(cluster_threshold_value, sufficient_digits), r=r, n=len(rndm_pts)), fontdict={'fontsize': 6})
        if X is not None:
            ax_rnd.imshow(X=X, interpolation='none', cmap=cmap_binary, extent=extent)
            non_valid_patch = _plt_Patch(facecolor=non_valid_area_color, label='Non-valid area', edgecolor='black')
            sample_patch = _plt_Patch(facecolor=sample_area_color, label='Sample area', edgecolor='black')
            ax_rnd.legend(handles=[non_valid_patch, sample_patch], loc='best', fontsize=5)
        plot_polygon(ax=ax_rnd, poly=non_valid_area, facecolor=non_valid_area_color, edgecolor='black', linewidth=sample_area_lw)
        sc = ax_rnd.scatter(x=_rnd_x, y=_rnd_y, c=random_vals[:,i], s=s, marker='.', norm=norm, cmap=cmap_scatter, linewidths=0.3)
        plot_polygon(ax=ax_rnd, poly=grid.sample_area, facecolor="none", edgecolor='black', linewidth=sample_area_lw)
        set_map_frame(ax=ax_rnd, xmin=map_xmin, xmax=map_xmax, ymin=map_ymin, ymax=map_ymax, r=r)

        # MAP — dataset points (second map panel)
        ax_pts = fig.add_subplot(bottom_gs[1])
        ax_pts.set_facecolor(sample_area_color)
        ax_pts.set_title(ax_titles.get('pts_map', "Dataset points (n={n})").format(col=colname, k=k, thresh=round(cluster_threshold_value, sufficient_digits), r=r, n=len(pts)), fontdict={'fontsize': 6})
        if X is not None:
            ax_pts.imshow(X=X, interpolation='none', cmap=cmap_binary, extent=extent)
        plot_polygon(ax=ax_pts, poly=non_valid_area, facecolor=non_valid_area_color, edgecolor='black', linewidth=sample_area_lw)
        sc_pts = ax_pts.scatter(x=_pts_x, y=_pts_y, c=pts_vals[:,i], s=s, marker='.', norm=norm, cmap=cmap_scatter, linewidths=0.3)
        set_map_frame(ax=ax_pts, xmin=map_xmin, xmax=map_xmax, ymin=map_ymin, ymax=map_ymax, r=r)

        # shared colorbar spanning the combined height of both map panels
        # make_axes_locatable respects set_aspect so the bar height matches the rendered maps
        if maps_side_by_side:
            _div = _make_axes_locatable(ax_pts)
            _cax = _div.append_axes("right", size="3%", pad=0.05)
            fig.colorbar(sc_pts, cax=_cax, extend='max')
        else:
            fig.colorbar(sc_pts, ax=[ax_rnd, ax_pts], extend='max', fraction=0.03, pad=0.02)

    if filename:
        fig.savefig(filename, dpi=300, bbox_inches="tight")
    if not show:
        _plt_close(fig)
    return fig
    #
#