package org.cbioportal.service.util;

import java.util.*;
import java.util.stream.Collectors;

import org.cbioportal.model.*;
import org.cbioportal.service.GenePanelService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
public class ProfiledCasesCounter {

    @Autowired
    private GenePanelService genePanelService;

    public void calculate(List<String> molecularProfileIds, List<String> sampleIds,
            List<? extends AlterationCountByGene> alterationCounts, boolean countByPatients) {

        List<GenePanelData> genePanelDataList = genePanelService
                .fetchGenePanelDataInMultipleMolecularProfiles(molecularProfileIds, sampleIds);
        Map<String, List<GenePanelData>> genePanelDataMap = genePanelDataList.stream()
                .filter(g -> g.getGenePanelId() != null).collect(Collectors.groupingBy(GenePanelData::getGenePanelId));
        List<GenePanel> genePanels = new ArrayList<>();
        if (!genePanelDataMap.isEmpty()) {
            genePanels = genePanelService.fetchGenePanels(new ArrayList<>(genePanelDataMap.keySet()), "DETAILED");
        }

        Map<Integer, List<GenePanel>> geneGenePanelMap = new HashMap<>();
        for (GenePanel genePanel : genePanels) {
            for (GenePanelToGene genePanelToGene : genePanel.getGenes()) {
                Integer entrezGeneId = genePanelToGene.getEntrezGeneId();
                if (geneGenePanelMap.containsKey(entrezGeneId)) {
                    geneGenePanelMap.get(entrezGeneId).add(genePanel);
                } else {
                    List<GenePanel> geneGenePanelList = new ArrayList<>();
                    geneGenePanelList.add(genePanel);
                    geneGenePanelMap.put(entrezGeneId, geneGenePanelList);
                }
            }
        }

        List<GenePanelData> profiled = genePanelDataList
                .stream()
                .filter(GenePanelData::getProfiled)
                .collect(Collectors.toList());

        Set<String> profiledCases = profiled
                .stream()
                .map(x -> countByPatients ? x.getPatientId() : x.getSampleId())
                .collect(Collectors.toSet());

        Set<String> casesWithoutPanelData = profiled
                .stream()
                .filter(g -> g.getGenePanelId() == null)
                .map(x -> countByPatients ? x.getPatientId() : x.getSampleId())
                .collect(Collectors.toSet());

        for (AlterationCountByGene alterationCountByGene : alterationCounts) {
            final Set<String> profiledCasesForGene = new HashSet<String>();
            Integer entrezGeneId = alterationCountByGene.getEntrezGeneId();
            List<GenePanel> allPanels = new ArrayList<>();
            
            if (geneGenePanelMap.containsKey(entrezGeneId)) {
                geneGenePanelMap.get(entrezGeneId).forEach(genePanel -> {
                    Set<String> casesWithPanelData = genePanelDataMap
                            .get(genePanel.getStableId())
                            .stream()
                            .map(x -> countByPatients ? x.getPatientId() : x.getSampleId())
                            .collect(Collectors.toSet());
                    profiledCasesForGene.addAll(casesWithPanelData);
                    allPanels.add(genePanel);
                });
                
                profiledCasesForGene.addAll(casesWithoutPanelData);
            } else {
                
                profiledCasesForGene.addAll(profiledCases);
            }
            
            alterationCountByGene.setMatchingGenePanelIds(
                    allPanels.stream().map(panel -> panel.getStableId()).collect(Collectors.toSet()));
            alterationCountByGene.setNumberOfProfiledCases(profiledCasesForGene.size());
        }
    }
}

@@ -2111,34 +2111,132 @@
             
-.map(x -%3E 
+// there can be duplicate patient or sample id, append study id%0A                .map(x -%3E x.getStudyId() + (
 countByPatie
@@ -2268,32 +2268,33 @@
 x.getSampleId())
+)
 %0A               
@@ -2461,16 +2461,96 @@
 = null)%0A
+                // there can be duplicate patient or sample id, append study id%0A
         
@@ -2559,32 +2559,50 @@
       .map(x -%3E 
+x.getStudyId() + (
 countByPatients 
@@ -2630,32 +2630,33 @@
 x.getSampleId())
+)
 %0A               
@@ -3284,32 +3284,124 @@
       .stream()%0A
+                            // there can be duplicate patient or sample id, append study id%0A
                 
@@ -3422,16 +3422,34 @@
 ap(x -%3E 
+x.getStudyId() + (
 countByP
@@ -3489,24 +3489,25 @@
 tSampleId())
+)
 %0A           

append study id while calculating profiled cases count