"""Stage 2: vectorize a probability raster into fracture centerlines (GeoPackage).

Ported from FraXtex's `fraxtex_tools/vectorize.py`, which is canonical for this
pipeline (see CLAUDE.md). Kept line-comparable with the plugin; the only changes are
the QGIS parameter plumbing and the output path, which the caller supplies.

Expects the inference convention: **high values mean fracture**.

All distances are given in pixels and multiplied by the raster resolution, so the
same parameter values behave the same on a georeferenced raster and on a plain PNG
(where the resolution is 1 and the units are pixels).
"""

import geopandas as gpd
import numpy as np
import rasterio
import shapely
import skimage.filters as skf
import skimage.morphology as skmp
from rasterio import features
from shapely.affinity import affine_transform

# Keep in the same order as the UI labels.
THRESH_METHODS = ["mean", "adaptive", "otsu", "manual"]
SKEL_METHODS = ["zhang", "lee", "medial_axis"]


def vectorize(
    raster_path,
    output_path,
    feedback,
    thresh_method="mean",
    manual_val=127,
    min_size=142,
    skel_method="zhang",
    buffer_px=3,
    min_branch=10,
    simplify_px=10,
):
    """Threshold -> skeletonize -> polygonize -> centerline -> GeoPackage.

    Returns `(output_path, feature_count)`.
    """
    with rasterio.open(raster_path) as src:
        arr = src.read(1)
        crs = src.crs
        transform = src.transform
        res_x = abs(transform.a)

    feedback.pushInfo(f"Shape: {arr.shape}, res: {res_x:.6f}")

    # 1. Threshold
    feedback.setProgress(5)
    if thresh_method == "otsu":
        threshold = skf.threshold_otsu(arr)
    elif thresh_method == "adaptive":
        threshold = skf.threshold_local(arr, block_size=51, method="gaussian")
    elif thresh_method == "mean":
        threshold = skf.threshold_mean(arr)
    else:
        threshold = manual_val

    # An adaptive threshold is a whole array, so summarise rather than print it.
    if np.ndim(threshold) == 0:
        feedback.pushInfo(f"Method: {thresh_method}, threshold: {threshold:.3f}")
    else:
        feedback.pushInfo(
            f"Method: {thresh_method}, per-pixel threshold "
            f"(mean {float(np.mean(threshold)):.3f})"
        )

    binary = (arr > threshold).astype(np.uint8)

    # 2. Remove small objects
    feedback.setProgress(15)
    binary = skmp.remove_small_objects(binary.astype(bool), min_size=min_size).astype(np.uint8)

    # 3. Skeletonize
    feedback.setProgress(25)
    if skel_method == "zhang":
        skel = skmp.skeletonize(binary, method="zhang")
    elif skel_method == "lee":
        skel = skmp.skeletonize(binary, method="lee")
    else:
        skel, _ = skmp.medial_axis(binary, return_distance=True)
    skel = skel.astype(np.uint8)

    # 4. Raster -> vector
    feedback.setProgress(40)
    polys = [shapely.geometry.shape(c) for c, v in features.shapes(skel, skel > 0.5)]
    if not polys:
        feedback.pushWarning(
            "No fractures detected. Try a lower threshold or a smaller minimum "
            "object size."
        )
        gpd.GeoDataFrame(geometry=[], crs=crs).to_file(output_path, driver="GPKG")
        return output_path, 0

    gdf = gpd.GeoDataFrame(geometry=polys, crs=crs)

    # 5. Affine transform to map coordinates
    aff = [transform.a, transform.b, transform.d, transform.e, transform.c, transform.f]
    gdf["geometry"] = gdf["geometry"].apply(lambda g: affine_transform(g, aff))

    # 6. Buffer -> dissolve -> explode
    feedback.setProgress(55)
    gdf["geometry"] = gdf.buffer(res_x * buffer_px, join_style="mitre")
    gdf = gdf.dissolve().explode()

    # 7. Centerline
    feedback.setProgress(70)
    try:
        import pygeoops
        gdf["geometry"] = pygeoops.centerline(
            gdf["geometry"],
            densify_distance=res_x,
            min_branch_length=res_x * min_branch,
            simplifytolerance=res_x,
        )
    except ImportError:
        feedback.pushInfo("pygeoops not available, skipping centerline.")

    # 8. Simplify
    feedback.setProgress(85)
    gdf = gdf.simplify(res_x * simplify_px, preserve_topology=False)
    gdf = gdf[~gdf.is_empty]
    gdf = gpd.GeoDataFrame(geometry=list(gdf), crs=crs)

    feedback.setProgress(95)
    gdf.to_file(output_path, driver="GPKG")
    feedback.pushInfo(f"Vectorize done: {len(gdf)} features.")
    feedback.setProgress(100)
    return output_path, len(gdf)