# Path: pydigree/io/vcf.py
# def read_vcf(filename, require_pass=False, freq_info=None):
#     """
#     Reads a VCF file and returns a Population object with the
#     individuals represented in the file
#     
#     Genotypes generated by this function will be sparse
# 
#     :param require_pass: only allow variants with PASS under FILTER
#     :type require_pass: bool
#     :param freq_info: INFO field to get allele frequency from
#     :param freq_info: string
# 
#     :returns: Individuals in the VCF
#     :rtype: Population
#     """
#     with smartopen(filename) as f:
# 
#         genotypes = []
# 
#         pop, inds = _vcf_parseheader(f)
# 
#         last_chrom = None
#         chromobj = None
# 
#         for line in f:
#             record = VCFRecord(line)
# 
#             if require_pass and not record.filter_passed:
#                 continue
# 
#             if record.chrom != last_chrom:
#                 if last_chrom is not None:
#                     pop.add_chromosome(chromobj)
#                 chromobj = ChromosomeTemplate(label=record.chrom)
# 
#             if freq_info is not None:
#                 freq = _vcf_get_infofreq(record.info, freq_info)
#             else:
#                 freq = 0
# 
#             genorow = record.genotypes()
#             genotypes.append(genorow)
# 
#             chromobj.add_genotype(bp=record.pos,
#                                   label=record.label,
#                                   frequency=freq)
# 
#             last_chrom = record.chrom
# 
#         pop.add_chromosome(chromobj)
#         pop.chromosomes.finalize()
#     for ind in inds:
#         # Initialize new genotypes
#         ind._init_genotypes(sparse=True)
# 
#     # Now actually sift through markers and assign them to individuals
#     final_indices = []
#     for chromidx, chromobj  in enumerate(pop.chromosomes):
#         indices = zip([chromidx]*chromobj.nmark(), range(chromobj.nmark()))
#         final_indices.extend(indices)
# 
#     raw_indices = range(len(genotypes))
# 
#     for raw, final in zip(raw_indices, final_indices):
#         chromidx, markidx = final
#         row = genotypes[raw]
#         assign_genorow(row, inds, chromidx, markidx)
# 
#         # Kill the row so we don't end up with the whole dataset in memory twice
#         genotypes[raw] = None
#     
#     return pop


import numpy as np
import os
import numpy as np
from pydigree.io.vcf import read_vcf
from pydigree.cydigree.vcfparse import vcf_allele_parser

testdir = os.path.dirname(os.path.abspath(__file__))
TESTDATA_DIR = os.path.join(testdir, 'test_data', 'vcf')

def test_vcf():
    testvcf = os.path.join(TESTDATA_DIR, 'test.vcf')
    
    pop = read_vcf(testvcf)
    assert len(pop.individuals) == 3
    assert len(pop.chromosomes) == 2

    # Test individual with good genotypes
    assert (pop['NA00001'].genotypes[1][0].todense() == np.array([0, 0, 1, 0, 0, 0, 1])).all()
    assert (pop['NA00001'].genotypes[1][1].todense() == np.array([0, 0, 2, 0, 1, 1, 1])).all()
    # assert not pop['NA00001'].genotypes[1][0].missing.all()

    # Test individual with bad genotypes.
    # assert pop['NA00003'].genotypes[1][0].missing.all()


    # Test for FILTER == PASS
    # pop = read_vcf(testvcf)
    # assert pop.chromosomes[1].nmark() == 6


# Path: pydigree/io/vcf.py
# def read_vcf(filename, require_pass=False, freq_info=None):
#     """
#     Reads a VCF file and returns a Population object with the
#     individuals represented in the file
#     
#     Genotypes generated by this function will be sparse
# 
#     :param require_pass: only allow variants with PASS under FILTER
#     :type require_pass: bool
#     :param freq_info: INFO field to get allele frequency from
#     :param freq_info: string
# 
#     :returns: Individuals in the VCF
#     :rtype: Population
#     """
#     with smartopen(filename) as f:
# 
#         genotypes = []
# 
#         pop, inds = _vcf_parseheader(f)
# 
#         last_chrom = None
#         chromobj = None
# 
#         for line in f:
#             record = VCFRecord(line)
# 
#             if require_pass and not record.filter_passed:
#                 continue
# 
#             if record.chrom != last_chrom:
#                 if last_chrom is not None:
#                     pop.add_chromosome(chromobj)
#                 chromobj = ChromosomeTemplate(label=record.chrom)
# 
#             if freq_info is not None:
#                 freq = _vcf_get_infofreq(record.info, freq_info)
#             else:
#                 freq = 0
# 
#             genorow = record.genotypes()
#             genotypes.append(genorow)
# 
#             chromobj.add_genotype(bp=record.pos,
#                                   label=record.label,
#                                   frequency=freq)
# 
#             last_chrom = record.chrom
# 
#         pop.add_chromosome(chromobj)
#         pop.chromosomes.finalize()
#     for ind in inds:
#         # Initialize new genotypes
#         ind._init_genotypes(sparse=True)
# 
#     # Now actually sift through markers and assign them to individuals
#     final_indices = []
#     for chromidx, chromobj  in enumerate(pop.chromosomes):
#         indices = zip([chromidx]*chromobj.nmark(), range(chromobj.nmark()))
#         final_indices.extend(indices)
# 
#     raw_indices = range(len(genotypes))
# 
#     for raw, final in zip(raw_indices, final_indices):
#         chromidx, markidx = final
#         row = genotypes[raw]
#         assign_genorow(row, inds, chromidx, markidx)
# 
#         # Kill the row so we don't end up with the whole dataset in memory twice
#         genotypes[raw] = None
#     
#     return pop


# Path: tests/test_vcf.py
import numpy as np
import os
import numpy as np
from pydigree.io.vcf import read_vcf
from pydigree.cydigree.vcfparse import vcf_allele_parser


testdir = os.path.dirname(os.path.abspath(__file__))
TESTDATA_DIR = os.path.join(testdir, 'test_data', 'vcf')

def test_vcf():
    testvcf = os.path.join(TESTDATA_DIR, 'test.vcf')
    
    pop = read_vcf(testvcf)
    assert len(pop.individuals) == 3
    assert len(pop.chromosomes) == 2

    # Test individual with good genotypes
    assert (pop['NA00001'].genotypes[1][0].todense() == np.array([0, 0, 1, 0, 0, 0, 1])).all()
    assert (pop['NA00001'].genotypes[1][1].todense() == np.array([0, 0, 2, 0, 1, 1, 1])).all()
    # assert not pop['NA00001'].genotypes[1][0].missing.all()

    # Test individual with bad genotypes.
    # assert pop['NA00003'].genotypes[1][0].missing.all()


    # Test for FILTER == PASS
    # pop = read_vcf(testvcf)
    # assert pop.chromosomes[1].nmark() == 6


    pop = read_vcf(testvcf, freq_info='AF')