/*
 * PairedPatterns.java
 *
 * SSGD: Serially-Sampled Genome Demographics
 *
 * Copyright (c) 2015 Arman Bilge <armanbilge@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package org.compevol.ssgd;

import dr.evolution.datatype.DataType;
import dr.evolution.util.Taxon;
import dr.evolution.util.TaxonList;
import dr.math.MathUtils;
import dr.util.Identifiable;

/**
 * @author Arman Bilge <armanbilge@gmail.com>
 */
public class PairedPatterns implements Identifiable {

    private static final long serialVersionUID = 0L;

    private final DataType dataType;
    private final int stateCount;
    private final TaxonList taxa;
    private final double[][] weights;

    private String id;

    public PairedPatterns(final DataType dataType, final TaxonList taxa) {
        this.dataType = dataType;
        stateCount = dataType.getStateCount();
        this.taxa = taxa;
        weights = new double[taxa.getTaxonCount() * (taxa.getTaxonCount() - 1) / 2][stateCount * stateCount];
    }

    public double getPatternWeight(final Taxon a, final int i, final Taxon b, final int j) {

        final int m = taxa.getTaxonIndex(a);
        final int n = taxa.getTaxonIndex(b);

        if (m == n)
            throw new IllegalArgumentException("The two taxa must be different.");
        else if (m > n)
            return getPatternWeight(b, j, a, i);

        return weights[m + n * (n - 1) / 2 ][stateCount * i + j];

    }

    public void addPattern(final Taxon a, final int i, final Taxon b, final int j) {
        addPattern(a, i, b, j, 1);
    }

    public void addPattern(final Taxon a, final int i, final Taxon b, final int j, final double w) {

        final int m = taxa.getTaxonIndex(a);
        final int n = taxa.getTaxonIndex(b);

        if (m == n) {
            throw new IllegalArgumentException("The two taxa must be different.");
        } else if (m > n) {
            addPattern(b, j, a, i, w);
            return;
        }

        final int x = m + n * (n - 1) / 2;
        final int y = stateCount * i + j;

        weights[x][y] += w;

    }

    public final DataType getDataType() {
        return dataType;
    }

    public final TaxonList getTaxa() {
        return taxa;
    }

    public final double[] getApproximateFrequencies() {

        final double[] freqs = new double[stateCount];

        for (int i = 0; i < stateCount; ++i)
            freqs[i] = weights[0][(stateCount+1) * i];

        final double sum = MathUtils.getTotal(freqs);

        for (int i = 0; i < stateCount; ++i)
            freqs[i] /= sum;

        return freqs;
    }

    public final double getTotalWeight() {
        double total = 0.0;
        for (final double[] weights : this.weights)
            total += MathUtils.getTotal(weights);
        return total;
    }

    @Override
    public String getId() {
        return id;
    }

    @Override
    public void setId(final String id) {
        this.id = id;
    }
}

@@ -1633,16 +1633,55 @@
 eCount;%0A
+    private final int transitionCount;%0A
     priv
@@ -1732,18 +1732,16 @@
 double%5B%5D
-%5B%5D
  weights
@@ -1923,16 +1923,67 @@
 ount();%0A
+        transitionCount = stateCount * stateCount;%0A
         
@@ -2086,28 +2086,21 @@
  / 2
-%5D%5BstateCount * state
+ * transition
 Coun
@@ -2499,48 +2499,28 @@
 hts%5B
-m + n * (n - 1) / 2 %5D%5BstateCount * i +
+getIndex(m, n, i,
  j
+)
 %5D;%0A%0A
@@ -3061,114 +3061,189 @@
     
-final int x = m + n * (n - 1) / 2;%0A        final int y = stateCount * i + j;%0A%0A        weights%5Bx%5D%5By%5D += w;%0A
+weights%5BgetIndex(m, n, i, j)%5D += w;%0A%0A    %7D%0A%0A    private int getIndex(final int m, int n, int i, int j) %7B%0A        return (m + n * (n - 1) / 2) * transitionCount + stateCount * i + j;
 %0A   
@@ -3580,11 +3580,8 @@
 hts%5B
-0%5D%5B
 (sta
@@ -3812,100 +3812,14 @@
     
-double total = 0.0;%0A        for (final double%5B%5D weights : this.weights)%0A            total +=
+return
  Mat
@@ -3848,30 +3848,8 @@
 s);%0A
-        return total;%0A
     

Use 1D-array for PairedPatterns