// ----------------------------------------------------------------------------
// SmartQuantileStretch.js
// Quantile-driven one-click stretch for linear astronomical images in PixInsight.
//
// Copyright (c) 2026 Peter Mrva
// SPDX-License-Identifier: MIT
// See LICENSE.txt for the full license text.
//
// The script:
// - analyzes the intensity distribution and dynamic range of linear data,
// - derives adaptive stretch parameters from robust image quantiles,
// - applies a nonlinear smooth hyperbolic transfer curve, and
// - converts the image from a linear to a nonlinear display state.
//
// Installation:
// - One-time run: Script > Execute Script File, then select this file.
// - Permanent menu entry: Script > Feature Scripts > Add, select the folder
//   containing this file, and run the feature search. The script will appear
//   under Script > IntensityTransformations > SmartQuantileStretch.
//
// Basic use:
// 1. Select a calibrated linear image.
// 2. Choose Full image, Stars only, or Starless.
// 3. Use linked RGB for color-calibrated data; use unlinked RGB only when
//    independent channel treatment is intended.
// 4. Keep Create a new image enabled for initial tests.
// 5. Click Analyze & Apply and inspect the output and console report.
// If an STF remains active on a modified target, reset it before judging the
// permanent stretch.
// ----------------------------------------------------------------------------

#engine v8

#feature-id    IntensityTransformations > SmartQuantileStretch
#feature-info  <b>Smart Quantile Stretch v0.8.2 beta</b><br/>\
               Quantile analysis and a smooth adaptive nonlinear hyperbolic stretch \
               for full, stars-only and starless linear images.

#include <pjsr/StdButton.jsh>
#include <pjsr/StdIcon.jsh>
#include <pjsr/TextAlign.jsh>
#include <pjsr/FrameStyle.jsh>

#define TITLE   "Smart Quantile Stretch"
#define VERSION "0.8.2 beta"

var MODE_FULL = 0;
var MODE_STARS = 1;
var MODE_STARLESS = 2;
var HISTOGRAM_BINS = 65536;
var MAX_SAMPLES = 500000;
var EPSILON = 1.0e-12;

// ----------------------------------------------------------------------------
// Histogram analysis

function clamp01( x )
{
   return Math.max( 0, Math.min( 1, x ) );
}

function HistogramAnalysis( bins, count, sampledMaximum )
{
   this.bins = bins;
   this.count = count;
   this.sampledMaximum = sampledMaximum;

   this.quantile = function( p )
   {
      p = clamp01( p );
      var target = Math.max( 1, Math.ceil( p*this.count ) );
      var sum = 0;
      for ( var i = 0; i < this.bins.length; ++i )
      {
         sum += this.bins[i];
         if ( sum >= target )
            return i/(this.bins.length - 1);
      }
      return 1;
   };

   this.fractionAt = function( x )
   {
      var last = Math.round( clamp01( x )*(this.bins.length - 1) );
      var sum = 0;
      for ( var i = 0; i <= last; ++i )
         sum += this.bins[i];
      return sum/this.count;
   };
}

function sampleIntensityHistogram( view, channel )
{
   var image = view.image;
   var width = image.width;
   var height = image.height;
   var color = image.numberOfNominalChannels >= 3;
   var separateChannel = color && channel >= 0;
   var step = Math.max( 1, Math.floor( Math.sqrt( width*height/MAX_SAMPLES ) ) );
   var x0 = Math.floor( step/2 );
   var y0 = Math.floor( step/2 );

   var bins = new Uint32Array( HISTOGRAM_BINS );
   var row0 = new Float32Array( width );
   var row1 = color && !separateChannel ? new Float32Array( width ) : null;
   var row2 = color && !separateChannel ? new Float32Array( width ) : null;
   var count = 0;
   var sampledMaximum = 0;

   var analysisName = separateChannel ?
      ["R channel", "G channel", "B channel"][channel] :
      (color ? "combined RGB intensity" : "K channel");
   console.writeln( format( "Analysis sampling (%s): step=%d, target <= %d pixels",
                            analysisName, step, MAX_SAMPLES ) );

   for ( var y = y0; y < height; y += step )
   {
      var r = new Rect( 0, y, width, y + 1 );
      image.getSamples( row0, r, separateChannel ? channel : 0 );
      if ( color && !separateChannel )
      {
         image.getSamples( row1, r, 1 );
         image.getSamples( row2, r, 2 );
      }

      for ( var x = x0; x < width; x += step )
      {
         // Linked mode uses arithmetic RGB intensity and makes no assumption
         // about a particular RGB working space. Unlinked mode samples the
         // selected nominal channel directly.
         var v = color && !separateChannel ?
            (row0[x] + row1[x] + row2[x])/3 : row0[x];
         v = clamp01( v );
         sampledMaximum = Math.max( sampledMaximum, v );
         ++bins[Math.round( v*(HISTOGRAM_BINS - 1) )];
         ++count;
      }

      if ( (count & 0x3ffff) == 0 )
         processEvents();
   }

   if ( count < 1 )
      throw new Error( "Could not obtain image samples." );

   return new HistogramAnalysis( bins, count, sampledMaximum );
}

function exactChannelMaximum( image, channel )
{
   var maximum = 0;
   image.pushSelections();
   try
   {
      image.resetSelections();
      var first = channel >= 0 ? channel : 0;
      var last = channel >= 0 ? channel + 1 : image.numberOfNominalChannels;
      for ( var c = first; c < last; ++c )
      {
         image.selectedChannel = c;
         maximum = Math.max( maximum, image.maximum() );
      }
   }
   catch ( x )
   {
      // A white point of 1 guarantees that no valid normalized sample clips.
      console.warningln( "Exact maximum unavailable; using white point 1.0." );
      maximum = 1;
   }
   finally
   {
      image.popSelections();
   }
   return clamp01( maximum );
}

// ----------------------------------------------------------------------------
// Smooth adaptive hyperbolic curve
//
//                 (1+k)x
//       f(x) = -------------
//                  x + k
//
// It maps 0 -> 0 and 1 -> 1, is smooth over the complete range and has no
// internal slope transitions that can produce rings around stellar profiles.
// The knee k is solved analytically from one robust data point and a desired
// output brightness. All low-level samples therefore receive a continuous
// linear gain, including nebulosity located just above the background median.

function fitHyperbolicCurve( xReference, yReference )
{
   xReference = Math.max( EPSILON, Math.min( 1 - EPSILON, xReference ) );
   yReference = Math.max( xReference + EPSILON,
                          Math.min( 1 - EPSILON, yReference ) );
   var k = xReference*(1 - yReference)/(yReference - xReference);
   if ( !isFinite( k ) || k <= 0 )
      throw new Error( "Could not fit a smooth hyperbolic stretch curve." );
   return { k: k };
}

function evaluateHyperbolicCurve( x, curve )
{
   x = clamp01( x );
   return (1 + curve.k)*x/(x + curve.k);
}

function normalizedPoint( x, blackPoint, whitePoint )
{
   return clamp01( (x - blackPoint)/Math.max( EPSILON, whitePoint - blackPoint ) );
}

function quantileSummary( H )
{
   return {
      q0001: H.quantile( 0.0001 ),
      q001:  H.quantile( 0.001 ),
      q01:   H.quantile( 0.01 ),
      q10:   H.quantile( 0.10 ),
      q16:   H.quantile( 0.16 ),
      q25:   H.quantile( 0.25 ),
      q50:   H.quantile( 0.50 ),
      q84:   H.quantile( 0.84 ),
      q90:   H.quantile( 0.90 ),
      q95:   H.quantile( 0.95 ),
      q99:   H.quantile( 0.99 ),
      q995:  H.quantile( 0.995 ),
      q999:  H.quantile( 0.999 ),
      q9995: H.quantile( 0.9995 ),
      q9999: H.quantile( 0.9999 )
   };
}

function modeName( mode )
{
   switch ( mode )
   {
   case MODE_STARS: return "Stars only";
   case MODE_STARLESS: return "Starless";
   default: return "Full image (stars + nebulosity)";
   }
}

function analyzeStretch( view, mode, channel )
{
   var H = sampleIntensityHistogram( view, channel );
   var q = quantileSummary( H );
   var dataMaximum = exactChannelMaximum( view.image, channel );
   // Keep the normalized white point at 1. The curve approaches white
   // smoothly and never rescales an unsaturated image maximum to clipping.
   var whitePoint = 1;
   var blackPoint = 0;
   var referencePoint;
   var referenceTarget;
   var dynamicHigh;

   if ( mode == MODE_STARS )
   {
      // Stars-only data may contain a large exactly-zero residual background.
      // Clip only that residual population and fit the smooth curve within the
      // detected stellar tail.
      var sigma = Math.max( 1/HISTOGRAM_BINS, 0.5*(q.q84 - q.q16) );
      var signalThreshold = q.q50 + 3*sigma;
      var signalStart = Math.min( 0.99, Math.max( 0.90,
         H.fractionAt( signalThreshold ) ) );

      blackPoint = H.quantile( Math.min( 0.65, signalStart - 0.05 ) );
      // Use the upper part of the detected stellar population. Anchoring the
      // curve near its lower edge produces excessive gain in faint PSF wings
      // and visibly inflates stars.
      referencePoint = H.quantile( signalStart + 0.90*(1 - signalStart) );
      dynamicHigh = H.quantile( signalStart + 0.99*(1 - signalStart) );
   }
   else if ( mode == MODE_STARLESS )
   {
      referencePoint = q.q50;
      dynamicHigh = q.q995;
   }
   else
   {
      referencePoint = q.q50;
      dynamicHigh = q.q999;
   }

   var xReference = normalizedPoint( referencePoint, blackPoint, whitePoint );
   var xHigh = normalizedPoint( dynamicHigh, blackPoint, whitePoint );
   if ( xReference <= 1/HISTOGRAM_BINS || xHigh <= xReference )
      throw new Error( "Insufficient nonzero tonal range for this profile." );

   // More dynamic stops imply a brighter stellar tail, so reduce the median
   // target slightly. This remains deliberately milder than a typical STF,
   // which places the background median around 0.25.
   var dynamicStops = Math.log( xHigh/xReference )/Math.LN2;
   if ( mode == MODE_STARLESS )
      referenceTarget = Math.max( 0.115, Math.min( 0.150,
         0.138 - 0.003*(dynamicStops - 5) ) );
   else if ( mode == MODE_FULL )
      referenceTarget = Math.max( 0.095, Math.min( 0.125,
         0.122 - 0.004*(dynamicStops - 5) ) );
   else
      referenceTarget = Math.max( 0.170, Math.min( 0.230,
         0.200 - 0.005*(dynamicStops - 4) ) );

   var curve = fitHyperbolicCurve( xReference, referenceTarget );

   return {
      histogram: H,
      quantiles: q,
      blackPoint: blackPoint,
      whitePoint: whitePoint,
      dataMaximum: dataMaximum,
      referencePoint: referencePoint,
      referenceTarget: referenceTarget,
      xReference: xReference,
      dynamicHigh: dynamicHigh,
      dynamicStops: dynamicStops,
      curve: curve
   };
}

function numberForPixelMath( x )
{
   return format( "%.17g", x );
}

function stretchExpression( analysis )
{
   var b = numberForPixelMath( analysis.blackPoint );
   var w = numberForPixelMath( analysis.whitePoint );
   var u = "max(0,min(1,($T-" + b + ")/(" + w + "-" + b + ")))";
   var k = numberForPixelMath( analysis.curve.k );
   return "(1+" + k + ")*(" + u + ")/((" + u + ")+" + k + ")";
}

function uniqueImageId( baseId )
{
   var base = baseId.replace( /[^A-Za-z0-9_]/g, "_" );
   var id = base;
   var n = 1;
   while ( !ImageWindow.windowById( id ).isNull )
      id = base + "_" + n++;
   return id;
}

function reportAnalysis( A, label )
{
   var q = A.quantiles;
   if ( label.length > 0 )
      console.writeln( "<br><b>" + label + "</b>" );
   console.writeln( format( "Samples analyzed: %d", A.histogram.count ) );
   console.writeln( format( "Quantiles: q01=%.8f, q50=%.8f, q99=%.8f, " +
                            "q99.9=%.8f", q.q01, q.q50, q.q99, q.q999 ) );
   console.writeln( format( "Black point: %.8f", A.blackPoint ) );
   console.writeln( format( "Exact maximum: %.8f", A.dataMaximum ) );
   console.writeln( format( "Protected normalization white point: %.8f",
                            A.whitePoint ) );
   console.writeln( format( "Analyzed dynamic range: %.3f stops",
                            A.dynamicStops ) );
   console.writeln( format( "Reference level: %.8f -> %.4f",
                            A.referencePoint, A.referenceTarget ) );
   console.writeln( format( "Smooth hyperbolic knee: %.10f", A.curve.k ) );
   console.writeln( format( "Initial faint-signal gain: %.1fx",
                            (1 + A.curve.k)/A.curve.k ) );
   console.writeln( format( "Predicted q84/q95/q99/q99.9: %.4f / %.4f / " +
                            "%.4f / %.4f",
      evaluateHyperbolicCurve( normalizedPoint( q.q84, A.blackPoint,
                                                 A.whitePoint ), A.curve ),
      evaluateHyperbolicCurve( normalizedPoint( q.q95, A.blackPoint,
                                                 A.whitePoint ), A.curve ),
      evaluateHyperbolicCurve( normalizedPoint( q.q99, A.blackPoint,
                                                 A.whitePoint ), A.curve ),
      evaluateHyperbolicCurve( normalizedPoint( q.q999, A.blackPoint,
                                                 A.whitePoint ), A.curve ) ) );
}

function exceptionText( x )
{
   if ( x === null )
      return "null";
   if ( x === undefined )
      return "undefined";
   if ( typeof x == "string" )
      return x;
   try
   {
      if ( typeof x.message == "string" && x.message.length > 0 )
         return x.message;
   }
   catch ( ignored )
   {
   }
   try
   {
      return String( x );
   }
   catch ( ignored )
   {
      return "Unknown exception";
   }
}

// ----------------------------------------------------------------------------
// Processing engine

function SmartQuantileStretchEngine()
{
   this.mode = MODE_FULL;
   this.linkedChannels = true;
   this.createNewImage = true;

   this.execute = function( view )
   {
      if ( view.isNull )
         throw new Error( "No target image has been selected." );
      if ( view.image.numberOfNominalChannels != 1 &&
           view.image.numberOfNominalChannels < 3 )
         throw new Error( "Unsupported image color space." );

      console.show();
      console.writeln( "<end><cbr><br><b>" + TITLE + " v" + VERSION + "</b>" );
      console.writeln( "Target: " + view.fullId );
      console.writeln( "Profile: " + modeName( this.mode ) );
      var color = view.image.numberOfNominalChannels >= 3;
      var linked = this.linkedChannels || !color;
      console.writeln( "RGB channels: " +
         (linked ? "linked (one common curve)" :
                   "unlinked (independent R, G and B curves)") );

      var analyses = [];
      if ( linked )
      {
         analyses.push( analyzeStretch( view, this.mode, -1 ) );
         reportAnalysis( analyses[0], color ? "Combined RGB intensity" : "K channel" );
      }
      else
      {
         for ( var c = 0; c < 3; ++c )
         {
            analyses.push( analyzeStretch( view, this.mode, c ) );
            reportAnalysis( analyses[c], ["R channel", "G channel", "B channel"][c] );
         }
      }

      var P;
      var pixelMathStage = "constructor";
      try
      {
         P = new PixelMath;
         pixelMathStage = "expression";
         P.expression = stretchExpression( analyses[0] );
         pixelMathStage = "expression1";
         P.expression1 = linked ? "" : stretchExpression( analyses[1] );
         pixelMathStage = "expression2";
         P.expression2 = linked ? "" : stretchExpression( analyses[2] );
         pixelMathStage = "expression3";
         P.expression3 = "";
         pixelMathStage = "useSingleExpression";
         P.useSingleExpression = linked;
         pixelMathStage = "symbols";
         P.symbols = "";
         pixelMathStage = "generateOutput";
         P.generateOutput = true;
         pixelMathStage = "singleThreaded";
         P.singleThreaded = false;
         pixelMathStage = "optimization";
         P.optimization = true;
         pixelMathStage = "use64BitWorkingImage";
         P.use64BitWorkingImage = false;
         pixelMathStage = "rescale";
         P.rescale = false;
         pixelMathStage = "rescaleLower";
         P.rescaleLower = 0;
         pixelMathStage = "rescaleUpper";
         P.rescaleUpper = 1;
         pixelMathStage = "truncate";
         P.truncate = true;
         pixelMathStage = "truncateLower";
         P.truncateLower = 0;
         pixelMathStage = "truncateUpper";
         P.truncateUpper = 1;
         pixelMathStage = "createNewImage";
         P.createNewImage = this.createNewImage;
         pixelMathStage = "showNewImage";
         P.showNewImage = true;
         pixelMathStage = "newImageId";
         P.newImageId = this.createNewImage ?
            uniqueImageId( view.id + "_SQS" ) : "";
         pixelMathStage = "newImageWidth";
         P.newImageWidth = 0;
         pixelMathStage = "newImageHeight";
         P.newImageHeight = 0;
         pixelMathStage = "newImageAlpha";
         P.newImageAlpha = false;
         pixelMathStage = "newImageColorSpace";
         // PixelMath enumeration values are signed integers. In the V8
         // runtime process constants are no longer exposed through
         // PixelMath.prototype. Zero is SameAsTarget for both properties.
         P.newImageColorSpace = 0;
         pixelMathStage = "newImageSampleFormat";
         P.newImageSampleFormat = 0;
      }
      catch ( x )
      {
         throw new Error( "PixelMath setup failed at '" + pixelMathStage +
                          "': " + exceptionText( x ) );
      }

      try
      {
         if ( !P.executeOn( view, false/*swapFile*/ ) )
            throw new Error( "executeOn() returned false." );
      }
      catch ( x )
      {
         throw new Error( "PixelMath execution failed: " + exceptionText( x ) );
      }

      console.writeln( "<end><cbr><b>Smart Quantile Stretch completed.</b>" );
   };
}

// ----------------------------------------------------------------------------
// Graphical interface

class SmartQuantileStretchDialog extends Dialog
{
   constructor( engine )
   {
      super();

   var labelWidth = this.font.width( "Image type:" ) + 16;

   this.helpLabel = new Label( this );
   this.helpLabel.frameStyle = FrameStyle_Box;
   this.helpLabel.margin = 8;
   this.helpLabel.wordWrapping = true;
   this.helpLabel.useRichText = true;
   this.helpLabel.text =
      "<b>Smart Quantile Stretch v" + VERSION + "</b><br/>" +
      "Analyzes linear astronomical data with robust quantiles, derives " +
      "adaptive parameters from the measured dynamic range, and applies a " +
      "smooth nonlinear hyperbolic curve. This converts the image from a " +
      "linear to a nonlinear state while lifting faint signal and smoothly " +
      "protecting highlights. RGB channels can use one linked curve or three " +
      "independently analyzed curves. Select the data type and click " +
      "<i>Analyze &amp; Apply</i>.";

   this.viewList = new ViewList( this );
   this.viewList.getAll();
   this.viewList.toolTip = "Linear target image.";
   if ( !ImageWindow.activeWindow.isNull )
      this.viewList.currentView = ImageWindow.activeWindow.currentView;

   this.targetLabel = new Label( this );
   this.targetLabel.text = "Target image:";
   this.targetLabel.textAlignment = TextAlign_Right|TextAlign_VertCenter;
   this.targetLabel.minWidth = labelWidth;

   this.targetSizer = new HorizontalSizer;
   this.targetSizer.spacing = 6;
   this.targetSizer.add( this.targetLabel );
   this.targetSizer.add( this.viewList, 100 );

   this.modeLabel = new Label( this );
   this.modeLabel.text = "Image type:";
   this.modeLabel.textAlignment = TextAlign_Right|TextAlign_VertCenter;
   this.modeLabel.minWidth = labelWidth;

   this.modeCombo = new ComboBox( this );
   this.modeCombo.addItem( "Full image — stars and nebulosity" );
   this.modeCombo.addItem( "Separated stars only" );
   this.modeCombo.addItem( "Starless image" );
   this.modeCombo.currentItem = engine.mode;
   this.modeCombo.toolTip =
      "Each profile uses a different robust reference level, target brightness " +
      "and dynamic-range estimate.";
   this.modeCombo.onItemSelected = function( itemIndex )
   {
      engine.mode = itemIndex;
   };

   this.modeSizer = new HorizontalSizer;
   this.modeSizer.spacing = 6;
   this.modeSizer.add( this.modeLabel );
   this.modeSizer.add( this.modeCombo, 100 );

   this.linkedChannels = new CheckBox( this );
   this.linkedChannels.text = "Link RGB channels (one common curve)";
   this.linkedChannels.checked = engine.linkedChannels;
   this.linkedChannels.toolTip =
      "Enabled: analyze combined RGB intensity and apply one curve to all " +
      "channels, preserving their relative balance. Disabled: analyze R, G " +
      "and B independently; useful for selected SHO/HOO workflows.";
   this.linkedChannels.onCheck = function( checked )
   {
      engine.linkedChannels = checked;
   };

   this.createNewImage = new CheckBox( this );
   this.createNewImage.text = "Create a new image";
   this.createNewImage.checked = engine.createNewImage;
   this.createNewImage.toolTip =
      "Recommended. Disable to modify the target image with Undo support.";
   this.createNewImage.onCheck = function( checked )
   {
      engine.createNewImage = checked;
   };

   var dialog = this;
   this.cancelButton = new PushButton( this );
   this.cancelButton.text = "Cancel";
   this.cancelButton.icon = this.scaledResource( ":/icons/cancel.png" );
   this.cancelButton.onClick = function()
   {
      dialog.cancel();
   };

   this.executeButton = new PushButton( this );
   this.executeButton.text = "Analyze && Apply";
   this.executeButton.icon = this.scaledResource( ":/icons/ok.png" );
   this.executeButton.defaultButton = true;
   this.executeButton.onClick = function()
   {
      dialog.ok();
   };

   this.buttonsSizer = new HorizontalSizer;
   this.buttonsSizer.spacing = 6;
   this.buttonsSizer.addStretch();
   this.buttonsSizer.add( this.cancelButton );
   this.buttonsSizer.add( this.executeButton );

   this.sizer = new VerticalSizer;
   this.sizer.margin = 8;
   this.sizer.spacing = 8;
   this.sizer.add( this.helpLabel );
   this.sizer.add( this.targetSizer );
   this.sizer.add( this.modeSizer );
   this.sizer.add( this.linkedChannels );
   this.sizer.add( this.createNewImage );
   this.sizer.addSpacing( 4 );
   this.sizer.add( this.buttonsSizer );

   this.windowTitle = TITLE + " v" + VERSION;
   this.adjustToContents();
      this.setFixedWidth( Math.max( this.width, 600 ) );
   }
}

// ----------------------------------------------------------------------------

function main()
{
   console.hide();
   var engine = new SmartQuantileStretchEngine;
   var dialog = new SmartQuantileStretchDialog( engine );

   for ( ;; )
   {
      if ( !dialog.execute() )
         break;
      try
      {
         engine.execute( dialog.viewList.currentView );
         break;
      }
      catch ( x )
      {
         var message = exceptionText( x );
         console.criticalln( "<end><cbr>*** Error: " + message );
         try
         {
            if ( x !== null && x !== undefined &&
                 typeof x.stack == "string" && x.stack.length > 0 )
               console.criticalln( x.stack );
         }
         catch ( ignored )
         {
         }
         (new MessageBox( message, TITLE, StdIcon_Error, StdButton_Ok )).execute();
      }
   }
}

main();

// ----------------------------------------------------------------------------
