package org.cbioportal.service.impl;

import org.apache.commons.lang.math.NumberUtils;
import org.apache.commons.math3.stat.correlation.PearsonsCorrelation;
import org.apache.commons.math3.stat.correlation.SpearmansCorrelation;
import org.cbioportal.model.CoExpression;
import org.cbioportal.model.Gene;
import org.cbioportal.model.GeneMolecularData;
import org.cbioportal.service.CoExpressionService;
import org.cbioportal.service.GeneService;
import org.cbioportal.service.MolecularDataService;
import org.cbioportal.service.exception.MolecularProfileNotFoundException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
public class CoExpressionServiceImpl implements CoExpressionService {
    
    @Autowired
    private MolecularDataService molecularDataService;
    @Autowired
    private GeneService geneService;
    
    private PearsonsCorrelation pearsonsCorrelation = new PearsonsCorrelation();
    private SpearmansCorrelation spearmansCorrelation = new SpearmansCorrelation();
    
    @Override
    public List<CoExpression> getCoExpressions(String molecularProfileId, String sampleListId, Integer entrezGeneId, 
                                              Double threshold) throws MolecularProfileNotFoundException {

        List<GeneMolecularData> molecularDataList = molecularDataService.getMolecularData(molecularProfileId, 
            sampleListId, null, "SUMMARY");
        
        return createCoExpressions(molecularDataList, entrezGeneId, threshold);
    }

    @Override
    public List<CoExpression> fetchCoExpressions(String molecularProfileId, List<String> sampleIds, Integer entrezGeneId, 
                                                Double threshold) throws MolecularProfileNotFoundException {

        List<GeneMolecularData> molecularDataList = molecularDataService.fetchMolecularData(molecularProfileId, 
            sampleIds, null, "SUMMARY");
        
        return createCoExpressions(molecularDataList, entrezGeneId, threshold);
    }
    
    private List<CoExpression> createCoExpressions(List<GeneMolecularData> molecularDataList, Integer queryEntrezGeneId,
                                                   Double threshold) {

        Map<Integer, List<GeneMolecularData>> molecularDataMap = molecularDataList.stream().filter(g -> 
            NumberUtils.isNumber(g.getValue())).collect(Collectors.groupingBy(GeneMolecularData::getEntrezGeneId));
        List<GeneMolecularData> queryMolecularDataList = molecularDataMap.remove(queryEntrezGeneId);
        
        Map<Integer, List<Gene>> genes = geneService.fetchGenes(molecularDataMap.keySet().stream()
            .map(String::valueOf).collect(Collectors.toList()), "ENTREZ_GENE_ID", "SUMMARY").stream()
            .collect(Collectors.groupingBy(Gene::getEntrezGeneId));
        
        double[] queryValues = queryMolecularDataList.stream().mapToDouble(g -> Double.parseDouble(g.getValue()))
            .toArray();
        
        List<CoExpression> coExpressionList = new ArrayList<>();
        for (Integer entrezGeneId : molecularDataMap.keySet()) {
            
            double[] values = molecularDataMap.get(entrezGeneId).stream().mapToDouble(g -> Double.parseDouble(
                g.getValue())).toArray();
            
            CoExpression coExpression = new CoExpression();
            coExpression.setEntrezGeneId(entrezGeneId);
            Gene gene = genes.get(entrezGeneId).get(0);
            coExpression.setCytoband(gene.getCytoband());
            coExpression.setHugoGeneSymbol(gene.getHugoGeneSymbol());
            
            double pearsonsValue = pearsonsCorrelation.correlation(queryValues, values);
            if (Double.isNaN(pearsonsValue) || Math.abs(pearsonsValue) < threshold) {
                continue;
            }
            coExpression.setPearsonsCorrelation(BigDecimal.valueOf(pearsonsValue));
            
            double spearmansValue = spearmansCorrelation.correlation(queryValues, values);
            if (Double.isNaN(spearmansValue) || Math.abs(spearmansValue) < threshold) {
                continue;
            }
            coExpression.setSpearmansCorrelation(BigDecimal.valueOf(spearmansValue));
            
            coExpressionList.add(coExpression);
        }
        
        return coExpressionList;
    }
}

@@ -2476,69 +2476,25 @@
 am()
-.filter(g -%3E %0A            NumberUtils.isNumber(g.getValue()))
+%0A                
 .col
@@ -2659,24 +2659,16 @@
 eneId);%0A
-        
 %0A       
@@ -2946,24 +2946,28 @@
         
-double%5B%5D
+List%3CString%3E
  queryVa
@@ -3012,77 +3012,55 @@
 .map
-ToDouble(g -%3E Double.parseDouble(g.getValue()))%0A            .toArray(
+(g -%3E g.getValue()).collect(Collectors.toList()
 );%0A 
@@ -3226,16 +3226,20 @@
     
-double%5B%5D
+List%3CString%3E
  val
@@ -3295,83 +3295,698 @@
 .map
-ToDouble(g -%3E Double.parseDouble(%0A                g.getValue())).toArray();
+(g -%3E g.getValue())%0A                .collect(Collectors.toList());%0A            List%3CString%3E queryValuesCopy = new ArrayList%3C%3E(queryValues);%0A%0A            List%3CInteger%3E valuesToRemove = new ArrayList%3C%3E();%0A            for (int i = 0; i %3C queryValuesCopy.size(); i++) %7B%0A                if (!NumberUtils.isNumber(queryValuesCopy.get(i)) %7C%7C !NumberUtils.isNumber(values.get(i))) %7B%0A                    valuesToRemove.add(i);%0A                %7D%0A            %7D%0A%0A            for (int i = 0; i %3C valuesToRemove.size(); i++) %7B%0A                int valueToRemove = valuesToRemove.get(i) - i;%0A                queryValuesCopy.remove(valueToRemove);%0A                values.remove(valueToRemove);%0A            %7D
 %0A   
@@ -4290,24 +4290,231 @@
 neSymbol());
+%0A%0A            double%5B%5D queryValuesNumber = queryValuesCopy.stream().mapToDouble(Double::parseDouble).toArray();%0A            double%5B%5D valuesNumber = values.stream().mapToDouble(Double::parseDouble).toArray();
 %0A           
@@ -4589,32 +4589,38 @@
 (queryValues
+Number
 , values
 );%0A         
@@ -4599,32 +4599,38 @@
 esNumber, values
+Number
 );%0A            i
@@ -4919,24 +4919,30 @@
 ryValues
+Number
 , values
 );%0A     
@@ -4933,16 +4933,22 @@
 , values
+Number
 );%0A     

Fix Co-Expression API when there are NA values