/*
 * Copyright (c) 2012 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 htsjdk.variant.vcf;

import htsjdk.tribble.TribbleException;
import htsjdk.variant.utils.GeneralUtils;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

/**
 * Manages header lines for standard VCF <pre>INFO</pre> and <pre>FORMAT</pre> fields.
 *
 * Provides simple mechanisms for
 *  1) registering standard lines,
 *  2) looking them up, and
 *  3) adding them to headers.
 *
 * @author Mark DePristo
 * @since 6/12
 */
public class VCFStandardHeaderLines {
    /**
     * Enabling this causes us to repair header lines even if only their descriptions differ.
     */
    private static final boolean REPAIR_BAD_DESCRIPTIONS = false;

    private static Standards<VCFFormatHeaderLine> formatStandards = new Standards<VCFFormatHeaderLine>();
    private static Standards<VCFInfoHeaderLine> infoStandards = new Standards<VCFInfoHeaderLine>();

    /**
     * Walks over the VCF header and repairs the standard VCF header lines in it, returning a freshly
     * allocated {@link VCFHeader} with standard VCF header lines repaired as necessary.
     */
    public static VCFHeader repairStandardHeaderLines(final VCFHeader oldHeader) {
        final Set<VCFHeaderLine> newLines = new LinkedHashSet<VCFHeaderLine>(
                oldHeader.getMetaDataInInputOrder().size());
        for (VCFHeaderLine line : oldHeader.getMetaDataInInputOrder()) {
            if (line instanceof VCFFormatHeaderLine) {
                line = formatStandards.repair((VCFFormatHeaderLine) line);
            } else if (line instanceof VCFInfoHeaderLine) {
                line = infoStandards.repair((VCFInfoHeaderLine) line);
            }

            newLines.add(line);
        }

        final VCFHeader repairedHeader = new VCFHeader(newLines, oldHeader.getGenotypeSamples());
        final VCFHeaderVersion oldHeaderVersion = oldHeader.getVCFHeaderVersion();
        if (oldHeaderVersion != null && oldHeaderVersion.isAtLeastAsRecentAs(VCFHeaderVersion.VCF4_3)) {
            // this needs to maintain version 4.3 (and not back-version to v4.2), so propagate
            // the old version only for v4.3
            repairedHeader.setVCFHeaderVersion(oldHeaderVersion);
        }
        return repairedHeader;
    }

    /**
     * Adds header lines for each of the format fields in IDs to header, returning the set of
     * {@code IDs} without standard descriptions, unless {@code throwErrorForMissing} is true, in which
     * case this situation results in a {@link TribbleException}
     */
    public static Set<String> addStandardFormatLines(
            final Set<VCFHeaderLine> headerLines, final boolean throwErrorForMissing, final Collection<String> IDs) {
        return formatStandards.addToHeader(headerLines, IDs, throwErrorForMissing);
    }

    /**
     * @see #addStandardFormatLines(java.util.Set, boolean, java.util.Collection)
     */
    public static Set<String> addStandardFormatLines(
            final Set<VCFHeaderLine> headerLines, final boolean throwErrorForMissing, final String... IDs) {
        return addStandardFormatLines(headerLines, throwErrorForMissing, Arrays.asList(IDs));
    }

    /**
     * Returns the standard format line for {@code ID}.
     * If none exists, return null or throw an exception, depending on {@code throwErrorForMissing}.
     */
    public static VCFFormatHeaderLine getFormatLine(final String ID, final boolean throwErrorForMissing) {
        return formatStandards.get(ID, throwErrorForMissing);
    }

    /**
     * Returns the standard format line for {@code ID}.
     * If none exists, throw an {@link TribbleException}
     */
    public static VCFFormatHeaderLine getFormatLine(final String ID) {
        return formatStandards.get(ID, true);
    }

    /**
     * Adds header lines for each of the info fields in {@code IDs} to header, returning the set of
     * IDs without standard descriptions, unless {@code throwErrorForMissing} is true, in which
     * case this situation results in a {@link TribbleException}.
     */
    public static Set<String> addStandardInfoLines(
            final Set<VCFHeaderLine> headerLines, final boolean throwErrorForMissing, final Collection<String> IDs) {
        return infoStandards.addToHeader(headerLines, IDs, throwErrorForMissing);
    }

    /**
     * @see #addStandardFormatLines(java.util.Set, boolean, java.util.Collection)
     */
    public static Set<String> addStandardInfoLines(
            final Set<VCFHeaderLine> headerLines, final boolean throwErrorForMissing, final String... IDs) {
        return addStandardInfoLines(headerLines, throwErrorForMissing, Arrays.asList(IDs));
    }

    /**
     * Returns the standard info line for {@code ID}.
     * If none exists, return {@code null} or throw a {@link TribbleException}, depending on {@code throwErrorForMissing}.
     */
    public static VCFInfoHeaderLine getInfoLine(final String ID, final boolean throwErrorForMissing) {
        return infoStandards.get(ID, throwErrorForMissing);
    }

    /**
     * Returns the standard info line for {@code ID}.
     * If none exists throw a {@link TribbleException}.
     */
    public static VCFInfoHeaderLine getInfoLine(final String ID) {
        return getInfoLine(ID, true);
    }

    private static void registerStandard(final VCFInfoHeaderLine line) {
        infoStandards.add(line);
    }

    private static void registerStandard(final VCFFormatHeaderLine line) {
        formatStandards.add(line);
    }

    //
    // VCF header line constants
    //
    static {
        // FORMAT lines
        registerStandard(new VCFFormatHeaderLine(VCFConstants.GENOTYPE_KEY, 1, VCFHeaderLineType.String, "Genotype"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.GENOTYPE_QUALITY_KEY, 1, VCFHeaderLineType.Integer, "Genotype Quality"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.DEPTH_KEY,
                1,
                VCFHeaderLineType.Integer,
                "Approximate read depth (reads with MQ=255 or with bad mates are filtered)"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.GENOTYPE_PL_KEY,
                VCFHeaderLineCount.G,
                VCFHeaderLineType.Integer,
                "Normalized, Phred-scaled likelihoods for genotypes as defined in the VCF specification"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.GENOTYPE_ALLELE_DEPTHS,
                VCFHeaderLineCount.R,
                VCFHeaderLineType.Integer,
                "Allelic depths for the ref and alt alleles in the order listed"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.GENOTYPE_FILTER_KEY,
                VCFHeaderLineCount.UNBOUNDED,
                VCFHeaderLineType.String,
                "Genotype-level filter"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.PHASE_SET_KEY,
                1,
                VCFHeaderLineType.Integer,
                "Phasing set (typically the position of the first variant in the set)"));
        registerStandard(new VCFFormatHeaderLine(
                VCFConstants.PHASE_QUALITY_KEY, 1, VCFHeaderLineType.Float, "Read-backed phasing quality"));

        // INFO lines
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.END_KEY, 1, VCFHeaderLineType.Integer, "Stop position of the interval"));
        registerStandard(new VCFInfoHeaderLine(VCFConstants.DBSNP_KEY, 0, VCFHeaderLineType.Flag, "dbSNP Membership"));
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.DEPTH_KEY,
                1,
                VCFHeaderLineType.Integer,
                "Approximate read depth; some reads may have been filtered"));
        registerStandard(
                new VCFInfoHeaderLine(VCFConstants.STRAND_BIAS_KEY, 1, VCFHeaderLineType.Float, "Strand Bias"));
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.ALLELE_FREQUENCY_KEY,
                VCFHeaderLineCount.A,
                VCFHeaderLineType.Float,
                "Allele Frequency, for each ALT allele, in the same order as listed"));
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.ALLELE_COUNT_KEY,
                VCFHeaderLineCount.A,
                VCFHeaderLineType.Integer,
                "Allele count in genotypes, for each ALT allele, in the same order as listed"));
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.ALLELE_NUMBER_KEY,
                1,
                VCFHeaderLineType.Integer,
                "Total number of alleles in called genotypes"));
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.MAPPING_QUALITY_ZERO_KEY,
                1,
                VCFHeaderLineType.Integer,
                "Total Mapping Quality Zero Reads"));
        registerStandard(new VCFInfoHeaderLine(
                VCFConstants.RMS_MAPPING_QUALITY_KEY, 1, VCFHeaderLineType.Float, "RMS Mapping Quality"));
        registerStandard(new VCFInfoHeaderLine(VCFConstants.SOMATIC_KEY, 0, VCFHeaderLineType.Flag, "Somatic event"));
    }

    private static class Standards<T extends VCFCompoundHeaderLine> {
        private final Map<String, T> standards = new HashMap<String, T>();

        public T repair(final T line) {
            final T standard = get(line.getID(), false);
            if (standard != null) {
                final boolean badCountType = line.getCountType() != standard.getCountType();
                final boolean badCount = line.isFixedCount() && !badCountType && line.getCount() != standard.getCount();
                final boolean badType = line.getType() != standard.getType();
                final boolean badDesc = !line.getDescription().equals(standard.getDescription());
                final boolean needsRepair = badCountType || badCount || badType || (REPAIR_BAD_DESCRIPTIONS && badDesc);

                if (needsRepair) {
                    if (GeneralUtils.DEBUG_MODE_ENABLED) {
                        System.err.println("Repairing standard header line for field " + line.getID() + " because"
                                + (badCountType
                                        ? " -- count types disagree; header has " + line.getCountType()
                                                + " but standard is " + standard.getCountType()
                                        : "")
                                + (badType
                                        ? " -- type disagree; header has " + line.getType() + " but standard is "
                                                + standard.getType()
                                        : "")
                                + (badCount
                                        ? " -- counts disagree; header has " + line.getCount() + " but standard is "
                                                + standard.getCount()
                                        : "")
                                + (badDesc
                                        ? " -- descriptions disagree; header has '" + line.getDescription()
                                                + "' but standard is '" + standard.getDescription() + "'"
                                        : ""));
                    }
                    return standard;
                } else {
                    return line;
                }
            } else {
                return line;
            }
        }

        public Set<String> addToHeader(
                final Set<VCFHeaderLine> headerLines,
                final Collection<String> IDs,
                final boolean throwErrorForMissing) {
            final Set<String> missing = new HashSet<String>();
            for (final String ID : IDs) {
                final T line = get(ID, throwErrorForMissing);
                if (line == null) missing.add(ID);
                else headerLines.add(line);
            }

            return missing;
        }

        public void add(final T line) {
            if (standards.containsKey(line.getID())) {
                throw new TribbleException("Attempting to add multiple standard header lines for ID " + line.getID());
            }
            standards.put(line.getID(), line);
        }

        public T get(final String ID, final boolean throwErrorForMissing) {
            final T x = standards.get(ID);
            if (throwErrorForMissing && x == null) {
                throw new TribbleException("Couldn't find a standard VCF header line for field " + ID);
            }
            return x;
        }
    }
}
