package org.broadinstitute.hellbender.tools.reference;

import org.broadinstitute.hellbender.CommandLineProgramTest;
import org.broadinstitute.hellbender.exceptions.UserException;
import org.broadinstitute.hellbender.testutils.IntegrationTestSpec;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class CheckReferenceCompatibilityIntegrationTest extends CommandLineProgramTest {

    private final String COMPARE_REFERENCES_TEST_FILES_DIRECTORY = toolsTestDir + "/reference/CompareReferences/";

    @DataProvider(name = "testReferenceCompatibilityBAMWithMD5sData")
    public Object[][] testReferenceCompatibilityBAMWithMD5sData() {
        return new Object[][]{
                // reference, dictionary for comparison, expected output
                new Object[]{
                        // exact match
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta"),
                        new File(getToolTestDataDir() + "reads_data_source_test1_withmd5s.bam"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithMD5s_exactmatch.table")
                },
                new Object[]{
                        // subset
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta"),
                        new File(getToolTestDataDir() + "reads_data_source_test1_withmd5s_missingchr1.bam"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithMD5s_subset.table")
                },
                new Object[]{
                        // not compatible
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_missingchr3.fasta"),
                        new File(getToolTestDataDir() + "reads_data_source_test1_withmd5s.bam"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithMD5s_notcompatible.table")
                },
        };
    }

    @Test(dataProvider = "testReferenceCompatibilityBAMWithMD5sData")
    public void testReferenceCompatibilityBAMWithMD5s(File ref1, File dict, File expectedOutput) throws IOException {
        final File output = createTempFile("testReferenceCompatibilityBAMsWithMD5s", ".table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-I", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    @DataProvider(name = "testReferenceCompatibilityBAMWithoutMD5sData")
    public Object[][] testReferenceCompatibilityBAMWithoutMD5sData() {
        return new Object[][]{
                // reference, dictionary for comparison, expected output
                new Object[]{
                        // exact match
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta"),
                        new File(getToolTestDataDir() + "reads_data_source_test1_withoutmd5s.bam"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithoutMD5s_compatible.table")
                },
                new Object[]{
                        // subset
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta"),
                        new File(getToolTestDataDir() + "reads_data_source_test1_withoutmd5s_missingchr1.bam"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithoutMD5s_subset.table")
                },
                new Object[]{
                        // not compatible
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_missingchr3.fasta"),
                        new File(getToolTestDataDir() + "reads_data_source_test1_withoutmd5s.bam"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithoutMD5s_notcompatible.table")
                },
        };
    }

    @Test(dataProvider = "testReferenceCompatibilityBAMWithoutMD5sData")
    public void testReferenceCompatibilityBAMWithoutMD5s(File ref1, File dict, File expectedOutput) throws IOException {
        final File output = createTempFile("testReferenceCompatibilityBAMsWithoutMD5s", ".table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-I", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    @Test(expectedExceptions = UserException.BadInput.class)
    public void testReferenceCompatibilityMultipleBAMs() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta");
        final File dict1 = new File(getToolTestDataDir() + "reads_data_source_test1_withmd5s.bam");
        final File dict2 = new File(getToolTestDataDir() + "reads_data_source_test1_withmd5s_missingchr1.bam");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-I", dict1.getAbsolutePath(), "-I", dict2.getAbsolutePath()};
        runCommandLine(args);
    }

    @Test
    public void testReferenceCompatibilityMultipleReferencesBAMWithMD5s() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta");
        final File ref2 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_1renamed.fasta");
        final File ref3 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_missingchr1.fasta");
        final File ref4 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_chr2snp.fasta");
        final File dict = new File(getToolTestDataDir() + "reads_data_source_test1_withmd5s_missingchr1.bam");
        final File output = createTempFile("testReferenceCompatibilityMultipleReferencesWithMD5s", ".table");
        final File expectedOutput = new File(getToolTestDataDir(), "expected.testReferenceCompatibilityMultipleReferencesBAMWithMD5s.table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath(), "-refcomp", ref2.getAbsolutePath(),
                "-refcomp", ref3.getAbsolutePath(), "-refcomp", ref4.getAbsolutePath(), "-I", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    // Note that hg19mini_chr2snp.fasta is NOT a valid reference, but tool renders them a match due to compatibility by name and length.
    // Included to demonstrate that this is the best the tool can do without MD5s present
    @Test
    public void testReferenceCompatibilityMultipleReferencesBAMWithoutMD5s() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta");
        final File ref2 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_1renamed.fasta");
        final File ref3 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_chr2snp.fasta");
        final File dict = new File(getToolTestDataDir() + "reads_data_source_test1_withoutmd5s.bam");
        final File output = createTempFile("testReferenceCompatibilityMultipleReferencesWithoutMD5s", ".table");
        final File expectedOutput = new File(getToolTestDataDir(), "expected.testReferenceCompatibilityMultipleReferencesBAMWithoutMD5s.table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath(), "-refcomp", ref2.getAbsolutePath(),
                "-refcomp", ref3.getAbsolutePath(), "-I", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    // References are not compatible since one has SNP on chr 2, but tool renders them seemingly a match
    // Included to demonstrate that this is the best the tool can do without MD5s present
    @Test
    public void testReferenceCompatibilityWithoutMD5sMismatch() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_chr2snp.fasta");
        final File dict = new File(getToolTestDataDir() + "reads_data_source_test1_withoutmd5s.bam");
        final File output = createTempFile("testReferenceCompatibilityVCFWithMD5s", ".table");
        final File expectedOutput = new File(getToolTestDataDir(), "expected.testReferenceCompatibilityWithoutMD5Mismatch.table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-I", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    // TODO: compatibility based on MD5 faulty since MD5s not in sequence dictionary (https://github.com/samtools/htsjdk/issues/730)
    @Test
    public void testReferenceCompatibilityVCFWithMD5s() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta");
        final File dict = new File(getToolTestDataDir() + "example_variants_withSequenceDict_withmd5.vcf");
        final File output = createTempFile("testReferenceCompatibilityVCFWithMD5s", ".table");
        final File expectedOutput = new File(getToolTestDataDir(), "expected.testReferenceCompatibilityVCFWithMD5s.table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-V", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    @DataProvider(name = "testReferenceCompatibilityVCFWithoutMD5sData")
    public Object[][] testReferenceCompatibilityVCFWithoutMD5sData() {
        return new Object[][]{
                // reference, dictionary for comparison, expected output
                new Object[]{
                        // exact match
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta"),
                        new File(getToolTestDataDir() + "example_variants_withSequenceDict_withoutmd5s.vcf"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityVCFWithoutMD5s_match.table")
                },
                new Object[]{
                        // subset
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini.fasta"),
                        new File(getToolTestDataDir() + "example_variants_withSequenceDict_withoutmd5s_missingchr1.vcf"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityVCFWithoutMD5s_subset.table")
                },
                new Object[]{
                        // not compatible
                        new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_missingchr3.fasta"),
                        new File(getToolTestDataDir() + "example_variants_withSequenceDict_withoutmd5s.vcf"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityVCFWithoutMD5s_notcompatible.table")
                },
        };
    }

    @Test(dataProvider = "testReferenceCompatibilityVCFWithoutMD5sData")
    public void testReferenceCompatibilityVCFWithoutMD5s(File ref1, File dict, File expectedOutput) throws IOException {
        final File output = createTempFile("testReferenceCompatibilityVCFWithoutMD5s", ".table");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-V", dict.getAbsolutePath(), "-O", output.getAbsolutePath()};
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expectedOutput);
    }

    @Test(expectedExceptions = UserException.BadInput.class)
    public void testReferenceCompatibilityBAMandVCF() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_chr2snp.fasta");
        final File bam = new File(getToolTestDataDir() + "reads_data_source_test1_withoutmd5s.bam");
        final File vcf = new File(getToolTestDataDir() + "example_variants_withSequenceDict_withoutmd5s.vcf");

        final String[] args = new String[] {"-refcomp", ref1.getAbsolutePath() , "-I", bam.getAbsolutePath(), "-V", vcf.getAbsolutePath()};
        runCommandLine(args);
    }

    @Test(expectedExceptions = UserException.BadInput.class)
    public void testReferenceCompatibilityNoBAMOrVCF() throws IOException {
        final File ref1 = new File(COMPARE_REFERENCES_TEST_FILES_DIRECTORY + "hg19mini_chr2snp.fasta");

        final String[] args = new String[]{"-refcomp", ref1.getAbsolutePath()};
        runCommandLine(args);
    }

    @DataProvider(name = "cloudInputData")
    public Object[][] cloudInputData() {
        String testBucket = getGCPTestInputPath() + "org/broadinstitute/hellbender/tools/reference/";
        return new Object[][]{
                // dict, list of refs, expected output, isBAM, isVCF
                new Object[]{testBucket + "reads_data_source_test1_withmd5s.bam",
                        Arrays.asList(testBucket + "hg19mini.fasta"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityBAMWithMD5s_exactmatch.table"),
                        true, false
                },
                new Object[]{ testBucket + "example_variants_withSequenceDict_withmd5.vcf",
                        Arrays.asList(testBucket + "hg19mini.fasta"),
                        new File(getToolTestDataDir(), "expected.testReferenceCompatibilityVCFWithMD5s.table"),
                        false, true
                },
        };
    }

    @Test(dataProvider = "cloudInputData", groups = "bucket")
    public void testReferenceCompatibilityCloudInputs(String dict, List<String> fastas, File expected, boolean isBAM, boolean isVCF) throws IOException{
        final File output = createTempFile("testCompareReferencesCloudInputs", ".table");

        List<String> arguments = new ArrayList<>();
        for(int i = 0; i < fastas.size(); i++){
            arguments.add("-refcomp");
            arguments.add(fastas.get(i));
        }

        if(isBAM){
            arguments.add("-I");
        } else if(isVCF){
            arguments.add("-V");
        }

        arguments.add(dict);
        arguments.add("-O");
        arguments.add(output.getAbsolutePath());

        final String[] args = arguments.toArray(new String[0]);
        runCommandLine(args);

        IntegrationTestSpec.assertEqualTextFiles(output, expected);
    }
}