"""
Method B — Per-pole robust Z-score detection.
==============================================

Ported from the R prototype ``program.R``.

For each pole point:
  1. **Pole sampling window** — a circle of radius ``pole_window_m / 2``
     around the pole.  Three radiance statistics are extracted:
       I_pole_max    = max(pole pixels)
       I_pole_median = median(pole pixels)
       I_pole_Q95    = 95th percentile(pole pixels)
  2. **Background ring (annulus)** — between ``bg_inner_m`` and
     ``bg_outer_m``.  Two statistics are extracted:
       I_bg_median = median(background pixels)
       MAD_bg      = median(|bg - I_bg_median|)
  3. **Robust Z-scores** (per-pole, using the *local* background):
       sigma_robust = max(1.4826 * MAD_bg, eps)
       Z_max    = (I_pole_max    - I_bg_median) / sigma_robust
       Z_median = (I_pole_median - I_bg_median) / sigma_robust
       Z_Q95    = (I_pole_Q95    - I_bg_median) / sigma_robust
  4. **Detection decisions** (Z >= T_z):
       det_max, det_median, det_Q95

Key difference from RAST / Method A
-----------------------------------
Method B uses a **local** background ring per pole (not global statistics)
and produces three detection metrics (max / median / Q95).

Produced attributes
-------------------
I_pole_max, I_pole_median, I_pole_Q95, I_bg_median, MAD_bg,
Z_max, Z_median, Z_Q95, det_max, det_median, det_Q95,
pole_window_m, bg_inner_m, bg_outer_m, T_z, method
"""

import math
import numpy as np
import geopandas as gpd
import rasterio
from rasterio.mask import mask as rasterio_mask


def run(tiff_path, vector_path, band=1, threshold=3.0,
        pole_window_m=30.0, bg_inner_m=40.0, bg_outer_m=80.0):
    """
    Run the Method B analysis.

    Parameters
    ----------
    tiff_path : str
        Path to the input GeoTIFF (night-time light raster).
    vector_path : str
        Path to the input vector file (pole points).
    band : int
        Raster band to use (1, 2 or 3).
    threshold : float
        Robust Z-score threshold ``T_z``.  A pole is *detected* when
        ``Z >= T_z``.
    pole_window_m : float
        Diameter of the pole sampling window in metres
        (buffer radius = ``pole_window_m / 2``).
    bg_inner_m : float
        Inner radius of the background ring in metres.
    bg_outer_m : float
        Outer radius of the background ring in metres.

    Returns
    -------
    geopandas.GeoDataFrame
        Original poles enriched with Method B attributes.
    """
    T_z = float(threshold)
    eps = 1e-6
    pole_radius = pole_window_m / 2.0

    # ------------------------------------------------------------------
    # Read vector data
    # ------------------------------------------------------------------
    gdf = gpd.read_file(vector_path)

    for col in gdf.columns:
        if str(gdf[col].dtype).startswith('datetime'):
            gdf[col] = gdf[col].apply(
                lambda x: x.isoformat()
                if x is not None and not (isinstance(x, float) and math.isnan(x))
                else None
            )

    # ------------------------------------------------------------------
    # Open raster and reproject vector
    # ------------------------------------------------------------------
    with rasterio.open(tiff_path) as src:
        raster_crs = src.crs
        gdf_raster_crs = gdf.to_crs(raster_crs)

        centroid = gdf_raster_crs.geometry.unary_union.centroid
        lon, lat = centroid.x, centroid.y
        if gdf_raster_crs.crs.is_geographic:
            utm_zone = int((lon + 180) / 6) + 1
            hemisphere = 'north' if lat >= 0 else 'south'
            utm_crs = f'+proj=utm +zone={utm_zone} +{hemisphere} +datum=WGS84'
        else:
            utm_crs = raster_crs

        gdf_proj = gdf_raster_crs.to_crs(utm_crs)

        # ------------------------------------------------------------------
        # Per-pole extraction
        # ------------------------------------------------------------------
        n = len(gdf_proj)
        I_pole_max_arr    = np.full(n, np.nan)
        I_pole_median_arr = np.full(n, np.nan)
        I_pole_Q95_arr    = np.full(n, np.nan)
        I_bg_median_arr   = np.full(n, np.nan)
        MAD_bg_arr        = np.full(n, np.nan)
        Z_max_arr         = np.full(n, np.nan)
        Z_median_arr      = np.full(n, np.nan)
        Z_Q95_arr         = np.full(n, np.nan)
        det_max_arr       = np.zeros(n, dtype=int)
        det_median_arr    = np.zeros(n, dtype=int)
        det_Q95_arr       = np.zeros(n, dtype=int)

        for i, (_idx, row) in enumerate(gdf_proj.iterrows()):
            point = row.geometry

            # --- 6.1 Pole sampling window ---
            pole_circle = point.buffer(pole_radius)
            pole_gdf = gpd.GeoDataFrame(
                geometry=[pole_circle], crs=utm_crs
            ).to_crs(raster_crs)
            pole_geom = [pole_gdf.geometry.iloc[0]]

            try:
                pole_data, _ = rasterio_mask(
                    src, pole_geom, crop=True, nodata=0, filled=True
                )
                pole_band = pole_data[band - 1]
                pole_vals = pole_band[pole_band != 0]
            except Exception:
                pole_vals = np.array([])

            if len(pole_vals) == 0:
                continue

            I_pole_max_arr[i]    = float(np.max(pole_vals))
            I_pole_median_arr[i] = float(np.median(pole_vals))
            I_pole_Q95_arr[i]    = float(np.percentile(pole_vals, 95))

            # --- 6.2 Background ring (annulus) ---
            bg_outer_circle = point.buffer(bg_outer_m)
            bg_inner_circle = point.buffer(bg_inner_m)
            bg_ring = bg_outer_circle.difference(bg_inner_circle)

            bg_gdf = gpd.GeoDataFrame(
                geometry=[bg_ring], crs=utm_crs
            ).to_crs(raster_crs)
            bg_geom = [bg_gdf.geometry.iloc[0]]

            try:
                bg_data, _ = rasterio_mask(
                    src, bg_geom, crop=True, nodata=0, filled=True
                )
                bg_band = bg_data[band - 1]
                bg_vals = bg_band[bg_band != 0]
            except Exception:
                bg_vals = np.array([])

            # Need at least 5 background pixels
            if len(bg_vals) < 5:
                continue

            bg_median = float(np.median(bg_vals))
            bg_mad = float(np.median(np.abs(bg_vals - bg_median)))
            sigma_robust = max(1.4826 * bg_mad, eps)

            I_bg_median_arr[i] = bg_median
            MAD_bg_arr[i]      = bg_mad

            # --- 6.3 Robust Z-scores ---
            Z_max_arr[i]    = (I_pole_max_arr[i]    - bg_median) / sigma_robust
            Z_median_arr[i] = (I_pole_median_arr[i] - bg_median) / sigma_robust
            Z_Q95_arr[i]    = (I_pole_Q95_arr[i]    - bg_median) / sigma_robust

            # --- 6.4 Detection decisions ---
            det_max_arr[i]    = int(Z_max_arr[i]    >= T_z)
            det_median_arr[i] = int(Z_median_arr[i] >= T_z)
            det_Q95_arr[i]    = int(Z_Q95_arr[i]    >= T_z)

    # ------------------------------------------------------------------
    # Build output GeoDataFrame
    # ------------------------------------------------------------------
    result_gdf = gdf.copy()
    result_gdf['I_pole_max']    = I_pole_max_arr
    result_gdf['I_pole_median'] = I_pole_median_arr
    result_gdf['I_pole_Q95']    = I_pole_Q95_arr
    result_gdf['I_bg_median']   = I_bg_median_arr
    result_gdf['MAD_bg']        = MAD_bg_arr
    result_gdf['Z_max']         = Z_max_arr
    result_gdf['Z_median']      = Z_median_arr
    result_gdf['Z_Q95']         = Z_Q95_arr
    result_gdf['det_max']       = det_max_arr
    result_gdf['det_median']    = det_median_arr
    result_gdf['det_Q95']       = det_Q95_arr
    result_gdf['pole_window_m'] = pole_window_m
    result_gdf['bg_inner_m']    = bg_inner_m
    result_gdf['bg_outer_m']    = bg_outer_m
    result_gdf['T_z']           = T_z
    result_gdf['method']        = 'Method B'

    return result_gdf