import os
import warnings

import matplotlib
import numpy as np
import rasterio
from sklearn.linear_model import LinearRegression, RANSACRegressor, HuberRegressor
from sklearn.preprocessing import PolynomialFeatures

try:

    from qgis.PyQt.QtCore import QVariant
    from qgis.core import (
        QgsCoordinateTransform,
        QgsCoordinateReferenceSystem,
        QgsFeature,
        QgsField,
        QgsGeometry,
        QgsProcessingException,
        QgsProject,
        QgsRasterLayer,
        QgsVectorLayer,
        QgsVectorFileWriter,
        QgsWkbTypes,
    )
except ImportError:
    QVariant = None
    QgsCoordinateTransform = None
    QgsCoordinateReferenceSystem = None
    QgsFeature = None
    QgsField = None
    QgsGeometry = None
    QgsProcessingException = Exception
    QgsProject = None
    QgsRasterLayer = None
    QgsVectorLayer = None
    QgsVectorFileWriter = None
    QgsWkbTypes = None


try:
    from ...infrastructure.vector_io import resolve_depth_field
except (ImportError, ValueError):
    from infrastructure.vector_io import resolve_depth_field


warnings.filterwarnings("ignore")
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

FILTER_MODES = ["Linear RANSAC", "LS Variance Fit", "Huber Variance Fit"]


def extract_and_calc_robust(ras_path, vec_layer, depth_fld, b_idx, g_idx, fb):
    rlayer = QgsRasterLayer(ras_path)
    
    # Ensure vector and raster CRS are valid
    v_crs = vec_layer.crs() if vec_layer and vec_layer.crs().isValid() else QgsCoordinateReferenceSystem("EPSG:4326")
    r_crs = rlayer.crs() if rlayer and rlayer.crs().isValid() else QgsCoordinateReferenceSystem("EPSG:4326")
    
    tr = QgsCoordinateTransform(v_crs, r_crs, QgsProject.instance())

    # Resolve actual depth field name dynamically
    actual_depth_fld = resolve_depth_field(vec_layer, depth_fld)

    total_pts = vec_layer.featureCount() if vec_layer else 0
    count_invalid_depth = 0
    count_out_of_bounds = 0
    count_nodata = 0
    count_valid = 0

    with rasterio.open(ras_path) as src:
        band_b = src.read(b_idx)
        band_g = src.read(g_idx)
        raw_b, raw_g, depths_list, feats_list = [], [], [], []

        for f in vec_layer.getFeatures():
            raw_d = f[actual_depth_fld] if actual_depth_fld in [field.name() for field in vec_layer.fields()] else None
            if raw_d is None or raw_d == "" or str(raw_d).strip() == "":
                count_invalid_depth += 1
                continue
            try:
                d = float(raw_d)
            except (ValueError, TypeError):
                count_invalid_depth += 1
                continue

            if d == 0.0:
                count_invalid_depth += 1
                continue

            geom = QgsGeometry(f.geometry())
            raw_pt = geom.asMultiPoint()[0] if geom.isMultipart() else geom.asPoint()
            
            pt = raw_pt
            in_bounds = False
            r, c = -1, -1

            # 1. Test if raw point coordinates are ALREADY inside raster bounds
            try:
                r_raw, c_raw = src.index(raw_pt.x(), raw_pt.y())
                if 0 <= r_raw < src.height and 0 <= c_raw < src.width:
                    in_bounds = True
                    r, c = r_raw, c_raw
            except Exception:
                in_bounds = False

            if not in_bounds:
                try:
                    geom_trans = QgsGeometry(geom)
                    geom_trans.transform(tr)
                    trans_pt = geom_trans.asMultiPoint()[0] if geom_trans.isMultipart() else geom_trans.asPoint()
                    r_trans, c_trans = src.index(trans_pt.x(), trans_pt.y())
                    if 0 <= r_trans < src.height and 0 <= c_trans < src.width:
                        in_bounds = True
                        r, c = r_trans, c_trans
                except Exception:
                    pass

            # PROJ 6+ Strict Axis Order Fallback (Lat/Lon vs Lon/Lat)
            if not in_bounds and 'trans_pt' in locals():
                try:
                    r_trans_swap, c_trans_swap = src.index(trans_pt.y(), trans_pt.x())
                    if 0 <= r_trans_swap < src.height and 0 <= c_trans_swap < src.width:
                        in_bounds = True
                        r, c = r_trans_swap, c_trans_swap
                except Exception:
                    pass

            if not in_bounds:
                count_out_of_bounds += 1
                continue

            if 0 <= r < src.height and 0 <= c < src.width:
                vb, vg = band_b[r, c], band_g[r, c]
                if vb > 0 and vg > 0:
                    raw_b.append(vb)
                    raw_g.append(vg)
                    depths_list.append(d)
                    feats_list.append(f)
                    count_valid += 1
                else:
                    count_nodata += 1
            else:
                count_out_of_bounds += 1

        fb.pushInfo(f"   [DEBUG] Total Points: {total_pts}")
        fb.pushInfo(f"   [DEBUG] Depth Field Used: '{actual_depth_fld}'")
        fb.pushInfo(f"   [DEBUG] Vector CRS: {v_crs.authid() or v_crs.userFriendlyIdentifier()} | Raster CRS: {r_crs.authid() or r_crs.userFriendlyIdentifier()}")
        fb.pushInfo(f"   [DEBUG] Invalid Depth (0/None): {count_invalid_depth}")
        fb.pushInfo(f"   [DEBUG] Out of Bounds: {count_out_of_bounds}")
        fb.pushInfo(f"   [DEBUG] NoData/Masked Pixels: {count_nodata}")
        fb.pushInfo(f"   [DEBUG] Candidate Points for Filtering: {count_valid}")

        if count_valid == 0:
            hint = (
                "Hint: Uncheck 'Enable Water Masking' in Phase 01 if 'NoData' is high."
                if count_nodata > 0
                else f"Hint: Check CRS alignment between vector ({v_crs.authid()}) and raster ({r_crs.authid()})."
            )
            raise QgsProcessingException(
                f"No valid training points found! (NoData={count_nodata}, OutBounds={count_out_of_bounds}). {hint}"
            )

        np_b, np_g = np.array(raw_b), np.array(raw_g)
        min_p5 = min(np.percentile(np_b, 5), np.percentile(np_g, 5))
        robust_scale = 10.0 / max(min_p5, 0.0001)
        lb = np.log(np.maximum(np_b * robust_scale, 1.1))
        lg = np.log(np.maximum(np_g * robust_scale, 1.1))
        ratios = lb / lg
        valid = (ratios > 0.1) & (ratios < 5.0)
        return (
            ratios[valid],
            np.array(depths_list)[valid],
            [feats_list[i] for i in range(len(valid)) if valid[i]],
            robust_scale,
        )


def save_subset_with_uncert(
    original_layer, all_features, mask, uncert, out_path, layer_name
):
    from qgis.core import QgsFields
    fields = QgsFields()
    for f in original_layer.fields():
        fields.append(f)
    uncert_idx = fields.indexOf("SDB_Uncert")
    if uncert_idx == -1:
        fields.append(QgsField("SDB_Uncert", QVariant.Double))
        uncert_idx = fields.count() - 1

    base, _ = os.path.splitext(out_path)
    for ext in [".shp", ".shx", ".dbf", ".prj", ".cpg", ".qpj"]:
        f_to_del = base + ext
        if os.path.exists(f_to_del):
            try:
                os.remove(f_to_del)
            except Exception:
                pass
    writer = QgsVectorFileWriter(
        out_path,
        "UTF-8",
        fields,
        QgsWkbTypes.Point,
        original_layer.sourceCrs(),
        "ESRI Shapefile",
    )
    if writer.hasError() == QgsVectorFileWriter.NoError:
        for i, m in enumerate(mask):
            if m:
                feat = QgsFeature(fields)
                feat.setGeometry(all_features[i].geometry())
                attrs = all_features[i].attributes()
                if len(attrs) <= uncert_idx:
                    attrs.append(float(uncert[i]))
                else:
                    attrs[uncert_idx] = float(uncert[i])
                feat.setAttributes(attrs)
                writer.addFeature(feat)
        del writer


def plot_physics_trend(X, y, mask, mode, path):
    plt.figure(figsize=(10, 6))
    
    n_acc = np.sum(mask)
    n_rej = np.sum(~mask)
    
    plt.scatter(X[mask], y[mask], c="dodgerblue", s=15, alpha=0.6, label=f"Accepted ({n_acc})")
    plt.scatter(
        X[~mask], y[~mask], c="gray", marker="x", s=15, alpha=0.4, label=f"Rejected ({n_rej})"
    )

    # Dynamic Outlier-Robust Axis Limiting (Percentile-based Zooming)
    x_valid = X[mask] if np.sum(mask) > 0 else X
    y_valid = y[mask] if np.sum(mask) > 0 else y

    x_pmin, x_pmax = np.percentile(x_valid, [0.5, 99.5])
    y_pmin, y_pmax = np.percentile(y_valid, [0.5, 99.5])
    
    x_span = x_pmax - x_pmin if x_pmax > x_pmin else 1.0
    y_span = y_pmax - y_pmin if y_pmax > y_pmin else 1.0

    # Trend line strictly inside valid domain to avoid shooting off to infinity
    x_rng = np.linspace(x_pmin - 0.05 * x_span, x_pmax + 0.05 * x_span, 100).reshape(-1, 1)
    if mode == 0:
        model = LinearRegression().fit(x_valid, y_valid)
        y_rng = model.predict(x_rng)
    else:
        poly = PolynomialFeatures(degree=2, include_bias=False)
        model = LinearRegression().fit(poly.fit_transform(x_valid), y_valid)
        y_rng = model.predict(poly.transform(x_rng))

    plt.plot(x_rng, y_rng, "k-", lw=2, label="Trend Model")
    plt.xlim(x_pmin - 0.1 * x_span, x_pmax + 0.1 * x_span)
    plt.ylim(y_pmin - 0.15 * y_span, y_pmax + 0.15 * y_span)
    
    plt.title("Plot 1: Band Ratio vs Depth")
    plt.xlabel("Log Ratio")
    plt.ylabel("Depth (m)")
    plt.legend()
    plt.grid(True, linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig(path, dpi=120)
    plt.close()


def plot_variance_analysis(depths, residuals, sigma_model, mode, path):
    if mode == 0 or sigma_model is None:
        return
    plt.figure(figsize=(10, 6))
    d_abs = np.abs(depths)
    res_sq = residuals**2
    
    d_pmin, d_pmax = np.percentile(d_abs, [0.5, 99.5])
    d_span = d_pmax - d_pmin if d_pmax > d_pmin else 1.0

    poly = PolynomialFeatures(degree=2, include_bias=False)
    d_rng = np.linspace(d_pmin, d_pmax + 0.05 * d_span, 100).reshape(-1, 1)
    var_pred = sigma_model.predict(poly.fit_transform(d_rng))

    plt.scatter(d_abs, res_sq, c="black", s=10, alpha=0.5, label="Residual Squares")
    plt.plot(d_rng, var_pred, "r-", lw=2.5, label=r"Variance Trend ($\sigma^2$)")
    
    # Robust Y-axis scaling to prevent extreme residual outliers (e.g. res^2 = 800) from squishing the variance trend
    y_max = max(np.percentile(res_sq, 98.0) * 1.5, np.max(var_pred) * 2.5)
    if np.isnan(y_max) or y_max <= 0:
        y_max = np.max(res_sq) if len(res_sq) > 0 else 100
        
    plt.xlim(max(0, d_pmin - 0.05 * d_span), d_pmax + 0.05 * d_span)
    plt.ylim(-0.02 * y_max, y_max)

    plt.title("Plot 2: Depth vs Residual Square (Variance Analysis)")
    plt.xlabel("Depth (m)")
    plt.ylabel(r"Residual² ($m^2$)")
    plt.legend()
    plt.grid(True, linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig(path, dpi=120)
    plt.close()


def plot_envelope_analysis(depths, residuals, sigmas, mask, multiplier, mode, path):
    plt.figure(figsize=(10, 6))
    
    n_acc = np.sum(mask)
    n_rej = np.sum(~mask)
    
    d_abs = np.abs(depths)
    sort_idx = np.argsort(d_abs)
    d_sorted = d_abs[sort_idx]
    sigma_sorted = sigmas[sort_idx]

    upper = multiplier * sigma_sorted
    lower = -multiplier * sigma_sorted

    plt.plot(d_sorted, upper, "r-", lw=2, label=rf"Upper (+{multiplier}$\sigma$)")
    plt.plot(d_sorted, lower, "r-", lw=2, label=rf"Lower (-{multiplier}$\sigma$)")
    plt.fill_between(d_sorted, lower, upper, color="red", alpha=0.1)
    plt.scatter(
        d_abs[mask], residuals[mask], c="dodgerblue", s=15, alpha=0.6, label=f"Accepted ({n_acc})"
    )
    plt.scatter(
        d_abs[~mask],
        residuals[~mask],
        c="gray",
        marker="x",
        s=15,
        alpha=0.4,
        label=f"Rejected ({n_rej})",
    )
    plt.axhline(0, c="black", lw=1)
    
    # Robust Envelope Limiting to zoom in on data & envelope
    res_valid = residuals[mask] if np.sum(mask) > 0 else residuals
    r_pmin, r_pmax = np.percentile(res_valid, [0.5, 99.5])
    
    env_max = max(np.max(upper), r_pmax)
    env_min = min(np.min(lower), r_pmin)
    r_span = env_max - env_min if env_max > env_min else 2.0
    
    d_pmin, d_pmax = np.percentile(d_abs, [0.5, 99.5])
    d_span = d_pmax - d_pmin if d_pmax > d_pmin else 1.0

    plt.xlim(max(0, d_pmin - 0.05 * d_span), d_pmax + 0.05 * d_span)
    plt.ylim(env_min - 0.15 * r_span, env_max + 0.15 * r_span)

    plt.title(f"Plot 3: Residuals & Uncertainty Envelope (Mode {mode})")
    plt.xlabel("Depth (m)")
    plt.ylabel("Residual (m)")
    plt.legend()
    plt.grid(True, linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig(path, dpi=120)
    plt.close()


def run_phase02_filtering(algorithm, parameters, context, feedback):
    out_dir = algorithm.parameterAsString(parameters, algorithm.OUTPUT_FOLDER, context)
    os.makedirs(out_dir, exist_ok=True)
    stack_path = algorithm.parameterAsRasterLayer(
        parameters, algorithm.INPUT_STACK, context
    ).source()
    points_layer = algorithm.parameterAsVectorLayer(
        parameters, algorithm.INPUT_POINTS, context
    )
    
    # Safely reproject points to match stack before proceeding
    import processing
    r_temp = QgsRasterLayer(stack_path, "temp")
    
    if r_temp.isValid() and points_layer and points_layer.crs().isValid():
        v_crs = points_layer.crs()
        if v_crs == r_temp.crs() and not r_temp.crs().isGeographic():
            bbox = points_layer.extent()
            if bbox.xMinimum() >= -180 and bbox.xMaximum() <= 180 and bbox.yMinimum() >= -90 and bbox.yMaximum() <= 90:
                v_crs = QgsCoordinateReferenceSystem("EPSG:4326")
                points_layer.setCrs(v_crs)
                feedback.pushInfo("   [DEBUG] Auto-detected geographic coordinates despite projected CRS. Forcing Vector CRS to EPSG:4326.")

    if r_temp.isValid() and points_layer and points_layer.crs().isValid() and points_layer.crs() != r_temp.crs():
        feedback.pushInfo(f"   [DEBUG] Reprojecting training points from {points_layer.crs().authid()} to {r_temp.crs().authid()}...")
        try:
            out_reproj = os.path.join(out_dir, "1_Reprojected_Points.gpkg")
            res = processing.run("native:reprojectlayer", {
                'INPUT': points_layer,
                'TARGET_CRS': r_temp.crs(),
                'OUTPUT': out_reproj
            }, context=context, feedback=feedback)
            points_layer = QgsVectorLayer(out_reproj, "reprojected_points", "ogr")
        except Exception as e:
            feedback.pushWarning(f"Failed to reproject points natively: {e}")

    depth_fld = algorithm.parameterAsString(parameters, algorithm.FIELD_DEPTH, context)
    b_idx = algorithm.parameterAsInt(parameters, algorithm.BLUE_BAND, context)
    g_idx = algorithm.parameterAsInt(parameters, algorithm.GREEN_BAND, context)

    mode_idx = algorithm.parameterAsEnum(parameters, algorithm.FILTER_MODE, context)
    user_val = algorithm.parameterAsDouble(
        parameters, algorithm.RESIDUAL_THRESHOLD, context
    )
    max_trials = algorithm.parameterAsInt(
        parameters, algorithm.RANSAC_MAX_TRIALS, context
    )

    feedback.pushInfo(
        f"\n>>> MODULE 02: STARTING FILTERING [{FILTER_MODES[mode_idx]}]..."
    )

    X_ratio, y_depth, features_list, _calc_scale = extract_and_calc_robust(
        stack_path, points_layer, depth_fld, b_idx, g_idx, feedback
    )

    if len(y_depth) < 20:
        raise QgsProcessingException("Not enough valid points found (<20)!")

    X = X_ratio.reshape(-1, 1)
    y = y_depth

    final_mask = None
    est_sigma = None
    sigma_model_for_plot = None

    if mode_idx == 0:
        feedback.pushInfo("   [Logic] Running Linear RANSAC...")
        if user_val <= 0:
            lr = LinearRegression().fit(X, y)
            residuals_raw = np.abs(y - lr.predict(X))
            mad = np.median(residuals_raw)
            thresh = 2.5 * mad
            feedback.pushInfo(f"   [Auto] Threshold: {thresh:.4f}m (MAD-based)")
        else:
            thresh = user_val
            feedback.pushInfo(f"   [User] Using Threshold: {thresh:.4f}m")

        ransac = RANSACRegressor(
            random_state=42,
            min_samples=0.1,
            residual_threshold=thresh,
            max_trials=max_trials,
        )
        ransac.fit(X, y)
        final_mask = ransac.inlier_mask_
        est_sigma = np.full_like(y, thresh / 3.0)
        residuals = y - ransac.predict(X)

    elif mode_idx == 1:
        feedback.pushInfo("   [Logic] Running LS Variance Fit...")
        poly = PolynomialFeatures(degree=2, include_bias=False)
        X_poly = poly.fit_transform(X)
        ls_model = LinearRegression().fit(X_poly, y)
        residuals = y - ls_model.predict(X_poly)
        mad = np.median(np.abs(residuals))
        sigma_val = 1.4826 * mad

        if user_val <= 0:
            multiplier = 3.0
            feedback.pushInfo(
                f"   [Auto] Global Sigma: {sigma_val:.4f}m | Multiplier: 3.0"
            )
        else:
            multiplier = user_val

        limit = multiplier * sigma_val
        final_mask = np.abs(residuals) <= limit
        est_sigma = np.full_like(y, sigma_val)

        res_sq = residuals**2
        depth_poly_feat = poly.fit_transform(y.reshape(-1, 1))
        sigma_model_for_plot = LinearRegression().fit(depth_poly_feat, res_sq)

    elif mode_idx == 2:
        feedback.pushInfo("   [Logic] Running Huber Variance Fit...")
        poly = PolynomialFeatures(degree=2, include_bias=False)
        X_poly = poly.fit_transform(X)
        ransac_trend = RANSACRegressor(
            random_state=42, min_samples=0.1, max_trials=max_trials
        )
        ransac_trend.fit(X_poly, y)

        residuals = y - ransac_trend.predict(X_poly)
        mask_trend = ransac_trend.inlier_mask_

        y_clean = y[mask_trend]
        res_sq_clean = residuals[mask_trend] ** 2
        depth_poly_feat = poly.fit_transform(y_clean.reshape(-1, 1))

        sigma_model = HuberRegressor(epsilon=1.35, max_iter=max_trials)
        sigma_model.fit(depth_poly_feat, res_sq_clean)
        sigma_model_for_plot = sigma_model

        all_depth_poly = poly.transform(y.reshape(-1, 1))
        est_sigma_sq = sigma_model.predict(all_depth_poly)
        est_sigma = np.sqrt(np.maximum(est_sigma_sq, 0.0025))

        if user_val <= 0:
            multiplier = 3.0
            feedback.pushInfo("   [Auto] Using Default 3-Sigma Envelope")
        else:
            multiplier = user_val

        final_mask = np.abs(residuals) <= (multiplier * est_sigma)

    clean_shp = os.path.join(out_dir, "2_Cleaned_Training_Data.shp")
    reject_shp = os.path.join(out_dir, "2_Outliers_Rejected.shp")

    save_subset_with_uncert(
        points_layer, features_list, final_mask, est_sigma, clean_shp, "Cleaned Data"
    )
    save_subset_with_uncert(
        points_layer, features_list, ~final_mask, est_sigma, reject_shp, "Rejected Data"
    )

    plot_mult = user_val if user_val > 0 else (2.5 if mode_idx == 0 else 3.0)
    plot_physics_trend(
        X, y, final_mask, mode_idx, os.path.join(out_dir, "2_Plot_1_Trend.png")
    )
    plot_variance_analysis(
        y,
        residuals,
        sigma_model_for_plot,
        mode_idx,
        os.path.join(out_dir, "2_Plot_2_Variance.png"),
    )
    plot_envelope_analysis(
        y,
        residuals,
        est_sigma,
        final_mask,
        plot_mult,
        mode_idx,
        os.path.join(out_dir, "2_Plot_3_Envelope.png"),
    )

    n_acc = int(np.sum(final_mask))
    n_rej = int(np.sum(~final_mask))
    pct_acc = (n_acc / len(final_mask)) * 100.0 if len(final_mask) > 0 else 0.0
    pct_rej = (n_rej / len(final_mask)) * 100.0 if len(final_mask) > 0 else 0.0

    feedback.pushInfo("   ══════════════════════════════════════════════════")
    feedback.pushInfo(f"   📊 [Phase 02 Summary - {FILTER_MODES[mode_idx]} Output]:")
    feedback.pushInfo(f"      • Initial Candidates: {len(final_mask)}")
    feedback.pushInfo(f"      • Accepted (Cleaned Inliers): {n_acc} ({pct_acc:.1f}%)")
    feedback.pushInfo(f"      • Rejected (Outliers Filtered): {n_rej} ({pct_rej:.1f}%)")
    feedback.pushInfo(f"      • Cleaned Output Layer: {clean_shp}")
    feedback.pushInfo("   ══════════════════════════════════════════════════")

    try:
        from Bathymetrix_AI.infrastructure.logging import log_module_completion
        primary_files = {
            "Cleaned Inliers": clean_shp,
            "Outliers Rejected": reject_shp,
            "Physics Trend Plot": os.path.join(out_dir, "2_Plot_1_Trend.png"),
            "Variance Analysis": os.path.join(out_dir, "2_Plot_2_Variance.png"),
            "Envelope Plot": os.path.join(out_dir, "2_Plot_3_Envelope.png"),
        }
        log_module_completion(
            module_title="Phase 02: Robust Bathymetric Filtering",
            out_dir=out_dir,
            primary_files=primary_files,
            log_path=os.path.join(out_dir, "02_Filtering_Log.txt"),
            feedback=feedback
        )
    except Exception:
        pass

    return {"OUTPUT_CLEAN_VEC": clean_shp}
