package org.opencb.hpg.bigdata.core.io;

import htsjdk.variant.variantcontext.VariantContext;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.models.variant.avro.VariantAvro;
import org.opencb.biodata.tools.variant.converters.avro.VariantContextToVariantConverter;
import org.opencb.commons.run.ParallelTaskRunner;
import org.opencb.hpg.bigdata.core.avro.VariantAvroAnnotator;
import org.opencb.hpg.bigdata.core.io.avro.AvroEncoder;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Predicate;

/**
 * Created by jtarraga on 04/10/17.
 */
public class ConvertTask implements ParallelTaskRunner.Task<VariantContext, VariantAvro> {
    private VariantContextToVariantConverter converter;
    private List<List<Predicate<VariantAvro>>> filters;
    private boolean annotate = false;

    protected final AvroEncoder<VariantAvro> encoder;

    public ConvertTask(VariantContextToVariantConverter converter, List<List<Predicate<VariantAvro>>> filters,
                       boolean annotate) {
        this.converter = converter;
        this.filters = filters;
        this.annotate = annotate;

        this.encoder = new AvroEncoder<>(VariantAvro.getClassSchema());
    }

    @Override
    public List<VariantAvro> apply(List<VariantContext> variantContexts) throws RuntimeException {
        List<VariantAvro> variantAvros;
        if (annotate) {
            // Annotate before converting to Avro
            // Duplicate code for efficiency purposes
            VariantAvroAnnotator variantAvroAnnotator = new VariantAvroAnnotator();
            List<VariantAvro> variants = new ArrayList<>(variantContexts.size());

            for (VariantContext vc : variantContexts) {
                try {
                    Variant variant = converter.convert(vc);
                    if (filter(variant.getImpl())) {
                        variants.add(variant.getImpl());
//                    statsCalculator.updateGlobalStats(variant);
                    }
                } catch (IllegalArgumentException e) {
                    e.printStackTrace();
                }
            }
            // Annotate variants and then write them to disk
            variantAvros = variantAvroAnnotator.annotate(variants);
        } else {
            // Convert to Avro without annotating
            variantAvros = new ArrayList<>(variantContexts.size());
            for (VariantContext vc : variantContexts) {
                Variant variant = converter.convert(vc);
                if (filter(variant.getImpl())) {
                    variantAvros.add(variant.getImpl());
                }
            }
        }

        // Return variants
        return variantAvros;
    }

    public boolean filter(VariantAvro record) {
        for (List<Predicate<VariantAvro>> list: filters) {
            if (list.size() == 1) {
                if (!list.get(0).test(record)) {
                    return false;
                }
            } else if (list.size() > 1) {
                boolean or = false;
                for (Predicate<VariantAvro> filter: list) {
                    if (filter.test(record)) {
                        or = true;
                        break;
                    }
                }
                if (!or) {
                    return false;
                }
            }
        }
        return true;
    }
}

@@ -196,16 +196,75 @@
 ntAvro;%0A
+import org.opencb.biodata.tools.variant.VariantNormalizer;%0A
 import o
@@ -885,24 +885,75 @@
 e = false;%0A%0A
+    protected VariantNormalizer variantNormalizer;%0A
     protecte
@@ -1248,24 +1248,99 @@
  annotate;%0A%0A
+        this.variantNormalizer = new VariantNormalizer(true, true, false);%0A
         this
@@ -1563,224 +1563,69 @@
 os;%0A
-        if (annotate) %7B%0A            // Annotate before converting to Avro%0A            // Duplicate code for efficiency purposes%0A            VariantAvroAnnotator variantAvroAnnotator = new VariantAvroAnnotator();%0A    
+%0A        // Convert and filter (from VariantContext to Avro)%0A
     
@@ -1632,36 +1632,32 @@
     List%3CVariant
-Avro
 %3E variants = new
@@ -1694,21 +1694,16 @@
 ize());%0A
-%0A    
         
@@ -1750,38 +1750,8 @@
 ) %7B%0A
-                try %7B%0A        
     
@@ -1807,32 +1807,24 @@
             
-        
 if (filter(v
@@ -1856,32 +1856,24 @@
             
-  
-      
 variants.add
@@ -1884,95 +1884,11 @@
 iant
-.getImpl());%0A//                    statsCalculator.updateGlobalStats(variant);%0A        
+);%0A
     
@@ -1909,200 +1909,40 @@
     
-  
-      %7D catch (IllegalArgumentException e) %7B%0A                    e.printStackTrace();%0A                %7D%0A            %7D%0A            // Annotate variants and then write them to disk%0A   
+%7D%0A%0A        // Normalize%0A
         
  var
@@ -1933,36 +1933,31 @@
 ize%0A        
- 
 variant
-Avro
 s = variantA
@@ -1959,113 +1959,37 @@
 iant
-AvroAnnotator.annotate(variants);%0A        %7D else %7B%0A            // Convert to Avro without annotating%0A    
+Normalizer.apply(variants);%0A%0A
     
@@ -2022,39 +2022,32 @@
 ayList%3C%3E(variant
-Context
 s.size());%0A     
@@ -2045,28 +2045,24 @@
 ));%0A        
-   
- 
 for (Variant
@@ -2057,35 +2057,32 @@
 for (Variant
-Context vc 
+ variant
 : variantCon
@@ -2070,39 +2070,32 @@
 variant: variant
-Context
 s) %7B%0A           
@@ -2099,50 +2099,42 @@
     
-   
- V
+v
 ariant
- variant = converter.convert(vc
+Avros.add(variant.getImpl()
 );%0A 
@@ -2132,32 +2132,35 @@
 mpl());%0A        
+%7D%0A%0A
         if (filt
@@ -2159,43 +2159,20 @@
 if (
-filter(variant.getImpl())) %7B%0A      
+annotate) %7B%0A
     
@@ -2171,35 +2171,33 @@
 ) %7B%0A            
-  v
+V
 ariantAvros.add(
@@ -2194,65 +2194,140 @@
 Avro
-s.add(variant.getImpl());%0A                %7D%0A            %7D
+Annotator variantAvroAnnotator = new VariantAvroAnnotator();%0A            variantAvros = variantAvroAnnotator.annotate(variantAvros);
 %0A   
@@ -2360,16 +2360,31 @@
 variants
+ (Avro objects)
 %0A       

normalize variant when converting VCF to Avro/Parquet #131