/*******************************************************************************
 * Copyright (c) 2013 Michael Kutschke.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     Michael Kutschke - initial API and implementation
 ******************************************************************************/
package bayesnet.jayes.inference.jtree;

import static bayesnet.jayes.util.Pair.newPair;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.ListIterator;

import bayesnet.internal.jayes.util.UnionFind;
import bayesnet.jayes.BayesNet;
import bayesnet.jayes.util.Graph;
import bayesnet.jayes.util.OrderIgnoringPair;
import bayesnet.jayes.util.Pair;

public class SepsetComputer {

    public List, List>> computeSepsets(JunctionTree junctionTree, BayesNet net) {
        final List, List>> candidates = enumerateCandidateSepSets(junctionTree
                .getClusters());
        Collections.sort(candidates, new SepsetComparator(net));
        return computeMaxSpanningTree(junctionTree.getGraph(), candidates);

    }

    private List, List>> enumerateCandidateSepSets(List> clusters) {
        final List, List>> sepSets = new ArrayList, List>>();
        final ListIterator> it = clusters.listIterator();
        while (it.hasNext()) {
            final List clique1 = it.next();
            final ListIterator> remainingIt = clusters.listIterator(it.nextIndex());
            while (remainingIt.hasNext()) { // generate sepSets
                final List clique2 = new ArrayList(remainingIt.next());
                clique2.retainAll(clique1);
                sepSets.add(newPair(new OrderIgnoringPair(it.nextIndex() - 1, remainingIt.nextIndex() - 1),
                        clique2));
            }
        }
        return sepSets;
    }

    private List, List>> computeMaxSpanningTree(Graph graph,
            final List, List>> sortedCandidateSepSets) {

        final ArrayDeque, List>> pq = new ArrayDeque, List>>(
                sortedCandidateSepSets);

        final int vertexCount = graph.numberOfVertices();
        final UnionFind[] sets = UnionFind.createArray(vertexCount);

        final List, List>> leftSepSets = new ArrayList, List>>();
        while (leftSepSets.size() < vertexCount - 1) {
            final Pair, List> sep = pq.poll();
            final boolean bothEndsInSameTree = sets[sep.getFirst().getFirst()].find() == sets[sep.getFirst()
                                                                                              .getSecond()].find();
            if (!bothEndsInSameTree) {
                sets[sep.getFirst().getFirst()].merge(sets[sep.getFirst().getSecond()]);
                leftSepSets.add(sep);
                graph.addEdge(sep.getFirst().getFirst(), sep.getFirst().getSecond());
            }
        }
        return leftSepSets;
    }

    private static final class SepsetComparator implements Comparator, List>> {

        private final BayesNet net;

        public SepsetComparator(BayesNet net) {
            this.net = net;
        }

        // heuristic: choose sepSet with most variables first,
        // if equal, choose the on with least table size
        @Override
        public int compare(final Pair, List> sepSet1,
                final Pair, List> sepSet2) {
            final int compareNumberOfVariables = compare(sepSet2.getSecond().size(), sepSet1.getSecond().size());
            if (compareNumberOfVariables != 0) {
                return compareNumberOfVariables;
            }
            final int tableSize1 = getTableSize(sepSet1.getSecond());
            final int tableSize2 = getTableSize(sepSet2.getSecond());
            return compare(tableSize1, tableSize2);

        }

        private int getTableSize(final List cluster) {
            int tableSize = 1;
            for (final int id : cluster) {
                tableSize *= net.getNode(id).getOutcomeCount();
            }
            return tableSize;
        }

        private int compare(final int i1, final int i2) {
            return i1 - i2;
        }
    }

}