package edu.cmu.sphinx.decoder.adaptation;

import java.io.File;
import java.io.PrintWriter;
import java.util.Scanner;

import org.apache.commons.math3.linear.Array2DRowRealMatrix;
import org.apache.commons.math3.linear.ArrayRealVector;
import org.apache.commons.math3.linear.DecompositionSolver;
import org.apache.commons.math3.linear.LUDecomposition;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.linear.RealVector;

import edu.cmu.sphinx.linguist.acoustic.tiedstate.Sphinx3Loader;

public class Transform {

    private float[][][][] As;
    private float[][][] Bs;
    private Sphinx3Loader loader;
    private int nrOfClusters;

    public Transform(Sphinx3Loader loader, int nrOfClusters) {
        this.loader = loader;
        this.nrOfClusters = nrOfClusters;
    }

    /**
     * Used for access to A matrix.
     * 
     * @return A matrix (representing A from A*x + B = C)
     */
    public float[][][][] getAs() {
        return As;
    }

    /**
     * Used for access to B matrix.
     * 
     * @return B matrix (representing B from A*x + B = C)
     */
    public float[][][] getBs() {
        return Bs;
    }

    /**
     * Writes the transformation to file in a format that could further be used
     * in Sphinx3 and Sphinx4.
     * 
     * @param filePath
     *            path to store transform matrix
     * @param index
     *            index of transform to store
     * @throws Exception
     *             if something went wrong
     */
    public void store(String filePath, int index) throws Exception {
        PrintWriter writer = new PrintWriter(filePath, "UTF-8");

        // nMllrClass
        writer.println("1");
        writer.println(loader.getNumStreams());

        for (int i = 0; i < loader.getNumStreams(); i++) {
            writer.println(loader.getVectorLength()[i]);

            for (int j = 0; j < loader.getVectorLength()[i]; j++) {
                for (int k = 0; k < loader.getVectorLength()[i]; ++k) {
                    writer.print(As[index][i][j][k]);
                    writer.print(" ");
                }
                writer.println();
            }

            for (int j = 0; j < loader.getVectorLength()[i]; j++) {
                writer.print(Bs[index][i][j]);
                writer.print(" ");

            }
            writer.println();

            for (int j = 0; j < loader.getVectorLength()[i]; j++) {
                writer.print("1.0 ");

            }
            writer.println();
        }
        writer.close();
    }

    /**
     * Used for computing the actual transformations (A and B matrices). These
     * are stored in As and Bs.
     */
    private void computeMllrTransforms(double[][][][][] regLs, double[][][][] regRs) {
        int len;
        DecompositionSolver solver;
        RealMatrix coef;
        RealVector vect, ABloc;

        for (int c = 0; c < nrOfClusters; c++) {
            this.As[c] = new float[loader.getNumStreams()][][];
            this.Bs[c] = new float[loader.getNumStreams()][];

            for (int i = 0; i < loader.getNumStreams(); i++) {
                len = loader.getVectorLength()[i];
                this.As[c][i] = new float[len][len];
                this.Bs[c][i] = new float[len];

                for (int j = 0; j < len; ++j) {
                    coef = new Array2DRowRealMatrix(regLs[c][i][j], false);
                    solver = new LUDecomposition(coef).getSolver();
                    vect = new ArrayRealVector(regRs[c][i][j], false);
                    ABloc = solver.solve(vect);

                    for (int k = 0; k < len; ++k) {
                        this.As[c][i][j][k] = (float) ABloc.getEntry(k);
                    }

                    this.Bs[c][i][j] = (float) ABloc.getEntry(len);
                }
            }
        }
    }

    /**
     * Read the transformation from a file
     * 
     * @param filePath
     *            file path to load transform
     * @throws Exception
     *             if something went wrong
     */
    public void load(String filePath) throws Exception {

        Scanner input = new Scanner(new File(filePath));
        int numStreams, nMllrClass;
        int[] vectorLength = new int[1];

        nMllrClass = input.nextInt();

        assert nMllrClass == 1;

        numStreams = input.nextInt();

        this.As = new float[nMllrClass][][][];
        this.Bs = new float[nMllrClass][][];

        for (int i = 0; i < numStreams; i++) {
            vectorLength[i] = input.nextInt();

            int length = vectorLength[i];

            this.As[0] = new float[numStreams][length][length];
            this.Bs[0] = new float[numStreams][length];

            for (int j = 0; j < length; j++) {
                for (int k = 0; k < length; ++k) {
                    As[0][i][j][k] = input.nextFloat();
                }
            }

            for (int j = 0; j < length; j++) {
                Bs[0][i][j] = input.nextFloat();
            }
        }
        input.close();
    }

    /**
     * Stores in current object a transform generated on the provided stats.
     * 
     * @param stats
     *            provided stats that were previously collected from Result
     *            objects.
     */
    public void update(Stats stats) {
        stats.fillRegLowerPart();
        As = new float[nrOfClusters][][][];
        Bs = new float[nrOfClusters][][];
        this.computeMllrTransforms(stats.getRegLs(), stats.getRegRs());
    }
}

@@ -4195,49 +4195,8 @@
 ass;
-%0A        int%5B%5D vectorLength = new int%5B1%5D;
 %0A%0A  
@@ -4344,16 +4344,26 @@
 rClass%5D%5B
+numStreams
 %5D%5B%5D%5B%5D;%0A 
@@ -4401,16 +4401,26 @@
 rClass%5D%5B
+numStreams
 %5D%5B%5D;%0A%0A  
@@ -4476,31 +4476,26 @@
         
-vectorL
+int l
 ength
-%5Bi%5D
  = input
@@ -4511,51 +4511,8 @@
 );%0A%0A
-            int length = vectorLength%5Bi%5D;%0A%0A
     
@@ -4525,24 +4525,27 @@
   this.As%5B0%5D
+%5Bi%5D
  = new float
@@ -4541,36 +4541,24 @@
 = new float%5B
-numStreams%5D%5B
 length%5D%5Bleng
@@ -4584,16 +4584,19 @@
 is.Bs%5B0%5D
+%5Bi%5D
  = new f
@@ -4600,28 +4600,16 @@
 w float%5B
-numStreams%5D%5B
 length%5D;
@@ -4701,19 +4701,19 @@
 length; 
+k
 ++
-k
 ) %7B%0A    
@@ -4788,33 +4788,32 @@
 %7D%0A            %7D%0A
-%0A
             for 
@@ -4898,32 +4898,172 @@
 ;%0A            %7D%0A
+            for (int j = 0; j %3C length; j++) %7B%0A                // Skip MLLR variance scale%0A                input.nextFloat();%0A            %7D%0A
         %7D%0A      

Load MLLR transforms with only cluster properly