package uk.ac.ebi.atlas.tsne;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Multimap;
import com.google.common.collect.TreeMultimap;

import javax.inject.Inject;
import javax.inject.Named;
import java.util.Map;
import java.util.stream.Collectors;

@Named
public class TSnePlotCollator {

    // TODO This class will change a lot (and very probably won’t be necessary) when all data is available in the DB
    // TODO and we can do a JOIN on the cell ID field

    private final TSnePlotDao tSnePlotDao;
    private final ClusterDao clusterDao;
    private final GeneExpressionDao geneExpressionDao;
    private final RandomTSnePlot randomTSnePlot = new RandomTSnePlot();

    @Inject
    public TSnePlotCollator(TSnePlotDao tSnePlotDao, ClusterDao clusterDao, GeneExpressionDao geneExpressionDao) {
        this.tSnePlotDao = tSnePlotDao;
        this.clusterDao = clusterDao;
        this.geneExpressionDao = geneExpressionDao;
    }

    public TreeMultimap<Integer, TSnePoint> getTSnePlotWithClusters(String experimentAccession, int k) {
        if ("E-RANDOM".equalsIgnoreCase(experimentAccession)) {
            return randomTSnePlot.getTSnePlotWithClusters(k);
        }

        Map<String, TSnePoint> points = tSnePlotDao.fetchTSnePlotPoints(experimentAccession);
        Multimap<Integer, String> clusters = clusterDao.fetchClusters(experimentAccession, k);

        TreeMultimap<Integer, TSnePoint> clusterPoints = TreeMultimap.create();
        clusters.entries().forEach(entry -> clusterPoints.put(entry.getKey(), points.get(entry.getValue())));

        return clusterPoints;
    }

    public ImmutableList<TSnePoint> getTSnePlotWithExpression(String experimentAccession, String geneId) {
        Map<String, TSnePoint> points = tSnePlotDao.fetchTSnePlotPoints(experimentAccession);
        Map<String, Double> expression = geneExpressionDao.fetchGeneExpression(experimentAccession, geneId);

        return ImmutableList.copyOf(
                points.values().stream()
                        .map(point ->
                                TSnePoint.create(point.x(), point.y(), expression.get(point.name()), point.name()))
                        .collect(Collectors.toList()));
    }

}

@@ -28,56 +28,8 @@
 e;%0A%0A
-import com.google.common.collect.ImmutableList;%0A
 impo
@@ -195,44 +195,8 @@
 Map;
-%0Aimport java.util.stream.Collectors;
 %0A%0A@N
@@ -916,32 +916,1246 @@
 eger, TSnePoint%3E
+ getTsnePlotWithClustersAndExpression(String experimentAccession,%0A                                                                                 int k,%0A                                                                                 String geneId) %7B%0A        TreeMultimap%3CInteger, TSnePoint%3E clusters = this.getTSnePlotWithClusters(experimentAccession, k);%0A        Map%3CString, Double%3E expression = geneExpressionDao.fetchGeneExpression(experimentAccession, geneId);%0A%0A        TreeMultimap%3CInteger, TSnePoint%3E clusterPointsWithExpression = TreeMultimap.create();%0A        clusters.entries().forEach(%0A                entry -%3E%0A                        clusterPointsWithExpression.put(%0A                                entry.getKey(),%0A                                TSnePoint.create(%0A                                        entry.getValue().x(),%0A                                        entry.getValue().y(),%0A                                        expression.get(entry.getValue().name()),%0A                                        entry.getValue().name()%0A                                )%0A                        )%0A        );%0A%0A        return clusterPointsWithExpression;%0A    %7D%0A%0A    private TreeMultimap%3CInteger, TSnePoint%3E
  getTSnePlotWith
@@ -2758,613 +2758,6 @@
   %7D%0A
-%0A    public ImmutableList%3CTSnePoint%3E getTSnePlotWithExpression(String experimentAccession, String geneId) %7B%0A        Map%3CString, TSnePoint%3E points = tSnePlotDao.fetchTSnePlotPoints(experimentAccession);%0A        Map%3CString, Double%3E expression = geneExpressionDao.fetchGeneExpression(experimentAccession, geneId);%0A%0A        return ImmutableList.copyOf(%0A                points.values().stream()%0A                        .map(point -%3E%0A                                TSnePoint.create(point.x(), point.y(), expression.get(point.name()), point.name()))%0A                        .collect(Collectors.toList()));%0A    %7D%0A%0A
 %7D%0A

Remove getTSnePlotWithExpression and expose only one public method Every t-SNE plot with expression levels will also be showing the clustered points, so this is not necessary anymore