class VCFReader(object):
    def __init__(self, fhand, compressed=None, filename=None,
                 min_calls_for_pop_stats=DEF_MIN_CALLS_FOR_POP_STATS):
        self.fhand = fhand
        self.pyvcf_reader = pyvcfReader(fsock=fhand, compressed=compressed,
                                        filename=filename)
        self.min_calls_for_pop_stats = min_calls_for_pop_stats
        self._snpcaller = None
        self._samples = self.pyvcf_reader.samples
        self._header_lines = self.pyvcf_reader._header_lines
        self._column_headers = self.pyvcf_reader._column_headers

    def parse_snvs(self):
        min_calls_for_pop_stats = self.min_calls_for_pop_stats
        last_snp = None
        try:
            counter =0
            for snp in self.pyvcf_reader:
                counter +=1
                snp = SNV(snp, reader=self,
                          min_calls_for_pop_stats=min_calls_for_pop_stats)
                last_snp = snp
                yield snp
        except Exception:
            from traceback import print_exception
            exc_type, exc_value, exc_traceback = sys.exc_info()

            print_exception(exc_type, exc_value, exc_traceback,
                            limit=20, file=sys.stderr)

            if last_snp is not None:
                chrom = str(last_snp.chrom)
                pos = last_snp.pos
                msg = 'Last parsed SNP was: {} {}\n'.format(chrom, pos + 1)
                sys.stderr.write(msg)
                raise

    def fetch_snvs(self, chrom, start, end=None):
        min_calls_for_pop_stats = self.min_calls_for_pop_stats
        try:
            snvs = self.pyvcf_reader.fetch(chrom, start + 1, end=end)
        except KeyError:
            snvs = []
        if snvs is None:
            snvs = []

        for snp in snvs:
            snp = SNV(snp, reader=self,
                      min_calls_for_pop_stats=min_calls_for_pop_stats)
            yield snp

    def sliding_windows(self, size, step=None, ref_fhand=None,
                        min_num_snps=DEF_MIN_NUM_SNPS_IN_WIN):
        random_snp_reader = VCFReader(open(self.fhand.name))
        sliding_window = _SNPSlidingWindow(snp_reader=random_snp_reader,
                                           win_size=size, win_step=step,
                                           min_num_snps=min_num_snps,
                                           ref_fhand=ref_fhand)
        for window in sliding_window.windows():
            yield window

    @property
    def snpcaller(self):
        if self._snpcaller is not None:
            return self._snpcaller

        metadata = self.pyvcf_reader.metadata
        if 'source' in metadata:
            if 'VarScan2' in metadata['source']:
                snpcaller = VARSCAN
            elif 'freebayes' in metadata['source'][0].lower():
                snpcaller = FREEBAYES
            else:
                snpcaller = GENERIC
        elif 'UnifiedGenotyper' in metadata:
            snpcaller = GATK
        else:
            snpcaller = GENERIC
        self._snpcaller = snpcaller
        return snpcaller

    @property
    def samples(self):
        return self._samples

    @property
    def filters(self):
        return self.pyvcf_reader.filters

    @property
    def infos(self):
        return self.pyvcf_reader.infos

    def _build_header(self, samples=None):
        if samples is None:
            samples = self.samples
        header = '\n'.join(self._header_lines)
        header += '\n#' + '\t'.join(self._column_headers)
        header += '\t' + '\t'.join(samples)
        return header

    @property
    def header(self):
        return self._build_header()

    def create_template_reader(self, samples=None):
        vcf = self._build_header(samples)
        fake_vcf = BytesIO(vcf)
        return VCFReader(fake_vcf) def compress_with_bgzip(in_fhand, compressed_fhand):
    '''It compresses the input fhand.

    The input fhand has to be a real file with a name.
    '''
    cmd = ['bgzip', '-c', in_fhand.name]
    check_call(cmd, stdout=compressed_fhand) def index_vcf_with_tabix(in_fpath):
    cmd = ['tabix', '-p', 'vcf', in_fpath]
    tabix = Popen(cmd, stdout=PIPE, stderr=PIPE)
    stdout, stderr = tabix.communicate()
    if tabix.returncode:
        msg = 'Some problem indexing with tabix.\n'
        msg += 'stdout: ' + stdout
        msg += 'stderr: ' + stderr
        raise RuntimeError(msg) class IlluminaWriter(object):
    '''It writes the SNPs in Illumina format

    ref_fpath should be in fasta format and it has to have a name attribute.
    min_maf controls the SNPs reported in the adjacent segments as IUPAC codes.
    '''

    # TODO add extra error classes
    # TODO include the error classes inside this class to easy access
    class NotEnoughAdjacentSequenceError(Exception):
        pass

    def __init__(self,  ref_fpath, out_fhand, length=60, vcf_fpath=None,
                 min_length=None):
        ''''It inits.

        The vcf will be used to replace in the reference sequence the SNPs
        around the SNP of interest with IUPAC codes
        '''
        self._sep = u'\t'
        self._len = length
        if min_length is None:
            min_length = length
        if min_length > length:
            msg = 'Minimum length must be smaller than required length'
            raise ValueError(msg)
        self._min_len = min_length

        self._ref_seqs = seq_index(ref_fpath, format='fasta')

        if vcf_fpath:
            self._snvs = Reader(filename=vcf_fpath)
        else:
            self._snvs = None
        self._out_fhand = out_fhand
        out_fhand.write(u'CHROM\tPOS\tID\tseq\n')
        self._prev_chrom = None

    def write(self, snv):
        chrom_name = snv.CHROM

        prev_chrom = self._prev_chrom
        if prev_chrom is None or prev_chrom.name != chrom_name:
            chrom = self._ref_seqs[chrom_name]
            self._prev_chrom = chrom
        else:
            chrom = prev_chrom

        length = self._len
        min_len = self._min_len

        snv_start = snv.start   # 0 based
        snv_end = snv.end       # 1 based
        desired_start = snv_start - length  # desired segment start
        end = snv_end + length      # desired segment end
        chrom_seq = chrom.seq
        first_segment = unicode(chrom_seq[desired_start:snv_start])

        if len(first_segment) < min_len:
            msg = "Not enough sequence in 3'. ID: %s, POS: %d, CHROM: %s"
            msg %= (snv.ID, snv.POS, snv.CHROM)
            raise self.NotEnoughAdjacentSequenceError(msg)

        if self._snvs:
            real_start = snv_start - len(first_segment)
            close_snvs = self._snvs.fetch(chrom.name, start=real_start,
                                          end=snv_start)
            first_segment = _replace_snvs_with_iupac(first_segment, close_snvs,
                                                     seq_offset=real_start)

        snv_segment = _build_snv_section(snv)
        second_segment = unicode(chrom_seq[snv_end:end])
        if len(second_segment) < min_len:
            msg = "Not enough sequence in 5'. ID: %s, POS: %d, CHROM: %s"
            msg %= (snv.ID, snv.POS, snv.CHROM)
            raise self.NotEnoughAdjacentSequenceError(msg)

        if self._snvs:
            real_end = snv_end + len(second_segment)
            close_snvs = self._snvs.fetch(chrom.name, start=snv_end,
                                          end=real_end)
            second_segment = _replace_snvs_with_iupac(second_segment,
                                                      close_snvs,
                                                      seq_offset=snv_end)

        out_fhand = self._out_fhand
        sep = self._sep
        out_fhand.write(unicode(snv.CHROM))
        out_fhand.write(sep)
        out_fhand.write(unicode(snv.POS))
        out_fhand.write(sep)
        snp_id = snv.ID
        if snp_id is None:
            snp_id = u'.'
        out_fhand.write(snp_id)
        out_fhand.write(sep)
        out_fhand.write(first_segment)
        out_fhand.write(snv_segment)
        out_fhand.write(second_segment)
        out_fhand.write(u'\n')

    def flush(self):
        self._out_fhand.flush()

    def close(self):
        self._out_fhand.close() def _replace_snvs_with_iupac(str_seq, snvs, seq_offset):
    # TODO add parameter that accepts a function to be run fore
    # every SNP, to select which ones get put as IUPAC, e.g. maf filter
    seq = list(str_seq)
    for snv in snvs:
        start = snv.start - seq_offset
        if start < 0:
            continue
        end = snv.end - seq_offset
        if end > len(seq):
            continue
        alleles = tuple(sorted(_get_unicode_alleles(snv)))
        if snv.is_snp:
            iupac_snv = IUPAC_CODES[alleles]
            seq[start] = iupac_snv
        elif snv.is_indel:
            if snv.is_deletion:
                for loc in range(start + 1, end):
                    seq[loc] = '-'
            elif len(alleles) > 2:
                # complex event
                seq[start] = '*'
            else:
                seq[start] = seq[start] + '^'
        else:
            seq[start] = seq[start] + '*'
    return u''.join(seq) class RQTLWriter(object):
    geno_coding = {('A', 'A'): 'AA',
                   ('A', '-'): 'A-',
                   ('-', 'A'): 'A-',
                   ('B', 'B'): 'BB',
                   ('B', '-'): 'B-',
                   ('-', 'B'): 'B-',
                   ('A', 'B'): 'AB',
                   ('B', 'A'): 'AB'}

    def __init__(self, out_fhand, phys_to_genet_dist=None):
        self._fhand = out_fhand
        self._snv_count = 0
        self.sep = ','
        self._first_snp = True
        self._phys_to_genet_dist = phys_to_genet_dist

    def _write_header_line(self, snv):
        line = ['phenot', '', '']
        line.extend(unicode(idx) for idx in range(1, len(snv.calls) + 1))
        self._fhand.write(self.sep.join(line))
        self._fhand.write(unicode('\n'))

    def _calls_to_ab(self, snv):
        if self._first_snp:
            self._write_header_line(snv)
            self._first_snp = False

        genotype = []
        ab_coding = {'.': '-'}

        for call in snv.record.samples:
            if not call.called:
                ab_geno = 'NA'
            else:
                geno = call.gt_alleles
                ab_geno = []
                for allele in geno:
                    if allele not in ab_coding:
                        # We do not know the phase of the parents so
                        # AB won't be related with the parent phases
                        ab_allele = chr(len(ab_coding) + 64)
                        ab_coding[allele] = ab_allele
                    else:
                        ab_allele = ab_coding[allele]
                    ab_geno.append(ab_allele)
                # we have to destroy the phase of the individual
                try:
                    ab_geno = self.geno_coding[tuple(ab_geno)]
                except KeyError:
                    msg = 'Until now, there is only support for bialleic SNPs'
                    raise NotImplementedError(msg)
            genotype.append(ab_geno)
        return genotype

    def write(self, snv):
        line = []
        self._snv_count += 1
        line.append(unicode(self._snv_count))
        line.append(unicode(snv.chrom))

        pos = snv.pos + 1
        if self._phys_to_genet_dist is not None:
            pos /= self._phys_to_genet_dist
        else:
            msg = 'The genetic map availability has not been implemented yet'
            raise NotImplementedError(msg)
        pos = '%.6e' % (pos,)
        line.append(pos)

        ab_genotypes = self._calls_to_ab(snv)
        line.extend(ab_genotypes)

        self._fhand.write(self.sep.join(line))
        self._fhand.write(unicode('\n'))

    def flush(self):
        self._fhand.flush()

    def close(self):
        self._fhand.close() DEF_PHYS_TO_GENET_DIST = 1000000 def write_parent_checker(vcf_fhand, parents_a, parents_b, genos_fhand,
                         phys_map_fhand=None,
                         coder_threshold=DEF_AB_CODER_THRESHOLD,
                         smooth_threshold=None, recomb_threshold=None):
    sep = '\t'
    coder = ABCoder(vcf_fhand, parents_a, parents_b,
                    parent_index_threshold=coder_threshold,
                    smooth_threhsold=smooth_threshold,
                    recomb_threshold=recomb_threshold)
    samples = coder.offspring
    snp_ids = []
    snp_genos = []
    coding = 'ascii'

    if phys_map_fhand is not None:
        phys_map_fhand.write('marker\tposition\tChromosome\n')

    for snp, genos in coder.recode_genotypes(samples):
        snp_id = get_or_create_id(snp).encode(coding)
        snp_ids.append(snp_id)
        geno_array = array('c')
        for geno in genos.values():
            geno_array.append(_code_to_one_letter(geno))
        snp_genos.append(geno_array)
        if phys_map_fhand is not None:
            phys_map_fhand.write(snp_id)
            phys_map_fhand.write(sep)
            phys_map_fhand.write(str(snp.POS))
            phys_map_fhand.write(sep)
            phys_map_fhand.write(snp.CHROM.encode(coding))
            phys_map_fhand.write('\n')

    genos_fhand.write('ID')
    genos_fhand.write(sep)
    genos_fhand.write(sep.join(snp_ids))
    genos_fhand.write('\n')

    for sample_idx, sample in enumerate(samples):
        genos_fhand.write(sample.encode(coding))
        genos_fhand.write(sep)
        to_write = sep.join(snp_genos[snp_idx][sample_idx] for snp_idx in range(len(snp_ids)))
        genos_fhand.write(to_write)
        genos_fhand.write('\n')

    return coder VCF_BIN_DIR = os.path.join(BIN_DIR, 'vcf') TEST_DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
                                             '..', 'test', 'test_data'))
import unittest
from os.path import join as pjoin
from io import StringIO
from io import BytesIO
from os import remove
from tempfile import NamedTemporaryFile
from subprocess import check_output, Popen, PIPE
from vcf import Reader
from crumbs.vcf.snv import VCFReader
from crumbs.utils.file_utils import compress_with_bgzip, index_vcf_with_tabix
from crumbs.vcf.writers import (IlluminaWriter, _replace_snvs_with_iupac,
                                RQTLWriter, DEF_PHYS_TO_GENET_DIST,
                                write_parent_checker)
from crumbs.utils.bin_utils import VCF_BIN_DIR
from crumbs.utils.test_utils import TEST_DATA_DIR
# Copyright 2012 Jose Blanca, Peio Ziarsolo, COMAV-Univ. Politecnica Valencia
# This file is part of ngs_crumbs.
# ngs_crumbs is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# ngs_crumbs is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR  PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with ngs_crumbs. If not, see <http://www.gnu.org/licenses/>.




# Method could be a function
# pylint: disable=R0201
# Too many public methods
# pylint: disable=R0904
# Missing docstring
# pylint: disable=C0111

VCF_PATH = pjoin(TEST_DATA_DIR, 'sample.vcf.gz')
REF_PATH = pjoin(TEST_DATA_DIR, 'sample_ref.fasta')
VCF_INDEL_PATH = pjoin(TEST_DATA_DIR, 'sample_indel.vcf.gz')


class BareBoneSNP(object):
    def __init__(self, alleles, start, end, is_sv=False):
        self.alleles = alleles
        self.POS = start + 1
        self.start = start
        self.end = end
        allele_lens = [len(allele) for allele in alleles]
        self.is_sv = is_sv
        if is_sv:
            # sv
            self.is_snp = False
            self.is_deletion = False
            self.is_indel = False
        elif len(set(allele_lens)) == 1:
            # snp
            self.is_snp = True
            self.is_deletion = False
            self.is_indel = False
        elif allele_lens[0] == 1:
            # insertion
            self.is_snp = False
            self.is_deletion = False
            self.is_indel = True
        else:
            # deletion
            self.is_snp = False
            self.is_deletion = True
            self.is_indel = True
        self.is_sv = is_sv


class IlluminaWriterTest(unittest.TestCase):
    def test_illumina_writer(self):
        # 01234
        # 1234567890
        # CCTGATTT-A
        # TAACGA
        #   -  C -A
        vcf = '''##fileformat=VCFv4.1
#CHROM POS ID REF ALT QUAL FILTER INFO
ref 1 . C T 10 PASS .
ref 2 . CT CA,C 10 PASS .
ref 3 . T A 10 PASS .
ref 4 . G C 10 PASS .
ref 5 . A G 10 PASS .
ref 6 . T A,C 10 PASS .
ref 7 . TT T 10 PASS .
ref 8 . T TA 10 PASS .
ref 10 . A C 10 PASS .
'''

        vcf = vcf.replace(' ', '\t')
        vcf_fhand = NamedTemporaryFile(suffix='.vcf')
        vcf_fhand.write(vcf)
        vcf_fhand.flush()
        vcf_compressed = NamedTemporaryFile(suffix='.vcf.gz')
        compress_with_bgzip(vcf_fhand, vcf_compressed)

        index_vcf_with_tabix(vcf_compressed.name)