#!/usr/bin/env python3

import pysam

import numpy as np

import argparse

import os

import itertools

import gzip

import pandas as pd

import multiprocessing

if __name__ == '__main__':

    argparser = argparse.ArgumentParser(

        formatter_class=argparse.ArgumentDefaultsHelpFormatter,

        description="""Convert genome FASTA file to convert alternative bases to Ns""")

    argparser.add_argument(

        'fasta',

        metavar='fastafile',

        type=str,

        help='Fasta file to mask')

    argparser.add_argument(

        'vcf',

        metavar='vcffile',

        type=str,

        help='VCF file(s) to extract variants from')

    argparser.add_argument(

        '-o',

        type=str,

        default='./masked.fasta.gz',

        help='output file')

    argparser.add_argument(

        '-t',

        type=int,

        default=10,

        help='Amount of contigs to process in parallel')

    args = argparser.parse_args()

    try:

        faIn = pysam.FastaFile(args.fasta)

    except ValueError as e:

        print('FAI index missing. Now creating...')

        os.system(f'samtools faidx {args.fasta}')

    except Exception as e:

        raise

    try:

        faIn = pysam.FastaFile(args.fasta)

    except Exception as e:

        raise

    if args.o.endswith('.gz'):

        outputHandle = gzip.open(args.o, 'wb')

    else:

        outputHandle = open(args.o, 'wb')

    referenceNames = faIn.references

    referenceLengths = dict(zip(referenceNames, faIn.lengths))

    totalMasked = 0

    def get_masked_bytearray( jargs ):

        totalMasked = 0

        (chrom, fasta_file_path, variant_file_path) = jargs

        with pysam.FastaFile(fasta_file_path) as faIn, pysam.VariantFile(variant_file_path, threads=4) as  variants:

            chrom_seq = bytearray(faIn.fetch(chrom), 'ascii')

            if chrom in variants.index:

                print(f'masking {chrom}')

                for rec in variants.fetch(chrom):

                    if rec.start >= (referenceLengths[rec.chrom]):

                        print(

                            f"WARNING: record {rec.chrom} {rec.pos} defines a variant outside the supplied fasta file!")

                        continue

                    if len(rec.alleles) == 1 and len(rec.alleles[0]) == 1:

                        chrom_seq[rec.pos-1] = 78  # ord N

                        totalMasked += 1

                    elif len(rec.alleles[0]) == 1 and len(rec.alleles[1]) == 1:

                        chrom_seq[rec.pos-1] = 78  # ord N

                        totalMasked += 1

            print(f'Masked {totalMasked} bases of {chrom}')

            return chrom_seq

    workers = multiprocessing.Pool(args.t)

    for chrom, chrom_seq in zip(

                referenceNames,

                workers.imap(

                    get_masked_bytearray,

                    ((chrom, args.fasta, args.vcf )

                    for chrom in referenceNames ))):

        # Write chromsome

        outputHandle.write(f'>{chrom}\n'.encode('ascii'))

        outputHandle.write(chrom_seq)

        outputHandle.write('\n'.encode('ascii'))

    outputHandle.close()