from pydigree.population import Population
from pydigree.individual import Individual
from pydigree.genotypes import ChromosomeTemplate
from pydigree.io import smartopen
from pydigree.exceptions import FileFormatError

from pydigree.cydigree.vcfparse import vcf_allele_parser, assign_genorow

class VCFRecord(object):
    ''' A class for parsing lines in VCF files '''
    def __init__(self, line):
        chromid, pos, varid, ref, alt, qual, filter_passed, info, field_format, data = line.strip(
        ).split(None, 9)
        self.chrom = chromid
        self.pos = int(pos)
        self.label = varid
        self.ref = ref
        self.alt = alt.split(',')
        self.qual = float(qual)
        self.filter_passed = (filter_passed == 'PASS')
        self._info = info
        self.field_format = field_format
        self.data = data

    @property
    def info(self):
        """
        A dictionary representing the INFO field in a VCF record

        :rtype dict:
        """
        infokv = [x.split('=') for x in self._info.split(';')]
        for flag in infokv:
            if len(flag) == 1:
                flag.append(True)

        return dict(infokv)

    def genotypes(self):
        """
        Extracts the genotypes from a VCF record 
        """
        alleles = vcf_allele_parser(self.data, self.field_format)

        return alleles
        
    def getitems(self, item):
        """
        Gets the value in each data field for the key specified in the
        format field.

        :param item: key to find
        :returns: values for each individual
        :rtype: list
        """
        format_ = self.field_format.split(':')
        idx = format_.index(item)
        return [x.split(':')[idx] for x in self.data]


def _vcf_parseheader(fileobj):
    pop = Population()
    for line in fileobj:

        if line.startswith('##'):
            continue

        elif line.startswith('#'):
            ind_ids = line.strip().split()[9:]
            inds = [Individual(pop, ind_id) for ind_id in ind_ids]
            for ind in inds:
                pop.register_individual(ind)

            return pop, inds

        else:
            raise FileFormatError("No header line in VCF")

def _vcf_get_infofreq(info, freq_info):
    if freq_info not in info:
        return 0.0

    freq = info[freq_info]
    if ',' in freq:
        freq = freq.split(',')[0]
    freq = float(freq)

    return freq 

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



