"""
Change analysis module: grid-based and transect-based methods
for quantifying riverbank erosion and accretion, with regime-
adaptive features including riparian buffer masking, MMU filter,
braiding index, connected component labelling, regime classification,
and grid-transect consistency index.
"""
import os
import numpy as np
import geopandas as gpd
from shapely.geometry import (
    box, Point, Polygon, MultiPolygon, GeometryCollection,
)
from shapely.ops import unary_union


# ======================================================================
# UTILITY
# ======================================================================

def extract_polygons(geom):
    """Extract only polygon/multipolygon parts from a geometry."""
    if geom is None or geom.is_empty:
        return None
    if isinstance(geom, (Polygon, MultiPolygon)):
        return geom
    if isinstance(geom, GeometryCollection):
        polys = [
            g for g in geom.geoms
            if isinstance(g, (Polygon, MultiPolygon))
        ]
        if not polys:
            return None
        if len(polys) == 1:
            return polys[0]
        parts = []
        for g in polys:
            if isinstance(g, MultiPolygon):
                parts.extend(g.geoms)
            else:
                parts.append(g)
        return MultiPolygon(parts)
    return None


# ======================================================================
# REGIME-ADAPTIVE PRE-PROCESSING
# ======================================================================

def compute_riparian_corridor(y1_gdf, y2_gdf, buffer_m):
    """
    Compute the riparian corridor as the buffered union of both
    epochs' waterbody polygons.
    """
    try:
        union = unary_union(
            list(y1_gdf.geometry.values)
            + list(y2_gdf.geometry.values)
        )
        if union is None or union.is_empty:
            return None
        corridor = union.buffer(buffer_m)
        return extract_polygons(corridor)
    except Exception:
        return None


def connected_component_label(gdf, min_area_ha=5):
    """
    Label disconnected waterbody clusters with a channel_id.
    Clusters below min_area_ha are excluded (non-riverine).
    """
    min_area_m2 = min_area_ha * 10000
    dissolved = gdf.geometry.unary_union
    if dissolved.is_empty:
        return gdf

    parts = []
    if dissolved.geom_type == 'MultiPolygon':
        parts = list(dissolved.geoms)
    else:
        parts = [dissolved]

    labeled = []
    for cid, geom in enumerate(parts, 1):
        area = geom.area
        if area >= min_area_m2:
            labeled.append({'channel_id': cid, 'geometry': geom})

    if not labeled:
        return gdf
    return gpd.GeoDataFrame(labeled, crs=gdf.crs)


def apply_mmu_filter(per_grid_gdf, min_cells=3, out_dir=None):
    """
    Apply Minimum Mapping Unit (MMU) filter: dissolve contiguous
    eroded/accreted cells and remove patches below threshold.
    """
    if min_cells < 1:
        return per_grid_gdf

    result = per_grid_gdf.copy()
    cell_area_m2 = (
        result.geometry.area.iloc[0]
        if len(result) > 0 else 0
    )
    mmu_area = min_cells * cell_area_m2

    for field, label in [('eros_m2', 'Erosion'), ('accr_m2', 'Accretion')]:
        if field not in result.columns:
            continue
        cells = result[result[field] > 0].copy()
        if cells.empty:
            continue
        dissolved = cells.geometry.unary_union
        if dissolved.is_empty:
            continue
        parts = []
        if dissolved.geom_type == 'MultiPolygon':
            parts = list(dissolved.geoms)
        else:
            parts = [dissolved]

        keep_ids = set()
        for patch in parts:
            if patch.area >= mmu_area:
                keep = cells[cells.geometry.intersects(patch)]
                keep_ids.update(keep.index.tolist())

        mask = result.index.isin(keep_ids)
        result.loc[~mask & (result[field] > 0), field] = 0.0
        result.loc[~mask & (result[field] > 0), field.replace('_m2', '_km2')] = 0.0

    if out_dir:
        result.to_file(os.path.join(out_dir, 'PerGrid_MMU_Filtered.shp'))

    return result


# ======================================================================
# FISHNET
# ======================================================================

def create_fishnet(bounds, cell_m, crs_authid):
    """Create a regular grid (fishnet) within given bounds."""
    minx, miny, maxx, maxy = [float(v) for v in bounds]
    xs = np.arange(minx, maxx, cell_m)
    ys = np.arange(miny, maxy, cell_m)
    polys = [
        box(x, y, x + cell_m, y + cell_m)
        for x in xs for y in ys
    ]
    return gpd.GeoDataFrame({'geometry': polys}, crs=crs_authid)


# ======================================================================
# PER-CELL CHANGE
# ======================================================================

def compute_per_grid_change(grid_gdf, y1_gdf, y2_gdf, out_dir,
                            crs_authid=None):
    """
    For each grid cell, compute erosion, accretion, and unchanged area.
    """
    if crs_authid:
        grid_gdf = grid_gdf.to_crs(crs_authid)
        y1_gdf = y1_gdf.to_crs(crs_authid)
        y2_gdf = y2_gdf.to_crs(crs_authid)

    results = []
    erosion_rows = []
    accretion_rows = []
    unchanged_rows = []

    for rec in grid_gdf[['grid_id', 'geometry']].itertuples(index=False):
        grid_id, g = rec
        y1_clip = y1_gdf[y1_gdf.intersects(g)]
        y2_clip = y2_gdf[y2_gdf.intersects(g)]

        if y1_clip.empty and y2_clip.empty:
            results.append((grid_id, 0.0, 0.0, 0.0))
            continue

        W1 = None
        W2 = None
        if not y1_clip.empty:
            inter = y1_clip.intersection(g).unary_union
            W1 = extract_polygons(inter)
        if not y2_clip.empty:
            inter = y2_clip.intersection(g).unary_union
            W2 = extract_polygons(inter)

        if (W1 is None or W1.is_empty) and \
           (W2 is None or W2.is_empty):
            results.append((grid_id, 0.0, 0.0, 0.0))
            continue

        unchanged_area = 0.0
        erosion_area = 0.0
        accretion_area = 0.0

        if W1 is not None and W2 is not None:
            inter = extract_polygons(W1.intersection(W2))
            if inter is not None:
                unchanged_area = inter.area
                unchanged_rows.append((grid_id, inter))

        if W2 is not None:
            diff = W2.difference(W1) if W1 is not None else W2
            er = extract_polygons(diff)
            if er is not None:
                erosion_area = er.area
                erosion_rows.append((grid_id, er))

        if W1 is not None:
            diff = W1.difference(W2) if W2 is not None else W1
            ac = extract_polygons(diff)
            if ac is not None:
                accretion_area = ac.area
                accretion_rows.append((grid_id, ac))

        results.append(
            (grid_id, erosion_area, accretion_area, unchanged_area)
        )

    per = gpd.GeoDataFrame(
        results,
        columns=['grid_id', 'eros_m2', 'accr_m2', 'unc_m2']
    )
    out = grid_gdf.merge(per, on='grid_id', how='left').fillna(0.0)
    out['eros_km2'] = out['eros_m2'] / 1e6
    out['accr_km2'] = out['accr_m2'] / 1e6
    out.to_file(os.path.join(out_dir, 'PerGrid_ErosionAccretion.shp'))

    def save_if_any(rows, fname, field, crs):
        if not rows:
            return
        gdf = gpd.GeoDataFrame(
            rows, columns=['grid_id', 'geometry'], crs=crs
        )
        gdf[field] = gdf.geometry.area
        gdf.to_file(os.path.join(out_dir, fname))

    save_if_any(erosion_rows, 'Erosion_Grids.shp',
                'eros_m2', out.crs)
    save_if_any(accretion_rows, 'Accretion_Grids.shp',
                'accr_m2', out.crs)
    save_if_any(unchanged_rows, 'Unchanged_Grids.shp',
                'unc_m2', out.crs)

    return out


# ======================================================================
# STRETCH AGGREGATION
# ======================================================================

def aggregate_by_stretch(per_grid_gdf, stretch_m, cell_m,
                         out_dir, top_n=15):
    """
    Aggregate per-grid results into river stretches and report Top-N.
    """
    grids_per_stretch = max(1, int(round(stretch_m / cell_m)))
    df = per_grid_gdf.copy()
    df['stretch_id'] = (
        (df['grid_id'] - 1) // grids_per_stretch
    ).astype(int)
    stretch = df.dissolve(
        by='stretch_id',
        aggfunc={'eros_km2': 'sum', 'accr_km2': 'sum'}
    )
    stretch.to_file(os.path.join(out_dir, 'Stretch_Summary.shp'))

    top_e = stretch.nlargest(top_n, 'eros_km2').reset_index()
    top_a = stretch.nlargest(top_n, 'accr_km2').reset_index()
    top_e.to_file(
        os.path.join(out_dir, f'Top{top_n}_Erosion_Stretches.shp')
    )
    top_a.to_file(
        os.path.join(out_dir, f'Top{top_n}_Accretion_Stretches.shp')
    )

    top_e[['stretch_id', 'eros_km2']].to_csv(
        os.path.join(out_dir, f'Top{top_n}_Erosion_Stretches.csv'),
        index=False
    )
    top_a[['stretch_id', 'accr_km2']].to_csv(
        os.path.join(out_dir, f'Top{top_n}_Accretion_Stretches.csv'),
        index=False
    )

    return stretch


# ======================================================================
# BRAIDING INDEX
# ======================================================================

def compute_braiding_index(stretch_gdf, y1_gdf, y2_gdf, stretch_m):
    """
    Compute braiding index (BI) per stretch from input polygons.
    BI = Total_wetted_perimeter / (2 * reach_length)
    """
    result = stretch_gdf.copy()
    bis = []
    reliability = []

    for idx in result.index:
        stretch_geom = result.loc[idx, 'geometry']
        if stretch_geom is None or stretch_geom.is_empty:
            bis.append(0)
            reliability.append('NODATA')
            continue

        # Overlap Y1 and Y2 water with this stretch
        y1_inter = y1_gdf[y1_gdf.intersects(stretch_geom)]
        y2_inter = y2_gdf[y2_gdf.intersects(stretch_geom)]

        combined = gpd.GeoDataFrame(
            {'geometry': list(y1_inter.geometry)
             + list(y2_inter.geometry)},
            crs=y1_gdf.crs
        )

        if combined.empty:
            bis.append(0)
            reliability.append('NODATA')
            continue

        try:
            merged = unary_union(combined.geometry.tolist())
        except Exception:
            bis.append(0)
            reliability.append('NODATA')
            continue

        if merged is None or merged.is_empty:
            bis.append(0)
            reliability.append('NODATA')
            continue

        total_perimeter = merged.length
        reach_length = stretch_geom.length if stretch_geom.length > 0 else stretch_m

        bi = total_perimeter / (2.0 * reach_length) if reach_length > 0 else 0

        if bi < 2:
            rel = 'HIGH'
        elif bi < 3:
            rel = 'MODERATE'
        else:
            rel = 'LOW'

        bis.append(round(bi, 2))
        reliability.append(rel)

    result['braiding_idx'] = bis
    result['cl_rel'] = reliability  # centerline reliability
    return result


# ======================================================================
# REGIME CLASSIFICATION
# ======================================================================

def classify_stretch_regime(stretch_gdf, y1_gdf, y2_gdf,
                            cell_m=30, stretch_m=3000,
                            user_river_type=None):
    """
    Classify each stretch as BRAIDED, MEANDERING, or TRANSITIONAL
    based on braiding index, width:length ratio, and component count.

    Uses the user-selected river morphological type as a prior to
    prevent misclassification of single-thread rivers (e.g., Narmada)
    whose water-mask fragments can artificially inflate component count.
    """
    result = stretch_gdf.copy()
    regimes = []

    # Thresholds per river type prior
    if user_river_type and user_river_type.lower() == "meandering":
        bi_braided = 10.0       # only override if extremely fragmented
        bi_meander = 5.0
    else:
        bi_braided = 3.0
        bi_meander = 2.0

    for idx in result.index:
        stretch_geom = result.loc[idx, 'geometry']
        if stretch_geom is None or stretch_geom.is_empty:
            regimes.append('NODATA')
            continue

        y1_inter = y1_gdf[y1_gdf.intersects(stretch_geom)]
        y2_inter = y2_gdf[y2_gdf.intersects(stretch_geom)]
        combined = gpd.GeoDataFrame(
            {'geometry': list(y1_inter.geometry)
             + list(y2_inter.geometry)},
            crs=y1_gdf.crs
        )

        if combined.empty:
            regimes.append('NODATA')
            continue

        try:
            merged = unary_union(combined.geometry.tolist())
        except Exception:
            regimes.append('NODATA')
            continue

        if merged is None or merged.is_empty:
            regimes.append('NODATA')
            continue

        n_components = 0
        if merged.geom_type == 'MultiPolygon':
            n_components = len(merged.geoms)
        elif merged.geom_type == 'Polygon':
            n_components = 1

        total_perimeter = merged.length
        reach_length = stretch_geom.length if stretch_geom.length > 0 else stretch_m
        bi = total_perimeter / (2.0 * reach_length) if reach_length > 0 else 0

        bounds = merged.bounds
        width = max(bounds[2] - bounds[0], bounds[3] - bounds[1])
        wl_ratio = width / reach_length if reach_length > 0 else 0

        if bi >= bi_braided and n_components > 3:
            regimes.append('BRAIDED')
        elif wl_ratio > 0.3 and n_components >= 2:
            regimes.append('BRAIDED')
        elif bi < bi_meander and n_components <= 2 and wl_ratio < 0.15:
            regimes.append('MEANDERING')
        else:
            regimes.append('TRANSITIONAL')

    result['regime'] = regimes
    return result[['regime', 'geometry']]


# ======================================================================
# CONSISTENCY INDEX
# ======================================================================

def compute_consistency_index(
    per_grid_gdf, transect_rates_gdf,
    transects_gdf, stretch_m, cell_m,
    out_dir=None,
):
    """
    Compute Grid-Transect Consistency Index per stretch.
    CI = sum(transect_erosion_area) / sum(grid_erosion_area)
    """
    if per_grid_gdf is None or per_grid_gdf.empty:
        return None
    if transect_rates_gdf is None or transect_rates_gdf.empty:
        return None
    if transects_gdf is None or transects_gdf.empty:
        return None

    import pandas as pd

    # Build transect erosion area
    transect_area_total = 0.0
    for _, row in transect_rates_gdf.iterrows():
        rate = row.get('avg_rate_yr', 0)
        rate = float(rate) if not isinstance(rate, (int, float)) else rate
        transect_area_total += rate * stretch_m

    # Build grid erosion area per stretch
    grids_per_stretch = max(1, int(round(stretch_m / cell_m)))
    df = per_grid_gdf.copy()
    df['stretch_id'] = ((df['grid_id'] - 1) // grids_per_stretch).astype(int)
    grid_area_total = float(df['eros_m2'].sum())

    _ = transect_area_total / (grid_area_total + 1e-10)

    rows = []
    for sid in df['stretch_id'].unique():
        g_area = float(df[df['stretch_id'] == sid]['eros_m2'].sum())
        t_area = transect_area_total / (df['stretch_id'].nunique() + 1e-10)
        s_ci = t_area / (g_area + 1e-10)

        if s_ci >= 0.85 and s_ci <= 1.15:
            conf = 'HIGH'
        elif s_ci >= 0.70 and s_ci <= 1.30:
            conf = 'MODERATE'
        else:
            conf = 'LOW'

        rows.append({
            'stretch_id': sid,
            'grid_erosion_m2': round(g_area, 2),
            'transect_erosion_m2': round(t_area, 2),
            'consistency_idx': round(s_ci, 3),
            'confidence': conf,
        })

    cdf = pd.DataFrame(rows)

    if out_dir:
        cdf.to_csv(
            os.path.join(out_dir, 'GridTransect_Consistency.csv'),
            index=False
        )

    return cdf


# ======================================================================
# CENTERLINE (Voronoi-enhanced)
# ======================================================================

def extract_centerline(waterbody_gdf, simplify_tolerance=30,
                       cell_size=30):
    """
    Extract river centerline using Voronoi skeletonisation with
    convex-hull fallback.
    """
    try:
        merged = unary_union(waterbody_gdf.geometry.tolist())
    except Exception:
        return None

    if merged is None or merged.is_empty:
        return None

    if isinstance(merged, MultiPolygon):
        # Buffer-dissolve to bridge sub-resolution gaps (narrow
        # segments, sandbars) that artificially fragment the
        # waterbody. Buffer = 1 cell_size: gaps narrower than
        # the analysis grid are likely artifacts, while wider
        # separations preserve genuine multi-channel morphology.
        merged = merged.buffer(cell_size).buffer(-cell_size)

    # Try Voronoi skeleton first
    try:
        from shapely.geometry import (
            LineString, MultiLineString, MultiPoint
        )
        from shapely.ops import voronoi_diagram

        coords = list(merged.exterior.coords)
        if len(coords) < 4:
            return None

        env = merged.envelope
        voronoi = voronoi_diagram(
            MultiPoint(coords), envelope=env
        )
        edges = []
        if hasattr(voronoi, 'geoms'):
            for g in voronoi.geoms:
                if g.within(merged) and isinstance(g, LineString):
                    edges.append(g)

        if not edges:
            raise ValueError("No interior edges found")

        merged_edges = unary_union(edges)
        if merged_edges.is_empty:
            raise ValueError("Empty edge union")

        if isinstance(merged_edges, MultiLineString):
            # Take the longest edge
            edge_list = sorted(merged_edges.geoms, key=lambda x: x.length, reverse=True)
            centerline = edge_list[0]
        else:
            centerline = merged_edges

        if simplify_tolerance > 0:
            centerline = centerline.simplify(simplify_tolerance)

        return gpd.GeoDataFrame(
            {'geometry': [centerline]}, crs=waterbody_gdf.crs
        )

    except Exception:
        # Fallback to convex hull method
        return _extract_centerline_fallback(merged, waterbody_gdf.crs,
                                            simplify_tolerance)


def _extract_centerline_fallback(merged, crs, simplify_tolerance=30):
    """Convex hull midpoint fallback for straight reaches."""
    try:
        from shapely.geometry import LineString
        bounds = merged.bounds
        width = max(bounds[2] - bounds[0], bounds[3] - bounds[1])
        spacing = width / 50.0

        hull = merged.convex_hull
        hull_pts = list(hull.exterior.coords)
        hull_line = LineString(hull_pts)
        centerline = []
        for i in range(int(hull_line.length / spacing) + 1):
            pt = hull_line.interpolate(i * spacing)
            centerline.append(pt)

        if len(centerline) < 2:
            centerline_pts = [
                (merged.centroid.x - width * 0.5, merged.centroid.y),
                (merged.centroid.x + width * 0.5, merged.centroid.y),
            ]
        else:
            centerline_pts = [(pt.x, pt.y) for pt in centerline]

        cl = LineString(centerline_pts)
        if simplify_tolerance > 0:
            cl = cl.simplify(simplify_tolerance)

        return gpd.GeoDataFrame(
            {'geometry': [cl]}, crs=crs
        )
    except Exception:
        return None


# ======================================================================
# TRANSECT GENERATION
# ======================================================================

def generate_transects(centerline_gdf, spacing=100,
                       transect_length=2000):
    """
    Generate perpendicular transects along a river centerline.
    """
    if centerline_gdf is None or centerline_gdf.empty:
        return None

    from shapely.geometry import LineString
    import math

    line = centerline_gdf.geometry.iloc[0]
    if line is None or line.is_empty:
        return None

    transects = []
    stations = []
    total_length = line.length
    num_transects = max(2, int(total_length / spacing))

    for i in range(num_transects + 1):
        dist = i * spacing
        if dist > total_length:
            break

        pt = line.interpolate(dist)
        angle = None

        try:
            idx = line.project(pt, normalized=False)
            segment_start = max(0, int(idx) - 1)
            segment_end = min(len(line.coords) - 1, segment_start + 1)
            if segment_start == segment_end:
                segment_end = segment_start
            p1 = line.coords[segment_start]
            p2 = line.coords[segment_end]
            dx = p2[0] - p1[0]
            dy = p2[1] - p1[1]
            angle = math.atan2(dy, dx)
        except Exception:
            continue

        if angle is None:
            continue

        perp_angle = angle + math.pi / 2
        half = transect_length / 2.0
        x1 = pt.x + half * math.cos(perp_angle)
        y1 = pt.y + half * math.sin(perp_angle)
        x2 = pt.x - half * math.cos(perp_angle)
        y2 = pt.y - half * math.sin(perp_angle)

        transects.append(LineString([(x1, y1), (x2, y2)]))
        stations.append(dist)

    if not transects:
        return None

    result = gpd.GeoDataFrame(
        {'station': stations, 'geometry': transects},
        crs=centerline_gdf.crs
    )
    return result


# ======================================================================
# TRANSECT RATES
# ======================================================================

def compute_transect_rates(transects_gdf, y1_gdf, y2_gdf,
                           years_interval, out_dir=None):
    """
    For each transect, compute bank positions and retreat rate.
    """
    result_rows = []

    for idx, row in transects_gdf.iterrows():
        transect = row.geometry
        station = row.station

        if transect is None or transect.is_empty:
            continue

        y1_inter = y1_gdf[y1_gdf.intersects(transect)]
        y2_inter = y2_gdf[y2_gdf.intersects(transect)]

        y1_crossings = _intersection_points(transect, y1_inter)
        y2_crossings = _intersection_points(transect, y2_inter)

        left_y1 = _left_bank(y1_crossings, transect)
        right_y1 = _right_bank(y1_crossings, transect)
        left_y2 = _left_bank(y2_crossings, transect)
        right_y2 = _right_bank(y2_crossings, transect)

        width_y1 = abs(left_y1 - right_y1) if left_y1 and right_y1 else 0
        width_y2 = abs(left_y2 - right_y2) if left_y2 and right_y2 else 0

        left_change = abs(left_y2 - left_y1) \
            if left_y1 and left_y2 else 0
        right_change = abs(right_y2 - right_y1) \
            if right_y1 and right_y2 else 0
        avg_change = (left_change + right_change) / 2.0

        left_rate = left_change / years_interval if years_interval else 0
        right_rate = right_change / years_interval if years_interval else 0
        avg_rate = avg_change / years_interval if years_interval else 0

        result_rows.append({
            'station': station,
            'width_y1': width_y1,
            'width_y2': width_y2,
            'left_change': left_change,
            'right_change': right_change,
            'avg_change': avg_change,
            'left_rate_yr': left_rate,
            'right_rate_yr': right_rate,
            'avg_rate_yr': avg_rate,
        })

    if not result_rows:
        return None

    result = gpd.GeoDataFrame(
        result_rows,
        geometry=[transects_gdf.geometry.iloc[0]] * len(result_rows)
        if len(result_rows) == 1
        else transects_gdf.geometry.iloc[:len(result_rows)],
    )

    if out_dir:
        result.to_file(os.path.join(out_dir, 'Transect_Rates.shp'))
        result[['station', 'avg_rate_yr']].to_csv(
            os.path.join(out_dir, 'Transect_Rates.csv'),
            index=False
        )

    return result


def _intersection_points(transect, water_gdf):
    points = []
    for _, row in water_gdf.iterrows():
        poly = row.geometry
        if poly is None or poly.is_empty:
            continue
        inter = transect.intersection(poly)
        if inter.is_empty:
            continue
        if isinstance(inter, Point):
            points.append(inter)
        elif hasattr(inter, 'geoms'):
            for g in inter.geoms:
                if isinstance(g, Point):
                    points.append(g)
    return points


def _left_bank(points, transect):
    """Distance from transect start to left bank intersection."""
    if not points:
        return None
    try:
        distances = [transect.project(p) for p in points]
        return min(distances)
    except Exception:
        return None


def _right_bank(points, transect):
    """Distance from transect start to right bank intersection."""
    if not points:
        return None
    try:
        distances = [transect.project(p) for p in points]
        return max(distances)
    except Exception:
        return None
