"""
Method A — Lumen-based anomaly scoring per lumen group.
========================================================

Enriches the input points with luminous-flux features and flags suspect
poles whose observed radiance is anomalously low compared to other poles
with the same lumen rating.

Required attributes
-------------------
  ARMATUR_GU  — armature power (e.g. "150 W")
  ARMATUR_AD  — armature count (0 / empty treated as 1)
  lumen / LUMEN / lumen_filled — optional, used if present

Pipeline
--------
  1. Raster extraction (same as RAST): inner circle → l_pole, annulus → l_bg
  2. Global background stats: mu_bg, sigma_bg, per-pole score s
  3. Lumen enrichment: estimate lumen for each pole from input / medians /
     reference table / fallback efficacy.
  4. Anomaly scoring per lumen group:
       anomaly_score = (nmf - group_median_nmf) / group_mad_nmf
     where nmf = img_radiance_near_minus_far
  5. Suspect flags: is_suspect (score < -2), is_strong_suspect (score < -3)

Produced attributes
-------------------
power_w, armatur_ad_used, total_powe, lm_per_armatur_used, lumen,
lumen_filled, lumen_source, rad_flux_w, up_flux_w, upward_ratio,
luminous_efficacy_lm_per_w, maintenance_factor, lumen_maintained,
l_pole, l_bg, mu_bg, sigma_bg, s, img_radiance_near_mean,
img_radiance_near_minus_far, group_label, group_median_nmf,
group_mad_nmf, anomaly_score, is_suspect, is_strong_suspect, method
"""

import math
import re
import json
from statistics import median as _median
from typing import Any, Dict, Iterable, List, Optional, Tuple

import numpy as np
import geopandas as gpd
import rasterio
from rasterio.mask import mask as rasterio_mask
from scipy.stats import median_abs_deviation

# ======================================================================
# Enrichment configuration
# ======================================================================

DEFAULT_LM_PER_ARMATUR: Dict[int, float] = {
    125: 11875.0,
    150: 14000.0,
    250: 25000.0,
    400: 38000.0,
}

DEFAULT_UPWARD_RATIO = 0.15
DEFAULT_LUMINOUS_EFFICACY = 683.0
DEFAULT_MAINTENANCE_FACTOR = 0.80
DEFAULT_FALLBACK_ETA = 100.0  # lm/W


# ======================================================================
# Enrichment helpers
# ======================================================================

def _safe_float(value: Any) -> Optional[float]:
    if value is None:
        return None
    if isinstance(value, bool):
        return float(value)
    if isinstance(value, (int, float)):
        if math.isnan(value) or math.isinf(value):
            return None
        return float(value)
    if isinstance(value, str):
        txt = value.strip().replace(",", ".")
        if txt == "":
            return None
        try:
            out = float(txt)
            if math.isnan(out) or math.isinf(out):
                return None
            return out
        except ValueError:
            return None
    return None


def _extract_power_w(text: Any) -> Optional[int]:
    if text is None:
        return None
    s = str(text)
    m = re.search(r"(\d+(?:\.\d+)?)", s)
    if not m:
        return None
    try:
        return int(round(float(m.group(1))))
    except ValueError:
        return None


def _normalize_armatur_count(value: Any) -> int:
    x = _safe_float(value)
    if x is None or x <= 0:
        return 1
    return int(round(x))


def _find_existing_lumen_key(properties: Dict[str, Any]) -> Optional[str]:
    candidates = [
        "lumen", "LUMEN", "lumen_filled", "LUMEN_FILLED",
        "lumen_recalc_used", "LUMEN_RECALC_USED",
    ]
    for key in candidates:
        if key in properties:
            return key
    return None


def _list_medians(values: Iterable[float]) -> Optional[float]:
    vals = [float(v) for v in values
            if v is not None and not math.isnan(v) and math.isfinite(v)]
    if not vals:
        return None
    return float(_median(vals))


def _build_training_stats(
    features: List[Dict[str, Any]]
) -> Tuple[Dict[Tuple[int, int], float], Dict[int, float]]:
    combo_values: Dict[Tuple[int, int], List[float]] = {}
    power_lm_per_arm_values: Dict[int, List[float]] = {}

    for feat in features:
        props = feat.get("properties", {}) or {}
        lum_key = _find_existing_lumen_key(props)
        if lum_key is None:
            continue
        lum = _safe_float(props.get(lum_key))
        power_w = _extract_power_w(props.get("ARMATUR_GU"))
        arm_count = _normalize_armatur_count(props.get("ARMATUR_AD"))
        if lum is None or lum <= 0 or power_w is None:
            continue
        combo_values.setdefault((power_w, arm_count), []).append(lum)
        power_lm_per_arm_values.setdefault(power_w, []).append(lum / arm_count)

    combo_median = {k: _list_medians(v) for k, v in combo_values.items()}
    power_median = {k: _list_medians(v) for k, v in power_lm_per_arm_values.items()}
    return combo_median, power_median


def _enrich_feature(
    feature: Dict[str, Any],
    combo_median: Dict[Tuple[int, int], float],
    power_median: Dict[int, float],
    lm_per_armatur_ref: Dict[int, float],
    upward_ratio: float,
    luminous_efficacy: float,
    maintenance_factor: float,
    fallback_eta: float,
) -> Dict[str, Any]:
    feat = {
        "type": feature.get("type", "Feature"),
        "geometry": feature.get("geometry"),
        "properties": dict(feature.get("properties", {}) or {}),
    }
    p = feat["properties"]

    power_w = _extract_power_w(p.get("ARMATUR_GU"))
    arm_count_original = _safe_float(p.get("ARMATUR_AD"))
    arm_count = _normalize_armatur_count(p.get("ARMATUR_AD"))
    total_powe = float(power_w * arm_count) if power_w is not None else None

    lum_key = _find_existing_lumen_key(p)
    input_lumen = _safe_float(p.get(lum_key)) if lum_key is not None else None

    lm_per_armatur_used = None
    lumen = None
    lumen_source = None

    if input_lumen is not None and input_lumen > 0:
        lumen = float(input_lumen)
        lm_per_armatur_used = lumen / arm_count
        lumen_source = f"input:{lum_key}"
    elif power_w is not None and (power_w, arm_count) in combo_median \
            and combo_median[(power_w, arm_count)] is not None:
        lumen = float(combo_median[(power_w, arm_count)])
        lm_per_armatur_used = lumen / arm_count
        lumen_source = "median_same_power_same_count"
    elif power_w is not None and power_w in power_median \
            and power_median[power_w] is not None:
        lm_per_armatur_used = float(power_median[power_w])
        lumen = lm_per_armatur_used * arm_count
        lumen_source = "median_same_power_lm_per_arm"
    elif power_w is not None and power_w in lm_per_armatur_ref:
        lm_per_armatur_used = float(lm_per_armatur_ref[power_w])
        lumen = lm_per_armatur_used * arm_count
        lumen_source = "reference_table"
    elif power_w is not None:
        lm_per_armatur_used = float(fallback_eta * power_w)
        lumen = lm_per_armatur_used * arm_count
        lumen_source = "fallback_eta"

    rad_flux_w = None
    up_flux_w = None
    lumen_maintained = None
    if lumen is not None:
        rad_flux_w = float(lumen / luminous_efficacy)
        up_flux_w = float(rad_flux_w * upward_ratio)
        lumen_maintained = float(lumen * maintenance_factor)

    p["power_w"] = power_w
    p["armatur_ad_original"] = arm_count_original
    p["armatur_ad_used"] = arm_count
    p["total_powe"] = total_powe
    p["lm_per_armatur_used"] = lm_per_armatur_used
    p["lumen"] = lumen
    p["lumen_filled"] = lumen
    p["lumen_source"] = lumen_source
    p["rad_flux_w"] = rad_flux_w
    p["up_flux_w"] = up_flux_w
    p["upward_ratio"] = upward_ratio
    p["luminous_efficacy_lm_per_w"] = luminous_efficacy
    p["maintenance_factor"] = maintenance_factor
    p["lumen_maintained"] = lumen_maintained

    return feat


def enrich_geojson(
    data: Dict[str, Any],
    lm_per_armatur_ref: Dict[int, float],
    upward_ratio: float,
    luminous_efficacy: float,
    maintenance_factor: float,
    fallback_eta: float,
) -> Dict[str, Any]:
    features = data.get("features", [])
    combo_median, power_median = _build_training_stats(features)
    new_features = [
        _enrich_feature(
            feature=f,
            combo_median=combo_median,
            power_median=power_median,
            lm_per_armatur_ref=lm_per_armatur_ref,
            upward_ratio=upward_ratio,
            luminous_efficacy=luminous_efficacy,
            maintenance_factor=maintenance_factor,
            fallback_eta=fallback_eta,
        )
        for f in features
    ]
    out = dict(data)
    out["features"] = new_features
    out.setdefault("name", "enriched_poles")
    return out


# ======================================================================
# Main entry point
# ======================================================================

def run(tiff_path, vector_path, band=1, threshold=3.0,
        inner_radius=5, outer_radius=15):
    """
    Run the Method A analysis.

    Parameters
    ----------
    tiff_path : str
        Path to the input GeoTIFF.
    vector_path : str
        Path to the input vector file (pole points).
    band : int
        Raster band to use (1, 2 or 3).
    threshold : float
        Observability threshold (used for the ``s`` score, kept for
        compatibility; Method A primarily uses anomaly_score).
    inner_radius : int
        Inner radius of the annulus in metres.
    outer_radius : int
        Outer radius of the annulus in metres.

    Returns
    -------
    geopandas.GeoDataFrame
        Enriched poles with anomaly-scoring attributes.
    """
    # ------------------------------------------------------------------
    # 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
            )

    # ------------------------------------------------------------------
    # Raster extraction (same as RAST)
    # ------------------------------------------------------------------
    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)

        l_pole_list = []
        l_bg_list = []
        near_mean_list = []
        near_minus_far_list = []

        for _idx, row in gdf_proj.iterrows():
            point = row.geometry
            inner_circle = point.buffer(inner_radius)
            outer_circle = point.buffer(outer_radius)
            annulus = outer_circle.difference(inner_circle)

            inner_gdf = gpd.GeoDataFrame(
                geometry=[inner_circle], crs=utm_crs).to_crs(raster_crs)
            annulus_gdf = gpd.GeoDataFrame(
                geometry=[annulus], crs=utm_crs).to_crs(raster_crs)

            inner_geom = [inner_gdf.geometry.iloc[0]]
            annulus_geom = [annulus_gdf.geometry.iloc[0]]

            try:
                inner_data, _ = rasterio_mask(
                    src, inner_geom, crop=True, nodata=0, filled=True)
                band_data_inner = inner_data[band - 1]
                valid_inner = band_data_inner[band_data_inner != 0]
                l_pole = float(np.max(valid_inner)) if len(valid_inner) > 0 else 0.0
                near_mean = float(np.mean(valid_inner)) if len(valid_inner) > 0 else 0.0
            except Exception:
                l_pole = 0.0
                near_mean = 0.0

            try:
                annulus_data, _ = rasterio_mask(
                    src, annulus_geom, crop=True, nodata=0, filled=True)
                band_data_annulus = annulus_data[band - 1]
                valid_annulus = band_data_annulus[band_data_annulus != 0]
                l_bg = float(np.max(valid_annulus)) if len(valid_annulus) > 0 else 0.0
                far_mean = float(np.mean(valid_annulus)) if len(valid_annulus) > 0 else 0.0
            except Exception:
                l_bg = 0.0
                far_mean = 0.0

            l_pole_list.append(l_pole)
            l_bg_list.append(l_bg)
            near_mean_list.append(near_mean)
            near_minus_far_list.append(near_mean - far_mean)

    # ------------------------------------------------------------------
    # Global statistics
    # ------------------------------------------------------------------
    l_bg_array = np.array(l_bg_list)
    mu_bg = float(np.median(l_bg_array))
    sigma_bg = float(median_abs_deviation(l_bg_array))

    s_list = []
    for l_pole in l_pole_list:
        s = (l_pole - mu_bg) / sigma_bg if sigma_bg > 0 else 0.0
        s_list.append(s)

    # ------------------------------------------------------------------
    # Step 1: Enrich the input GeoJSON with lumen-based features
    # ------------------------------------------------------------------
    raw_geojson = json.loads(gdf.to_json())
    enriched_geojson = enrich_geojson(
        data=raw_geojson,
        lm_per_armatur_ref=DEFAULT_LM_PER_ARMATUR,
        upward_ratio=DEFAULT_UPWARD_RATIO,
        luminous_efficacy=DEFAULT_LUMINOUS_EFFICACY,
        maintenance_factor=DEFAULT_MAINTENANCE_FACTOR,
        fallback_eta=DEFAULT_FALLBACK_ETA,
    )

    enriched_gdf = gpd.GeoDataFrame.from_features(
        enriched_geojson['features'], crs=gdf.crs)

    # Carry over the radiance columns
    enriched_gdf['l_pole'] = l_pole_list
    enriched_gdf['l_bg'] = l_bg_list
    enriched_gdf['mu_bg'] = mu_bg
    enriched_gdf['sigma_bg'] = sigma_bg
    enriched_gdf['s'] = s_list
    enriched_gdf['img_radiance_near_mean'] = near_mean_list
    enriched_gdf['img_radiance_near_minus_far'] = near_minus_far_list

    # ------------------------------------------------------------------
    # Step 2: Anomaly scoring per lumen group
    # ------------------------------------------------------------------
    lumen_filled = enriched_gdf['lumen_filled'].values
    nmf = enriched_gdf['img_radiance_near_minus_far'].values

    valid_mask = np.array([
        (lf is not None and not math.isnan(float(lf)) if lf is not None else False)
        and not math.isnan(float(n))
        for lf, n in zip(lumen_filled, nmf)
    ])

    group_median_arr = np.full(len(enriched_gdf), np.nan)
    group_mad_arr = np.full(len(enriched_gdf), np.nan)
    anomaly_score_arr = np.full(len(enriched_gdf), np.nan)
    group_label_arr = np.full(len(enriched_gdf), '', dtype=object)

    min_group_size = 5
    valid_indices = np.where(valid_mask)[0]

    if len(valid_indices) > 0:
        valid_lumen = np.array([float(lumen_filled[i]) for i in valid_indices])
        valid_nmf = np.array([float(nmf[i]) for i in valid_indices])
        unique_lumens = np.unique(valid_lumen)

        for lum_val in unique_lumens:
            group_idx_in_valid = np.where(valid_lumen == lum_val)[0]
            if len(group_idx_in_valid) < min_group_size:
                continue

            global_indices = valid_indices[group_idx_in_valid]
            group_nmf = valid_nmf[group_idx_in_valid]

            med_nmf = float(np.median(group_nmf))
            mad_nmf = float(median_abs_deviation(group_nmf))
            label = str(round(lum_val))

            for gi, nmf_val in zip(global_indices, group_nmf):
                group_median_arr[gi] = med_nmf
                group_mad_arr[gi] = mad_nmf
                group_label_arr[gi] = label
                if mad_nmf < 1e-10:
                    anomaly_score_arr[gi] = nmf_val - med_nmf
                else:
                    anomaly_score_arr[gi] = (nmf_val - med_nmf) / mad_nmf

    enriched_gdf['group_label'] = group_label_arr
    enriched_gdf['group_median_nmf'] = group_median_arr
    enriched_gdf['group_mad_nmf'] = group_mad_arr
    enriched_gdf['anomaly_score'] = anomaly_score_arr
    enriched_gdf['is_suspect'] = anomaly_score_arr < -2
    enriched_gdf['is_strong_suspect'] = anomaly_score_arr < -3
    enriched_gdf['method'] = 'Method A'

    return enriched_gdf