from __future__ import print_function
from rgt.GenomicVariant import GenomicVariant
from rgt.GenomicRegionSet import GenomicRegionSet
import vcf

"""
Represent list of GenomicVariant.

Authors: Manuel Allhoff

"""

class GenomicVariantSet(GenomicRegionSet):

    def __init__(self, vcf_path = None, name='GenomicVariantSet'):
        """Initialize"""
        GenomicRegionSet.__init__(self, name = name)
        if vcf_path:
            self.read_vcf(vcf_path)
    
    def sort(self):
        """Sort Elements by criteria defined by a GenomicRegion"""
        self.sequences.sort(cmp = GenomicVariant.__cmp__)
        self.sorted = True
         
    def read_vcf(self, vcf_path):
        """Read SNPs and InDels form VCF file"""
        self.reader = vcf.Reader(open(vcf_path), 'r')
         
        self.metadata = self.reader.metadata
        self.infos = self.reader.infos
        self.filters = self.reader.filters
        self.formats = self.reader.formats
        self.samples = self.reader.samples
         
        for v in self.reader:
            variant = GenomicVariant(v.CHROM, v.POS, v.REF, v.ALT, v.QUAL, filter = v.FILTER, id = v.ID, \
                                     info = v.INFO, format = v.FORMAT, genotype = v.genotype, samples = v.samples)
            self.add(variant)
    
    def write_vcf(self, vcf_path):
        """Write VCF file"""
        if not self.reader:
            raise Exception("No data available")
         
        writer = vcf.Writer(open(vcf_path, 'w'), self.reader)
        for v in self.sequences:
            record = vcf.model._Record(v.chrom, v.pos, v.id, v.ref, v.alt, v.qual, [], v.info, v.format, [], v.samples)
            writer.write_record(record)
 
    def filter_dbSNP(self):
        """Filter for dbSNP"""
        self.sequences = filter(lambda x: 'DB' not in x.info.keys(), self.sequences)
    
    def filter(self, at, op, t):
        """Filter for Attributes evaluate( at = <MQ, DP> <operator> <t> )"""
        self.sequences = filter(lambda x: eval(str(x.info[at]) + op + str(t)), self.sequences)
    
    def _reconstruct_info(self, GenomicRegionSet):
        """Reconstruct all information for GenomicVariantSet that
        get lost when using a GenomicRegionSet method"""
        tmp_sequences = []
        for genomic_region in GenomicRegionSet:
            c, p = genomic_region.chrom, genomic_region.initial
            for var in self.sequences:
                if var.chrom == c and var.pos == p: #1:1 mapping
                    tmp_sequences.append(var)
                    break
        
        return tmp_sequences
    
    def subtract(self, x):
        tmp = GenomicRegionSet.subtract(self, x, whole_region=False)
        self.sequences = self._reconstruct_info(tmp)
    
    def intersect(self, y):
        tmp = self._intersect(y)
        self.sequences = self._reconstruct_info(tmp)
    
    def _intersect(self, y, rm_duplicates=False):
        """Return the overlapping regions with three different modes.
                   
        (mode = OverlapType.ORIGINAL)
        Return the regions of original GenomicRegionSet which have any intersections with y.
        
            Keyword arguments:
            y -- the GenomicRegionSet which to compare with
            
            Return:
            z -- the regions of original GenomicRegionSet which have any intersections with y
            
            Graphical explanation:
            self       ----------              ------
            y              ----------                    ----
            Result     ----------
            
        """
        a = self
        b = y
        
        z = GenomicRegionSet(a.name + ' + ' + b.name)
        # XXX - someone putted an special symbol and spaces in the name! this is used as file name, never use strange characters.
        if len(a) == 0 or len(b) == 0: return z
        
        else:
            # If there is overlap within self or y, they should be merged first. 
            if a.sorted == False: a.sort()
            if b.sorted == False: b.sort()
            
            iter_a = iter(a)
            s = iter_a.next()
            last_j = len(b)-1
            j = 0
            cont_loop = True
                 
            ########################### OverlapType.ORIGINAL ###################################
            while cont_loop:
                #print(str(s),"\t",str(b[j]))
                # When the regions overlap
                if s.overlap(b[j]):
                    z.add(s)
                    
                    
                    try: s = iter_a.next()
                    except: cont_loop = False
                    
                elif s < b[j]:
                    try: s = iter_a.next()
                    except: cont_loop = False
                elif s > b[j]:
                    if j == last_j: cont_loop = False
                    else: j = j + 1
                else:
                    try: s = iter_a.next()
                    except: cont_loop = False
                        
            if rm_duplicates: 
                z.remove_duplicates()
            
            return z
    
    
if __name__ == '__main__':
    s = GenomicVariantSet('/home/manuel/data/humangenetics/01_S1_L001_R1_001.filtered.vcf')
    b = GenomicVariantSet('/home/manuel/data/humangenetics/K28_S8_L001_R1_001.filtered.vcf')
    print(len(s))
    s.subtract(b)
    
    print(len(s))
    #print(b.sequences[:10])
    #print(c.sequences[:10])
    print(len(s))
    s.filter('MQ', '>=', 40.0)
    print(len(s))
    #s.filter_DB()
    #print(len(s))
    #s.write_vcf('/home/manuel/1.vcf')
    