import os
import gzip
import multiprocessing
import itertools

import pysam
import numpy as np

import handygenome.tools as tools
import handygenome.refgenome.refgenome as refgenome
import handygenome.interval as libinterval
import handygenome.logutils as logutils
from handygenome.genomedf.genomedf import GenomeDataFrame as GDF


# pysam VariantFile mode related ones

"""
pysam file mode string

- Valid mode string pattern : ^[rwa]([bzu]?[0-9]?|[0-9]?[bzu]?)$
- Conflicting characters: (b,z,u)

- 'wb' : Compressed BCF (level 6), regardless of file name. 
- 'wb[0-9]' : BCF with level of compression indicated by the number, regardless of file name. 

- 'wz' : Compressed VCF (level 6), regardless of file name.
- 'wz[0-9]' : VCF with level of compression indicated by the number, regardless of file name. 

- 'wu' : Uncompressed VCF, regardless of file name.
- 'wu[0-9]' : Uncompressed VCF, regardless of file name.

- 'w[0-9]' : Uncompressed VCF, regardless of file name. 

- 'w' :
    *.vcf : uncompressed VCF
    *.vcf.gz : compressed VCF (level 6)
    *.bcf : compressed BCF (level 6)
    *.bcf.gz : compressed VCF (level 6)
"""

CYVCF2_FORMAT_DICT = {'v':'w', 'z':'wz', 'u':'wbu', 'b':'wb'}  
    # keys: bcftools format strings
    # values: cyvcf2 format strings
PYSAM_MODE_DICT = {'v':'wz0', 'z':'wz', 'u':'wb0', 'b':'wb'}
    # keys: bcftools format strings
    # values: pysam format strings
#DEFAULT_MODE_BCFTOOLS = 'z'


def write_mode_arghandler(mode_bcftools, mode_pysam):
    """
    Args:
        mode_bcftools: bcftools style mode string
        mode_pysam: pysam style mode string

    If 'mode_pysam' is set as a non-None value,
    it overrides 'mode_bcftools'.
    """
    if mode_pysam is None:
        return PYSAM_MODE_DICT[mode_bcftools]
    else:
        return mode_pysam

###

def get_vcf_format(vcf_path):
    with pysam.VariantFile(vcf_path) as vcf:
        is_vcf = vcf.is_vcf
        #is_bcf = vcf.is_bcf
        comp = vcf.compression

    if comp not in ('NONE', 'GZIP', 'BGZF'):
        raise Exception(f'Unexpected compression type')

    is_bgzf = (comp == 'BGZF')

    return is_vcf, comp, is_bgzf


def check_has_index(vcf_path):
    with pysam.VariantFile(vcf_path) as vcf:
        has_index = (vcf.index is not None)
    return has_index


def get_indexfile_path(vcf_path):
    if os.path.exists(f'{vcf_path}.csi'):
        return f'{vcf_path}.csi'
    elif os.path.exists(f'{vcf_path}.tbi'):
        return f'{vcf_path}.tbi'
    else:
        return None


def make_index(vcf_path):
    index_filename = f'{vcf_path}.csi'
    pysam.tabix_index(vcf_path, preset='vcf', csi=True, index=index_filename)
    return index_filename


###################################################
# VCF position extraction and fetch - new version #
###################################################

'''pysam.VariantFile.fetch regards span of a VCF record as an interval:
start0 = POS - 1
end0 = start0 + len(REF)
'''

def get_vcf_positions_new(vcf_path, split=None, nproc=1):
    if split is None:
        split = nproc * 10

    refver = refgenome.infer_refver_vcfpath(vcf_path)
    if not check_has_index(vcf_path):
        make_index(vcf_path) 

    allregion_gdf = GDF.all_regions(refver, assembled_only=False)
    allregion_split = allregion_gdf.equal_length_split(n=split)

    with multiprocessing.Pool(nproc) as pool:
        args = (
            (vcf_path, gdf) 
            for gdf in allregion_split
        )
        pool_result = pool.starmap(
            get_vcf_positions_new_target,
            args,
        )

    all_vr_intervals = list(itertools.chain.from_iterable(pool_result))
    chroms, start0s, end0s = zip(*all_vr_intervals)
    result_gdf = GDF.from_data(refver=refver, chroms=chroms, start0s=start0s, end0s=end0s)
    result_gdf = result_gdf.merge()

    return result_gdf


def get_vcf_positions_new_target(vcf_path, gdf):
    result = list()
    with pysam.VariantFile(vcf_path) as vcf:
        for chrom, start0, end0 in gdf.iter_coords():
            for vr in vcf.fetch(chrom, start0, end0):
                result.append(
                    (vr.chrom, vr.pos - 1, vr.pos - 1 + len(vr.ref))
                )
    return result


def get_vcf_fetchregions_new(vcf_path, n, nproc=1):
    refver = refgenome.infer_refver_vcfpath(vcf_path)
    position_gdf = get_vcf_positions_new(vcf_path, nproc=nproc)
    fetchregion_gdfs = list()
    for subgdf in position_gdf.equal_nrow_split(n=n):
        chrom_left = subgdf.chromosomes[0]
        start0_left = subgdf.starts[0]
        chrom_right = subgdf.chromosomes[-1]
        end0_right = subgdf.ends[-1]
        fetchregion_gdfs.append(
            GDF.from_margins(
                refver, chrom_left, start0_left, chrom_right, end0_right,
            )
        )

    return fetchregion_gdfs


def get_vr_fetcher_new(vcf, refver, chrom=None, start0=None, end0=None):
    if chrom is None:
        for vr in vcf.fetch():
            yield vr
    else:
        if start0 is None:
            start0 = 0
        if end0 is None:
            end0 = refgenome.get_chromdict(refver)[chrom]
        
        for vr in vcf.fetch(contig=chrom, start=start0, stop=end0):
            yield vr


###################################################
# VCF position extraction and fetch - old version #
###################################################

def get_vcf_positions(vcf_path, as_iter=False, verbose=False):
    """Returns:
        as_iter==True: A list of tuples (chrom, start0, end0) of each variant record
        as_iter==False: np.ndarray
    """
    is_vcf, comp, is_bgzf = get_vcf_format(vcf_path)
    if is_vcf:
        line_iter = get_line_iter(vcf_path, comp, verbose)
        result_iter = parse_line_iter(line_iter)
    else:  # BCF
        result_iter = get_vcf_positions_pysam(vcf_path)

    if as_iter:
        return result_iter
    else:
        #return list(result_iter)
        return np.array(
            list(result_iter), 
            dtype=[('Chromosome', object), ('Start', int), ('End', int)],
        )


def split_vcf_positions(vcf_positions, n):
    """Splits such that adjacent sub-positions do not meet each other with
    identical POS values
    """
    diffresult = np.diff(vcf_positions['Start'])
    nondiff_indexes = np.where(diffresult == 0)[0] + 1
    spittable_indexes = np.setdiff1d(np.arange(vcf_positions.shape[0]), nondiff_indexes)
    spittable_indexes_split = [
        x for x in np.array_split(spittable_indexes, n) 
        if x.shape[0] != 0
    ]
    splitpoints = [x[0] for x in spittable_indexes_split if x[0] != 0]
    split_positions = np.split(vcf_positions, splitpoints)
        # this works even when "splitpoints" is an empty list

    return split_positions


def get_fetchargs_from_vcf_positions(vcf_positions, refver):
    chrom_left = vcf_positions[0]['Chromosome']
    start0_left = vcf_positions[0]['Start']
    chrom_right = vcf_positions[-1]['Chromosome']
    end0_right = vcf_positions[-1]['Start'] + 1  
        # only the 1-base position indicated by POS value is used
    chromdict = refgenome.get_chromdict(refver)

    intvlist = libinterval.IntervalList.from_margin(
        chromdict, chrom_left, start0_left, chrom_right, end0_right,
    )
    fetchargs_list = [
        (intv.chrom, intv.start0, intv.end0) for intv in intvlist
    ]
    return fetchargs_list


def get_vcf_fetchregions(vcf_path, n, refver, verbose=False):
    all_position_info = get_vcf_positions(vcf_path, as_iter=False, verbose=verbose)
    split_position_info = split_vcf_positions(all_position_info, n)
    result = [
        get_fetchargs_from_vcf_positions(position_info, refver)
        for position_info in split_position_info
    ]

    return result


def get_vr_fetcher(vcf, refver, chrom=None, start0=None, end0=None, respect_refregion=False):
    if chrom is None:
        for vr in vcf.fetch():
            yield vr
    else:
        if start0 is None:
            start0 = 0
        if end0 is None:
            end0 = refgenome.get_chromdict(refver)[chrom]
        
        fetcher = vcf.fetch(contig=chrom, start=start0, stop=end0)

        if respect_refregion:
            def check_vr_inclusion(vr, chrom, start0, end0):
                vr_start0 = vr.pos - 1
                vr_end0 = vr_start0 + len(vr.ref)
                return (
                    (vr.contig == chrom)
                    and (vr_end0 > start0)
                    and (vr_start0 < end0)
                )
        else:
            def check_vr_inclusion(vr, chrom, start0, end0):
                vr_start0 = vr.pos - 1
                vr_end0 = vr.pos
                return (
                    (vr.contig == chrom)
                    and (vr_end0 > start0)
                    and (vr_start0 < end0)
                )

        for vr in fetcher:
            if check_vr_inclusion(vr, chrom, start0, end0):
                yield vr

#################
# miscellaneous #
#################

def get_vcf_lineno(vcf_path, verbose=False):
    is_vcf, comp, is_bgzf = get_vcf_format(vcf_path)
    if is_vcf:
        line_iter = get_line_iter(vcf_path, comp, verbose)
        return count_iterator(line_iter)
    else:  # BCF
        return get_vcf_lineno_pysam(vcf_path)


def count_iterator(iterator):
    idx = -1
    for idx, x in enumerate(iterator):
        pass
    return idx + 1


def get_line_iter(vcf_path, comp, verbose):
    if comp in ('BGZF', 'GZIP'):
        raw_iter = gzip.open(vcf_path, 'rb')
    elif comp == 'NONE':
        raw_iter = open(vcf_path, 'rb')

    if verbose:
        for idx, line in enumerate(raw_iter):
            if idx % 100000 == 0:
                print(idx)
            if not line.startswith(b'#'):
                yield line

    else:
        for line in raw_iter:
            if not line.startswith(b'#'):
                yield line


def parse_line_iter(line_iter):
    """Accepts byte mode iterator"""
    for line in line_iter:
        linesp = line.split(b'\t')
        pos0 = int(linesp[1]) - 1
        yield (linesp[0].decode(), pos0, pos0 + len(linesp[3]))


def get_vcf_positions_pysam(vcf_path):
    with pysam.VariantFile(vcf_path) as vcf:
        for vr in vcf.fetch():
            yield (vr.contig, vr.pos - 1, vr.pos - 1 + len(vr.ref))


def get_vcf_lineno_pysam(vcf_path):
    with pysam.VariantFile(vcf_path) as vcf:
        lineno = count_iterator(vcf.fetch())
    return lineno


# this does not work
def get_vcf_noerr(*args, **kwargs):
    with contextlib.redirect_stderr(io.StringIO()) as err, \
            contextlib.redirect_stdout(io.StringIO()) as out:
        vcf = pysam.VariantFile(*args, **kwargs)

    for buf in (err, out):
        msg = buf.getvalue()
        if not msg.startswith('[E::idx_find_and_load] '
                              'Could not retrieve index file for'):
            print(msg, end='', file=sys.stderr)

    return vcf


