/*
 *  $Id: linecorrect.c 29972 2026-05-21 17:17:25Z yeti-dn $
 *  Copyright (C) 2003-2026 David Necas (Yeti), Petr Klapetek, Luke Somers.
 *  E-mail: yeti@gwyddion.net, klapetek@gwyddion.net, lsomers@sas.upenn.edu.
 *
 *  This program is free software; you can redistribute it and/or modify it under the terms of the GNU General Public
 *  License as published by the Free Software Foundation; either version 2 of the License, or (at your option) any
 *  later version.
 *
 *  This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied
 *  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 *  details.
 *
 *  You should have received a copy of the GNU General Public License along with this program; if not, write to the
 *  Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include "config.h"
#include <glib/gi18n-lib.h>
#include <string.h>
#include <gwy.h>

#define RUN_MODES GWY_RUN_IMMEDIATE

typedef struct {
    gdouble *a;
    gdouble *b;
    guint n;
} MedianLineData;

static gboolean module_register    (void);
static void     line_correct_step  (GwyFile *data,
                                    GwyRunModeFlags mode);
static void     mark_inverted_lines(GwyFile *data,
                                    GwyRunModeFlags mode);

static GwyModuleInfo module_info = {
    GWY_MODULE_ABI_VERSION,
    &module_register,
    N_("Corrects line defects (mostly experimental algorithms)."),
    "Yeti <yeti@gwyddion.net>, Luke Somers <lsomers@sas.upenn.edu>",
    "2.0",
    "David Nečas (Yeti) & Petr Klapetek & Luke Somers",
    "2004",
};

GWY_MODULE_QUERY2(module_info, linecorrect)

static gboolean
module_register(void)
{
    gwy_image_func_register("line_correct_step",
                            &line_correct_step,
                            N_("/_Correct Data/Ste_p Line Correction"),
                            GWY_ICON_LINE_LEVEL,
                            RUN_MODES,
                            GWY_MENU_FLAG_IMAGE,
                            N_("Correct steps in lines"));
    gwy_image_func_register("line_correct_inverted",
                            &mark_inverted_lines,
                            N_("/_Correct Data/Mark _Inverted Rows"),
                            NULL,
                            RUN_MODES,
                            GWY_MENU_FLAG_IMAGE,
                            N_("Mark lines with inverted sign"));

    return TRUE;
}

static void
calculate_segment_correction(const gdouble *drow,
                             gdouble *output,
                             gint xres,
                             gint len)
{
    const gint min_len = 4;

    if (len >= min_len) {
        gdouble corr = 0.0;
        for (gint j = 0; j < len; j++)
            corr += (drow[j] + drow[2*xres + j])/2.0 - drow[xres + j];
        corr /= len;
        for (gint j = 0; j < len; j++)
            output[j] = (3*corr + (drow[j] + drow[2*xres + j])/2.0 - drow[xres + j])/4.0;
    }
    /* output is expected to be zero-initialised. */
}

static void
line_correct_step_iter(GwyField *field,
                       GwyField *tmp)
{
    const gdouble threshold = 3.0;

    gwy_field_clear(tmp);
    gint yres = gwy_field_get_yres(field), xres = gwy_field_get_xres(field);
    const gdouble *d = gwy_field_get_data_const(field);
    gdouble *b = gwy_field_get_data(tmp);

    gdouble w = 0.0;
    for (gint i = 0; i < yres-1; i++) {
        const gdouble *drow = d + i*xres;
        for (gint j = 0; j < xres; j++) {
            gdouble v = drow[j + xres] - drow[j];
            w += v*v;
        }
    }
    w = w/(yres-1)/xres;

    gint *mrow = g_new(gint, xres);
    for (gint i = 0; i < yres-2; i++) {
        const gdouble *drow = d + i*xres;
        gdouble *brow = b + i*xres;
        gwy_clear(mrow, xres);

        for (gint j = 0; j < xres; j++) {
            gdouble u = drow[xres + j];
            gdouble v = (u - drow[j])*(u - drow[j + 2*xres]);
            if (G_UNLIKELY(v > threshold*w)) {
                if (2*u - drow[j] - drow[j + 2*xres] > 0)
                    mrow[j] = 1;
                else
                    mrow[j] = -1;
            }
        }

        gint len = 1, j;
        for (j = 1; j < xres; j++) {
            if (mrow[j] == mrow[j-1])
                len++;
            else {
                if (mrow[j-1])
                    calculate_segment_correction(drow + j-len, brow + j-len, xres, len);
                len = 1;
            }
        }
        if (mrow[j-1])
            calculate_segment_correction(drow + j-len, brow + j-len, xres, len);
    }

    gwy_field_sum_fields(field, field, tmp);
}

static void
line_correct_step(GwyFile *data,
                  GwyRunModeFlags mode)
{
    GwyField *field;
    GQuark dquark;
    gint id;

    g_return_if_fail(mode & RUN_MODES);
    gwy_data_browser_get_current(GWY_APP_FIELD, &field,
                                 GWY_APP_FIELD_KEY, &dquark,
                                 GWY_APP_FIELD_ID, &id,
                                 0);
    g_return_if_fail(field && dquark);
    gwy_app_undo_qcheckpointv(GWY_DICT(data), 1, &dquark);

    gdouble avg = gwy_field_mean(field);
    GwyLine *shifts = gwy_field_find_row_shifts_trimmed_mean(field, NULL, GWY_MASK_IGNORE, 0.5, 0);
    gwy_field_subtract_row_shifts(field, shifts);
    g_object_unref(shifts);

    GwyField *tmp = gwy_field_new_alike(field, FALSE);
    line_correct_step_iter(field, tmp);
    line_correct_step_iter(field, tmp);
    g_object_unref(tmp);

    gwy_field_filter_conservative(field, 5);
    gwy_field_add(field, avg - gwy_field_mean(field));
    gwy_field_data_changed(field);
    gwy_log_add(data, GWY_FILE_IMAGE, id, id);
}

static gdouble
row_correlation(const gdouble *row1, gdouble avg1, gdouble rms1,
                const gdouble *row2, gdouble avg2, gdouble rms2,
                guint n, gdouble total_rms)
{
    gdouble s = 0.0;
    while (n--) {
        s += (*row1 - avg1)*(*row2 - avg2);
        row1++;
        row2++;
    }
    s /= (rms1*rms2 + total_rms*total_rms);
    return s;
}

static void
mark_inverted_lines(GwyFile *data,
                    GwyRunModeFlags mode)
{
    GwyField *field;
    GwyNield *existing_mask;
    GQuark mquark;
    gint id;

    g_return_if_fail(mode & RUN_MODES);
    gwy_data_browser_get_current(GWY_APP_FIELD, &field,
                                 GWY_APP_MASK_FIELD, &existing_mask,
                                 GWY_APP_MASK_FIELD_KEY, &mquark,
                                 GWY_APP_FIELD_ID, &id,
                                 0);
    g_return_if_fail(field && mquark);

    gdouble total_rms = gwy_field_rms(field);
    gint xres = gwy_field_get_xres(field), yres = gwy_field_get_yres(field);
    if (total_rms <= 0.0 || yres < 3 || xres < 3)
        return;

    GwyLine *avgline = gwy_line_new(yres, yres, FALSE);
    gwy_field_line_stats(field, avgline, GWY_LINE_STAT_MEAN, GWY_ORIENTATION_HORIZONTAL);
    gdouble *avg = gwy_line_get_data(avgline);

    GwyLine *rmsline = gwy_line_new(yres, yres, FALSE);
    gwy_field_line_stats(field, rmsline, GWY_LINE_STAT_RMS, GWY_ORIENTATION_HORIZONTAL);
    gdouble *rms = gwy_line_get_data(rmsline);

    gboolean have_negative = FALSE;
    const gdouble *d = gwy_field_get_data_const(field);
    /* Calculate neighbour row correlations. */
    gdouble *weight = g_new0(gdouble, yres-1);
    for (gint i = 0; i < yres-1; i++) {
        weight[i] = row_correlation(d + i*xres, avg[i], rms[i],
                                    d + (i+1)*xres, avg[i+1], rms[i+1],
                                    xres, total_rms);
        if (weight[i] < 0.0)
            have_negative = TRUE;
    }
    if (!have_negative) {
        g_object_unref(avgline);
        g_object_unref(rmsline);
        g_free(weight);
        return;
    }

    /* Find the most positively correlated block. */
    gint from = 0;
    for (gint i = 0; i < yres-2; i++) {
        if (weight[i]*weight[i+1] < 0.0) {
            gdouble s = 0.0;
            for (gint j = from; j <= i; j++)
                s += weight[j];
            gwy_math_fill(weight + from, i+1 - from, s);
            from = i+1;
        }
    }
    gdouble s = 0.0;
    for (gint j = from; j < yres-1; j++)
        s += weight[j];
    for (gint j = from; j < yres-1; j++)
        weight[j] = s;

    gdouble wmax = 0.0;
    from = 0;
    for (gint i = 0; i < yres-1; i++) {
        if (weight[i] > wmax) {
            wmax = weight[i];
            from = i;
        }
    }

    g_object_unref(avgline);
    g_object_unref(rmsline);

    GwyNield *mask = gwy_field_new_nield_alike(field);

    /* Extend the sign of this block downwards.  Note weight[i] represents the sign between rows i and i+1. */
    gboolean inverted = FALSE;
    for (gint i = from; i < yres-1; i++) {
        if (weight[i] < 0.0)
            inverted = !inverted;
        if (inverted)
            gwy_nield_area_fill(mask, NULL, GWY_MASK_IGNORE, 0, i+1, xres, 1, 1);
    }

    /* Extend the sign of this block upwards.  Note weight[i] represents the sign between rows i and i+1. */
    inverted = FALSE;
    for (gint i = from; i >= 0; i--) {
        if (weight[i] < 0.0)
            inverted = !inverted;
        if (inverted)
            gwy_nield_area_fill(mask, NULL, GWY_MASK_IGNORE, 0, i, xres, 1, 1);
    }

    g_free(weight);

    gint nonempty = gwy_nield_count(mask);
    if (!existing_mask && !nonempty) {
        g_object_unref(mask);
        return;
    }
    gwy_app_undo_qcheckpointv(GWY_DICT(data), 1, &mquark);
    gwy_file_set_image_mask(data, id, nonempty ? mask : NULL);
    g_object_unref(mask);
    gwy_log_add(data, GWY_FILE_IMAGE, id, id);
}

/* vim: set cin columns=120 tw=118 et ts=4 sw=4 cino=>1s,e0,n0,f0,{0,}0,^0,\:1s,=0,g1s,h0,t0,+1s,c3,(0,u0 : */
