/*
 *  $Id: xyz_level.c 30047 2026-06-11 08:30:45Z yeti-dn $
 *  Copyright (C) 2016-2026 David Necas (Yeti).
 *
 *  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 <gtk/gtk.h>
#include <gwy.h>

#define RUN_MODES (GWY_RUN_INTERACTIVE | GWY_RUN_IMMEDIATE)

enum {
    PARAM_DATA,
    PARAM_UPDATE_ALL,
};

typedef struct {
    GwyParams *params;
    GwyScatter *scatter;
    gint idata;
} ModuleArgs;

typedef struct {
    ModuleArgs *args;
    GtkWidget *dialog;
    GwyParamTable *table;
} ModuleGUI;

static gboolean         module_register     (void);
static GwyParamDef*     define_module_params(gboolean have_all);
static void             level_execute       (ModuleArgs *args);
static void             level_rotate_execute(ModuleArgs *args);
static void             setzero_execute     (ModuleArgs *args,
                                             gboolean to_mean);
static GwyDialogOutcome run_gui             (ModuleArgs *args,
                                             const gchar *funcname);
static void             param_changed_rotate(ModuleGUI *gui,
                                             gint id);
static void             xyz_level           (GwyFile *data,
                                             GwyRunModeFlags mode);
static void             xyz_level_rotate    (GwyFile *data,
                                             GwyRunModeFlags mode);
static void             level_rotate_xyz    (GwyXY *xy,
                                             gdouble *data,
                                             gsize n,
                                             gdouble bx,
                                             gdouble by,
                                             const GwyXYZ *c);
static void             rotate_xyz          (GwyXY *xy,
                                             gdouble *data,
                                             gsize n,
                                             const GwyXYZ *u,
                                             const GwyXYZ *c,
                                             gdouble phi);
static void             find_plane_coeffs   (const GwyXY *xy,
                                             const gdouble *data,
                                             gsize n,
                                             gdouble *a,
                                             gdouble *bx,
                                             gdouble *by,
                                             GwyXYZ *c);
static gboolean         unit_equal_to_xy    (GwyScatter *scatter,
                                             gint i);
static void             sanitise_params     (ModuleArgs *args);

static GwyModuleInfo module_info = {
    GWY_MODULE_ABI_VERSION,
    &module_register,
    N_("Simple XYZ data leveling."),
    "Yeti <yeti@gwyddion.net>",
    "3.1",
    "David Nečas (Yeti)",
    "2016",
};

GWY_MODULE_QUERY2(module_info, xyz_level)

static gboolean
module_register(void)
{
    gwy_xyz_func_register("xyz_fix_zero",
                          xyz_level,
                          N_("/_Level/Fix _Zero"),
                          GWY_ICON_FIX_ZERO,
                          RUN_MODES,
                          GWY_MENU_FLAG_XYZ,
                          N_("Shift minimum data value to zero"));
    gwy_xyz_func_register("xyz_zero_mean",
                          xyz_level,
                          N_("/_Level/Zero _Mean Value"),
                          GWY_ICON_ZERO_MEAN,
                          RUN_MODES,
                          GWY_MENU_FLAG_XYZ,
                          N_("Shift mean data value to zero"));
    gwy_xyz_func_register("xyz_level",
                          xyz_level,
                          N_("/_Level/Plane _Level"),
                          GWY_ICON_LEVEL,
                          RUN_MODES,
                          GWY_MENU_FLAG_XYZ,
                          N_("Level data by mean plane subtraction"));
    gwy_xyz_func_register("xyz_level_rotate",
                          xyz_level_rotate,
                          N_("/_Level/Level _Rotate"),
                          NULL,
                          RUN_MODES,
                          GWY_MENU_FLAG_XYZ,
                          N_("Automatically level data by plane rotation"));

    return TRUE;
}

static GwyParamDef*
define_module_params(gboolean have_all)
{
    static GwyParamDef *paramdef = NULL;

    if (paramdef)
        return paramdef;

    paramdef = gwy_param_def_new();
    gwy_param_def_set_function_name(paramdef, gwy_xyz_func_current());
    gwy_param_def_add_scatter_data(paramdef, PARAM_DATA, "data", _("Value to level"));
    if (have_all)
        gwy_param_def_add_boolean(paramdef, PARAM_UPDATE_ALL, "update_all", _("Apply to all values"), FALSE);

    return paramdef;
}

/* FIXME: We should use weighting by area.  But that must wait until we can do such thing... */
static void
xyz_level(GwyFile *data, GwyRunModeFlags mode)
{
    ModuleArgs args;
    gint id;

    gwy_data_browser_get_current(GWY_APP_SCATTER, &args.scatter,
                                 GWY_APP_SCATTER_ID, &id,
                                 0);
    g_return_if_fail(GWY_IS_SCATTER(args.scatter));

    gint nvalues = gwy_scatter_get_n_values(args.scatter);
    if (!nvalues)
        return;

    args.params = gwy_params_new_from_settings(define_module_params(TRUE));
    sanitise_params(&args);

    const gchar *funcname = gwy_xyz_func_current();
    GwyDialogOutcome outcome = GWY_DIALOG_PROCEED;
    if (nvalues > 1 && (mode == GWY_RUN_INTERACTIVE)) {
        outcome = run_gui(&args, funcname);
        gwy_params_save_to_settings(args.params);
        if (outcome == GWY_DIALOG_CANCEL)
            goto end;
        args.idata = gwy_params_get_int(args.params, PARAM_DATA);
    }
    else
        args.idata = MAX(args.idata, 0);

    GQuark quark = gwy_file_key_xyz(id);
    gwy_app_undo_qcheckpointv(GWY_DICT(data), 1, &quark);

    if (outcome == GWY_DIALOG_PROCEED) {
        if (gwy_strequal(funcname, "xyz_fix_zero"))
            setzero_execute(&args, FALSE);
        else if (gwy_strequal(funcname, "xyz_zero_mean"))
            setzero_execute(&args, TRUE);
        else if (gwy_strequal(funcname, "xyz_level"))
            level_execute(&args);
        else {
            g_assert_not_reached();
        }
    }

    gwy_log_add(data, GWY_FILE_XYZ, id, id);
    gwy_scatter_data_changed(args.scatter);

end:
    g_object_unref(args.params);
}

static void
xyz_level_rotate(GwyFile *data, GwyRunModeFlags mode)
{
    GwyDialogOutcome outcome = GWY_DIALOG_PROCEED;
    ModuleArgs args;
    gint id;

    g_return_if_fail(mode & RUN_MODES);

    gwy_data_browser_get_current(GWY_APP_SCATTER, &args.scatter,
                                 GWY_APP_SCATTER_ID, &id,
                                 0);
    g_return_if_fail(GWY_IS_SCATTER(args.scatter));

    gint nvalues = gwy_scatter_get_n_values(args.scatter);
    if (!nvalues)
        return;

    args.params = gwy_params_new_from_settings(define_module_params(FALSE));
    sanitise_params(&args);

    if (mode == GWY_RUN_INTERACTIVE) {
        outcome = run_gui(&args, "xyz_level_rotate");
        gwy_params_save_to_settings(args.params);
        if (outcome == GWY_DIALOG_CANCEL)
            goto end;
        args.idata = gwy_params_get_int(args.params, PARAM_DATA);
    }
    else {
        if (args.idata < 0 || !unit_equal_to_xy(args.scatter, args.idata)) {
            for (gint i = 0; i < nvalues; i++) {
                if (unit_equal_to_xy(args.scatter, i)) {
                    args.idata = i;
                    break;
                }
            }
        }
    }
    if (args.idata < 0)
        goto end;

    GQuark quark = gwy_file_key_xyz(id);
    gwy_app_undo_qcheckpointv(GWY_DICT(data), 1, &quark);
    if (outcome == GWY_DIALOG_PROCEED)
        level_rotate_execute(&args);

    gwy_log_add(data, GWY_FILE_XYZ, id, id);
    gwy_scatter_xy_changed(args.scatter);
    gwy_scatter_data_changed(args.scatter);

end:
    g_object_unref(args.params);
}

static GwyDialogOutcome
run_gui(ModuleArgs *args, const gchar *funcname)
{
    ModuleGUI gui;
    GwyDialog *dialog;
    GwyParamTable *table;
    const gchar *title;

    if (gwy_strequal(funcname, "xyz_level"))
        title = _("Plane Level");
    else if (gwy_strequal(funcname, "xyz_level_rotate"))
        title = _("Level Rotate");
    else if (gwy_strequal(funcname, "xyz_fix_zero"))
        title = _("Fix Zero");
    else if (gwy_strequal(funcname, "xyz_zero_mean"))
        title = _("Zero Mean Value");
    else {
        g_assert_not_reached();
    }

    gwy_clear1(gui);
    gui.args = args;
    gui.dialog = gwy_dialog_new(title);
    dialog = GWY_DIALOG(gui.dialog);
    gwy_dialog_add_stock_buttons(dialog, GWY_STOCK_CANCEL, GWY_STOCK_OK, 0);

    table = gui.table = gwy_param_table_new(args->params);
    gwy_param_table_append_scatter_data(table, PARAM_DATA, args->scatter);
    if (!gwy_strequal(funcname, "xyz_level_rotate"))
        gwy_param_table_append_checkbox(table, PARAM_UPDATE_ALL);
    gwy_dialog_add_content(dialog, gwy_param_table_widget(table), FALSE, TRUE, 0);

    if (gwy_strequal(funcname, "xyz_level_rotate"))
        g_signal_connect_swapped(table, "param-changed", G_CALLBACK(param_changed_rotate), &gui);

    return gwy_dialog_run(dialog);
}

static void
param_changed_rotate(ModuleGUI *gui, G_GNUC_UNUSED gint id)
{
    ModuleArgs *args = gui->args;

    gint i = gwy_params_get_int(args->params, PARAM_DATA);
    gtk_dialog_set_response_sensitive(GTK_DIALOG(gui->dialog), GTK_RESPONSE_OK,
                                      unit_equal_to_xy(args->scatter, i));
}

static void
setzero_execute(ModuleArgs *args, gboolean to_mean)
{
    GwyScatter *scatter = args->scatter;
    gboolean all_data = gwy_params_get_boolean(args->params, PARAM_UPDATE_ALL);
    gint ifrom = all_data ? 0 : args->idata;
    gint ito = all_data ? gwy_scatter_get_n_values(scatter) : args->idata+1;
    gsize n = gwy_scatter_get_n_points(scatter);

    for (gint i = ifrom; i < ito; i++) {
        gdouble *data = gwy_scatter_get_data(scatter, i);
        gdouble z0;
        if (to_mean)
            z0 = gwy_math_mean(data, n);
        else {
            gdouble zmax;
            gwy_scatter_min_max(scatter, i, &z0, &zmax);
        }

        for (gsize k = 0; k < n; k++)
            data[k] -= z0;
    }
}

static void
level_execute(ModuleArgs *args)
{
    GwyScatter *scatter = args->scatter;
    gboolean all_data = gwy_params_get_boolean(args->params, PARAM_UPDATE_ALL);
    gint ifrom = all_data ? 0 : args->idata;
    gint ito = all_data ? gwy_scatter_get_n_values(scatter) : args->idata+1;
    const GwyXY *xy = gwy_scatter_get_xy_const(scatter);
    gsize n = gwy_scatter_get_n_points(scatter);

    for (gint i = ifrom; i < ito; i++) {
        gdouble *data = gwy_scatter_get_data(scatter, i);
        gdouble a, bx, by;
        GwyXYZ c;
        find_plane_coeffs(xy, data, n, &a, &bx, &by, &c);
        gwy_debug("b %g %g", bx, by);
        for (gsize k = 0; k < n; k++)
            data[k] -= a + bx*xy[k].x + by*xy[k].y;
    }
}

static gboolean
accelerate_convergence(const gdouble *seq, gdouble *x)
{
    gdouble q, qdiff;

    /* If we are converged do not prevent convergence in other components. */
    if (seq[2] == 0.0) {
        *x = 0.0;
        return TRUE;
    }

    if (seq[0]*seq[1] <= 0.0 || seq[1]*seq[2] <= 0.0)
        return FALSE;

    q = seq[2]/seq[1];
    if (q > 0.5)
        return FALSE;

    qdiff = fabs(log(q*seq[0]/seq[1]));
    gwy_debug("q %g (diff %g)", q, qdiff);
    if (qdiff > 0.5*q)
        return FALSE;

    *x = seq[2]/(1.0 - q);
    return TRUE;
}

static void
level_rotate_execute(ModuleArgs *args)
{
    GwyScatter *scatter = args->scatter;
    GwyXY *xy = gwy_scatter_get_xy(scatter);
    gdouble *data = gwy_scatter_get_data(scatter, args->idata);
    gsize n = gwy_scatter_get_n_points(scatter);
    gdouble bxseq[3], byseq[3];

    /* find_plane_coeffs() calculates the mean plane in ordinary least-squares sense.  But this is not
     * self-consistent with rotation that should use total least squares.  Perform a few iterations. */
    gint ns = 0;
    for (gint k = 0; k < 20; k++) {
        gdouble a, bx, by;
        GwyXYZ c;

        find_plane_coeffs(xy, data, n, &a, &bx, &by, &c);
        gwy_debug("[%u] b %g %g", k, bx, by);
        if (ns < 3) {
            bxseq[ns] = bx;
            byseq[ns] = by;
            ns++;
        }
        else {
            bxseq[0] = bxseq[1];
            bxseq[1] = bxseq[2];
            bxseq[2] = bx;
            byseq[0] = byseq[1];
            byseq[1] = byseq[2];
            byseq[2] = by;
        }
        if (ns == 3) {
            gdouble bx_acc, by_acc;
            if (accelerate_convergence(bxseq, &bx_acc) && accelerate_convergence(byseq, &by_acc)) {
                bx = bx_acc;
                by = by_acc;
                ns = 0;
                gwy_debug("acceleratied b %g %g", bx, by);
            }
        }

        level_rotate_xyz(xy, data, n, bx, by, &c);
        if (k > 1 && sqrt(bx*bx + by*by) < 1e-15)
            break;
    }
}

static void
level_rotate_xyz(GwyXY *xy, gdouble *data, gsize n,
                 gdouble bx, gdouble by, const GwyXYZ *c)
{
    gdouble b = sqrt(bx*bx + by*by);
    if (!b)
        return;

    GwyXYZ u = { .x = -by/b, .y = bx/b, .z = 0.0 };
    rotate_xyz(xy, data, n, &u, c, atan2(b, 1.0));
}

static void
rotate_xyz(GwyXY *xy, gdouble *data, gsize n,
           const GwyXYZ *u, const GwyXYZ *c, gdouble phi)
{
    gdouble cphi = cos(phi), sphi = sin(phi);
    gdouble axx = cphi + u->x*u->x*(1.0 - cphi);
    gdouble axy = u->x*u->y*(1.0 - cphi) - u->z*sphi;
    gdouble axz = u->x*u->z*(1.0 - cphi) + u->y*sphi;
    gdouble ayx = u->y*u->x*(1.0 - cphi) + u->z*sphi;
    gdouble ayy = cphi + u->y*u->y*(1.0 - cphi);
    gdouble ayz = u->y*u->z*(1.0 - cphi) - u->x*sphi;
    gdouble azx = u->z*u->x*(1.0 - cphi) - u->y*sphi;
    gdouble azy = u->z*u->y*(1.0 - cphi) + u->x*sphi;
    gdouble azz = cphi + u->z*u->z*(1.0 - cphi);

    for (gsize k = 0; k < n; k++) {
        gdouble x = xy[k].x - c->x;
        gdouble y = xy[k].y - c->y;
        gdouble z = data[k] - c->z;

        xy[k].x = axx*x + axy*y + axz*z + c->x;
        xy[k].y = ayx*x + ayy*y + ayz*z + c->y;
        data[k] = azx*x + azy*y + azz*z + c->z;
    }
}

static void
find_plane_coeffs(const GwyXY *xy, const gdouble *data, gsize n,
                  gdouble *a, gdouble *bx, gdouble *by, GwyXYZ *c)
{
    gdouble sx, sy, sz, sxx, sxy, syy, sxz, syz, D;

    sx = sy = sz = 0.0;
    for (gsize k = 0; k < n; k++) {
        sx += xy[k].x;
        sy += xy[k].y;
        sz += data[k];
    }
    sx /= n;
    sy /= n;
    sz /= n;

    sxx = sxy = syy = sxz = syz = 0.0;
    for (gsize k = 0; k < n; k++) {
        gdouble x = xy[k].x - sx;
        gdouble y = xy[k].y - sy;
        gdouble z = data[k];

        sxx += x*x;
        syy += y*y;
        sxy += x*y;
        sxz += x*z;
        syz += y*z;
    }

    D = sxx*syy - sxy*sxy;
    *bx = (syy*sxz - sxy*syz)/D;
    *by = (sxx*syz - sxy*sxz)/D;
    *a = -(sx*(*bx) + sy*(*by));
    c->x = sx;
    c->y = sy;
    c->z = sz;
}

static gboolean
unit_equal_to_xy(GwyScatter *scatter, gint i)
{
    gint nvalues = gwy_scatter_get_n_values(scatter);
    g_return_val_if_fail(i >= 0 && i < nvalues, FALSE);
    return gwy_unit_equal(gwy_scatter_get_unit_xy(scatter), gwy_scatter_get_unit_data(scatter, i));
}

static void
sanitise_params(ModuleArgs *args)
{
    args->idata = gwy_scatter_find_data_by_label(args->scatter, gwy_params_get_string(args->params, PARAM_DATA));
}

/* 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 : */
