package org.cbioportal.service.impl;

import org.cbioportal.model.AlterationEnrichment;
import org.cbioportal.model.MolecularProfile;
import org.cbioportal.model.Mutation;
import org.cbioportal.model.MutationCountByGene;
import org.cbioportal.model.Sample;
import org.cbioportal.service.MolecularProfileService;
import org.cbioportal.service.MutationEnrichmentService;
import org.cbioportal.service.MutationService;
import org.cbioportal.service.SampleService;
import org.cbioportal.service.exception.MolecularProfileNotFoundException;
import org.cbioportal.service.util.AlterationEnrichmentUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Service
public class MutationEnrichmentServiceImpl implements MutationEnrichmentService {

    @Autowired
    private MutationService mutationService;
    @Autowired
    private SampleService sampleService;
    @Autowired
    private MolecularProfileService molecularProfileService;
    @Autowired
    private AlterationEnrichmentUtil alterationEnrichmentUtil;

    @Override
    public List<AlterationEnrichment> getMutationEnrichments(String molecularProfileId, List<String> alteredIds,
                                                             List<String> unalteredIds, List<Integer> queryGenes, 
                                                             String enrichmentType)
        throws MolecularProfileNotFoundException {

        List<String> allIds = new ArrayList<>(alteredIds);
        allIds.addAll(unalteredIds);
        List<MutationCountByGene> mutationCountByGeneList;
        List<Mutation> mutations;
        
        if (enrichmentType.equals("SAMPLE")) {
            mutationCountByGeneList = mutationService.getSampleCountByEntrezGeneIdsAndSampleIds(molecularProfileId, 
                allIds, null, false);
            mutations = mutationService.fetchMutationsInMolecularProfile(molecularProfileId, alteredIds, null, null, 
                "ID", null, null, null, null);
        } else {
            mutationCountByGeneList = mutationService.getPatientCountByEntrezGeneIdsAndSampleIds(molecularProfileId,
                allIds, null);
            MolecularProfile molecularProfile = molecularProfileService.getMolecularProfile(molecularProfileId);
            List<Sample> sampleList = sampleService.getAllSamplesOfPatientsInStudy(
                molecularProfile.getCancerStudyIdentifier(), alteredIds, "ID");
            mutations = mutationService.fetchMutationsInMolecularProfile(molecularProfileId,
                sampleList.stream().map(Sample::getStableId).collect(Collectors.toList()), null, null, "ID", null, null,
                null, null);
        }

        mutationCountByGeneList.removeIf(m -> queryGenes.contains(m.getEntrezGeneId()));
        return alterationEnrichmentUtil.createAlterationEnrichments(alteredIds.size(), unalteredIds.size(),
            mutationCountByGeneList, mutations, enrichmentType);
    }
}

@@ -1683,16 +1683,24 @@
 GeneList
+FromRepo
 ;%0A      
@@ -1810,32 +1810,40 @@
 nCountByGeneList
+FromRepo
  = mutationServi
@@ -2163,16 +2163,24 @@
 GeneList
+FromRepo
  = mutat
@@ -2815,16 +2815,157 @@
      %7D%0A%0A
+        List%3CMutationCountByGene%3E mutationCountByGeneList =%0A            new ArrayList%3CMutationCountByGene%3E(mutationCountByGeneListFromRepo);%0A
         

Fix array modification problem in Mutation Enrichments API