#!/usr/bin/env python

# -*- coding: utf-8 -*-

from singlecellmultiomics.molecule import MoleculeIterator, CHICMolecule

from singlecellmultiomics.fragment import CHICFragment

from singlecellmultiomics.bamProcessing.bamFunctions import get_reference_path_from_bam, get_contigs

from collections import defaultdict

import pyBigWig

import pysam

from singlecellmultiomics.bamProcessing.bamBinCounts import get_binned_counts

import pandas as pd

import argparse

from singlecellmultiomics.bamProcessing import get_contig_sizes

from singlecellmultiomics.utils import is_autosome, pool_wrapper

import numpy as np

from multiprocessing import Pool

from more_itertools import windowed

from glob import glob

import pyBigWig

from scipy.ndimage import gaussian_filter1d

from scipy.signal import find_peaks

import argparse

from itertools import product

import asyncio

def get_aligned_chic_len(molecule):

    for f in molecule:

        f.R2_primer_length = 6

    mlen = molecule.estimated_max_length

    if mlen is None:

        return None

    return mlen + 2 # 2 extra bases because 1 base is potentially lost on both sides

def _generate_molecule_coordinate_dict(args, max_fragment_size=800):

    bam_path,  sel_contig= args

    molecule_coordinate_dict = defaultdict(list) # cell->[ (start,end), (start,end) ..]

    with pysam.AlignmentFile(bam_path) as alignments:

       # pysam.FastaFile(get_reference_path_from_bam(alignments)) as reference:

        for molecule in MoleculeIterator(alignments,

                                         contig=sel_contig,

                                         molecule_class=CHICMolecule,

                                         fragment_class = CHICFragment,

                                         #molecule_class_args = {'reference':reference}

                                        ):

            contig, start, strand = molecule.get_cut_site()

            if not molecule.is_completely_matching:

                continue

            if molecule.get_mean_mapping_qual()<55:

                continue

            slen = get_aligned_chic_len(molecule)

            if slen is None or slen>max_fragment_size:

                continue

            if strand:

                end = start - slen

            else:

                end = start + slen

            molecule_coordinate_dict[molecule.sample].append( (start,end) )

    for sample in molecule_coordinate_dict:

        molecule_coordinate_dict[sample] =  sorted(molecule_coordinate_dict[sample])

    return sel_contig, molecule_coordinate_dict

def generate_molecule_coordinate_dict(bams, n_threads=None):

    molecule_coordinate_dicts = {}

    contigs = get_contigs(bams[0])

    with Pool(n_threads) as workers:

        for contig,d in workers.imap( _generate_molecule_coordinate_dict,

                                     product(bams,filter(is_autosome, contigs))):

            if not contig in molecule_coordinate_dicts:

                molecule_coordinate_dicts[contig] = d

            else:

                molecule_coordinate_dicts[contig].update( d )

            print(f'Finished {contig}')

    return molecule_coordinate_dicts

    #bin_size = 100_000

#distances_per_bin = defaultdict(lambda: defaultdict(list)  )

def contig_coordinate_dict_to_votes(d, n, min_total_cuts_per_cell=3, p_nuc_bin_size=5, max_linker_length=90, linker_vote_radius=25):

    #molecule_coordinate_dicts[contig] = d

    # n = size of contig

    vote_nucleosome = np.zeros(int( n/p_nuc_bin_size) )

    vote_linker = np.zeros(int( n/p_nuc_bin_size) )

    for cell in d.keys():

        if len(d[cell])<min_total_cuts_per_cell:

            continue

        for (start_a, end_a),(start_b,end_b) in windowed( d[cell],2):

            o_start_a = start_a

            #Check if the fragments are not overlapping:

            (start_a, end_a) = min(start_a, end_a), max(start_a, end_a)

            (start_b, end_b) = min(start_b, end_b), max(start_b, end_b)

            # The locations of the nucleosomes are constrained:

            # \\ start_a  -----  end_a \\     \\ start_b  ---- end_b \\

            if end_a - start_a >= 147 and end_a - start_a < (147*2+10): # it could contain a single nucleosome, but not more

                s = int(((start_a + 70 ))/p_nuc_bin_size)

                e = int(((end_a - 70))/p_nuc_bin_size)

                vote_nucleosome[ s:(e+1) ] += 1/(e-s)

            # Any starting point of a molecule is _always_ part of a linker

            # lets say += 25bp

            c=o_start_a

            s = int(((c-linker_vote_radius))/p_nuc_bin_size)

            e = int(((c+linker_vote_radius))/p_nuc_bin_size)+1

            vote_linker[s:(e+1)] += 1/(e-s)

            if end_a > start_b:

                continue

            if start_b - end_a  > max_linker_length: # The distance is larger than 90bp, which is a very long linker. Skip the linker vote

                continue

            s = int(((end_a))/p_nuc_bin_size)

            e = int(((start_b))/p_nuc_bin_size)

            vote_linker[s:(e+1)] += 1/(e-s+1)

    return vote_nucleosome, vote_linker

async def write_to_bw(handle, starts, ends,  values, contig,size=None):

    if size is not None:

        print( ends[ends>=size] )

    handle.addEntries(

        [contig]*len(starts), #Contig

        starts.astype(np.int64) , #Start

        ends= ends.astype(np.int64) ,  #end

        values=  values.astype(np.float64)

    )

async def coordinate_dicts_to_nucleosome_position_files(bams, molecule_coordinate_dicts, p_nuc_bin_size = 5, alias='npos', n_threads=None ):

    contigs = get_contigs(bams[0])

    sizes = get_contig_sizes(bams[0])

    # Create output bigwig file handles and worker pool

    linker_vote_write_path= f'{alias}_linkers.bw'

    nuc_vote_write_path= f'{alias}_nucleosomes.bw'

    merged_vote_write_path= f'{alias}_nucleosomes_min_linkers.bw'

    centroids = f'{alias}_nucleosome_centers.bed'

    with Pool(n_threads) as workers, \

        pyBigWig.open(linker_vote_write_path,'w') as linkers_out, \

        pyBigWig.open(nuc_vote_write_path,'w') as nuc_out, \

        pyBigWig.open(merged_vote_write_path,'w') as merged_out, \

        open(centroids,'w') as centroids_out:

        # Add headers for all output bigwigfiles

        for h in (linkers_out, nuc_out, merged_out):

            # Write header

            h.addHeader(list(zip(contigs, [sizes[c] for c in contigs])))

        # Obtain nucleosome and linker votes for each contig

        for ret_contig, (vote_nucleosome, vote_linker) in zip(

                    molecule_coordinate_dicts.keys(),

                    workers.imap(pool_wrapper,

            (

            (contig_coordinate_dict_to_votes,

            {

                'd':d,

                'n':sizes[contig],

                'min_total_cuts_per_cell':3,

                'p_nuc_bin_size':5,

                'max_linker_length':90,

                'linker_vote_radius':25

            })

                for contig, d in molecule_coordinate_dicts.items()

            ))):

            print(f'Writing votes for {ret_contig}')

            contig_len = sizes[ret_contig]

            smooth_linkers = gaussian_filter1d(vote_linker,2)

            smooth_nuc = gaussian_filter1d(vote_nucleosome,2)

            # Write nucleosome vote track:

            starts = np.array( np.linspace(0, contig_len-p_nuc_bin_size-1, len(vote_nucleosome)) )

            size = sizes[ret_contig]

            # Check if all numbers are preceding...

            d = np.diff(starts)

            d = d[d<1]

            if len(d):

                print(d)

            else:

                print(f'{len(starts)} Coordinates are in correct order')

            await asyncio.gather(

                write_to_bw(nuc_out, starts, starts+p_nuc_bin_size,  np.nan_to_num(smooth_nuc), ret_contig, size),

                write_to_bw(linkers_out, starts, starts+p_nuc_bin_size,  np.nan_to_num(smooth_linkers), ret_contig, size),

                write_to_bw(merged_out, starts, starts+p_nuc_bin_size,  np.nan_to_num(smooth_nuc - smooth_linkers), ret_contig, size=size)

            )

            # Use the find peak function to select properly spaced nucleosome positions.

            print(f'Writing centroids for {ret_contig}')

            mindist = 147

            min_dist_bins = int( mindist / p_nuc_bin_size )

            estimated_nucleosome_positions = find_peaks( gaussian_filter1d(smooth_nuc - smooth_linkers,2),distance=min_dist_bins)[0]

            for nucpos, score in zip( estimated_nucleosome_positions*p_nuc_bin_size, smooth_nuc[estimated_nucleosome_positions]):

                centroids_out.write(f'{ret_contig}\t{nucpos}\t{nucpos+p_nuc_bin_size}\t{score}\n')

    #contig_coordinate_dict_to_votes(d, n, min_total_cuts_per_cell=3, p_nuc_bin_size=5, max_linker_length=90, linker_vote_radius=25)

if __name__ == '__main__':

    argparser = argparse.ArgumentParser(

        formatter_class=argparse.ArgumentDefaultsHelpFormatter,

        description='Estimate nucleosome positions from scCHiC seq bam file(s)')

    argparser.add_argument('alignmentfiles', type=str, nargs='+')

    argparser.add_argument('-o', type=str, required=True, help='output prefix')

    argparser.add_argument('-bin_size', type=int,  default=3, help='Nucleosome position precision (bp), increasing this value increases memory consumption linearly')

    argparser.add_argument('-n_threads', type=int, help='Amount of worker threads')

    args = argparser.parse_args()

    async def main(args):

        # Run the analysis:

        print('Creating molecule coordinate database')

        d = generate_molecule_coordinate_dict(args.alignmentfiles, args.n_threads)

        print('Estimating nucleosome positions')

        await coordinate_dicts_to_nucleosome_position_files(args.alignmentfiles,

                    d,

                    p_nuc_bin_size = args.bin_size,

                    n_threads=args.n_threads,

                    alias=args.o)

    loop = asyncio.get_event_loop()

    loop.run_until_complete(main(args))

    #asyncio.run(main(args)) # python >= 3.7