package yi.stat.em;

import java.util.*;
import org.jmat.data.*;
import java.io.*;
import java.awt.*;
import yi.ds.*;
import yi.stat.hist.*;
import yi.util.io.*;
import org.jmat.data.matrixDecompositions.*;

/**
 * EMCumulator implements a Gaussian Model used in EM method.  The model is not
 * represented directly by (\mu, \Sigma), but \mu and \Sigma can be calculated
 * from the data added in.  Each added in data comes with a probability.  When
 * required, \mu and \Sigma will be calculated on fly.  The code is tuned that
 * \mu and \Sigma will only be calculated when necessary.
 *
 * \mu = \sum(p*x) / \sum(p);
 * \Sigma = \sum(p*(x-\mu)*(x-\mu)^T) / \sum(p)
 */

public class GaussianComponent  implements MixtureComponential, Serializable{
  static final long serialVersionUID = -1111502538601444719L;
  /**
   * Save the data
   */
  transient LinkedList m_data = new LinkedList();

  /**
   * Decide whether to use the cached data (mean)
   */
  boolean m_meanIsFreash = false;

  /**
   * Decide whether to use the cached data (covariance)
   */
  boolean m_covarianceIsFreash = false;

  /**
   * Cache the mean
   */
  double[] m_mean = null;

  /**
   * Cache the covariance matrix
   */
  double[][] m_covariance = null;

  /**
   * Cache the inverse of the covariance matrix
   */
  AbstractMatrix m_covarianceInverse = null;
  /**
   * Cache the denominator of the density function
   */
  double m_densityDenominator;
  boolean m_covarianceInverseIsFresh = false;
  boolean m_densityDenominatorIsFresh = false;

  SingularGaussianModelFixable m_singularFixer = null;

  /**
   * The covariance matrix determinant should be greater than 0.  However, when
   * it is very close to 0, the covariance matrix may have none inverse already.
   * So when the determiant is less than COV_MRT_LOW_LIMIT, it will be treated
   * as 0.
   */
//  public static double COV_MRT_LOW_LIMIT = 1e-30;
  public static double COV_MRT_LOW_LIMIT = 1e-100;

  public GaussianComponent() {
  }

  public void setSigularFixer(SingularGaussianModelFixable fixer) {
    m_singularFixer = fixer;
  }

  public void add(double[] data, double p) {
    m_data.add(new ProbabilityData( (double[]) data, p));
    m_meanIsFreash = m_covarianceIsFreash = false;
    m_covarianceInverseIsFresh = m_densityDenominatorIsFresh = false;
  }

  int getDataLength() {
    if (m_data == null || m_data.size() <= 0) {
      if (m_mean != null) return m_mean.length;
      if (m_covariance != null) return m_covariance.length;
      return 0;
    }
    else {
      ProbabilityData pdata = (ProbabilityData) m_data.getFirst();
      return pdata.data.length;
    }
  }

  /**
   * Learn the model.  For a Gaussian model, there is no need to learn at all.
   * Just to calculate the means and covariance matrixes.  However, should check
   * the integrity of the component.
   * @throws ModelSingularException
   */
  public void learnTheModel() throws ModelSingularException{
    double covinv = 0;
    try {
      getMeans();
      getCovariance();
      covinv = getDensityDenominator();
      getCovarianceInverse();
    } catch (ModelSingularException e) {
//      Iterator i = m_data.iterator();
//      while (i.hasNext()) {
//        ProbabilityData d = (ProbabilityData)i.next();
//        if (d.probability < 1e-6) continue;
//        else {
//          for (int j = 0; j < d.data.length; j++) {
//            if (j != 0)
//              System.err.print(' ');
//            System.err.print(d.data[j]);
//          }
//          System.err.println(" : ["+ d.probability + "]");
//        }
//      }
      throw e ;
    }
  }

  /**
   * For debug only
   */
  public double m_probSum = 0;

  /**
   * Set the means of the Gaussian.
   * @param means
   * @return
   */
  public void setMeans(double[] mean) {
    m_mean = mean;
    m_meanIsFreash = true;
  }

  public double[] getMeans() throws GaussianModelSingularException {

    if (m_meanIsFreash)
      return m_mean;

    int data_length = getDataLength();

    double[] means = new double[data_length];
    double prob_sum = 0;
    Iterator iter = m_data.iterator();
    while (iter.hasNext()) {
      ProbabilityData pdata = (ProbabilityData) iter.next();
      for (int i = 0; i < data_length; i++) {
        means[i] += pdata.data[i] * pdata.probability;
      }
      prob_sum += pdata.probability;
    }

    if (prob_sum == 0) {
      throw new GaussianModelSingularException("probability sum = 0");
    }
    m_probSum = prob_sum;

    for (int i = 0; i < data_length; i++) {
      means[i] /= prob_sum;
      if (Double.isNaN(means[i])) {
        double[] means0 = new double[data_length];
        double prob_sum0 = 0;
        Iterator iter0 = m_data.iterator();
        while (iter0.hasNext()) {
          ProbabilityData pdata = (ProbabilityData) iter0.next();
          for (int i0 = 0; i0 < data_length; i0++) {
            means0[i] += pdata.data[i0] * pdata.probability;
            System.err.print(" " + pdata.data[i0]);
          }
          prob_sum0 += pdata.probability;
          System.err.println(" /" + pdata.probability + ", prob_sum:" + prob_sum0);
        }
      }
    }

    // cache the result
    m_mean = (double[]) means.clone();
    m_meanIsFreash = true;

    return means;
  }

  /**
   * Return the max log(prob) on this Gaussian.  density on "mean" is the answer
   * @return
   * @throws GaussianModelSingularException
   */
  public double getMaxProb() throws GaussianModelSingularException{
    return getDensity(getMeans());
  }

  /**
   * Return the min log(prob) on this Gaussian.  Lower than that will be
   * treated as prob=0
   * @param portion how much the portion of the training data are considered
   *        as belonging to this Gaussian.  The prob(component) is a good
   *        estimation for this number
   * @return
   * @throws GaussianModelSingularException
   */
  public double getMinProb(double portion) throws GaussianModelSingularException{
    double[] probs = new double[m_data.size()];

    int i_index = 0;
    for (Iterator iter = m_data.iterator(); iter.hasNext(); i_index++) {
      ProbabilityData pdata = (ProbabilityData) iter.next();
      probs[i_index] = getDensity(pdata.data);
    }
    Arrays.sort(probs);
    int good_candidate = (int)(m_data.size() * portion);
    if (good_candidate == 0) good_candidate = 1; // at least one good
    return probs[probs.length - good_candidate];
  }

  public int getCount() {
    return m_data.size();
  }

  /**
   * Set the covariance
   * @param covariance
   */
  public void setCovariance(double[][] covariance) {
    m_covariance = covariance;
    m_covarianceIsFreash = true;
  }

  protected boolean fixCovarinceMatrix(){

    if (m_singularFixer == null) return true;

    // fix zero entry in the covariance matrix
    double[][] data_to_model = new double[m_data.size()][];
    int i_index = 0;
    for (Iterator i= m_data.iterator(); i.hasNext(); i_index++) {
      ProbabilityData pdata = (ProbabilityData) i.next();
      data_to_model[i_index] = pdata.data;
    }
    return m_singularFixer.fix(m_covariance, data_to_model, getProbabilitySum());
  }

  public double[][] getCovariance() throws GaussianModelSingularException {
    if (m_covarianceIsFreash)
      return m_covariance;

    int data_length = getDataLength();
    double[] means = getMeans();

    Matrix matrix_means = new Matrix(means, means.length);

    // accumulator
    double prob_sum = 0;
    Matrix matrix_sum = new Matrix(data_length, data_length);

    // loop through each vector
    Iterator iter = m_data.iterator();
    while (iter.hasNext()) {
      ProbabilityData pdata = (ProbabilityData) iter.next();
      // the data
      Matrix matrix_data = new Matrix(pdata.data, data_length);
      // data - mean
      AbstractMatrix matrix_diff = matrix_data.minus(matrix_means);
      // (data - mean)'
      AbstractMatrix matrix_diff_transpose = matrix_diff.transpose();
      // (data - mean)(data - mean)'prob
      AbstractMatrix matrix_times = matrix_diff.times(matrix_diff_transpose).times(
          pdata.probability);
      // acuumulators
      matrix_sum.plusEquals(matrix_times);
      prob_sum += pdata.probability;
    }

    if (prob_sum == 0)
      throw new GaussianModelSingularException("probability sum = 0");
    // get the average of accumulator
    AbstractMatrix matrix_covariance = matrix_sum.divide(prob_sum);

    m_covariance = matrix_covariance.getArrayCopy();

    // fix zero entries
    if (!fixCovarinceMatrix()) {
      throw new GaussianModelSingularException(
          "zero entry in covariance matrix can't be fixed");
    }

    m_covarianceIsFreash = true;
    return m_covariance;
  }

  /**
   * Get sqrt(Sigma) which is useful for standandlize the distribution
   * @return
   * @throws GaussianModelSingularException
   */
  public double[][] getSquareRootOfCovariance() throws GaussianModelSingularException {

    double[][] cov = getCovariance();

    // eigenvalue decomposition
    Matrix m = new Matrix(cov);
    EigenvalueDecomposition e = m.eig();
    AbstractMatrix ed = e.getD();
    AbstractMatrix ev = e.getV();

    // squared root of eigenvalues
    AbstractMatrix root_ed = ed.ebePower(0.5);
    AbstractMatrix root_cov = ev.times(root_ed.times(ev.transpose()));
    return root_cov.getArrayCopy();
  }

  /**
   * Standardize the data (so that the mean of distribution is 0 and the
   * dev is 1)
   * @param data
   * @return
   */
  public double[] standardize(double[] data) throws GaussianModelSingularException{
    double[] mean = getMeans();

    // make zero centered
    double[] zero_mean = new double[data.length];
    for (int i=0; i<data.length; i++) {
      zero_mean[i] = data[i] - mean[i];
    }

    Matrix root_cov = new Matrix(getSquareRootOfCovariance());
    double[] norm_data = (new Matrix(zero_mean, 1)).divide(root_cov).getRowArrayCopy(
        0);
    return norm_data;
  }


//  public final static double SCRAMBLE_SMALL_NUMBER = 1e-5;
//  /**
//   * Sometimes, the rank of the covariance matrix is not full, so inverse is not
//   * avaaible.  (in other words, the features are not independent)  In that
//   * case, scramble the covariance matrix a little to make them independent
//   * @param covariance
//   */
//  private void fixCovarianceMatrixRank(Matrix covariance) {
//    int size = covariance.getColumnDimension();
//    Random r = new Random(0);
//    for (int i=0; i<size; i++) {
//      covariance.set(i, i, covariance.get(i, i) + SCRAMBLE_SMALL_NUMBER);
//    }
//  }

  AbstractMatrix getCovarianceInverse() throws GaussianModelSingularException {
    if (m_covarianceInverseIsFresh)
      return m_covarianceInverse;
    double[][] covariance = getCovariance();
    AbstractMatrix mx_covariance = new Matrix(covariance);
    m_covarianceInverse = null;
    try {
      m_covarianceInverse = mx_covariance.inverse();
    }
    catch (Exception e) {
      System.err.println(dataToString());
      throw new GaussianModelSingularException(e.toString());
    }
    m_covarianceInverseIsFresh = true;
    return m_covarianceInverse;
  }

  public double getDensityDenominator() throws GaussianModelSingularException {
    if (m_densityDenominatorIsFresh)
      return m_densityDenominator;
    int data_length = getDataLength();
    double[][] covariance = getCovariance();
    Matrix mx_covariance = new Matrix(covariance);
    double covariance_determinant = mx_covariance.determinant();

    if (covariance_determinant < COV_MRT_LOW_LIMIT) {
      throw new GaussianModelSingularException("Covariance Matrix Determinant = " + covariance_determinant);
    }

    m_densityDenominator = Math.sqrt(Math.pow(2 * Math.PI, data_length) *
                                     covariance_determinant);
    m_densityDenominatorIsFresh = true;
    return m_densityDenominator;
  }

  /**
   * Calcuate the model parameters.  Since those parameters will be cached, so
   * this calculation is not wasting time
   */
  public void CalculateModelParameters() throws GaussianModelSingularException {
    getMeans();
    getCovariance();
    getCovarianceInverse();
    getDensityDenominator();
  }

  /**
   * Return the probability of the data given the model represented by the
   * EMCumulator.
   * @param data
   * @return
   */
  public double getDensity(double[] data) throws GaussianModelSingularException {
    int data_length = data.length;
    Matrix mx_data = new Matrix(data, data_length);

    double[] mean = getMeans();
    Matrix mx_mean = new Matrix(mean, data_length);

    double[][] covariance = getCovariance();
    if (covariance == null) {
      System.err.println("covariance == null.  Probability is assigned to 1");
      return 0;
    }
    Matrix mx_covariance = new Matrix(covariance);

    AbstractMatrix mx_diff = mx_data.minus(mx_mean);
    AbstractMatrix mx_diff_trans = mx_diff.transpose();
    AbstractMatrix mx_covariance_inv = getCovarianceInverse();
    AbstractMatrix mx_up = mx_diff_trans.times(mx_covariance_inv).times(mx_diff);
    double up = -mx_up.toDouble() / 2;
    double densityDenominator = getDensityDenominator();
    if (up-Math.log(densityDenominator) > Math.log(Double.MAX_VALUE)) return Double.MAX_VALUE;
    return Math.exp(up) / densityDenominator;
  }

  /**
   * Construct the histogram on one dimension
   * @param dim
   * @param min
   * @param max
   * @param bin_num
   * @return
   */
  public double[] getHistogram(int dim, double min, double max, int bin_num) {
    WeightedHistogram hist = new WeightedHistogram(bin_num);
    double bin_width = (max-min)/bin_num;
    for (Iterator i= m_data.iterator(); i.hasNext(); ) {
      ProbabilityData dp = (ProbabilityData)i.next();
      double d = dp.data[dim];
      double p = dp.probability;
      int bin_index = (int)((d - min) / bin_width);
      if (bin_index < 0) bin_index = 0;
      if (bin_index >= bin_num) bin_index = bin_num - 1;
      hist.putToBin(bin_index, p);
    }
    return hist.getWeightedHistogram();
  }

  /**
   * Return the probabilities of the data under the different models represented
   * by the EMCumulators.  This is the E-Step in EM algorithm
   */
  public static HashMap getProbabilities(HashMap cumulators,
                                         HashMap model_probs, double[] data) {
    Iterator iter = cumulators.keySet().iterator();
    HashMap probs = new HashMap();
    double prob_sum = 0;
    while (iter.hasNext()) {
      Object key = iter.next();
      GaussianComponent cumulator = (GaussianComponent) cumulators.get(key);
      try {
        double prob = cumulator.getDensity(data);
        double model_prob = ( (Double) model_probs.get(key)).doubleValue();
        probs.put(key, new Double(prob * model_prob));
        prob_sum += prob * model_prob;
      }
      catch (GaussianModelSingularException e) {
        e.printStackTrace();
      }
    }

    // normalize
    iter = probs.keySet().iterator();
    while (iter.hasNext()) {
      Object key = iter.next();
      double prob = ( (Double) probs.get(key)).doubleValue();
      probs.put(key, new Double(prob / prob_sum));
    }
    return probs;
  }

  /**
   * Create EM cumulators for an array of data.  data[i] is the i-th data and
   * labels[i] is the corresponding index.  This function is useful for
   * construct multiple EM cumulators for a segmented image in which each pixel
   * has its corresponding index and the data is its color feature or texture
   * feature.
   * @param data
   * @param labels
   * @return
   */
  public static GaussianComponent[] createEMCumulators(double[][] data, int[] labels, SingularGaussianModelFixable fixer) {
    int label_num = yi.stat.StatUtil.maximum(labels) + 1;
    GaussianComponent[] cumulators = new GaussianComponent[label_num];
    for (int i = 0; i < label_num; i++) {
      cumulators[i] = new GaussianComponent();
      cumulators[i].setSigularFixer(fixer);
    }

    // calculate the sum and count for every region
    for (int i = 0; i < data.length; i++) {
      if (labels[i] < 0)
        continue;
      cumulators[labels[i]].add(data[i], 1);
    }

    return cumulators;
  }


  /**
   * Return the mean of the probability of the data
   * @return
   */
  public double getProbabilitySum(){
    double prob_sum = 0;
    Iterator iter = m_data.iterator();
    while (iter.hasNext()) {
      ProbabilityData pdata = (ProbabilityData) iter.next();
      prob_sum += pdata.probability;
    }
    if (m_data.size() == 0) {
      return 0;
    }
    return prob_sum;

  }


  public String toString() {
    StringBuffer str = new StringBuffer();

    // output means
    try {
      double[] means = getMeans();
      Matrix mx_means = new Matrix(means, 1);
      str.append("Mean:\n" + mx_means);

      /** @todo remove the output color: this is for debug color clustering only */
      if (means.length >=3) {
        Color color = new Color(Color.HSBtoRGB( (float) means[0],
                                               (float) means[1],
                                               (float) means[2]));
        str.append("RGB(" + color.getRed() + "," + color.getGreen() + "," +
                   color.getBlue() + ")\n");
      }
    }
    catch (GaussianModelSingularException e) {
      str.append("Mean not available: " + e.toString() + "\n");
    }

    // output covariance
    try {
      double[][] covariance = getCovariance();
      Matrix mx_covariance = new Matrix(covariance);
      str.append("Covariance:\n" + mx_covariance);

    }
    catch (GaussianModelSingularException e) {
      str.append("Covariance not available: " + e.toString() + "\n");
    }

    return str.toString();
  }

  public String dataToString() {
    StringBuffer str = new StringBuffer();
    Iterator i = m_data.iterator();
    while (i.hasNext()) {
      str.append(i.next().toString() + '\n');
    }
    return str.toString();
  }

  /**
   * Play with Gaussian distribution
   */
  public static class GaussianExperiment {
    public static void main(String[] args) {
      double[][] data = {
          {0,0},
          {1,1.5},
          {0.5,2},
          {-0.5,-3},
          {-1,-4},
      };
      int[] labels = new int[data.length];
      Arrays.fill(labels, 0);
      GaussianComponent[] gs = createEMCumulators(data, labels, null);
      GaussianComponent g=gs[0];
//
//      AlignedGaussianComponent g = new AlignedGaussianComponent();
//      for (int i=0; i<data.length; i++) {
//        g.add(data[i], 1);
//      }
      try {
        double[] mean = g.getMeans();
        OutputResult.print(System.out, mean);

        double[][] cov = g.getCovariance();
        OutputResult.print(System.out, cov);

        // eigenvalue decomposition
        Matrix m = new Matrix(cov);
        EigenvalueDecomposition e = m.eig();
        AbstractMatrix ed = e.getD();
        AbstractMatrix ev = e.getV();
        System.out.println(ed);

        System.out.println(ev);

        // squared root of eigenvalues
        AbstractMatrix root_ed = ed.ebePower(0.5);
        AbstractMatrix root_cov = ev.times(root_ed.times(ev.transpose()));

        System.out.println(root_cov);

        double[][] norm_data = new double[data.length][];
        for (int i=0; i<data.length; i++) {
          norm_data[i] = g.standardize(data[i]);
        }
        OutputResult.print(System.out, norm_data);

        gs = createEMCumulators(norm_data, labels, null);
        GaussianComponent norm_g=gs[0];
//        AlignedGaussianComponent norm_g = new AlignedGaussianComponent();
//        for (int i=0; i<data.length; i++) {
//          norm_g.add(norm_data[i], 1);
//        }

        double[] norm_mean = norm_g.getMeans();
        OutputResult.print(System.out, norm_mean);

        cov = norm_g.getCovariance();
        OutputResult.print(System.out, cov);

        System.out.println(g.getDensity(mean));
        System.out.println(norm_g.getDensity(norm_mean));
        System.out.println(g.getDensity(mean)/norm_g.getDensity(norm_mean));
        for (int i=0; i<data.length; i++) {
          System.out.println(g.getDensity(data[i]));
          System.out.println(norm_g.getDensity(norm_data[i]));
          System.out.println(g.getDensity(data[i])/norm_g.getDensity(norm_data[i]));
        }

        // combine to 1-D data
        {
          GaussianComponent g1 = new GaussianComponent();
          double[] mean1 = {
              0, 0};
          g1.setMeans(mean1);
          double[][] cov1 = {
              {
              1, 0}
              , {
              0, 1}
          };
          g1.setCovariance(cov1);
          double[] d0 = {
              0, 0};
          System.out.println("p(0,0) =" + g1.getDensity(d0));
          double[] d1 = {
              1, 0};
          System.out.println("p(1,0) =" + g1.getDensity(d1) + ", ratio = " + (g1.getDensity(d1)/g1.getDensity(d0)));
          double[] d2 = {
              0, 1};
          System.out.println("p(0,1) =" + g1.getDensity(d2) + ", ratio = " + (g1.getDensity(d2)/g1.getDensity(d0)));
          double[] d3 = {
              1, 1};
          System.out.println("p(1,1) =" + g1.getDensity(d3) + ", ratio = " + (g1.getDensity(d3)/g1.getDensity(d0)));
        }

        {
          GaussianComponent g1 = new GaussianComponent();
          double[] mean1 = {
              0};
          g1.setMeans(mean1);
          double[][] cov1 = {
              {1}};
          g1.setCovariance(cov1);
          double[] d0 = {
              0};
          System.out.println("p(0) =" + g1.getDensity(d0));
          double[] d1 = {
              Math.sqrt(2)};
          System.out.println("p(sqrt(2)) =" + g1.getDensity(d1) + ", ratio = " + (g1.getDensity(d1)/g1.getDensity(d0)));
          double[] d3 = {
              1};
          System.out.println("p(1) =" + g1.getDensity(d3) + ", ratio = " + (g1.getDensity(d3)/g1.getDensity(d0)));
        }
      } catch (Exception e) {
        e.printStackTrace();
      }


    }
  }

}