/*
 * Copyright (c) 2009 The Broad Institute
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

package org.broadinstitute.sting.gatk.contexts;

import org.broadinstitute.sting.gatk.datasources.sample.Sample;
import org.broadinstitute.sting.utils.exceptions.ReviewedStingException;
import org.broadinstitute.sting.utils.GenomeLoc;
import org.broadinstitute.sting.utils.exceptions.UserException;
import org.broadinstitute.sting.utils.pileup.*;

import java.util.*;

/**
 * Useful utilities for storing different AlignmentContexts
 * User: ebanks
 */
public class StratifiedAlignmentContext {

    // Definitions:
    //   COMPLETE = full alignment context
    //   FORWARD  = reads on forward strand
    //   REVERSE  = reads on forward strand
    //
    public enum StratifiedContextType { COMPLETE, FORWARD, REVERSE }

    private StratifiedAlignmentContext() {
        // cannot be instantiated
    }

    /**
     * Returns a potentially derived subcontext containing only forward, reverse, or in fact all reads
     * in alignment context context.
     *
     * @param context
     * @param type
     * @return
     */
    public static AlignmentContext stratify(AlignmentContext context, StratifiedContextType type) {
        switch(type) {
            case COMPLETE:
                return context;
            case FORWARD:
                return new AlignmentContext(context.getLocation(),context.getPileup().getPositiveStrandPileup());
            case REVERSE:
                return new AlignmentContext(context.getLocation(),context.getPileup().getNegativeStrandPileup());
            default:
                throw new ReviewedStingException("Unable to get alignment context for type = " + type);
        }
    }

    public static Map<String, AlignmentContext> splitContextBySampleName(AlignmentContext context) {
        return splitContextBySampleName(context, null);
    }

    public static Map<Sample, AlignmentContext> splitContextBySample(AlignmentContext context) {
        Map<Sample, AlignmentContext> m = new HashMap<Sample, AlignmentContext>();
        for ( Map.Entry<String, AlignmentContext> entry : splitContextBySampleName(context, null).entrySet() ) {
            m.put(new Sample(entry.getKey()), entry.getValue());
        }
        return m;
    }

    /**
     * Splits the given AlignmentContext into a StratifiedAlignmentContext per sample, but referencd by sample name instead
     * of sample object.
     *
     * @param context                the original pileup
     *
     * @return a Map of sample name to StratifiedAlignmentContext
     *
     **/
    public static Map<String, AlignmentContext> splitContextBySampleName(AlignmentContext context, String assumedSingleSample) {
        GenomeLoc loc = context.getLocation();
        HashMap<String, AlignmentContext> contexts = new HashMap<String, AlignmentContext>();

        for(String sample: context.getPileup().getSampleNames()) {
            ReadBackedPileup pileupBySample = context.getPileup().getPileupForSampleName(sample);

            // Don't add empty pileups to the split context.
            if(pileupBySample.size() == 0)
                continue;

            if(sample != null)
                contexts.put(sample, new AlignmentContext(loc, pileupBySample));
            else {
                if(assumedSingleSample == null) {
                    throw new UserException.ReadMissingReadGroup(pileupBySample.iterator().next().getRead());
                }
                contexts.put(assumedSingleSample,new AlignmentContext(loc, pileupBySample));
            }
        }

        return contexts;
    }

    public static Map<String, AlignmentContext> splitContextBySampleName(ReadBackedPileup pileup, String assumedSingleSample) {
        return splitContextBySampleName(new AlignmentContext(pileup.getLocation(), pileup));
    }


    public static AlignmentContext joinContexts(Collection<AlignmentContext> contexts) {
        // validation
        GenomeLoc loc = contexts.iterator().next().getLocation();
        boolean isExtended = contexts.iterator().next().basePileup instanceof ReadBackedExtendedEventPileup;
        for(AlignmentContext context: contexts) {
            if(!loc.equals(context.getLocation()))
                throw new ReviewedStingException("Illegal attempt to join contexts from different genomic locations");
            if(isExtended != (context.basePileup instanceof ReadBackedExtendedEventPileup))
                throw new ReviewedStingException("Illegal attempt to join simple and extended contexts");
        }

        AlignmentContext jointContext;
        if(isExtended) {
            List<ExtendedEventPileupElement> pe = new ArrayList<ExtendedEventPileupElement>();
            for(AlignmentContext context: contexts) {
                for(PileupElement pileupElement: context.basePileup)
                    pe.add((ExtendedEventPileupElement)pileupElement);
            }
            jointContext = new AlignmentContext(loc, new ReadBackedExtendedEventPileupImpl(loc,pe));
        }
        else {
            List<PileupElement> pe = new ArrayList<PileupElement>();
            for(AlignmentContext context: contexts) {
                for(PileupElement pileupElement: context.basePileup)
                    pe.add(pileupElement);
            }
            jointContext = new AlignmentContext(loc, new ReadBackedPileupImpl(loc,pe));
        }

        return jointContext;
    }
}