import java.io.*;
import java.util.*;
import java.text.DecimalFormat;

/// Sets up noisy data pixels, and uses Levenberg-Marquardt to 
/// fit a 2-D Gaussian: intensity, xcentroid, ycentroid.
/// Assumes background is known = 10000 units /pixel. 
/// Assumes Gaussian RMS breadth is known = 4.00 pixels. 
/// Iterates with fresh noise to get some statistics on the fitted parms.
/// M.Lampton UCB SSL 2006; 2013



interface Constants
// includes all constants that are to be shared among classes
{
    final static int NPARMS = 3;   // three parms: intensity, x, y.
    final static int NX     = 50;  // image width pixels
    final static int NY     = 50;  // image height pixels 
    final static int NPTS   = NX*NY;  
}






public class Gaus2D implements Constants
{
    final static int NREPS      = 9;
    static double trueparms[]   = new double[NPARMS]; 
    static double fittedparms[] = new double[NPARMS]; 
    static double cleandata[]   = new double[NPTS];
    static double errors[]      = new double[NPTS]; 
    static double noisydata[]   = new double[NPTS]; 
    static double stats[][]     = new double[2][NPARMS]; 
    static double parmerrors[]  = new double[NPARMS]; 


    public static void main (String args[])
    {
        showHeader(); 
        setTrueParms(trueparms);  
        showParms("True  ", trueparms); 
        setCleanData(trueparms, cleandata);
        setErrors(cleandata, errors);  
        zeroStats(stats); 
        for (int n=1; n<=NREPS; n++)
        {
            setNoisyData(cleandata, errors, noisydata); 
            initParms(trueparms, fittedparms); 
            int niter = solve(fittedparms, noisydata, errors); 
            String title = "Fit  "+n; 
            showParmsIter(title, fittedparms, niter); 
            addStats(fittedparms, stats); 
        } 
        finishStats(stats); 
        showParms("Means ", stats[0]);   // the parameter means
        showParms("RMSerr", stats[1]);   // the rms errors on the parms.
    }

    static void showHeader()
    {
       System.out.println("          Signal       <X>pix      <Y>pix  niters"); 
       System.out.println("         ----------------------------------------"); 
    }

    static void setTrueParms(double p[])
    {
        p[0] = 1E4;     // total signal 
        p[1] = 22.0;    // x centroid, pixels 
        p[2] = 26.5;    // y centroid, pixels
    } 

    static void setCleanData(double p[], double d[])
    {
        for (int i=0; i<NPTS; i++)
          d[i] = F.funci(i, p); 
    }

    static void setErrors(double d[], double e[])
    // Installs Poisson errors based on model data
    {
        for (int i=0; i<NPTS; i++)
          e[i] = Math.sqrt(d[i]); 
    }      

    static void showParms(String title, double p[])
    {
        System.out.print(title); 
        for (int i=0; i<NPARMS; i++)
          System.out.print(U.fwd(p[i], 12, 3)); 
        System.out.println();
    }

    static void showParmsIter(String title, double p[], int n)
    {
        System.out.print(title); 
        for (int i=0; i<NPARMS; i++)
          System.out.print(U.fwd(p[i], 12, 3)); 
        System.out.println("  "+n);
    }

    static void zeroStats(double s[][])
    {
        for (int j=0; j<2; j++)
          for (int k=0; k<NPARMS; k++)
            s[j][k] = 0.0; 
    }

    static void setNoisyData(double c[], double e[], double n[])
    // Adds Gaussian random errors to create noisy data.
    {
        for (int i=0; i<NPTS; i++)
          n[i] = c[i] + e[i]*U.gauss();
    } 

    static void initParms(double p[], double q[])
    // Start fitting process with true parms
    {
        for (int j=0; j<NPARMS; j++)
          q[j] = p[j];
    }

    static int solve(double parms[], double data[], double errors[])
    // Starts with given parms, exits with best fit parms. 
    {
        LM myLM = new LM(NPTS, NPARMS, data, errors, parms); 
        int n = myLM.lmfit();
        myLM.getparms(parms); 
        return n; 
    }

    static void addStats(double p[], double s[][])
    // Adds given parms and their squares to build statistics
    {
        for (int j=0; j<NPARMS; j++)
        {
           s[0][j] += p[j]; 
           s[1][j] += p[j]*p[j]; 
        }
    }

    static void finishStats(double s[][])
    {
        for (int j=0; j<NPARMS; j++)
        {
            double sum = s[0][j]; 
            double sumsq = s[1][j]; 
            s[0][j] = sum/NREPS;      // put average on row zero.
            double var = sumsq/(NREPS-1) - sum*sum/(NREPS*(NREPS-1)); 
            s[1][j] = Math.sqrt(var); // put rms on row one.
        }
    } 
}




class F  implements Constants
//-----------------function model----------------------
//----must remain visible to caller and LM classes-----
{
    static double funci(int i, double p[])
    // samples funcxy() for pixel "i"
    {
        double x = i % NX;   // convert 1D pixel number to x
        double y = i / NY;   // convert 1D pixel number to y
        return funcxy(x,y, p); 
    }

    static double funcxy(double x, double y, double p[])
    // Returns the model prediction for a given pixel & model parms.
    //  p[0] = total integrated signal level
    //  p[1] = Xcentroid of Gaussian
    //  p[2] = Ycentroid of Gaussian 
    // x,y is the center of a pixel of interest.
    // RMS breadth = 4.0 held fixed not fitted.  But it could be. 
    // background = 10000 held fixed not fitted.  But it could be.
    {
        double background = 10000;   
        double rmswidth = 4.0; 
        double numer = (x-p[1])*(x-p[1]) + (y-p[2])*(y-p[2]); 
        double denom = rmswidth*rmswidth;
        double ratio = numer/denom; 
        if (ratio < 30.0)  // near the peak of the Gaussian?
          return background + (0.39894*p[0]/rmswidth)*Math.exp(-0.5*ratio); 
        else               // far from the peak of the Gaussian?
          return background; 
    }
}








class LM implements Constants  /// Levenberg-Marquardt fitting routine
{
    static final double DELTAP     = 1E-6;     // step size for num diff
    static final int    LMITER     =   20;     // max number of L-M iterations
    static final double LMBOOST    =  2.0;     // damping increase per failed step
    static final double LMSHRINK   = 0.10;     // damping decrease per successful step
    static final double LAMBDAZERO =  0.01;    // initial damping
    static final double LAMBDAMAX  =  100;     // max damping
    static final double LMTOL      = 1E-6;     // end point exit tolerance

    private int NPTS, NPARMS, NADJ; 
    private int which[]; 
    private double data[];
    private double error[];
    private double parms[];

    void getparms(double p[])
    // delivers final parms to user
    {
       for (int i=0; i<NPARMS; i++)
         p[i] = parms[i]; 
    }

    public LM(int gnpts, int gnparms, double gdata[], double gerror[], double gparms[])
    // this constructor brings all client info into new object.
    {
        // System.out.println("LM constructor 1 starting"); 
        NPTS = gnpts; 
        NADJ = NPARMS = gnparms; 
        data = new double[NPTS];
        for (int i=0; i<NPTS; i++)
          data[i] = gdata[i]; 
        error = new double[NPTS]; 
        for (int i=0; i<NPTS; i++)
          error[i] = gerror[i];
        parms = new double[NPARMS]; 
        for (int i=0; i<NPARMS; i++)
          parms[i] = gparms[i];
    }

    double getsos(double parms[] )
    // returns sum of squares for any parm vector
    {
        double sum=0.0, term;
        for (int i=0; i<NPTS; i++)
        {
           term = (data[i] - F.funci(i, parms)) / error[i];
           sum += term*term;
        }
        return sum;
    }

    void showparms(int niter, double sos)
    {
        System.out.print(U.fwi(niter, 6)+"  sos="+U.fwd(sos,16,4)+"   ");
        for (int i=0; i<NPARMS; i++)
          System.out.print(U.fwd(parms[i],12,4));
        System.out.println(); 
    }

    int lmfit( )
    // global NPARMS, NADJ, NPTS, parms[].
    // F.func(i,parms) is the model, compare against data[] with error[]
    // Outer loop continues until improvements level off;
    // Inner loop adjusts damping to assure downhill step.
    // Returns exit value of sos.
    // Ref: M.Lampton, Computers in Physics v.11 pp.110-115 1997.
    {
        int    i, j, k, iter=0;
        double sos, sosprev, rise, lambda=LAMBDAZERO;
        double[] testp   = new double[NPARMS];       // trial parm
        double[] beta    = new double[NADJ];         // downhill gradient
        double[][] alpha = new double[NADJ][NADJ];
        double[][] a     = new double[NADJ][NADJ];        

        sos = getsos(parms);

        do  //// succession of improved locations
        {
           iter++;
           sosprev = sos;                  // save previous sos for comparison
           // showparms(iter, sos); 
           for (j=0; j<NADJ; j++)          // get downhill gradient
           {
              beta[j] = 0.0;
              for (i=0; i<NPTS; i++)
                beta[j] -= (F.funci(i,parms)-data[i])*dfdp(i,j,parms)/U.SQR(error[i]);
           }
           for (j=0; j<NADJ; j++)          // get alpha
             for (k=0; k<NADJ; k++)
             {
                alpha[j][k] = 0.0;
                for (i=0; i<NPTS; i++)
                  alpha[j][k] += dfdp(i,j,parms)*dfdp(i,k,parms)/U.SQR(error[i]);
             }

           do  /// loop adjusts damping to get a downhill step
           {
              // System.out.println("trying lambda = "+U.fwd(lambda, 18, 4)); 
              for (j=0; j<NADJ; j++)       // copy and damp it
                for (k=0; k<NADJ; k++)
                  a[j][k] = alpha[j][k] + ((j==k) ? lambda : 0.0);
              gaussj(a, NADJ);             // invert
              for(j=0; j<NPARMS; j++)      // compute testparms
                testp[j] = parms[j];
              for (j=0; j<NADJ; j++)       // modify only the adjustables
                for (k=0; k<NADJ; k++)
                  testp[j] += a[j][k]*beta[k];

              sos = getsos(testp);         // try it out.
              rise = sos - sosprev;

              if (rise <= 0.0)             // good step!
              {
                 for (j=0; j<NPARMS; j++)  // adopt testp[]
                   parms[j] = testp[j];
                 lambda *= LMSHRINK;       // shrink lambda
                 break;
              }
              if (rise < LMTOL)            // finished but keep prev parms
                break;                  
              lambda *= LMBOOST;           // else try more damping.
           } while ((lambda < LAMBDAMAX) && (iter<LMITER));
        } while ((rise<-LMTOL) && (iter<LMITER));
        return iter;
    }

    double dfdp(int i, int j, double parms[])
    // Two sided derivative of F.func(i,p) w.r.t. any given parm[j]
    // Gives one jacobian matrix element dfi/dpj at given parmvector.
    {
        double pminus[] = new double[NPARMS];
        double pplus[]  = new double[NPARMS];
        for (int k=0; k<NPARMS; k++)
          pminus[k] = pplus[k] = parms[k];
        pminus[j] -= DELTAP;
        pplus[j] += DELTAP;
        return (F.funci(i,pplus) - F.funci(i,pminus))/(2.0*DELTAP);
    }

    double gaussj( double a[][], int N )
    // inverts the double array a[N][N] by Gauss-Jordan method
    {
        double det = 1.0, big, save;
        int i,j,k,L;
        int[] ik = new int[100];
        int[] jk = new int[100];
        for (k=0; k<N; k++)
        {
            big = 0.0;
            for (i=k; i<N; i++)
              for (j=k; j<N; j++)          // find biggest element
                if (Math.abs(big) <= Math.abs(a[i][j]))
                {
                    big = a[i][j];
                    ik[k] = i;
                    jk[k] = j;
                }
            if (big == 0.0) return 0.0;
            i = ik[k];
            if (i>k)
              for (j=0; j<N; j++)          // exchange rows
              {
                  save = a[k][j];
                  a[k][j] = a[i][j];
                  a[i][j] = -save;
              }
            j = jk[k];
            if (j>k)
              for (i=0; i<N; i++)
              {
                  save = a[i][k];
                  a[i][k] = a[i][j];
                  a[i][j] = -save;
              }
            for (i=0; i<N; i++)            // build the inverse
              if (i != k)
                a[i][k] = -a[i][k]/big;
            for (i=0; i<N; i++)
              for (j=0; j<N; j++)
                if ((i != k) && (j != k))
                  a[i][j] += a[i][k]*a[k][j];
            for (j=0; j<N; j++)
              if (j != k)
                a[k][j] /= big;
            a[k][k] = 1.0/big;
            det *= big;                    // bomb point
        }                                  // end k loop
        for (L=0; L<N; L++)
        {
            k = N-L-1;
            j = ik[k];
            if (j>k)
              for (i=0; i<N; i++)
              {
                  save = a[i][k];
                  a[i][k] = -a[i][j];
                  a[i][j] = save;
              }
            i = jk[k];
            if (i>k)
              for (j=0; j<N; j++)
              {
                  save = a[k][j];
                  a[k][j] = -a[i][j];
                  a[i][j] = save;
              }
        }
        return det;
    }
}









class U  /////utilities ////////////////////
{
    static double SQR(double x)
    {
        return x*x;
    }

    static double gauss()
    {
        double sum = 0.0; 
        for (int i=0; i<12; i++)
           sum += Math.random(); 
        return sum - 6.0;
    }

    static String fwi(int n, int w)
    // converts an int to a string with given width.
    {
        String s = Integer.toString(n); 
        while (s.length() < w)
           s = " " + s; 
        return s;
    }


    static String fwd(double x, int w, int d)
    // converts a double to a string with given width and decimals.
    {
        java.text.DecimalFormat df = new DecimalFormat();
        df.setMaximumFractionDigits(d); 
        df.setMinimumFractionDigits(d);
        df.setGroupingUsed(false);   
        String s = df.format(x); 
        while (s.length() < w)
          s = " " + s;  
        if (s.length() > w)
        {
            s = "";
            for (int i=0; i<w; i++)
              s = s + "-";
        }  
        return s; 
    }

    static String fwe(double x)
    // converts a double into an E notation string
    {
        java.text.NumberFormat ef = new DecimalFormat("0.0E0"); 
        String s = ef.format(x);    
        while (s.length() < 8)
          s = " " + s; 
        return s; 
    }

    static String fws(String s, int len)
    // enlarges a String so that it occupies "len" characters
    {
        while (s.length() < len)
          s = s + " ";
        return s; 
    }

    static double log10(double arg)
    {
        if (arg>0.0)
          return Math.log(arg)/2.302585092994046;
        return -999.9;
    }

    static double dex(double arg)
    {
        if (arg > 300.0)
           arg = 300.0; 
        if (arg < -300.0)
           arg = -300.0; 
        return Math.pow(10.0, arg); 
    }

    static void beep()
    {
        char c = 7; 
        String s = ""+c; 
        System.out.print(s); 
    }
}