#MIT License
#
#Copyright (c) 2019 Iñigo Prada Luengo
#
#Permission is hereby granted, free of charge, to any person obtaining a copy
#of this software and associated documentation files (the "Software"), to deal
#in the Software without restriction, including without limitation the rights
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the Software is
#furnished to do so, subject to the following conditions:
#
#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#SOFTWARE.

import numpy as np
import os
import pysam as ps
from Bio import SeqIO
from io import StringIO
import random as rd
from Bio.Seq import Seq
from Bio.Alphabet import generic_dna
import sys
import pybedtools as bt
import time
import subprocess as sp
import warnings


def sim_ecc_reads(genome_fasta,read_length,directory,reads,exclude_regions,fastq,insert_size,errors,mean_cov,locker,
                process,sim_circles,paired_end_fastq_1,paired_end_fastq_2,skipped,correct,ins_rate1,ins_rate2,del_rate1,
                del_rate2,sim_pid):
    """Function that takes as arguments a genome fasta file, weights each chromosome based on the length
    and simulates single end eccDNA reads
    """


    # Get the length of the chromosomes and store them in a sequence dictionary
    chromosomes = {}
    whole_genome_len = 0
    for rec in SeqIO.parse(genome_fasta, 'fasta'):
        name = rec.id
        seqLen = len(rec)
        whole_genome_len += seqLen
        chromosomes[name] = seqLen
    #chromosome sampling probability weighted based on its length
    weighted_chromosomes = {}
    for contigs in chromosomes:
        weighted_chromosomes[contigs] = chromosomes[contigs]/whole_genome_len




    contig_list = []
    weights = []
    for contigs, value in weighted_chromosomes.items():
        weights.append(value)
        contig_list.append(contigs)

    #Simulate the reads:



    os.chdir(directory)
    circle_bed = []




    set_of_reads = []
    set_of_left_reads = []
    set_of_right_reads = []

    circle_number = 0
    #reads simulated by a process
    n_of_reads = 0
    #
    n_of_reads_it = 0

    begin = time.time()
    #simulated reads
    while n_of_reads < reads + 1:


        #sample weighted chromosome
        #set random seed, important for paralell
        np.random.seed()
        chr = np.random.choice(contig_list, p=weights)

        # decide ecDNA length

        #sample circle length
        circle_length = rd.randint(150,350)


        # linear decrease in coverage based on circle length





        # compute circles sequencing coverage

        rounds_of_sim = (circle_length * mean_cov)/(read_length*2)







        # take in to account short length contigs
        #start position can't be bigger than (chr_length-circle_length)
        chr_pos_start =  rd.randint(0,(chromosomes[chr] - circle_length))
        #set end
        if chromosomes[chr] == (chromosomes[chr] - circle_length):
            chr_pos_end = chromosomes[chr]
        else:
            chr_pos_end = chr_pos_start + circle_length

        #if user of provides regions to exclude, check within it is on the region. and skip it
        if exclude_regions != None and bt.BedTool(exclude_regions).sort().any_hits(bt.Interval(chr,chr_pos_start,chr_pos_end)) != 0:
            #hit in a gap region
            # shared memory object between processes. It is use to track the number of skipped circles
            with skipped.get_lock():
                skipped.value+=1
            continue
        else:
            #shared memory object between processes. It is use to track the number of correctly simulated circles
            with correct.get_lock():
                correct.value+=1
        #save each circle positions, so that then I can check true circles


            first_line = [chr, chr_pos_start, chr_pos_end]

            #create class object outside the loop
            new_read = sim_paired_end(n_of_reads, insert_size, genome_fasta, chr, chr_pos_start,
                                    chr_pos_end, read_length, circle_number,process)

            #simulation rounds
            for each_sim in range(0,round(int(rounds_of_sim))):


                if errors == True:



                    if (n_of_reads_it+1) !=1000:

                        # sim the read
                        get_seq = new_read.simulate_read()
                        # put it in fastq format
                        simulated_reads = sim_paired_end.simulate_read_with_errors(new_read, get_seq[0], get_seq[1],
                                                                                get_seq[2],ins_rate1,ins_rate2,del_rate1,
                                                                                del_rate2,sim_pid)
                        if simulated_reads != None:
                            # save the read
                            assert len(set_of_left_reads) == len(set_of_right_reads)
                            set_of_left_reads.append(simulated_reads[0])
                            set_of_right_reads.append(simulated_reads[1])
                            assert len(set_of_left_reads) == len(set_of_right_reads)
                            n_of_reads += 1
                            n_of_reads_it += 1

                        else:
                            continue




                    else:

                        # simulate reads and save to disk
                        get_seq = new_read.simulate_read()
                        simulated_reads = sim_paired_end.simulate_read_with_errors(new_read, get_seq[0], get_seq[1],
                                                                                get_seq[2],ins_rate1,ins_rate2,del_rate1,
                                                                                del_rate2,sim_pid)
                        set_of_left_reads.append(simulated_reads[0])
                        set_of_right_reads.append(simulated_reads[1])

                        # save to disk
                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        locker.acquire()
                        print("Process %s: writting to disk 10000 reads" % process )
                        fastq_1 = open(paired_end_fastq_1, "a")
                        SeqIO.write(set_of_left_reads,fastq_1, "fastq")
                        fastq_1.close()
                        fastq_2 = open(paired_end_fastq_2, "a")
                        SeqIO.write(set_of_right_reads,fastq_2, "fastq")
                        fastq_2.close()
                        locker.release()
                        assert len(set_of_left_reads) == len(set_of_right_reads)

                        n_of_reads += 1
                        n_of_reads_it += 1

                        # sim the first read of the list
                        new_read = sim_paired_end(n_of_reads, insert_size, genome_fasta, chr, chr_pos_start,
                                                chr_pos_end, read_length, circle_number,process)
                        get_seq = new_read.simulate_read()
                        simulated_reads = sim_paired_end.simulate_read_with_errors(new_read, get_seq[0], get_seq[1],
                                                                                get_seq[2],ins_rate1,ins_rate2,del_rate1,
                                                                                del_rate2,sim_pid)
                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        set_of_left_reads = [simulated_reads[0]]
                        set_of_right_reads = [simulated_reads[1]]
                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        n_of_reads += 1
                        n_of_reads_it = 1


                else:

                    if (n_of_reads_it+1) != 10000:

                        #sim the read
                        get_seq = new_read.simulate_read()
                        #put it in fastq format
                        simulated_reads = sim_paired_end.simulate_perfect_read(new_read,get_seq[0], get_seq[1], get_seq[2])
                        #save the read
                        set_of_left_reads.append(simulated_reads[0])
                        set_of_right_reads.append(simulated_reads[1])
                        n_of_reads +=1
                        n_of_reads_it += 1



                    else:
                        #simulate reads and save to disk
                        get_seq = new_read.simulate_read()
                        simulated_reads = sim_paired_end.simulate_perfect_read(new_read, get_seq[0], get_seq[1],
                                                                                get_seq[2])
                        set_of_left_reads.append(simulated_reads[0])
                        set_of_right_reads.append(simulated_reads[1])

                        #save to disk
                        locker.acquire()
                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        print("Process %s: writting to disk 10000 reads" % process)
                        fastq_1 = open(paired_end_fastq_1, "a")
                        SeqIO.write(set_of_left_reads, fastq_1, "fastq")
                        fastq_1.close()
                        fastq_2 = open(paired_end_fastq_2, "a")
                        SeqIO.write(set_of_right_reads, fastq_2, "fastq")
                        fastq_2.close()
                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        locker.release()



                        n_of_reads += 1
                        n_of_reads_it += 1

                        #sim the first read of the list
                        new_read = sim_paired_end(n_of_reads, insert_size, genome_fasta, chr, chr_pos_start,
                                                chr_pos_end, read_length, circle_number,process)
                        get_seq = new_read.simulate_read()
                        simulated_reads = sim_paired_end.simulate_perfect_read(new_read, get_seq[0], get_seq[1],
                                                                                get_seq[2])

                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        set_of_left_reads = [simulated_reads[0]]
                        set_of_right_reads = [simulated_reads[1]]
                        assert len(set_of_left_reads) == len(set_of_right_reads)
                        n_of_reads +=1
                        n_of_reads_it = 1





            circle_bed.append(first_line)

    # last save to disk

    locker.acquire()
    fastq_1 = open(paired_end_fastq_1, "a")
    SeqIO.write(set_of_left_reads, fastq_1, "fastq")
    fastq_1.close()
    fastq_2 = open(paired_end_fastq_2, "a")
    SeqIO.write(set_of_right_reads, fastq_2, "fastq")
    fastq_2.close()
    locker.release()




    #shared memory between the processes.This is a list that every process will rate the simulated circles
    for element in circle_bed:
        sim_circles.append(element)




class sim_paired_end:

    #init the class
    def __init__(self,read_number,insert_size,genome_fa,chr,chr_pos_start,chr_pos_end,read_length,circle_id,process):
        self.read_number = read_number
        self.insert_size = insert_size
        self.genome_fa = genome_fa
        self.chr = chr
        self.chr_pos_start = chr_pos_start
        self.chr_pos_end = chr_pos_end
        self.read_length = read_length
        self.circle_id = circle_id
        self.process = process

    def simulate_read(self):
        """Function that simulates perfect paired-end reads"""

        fastafile = ps.FastaFile(self.genome_fa)
        # left split read

        insert = int(np.random.normal(self.insert_size, (self.insert_size / 12), 1))
        start = int(np.random.randint(self.chr_pos_start, (self.chr_pos_end + 1)))
        left_end = start + self.read_length
        total_end = start + int(np.round(insert))
        right_start = total_end - self.read_length
        if total_end > self.chr_pos_end:
            # split read scenario or insert spanning split read scenario
            if left_end > self.chr_pos_end:
                # left read spanning split read scenario
                # left_read
                left_dntps = self.chr_pos_end - start
                right_dntps = self.read_length - left_dntps

                # the error could be here
                left_split_read = fastafile.fetch(self.chr, start, self.chr_pos_end)
                right_split_read = fastafile.fetch(self.chr, self.chr_pos_start, (self.chr_pos_start + right_dntps))
                left_read = left_split_read + right_split_read

                # right_read
                right_start = self.chr_pos_start + int(round(self.insert_size - left_dntps - self.read_length))
                right_read = fastafile.fetch(self.chr, right_start, (right_start + self.read_length))

                # assertion to check the error here

                common_id = "%s|%s|%s:%s-%s:%s|%s:%s|1|%s" % (
                self.read_number, self.chr, start, self.chr_pos_end, self.chr_pos_start, (self.chr_pos_start + right_dntps), right_start,
                (right_start + self.read_length), self.circle_id)




            else:
                if right_start > self.chr_pos_end:
                    # insert spanning split read scenario
                    left_read = fastafile.fetch(self.chr, start, (start + self.read_length))
                    right_start = self.chr_pos_start + (right_start - self.chr_pos_end)
                    right_read = fastafile.fetch(self.chr, right_start, (right_start + self.read_length))
                    common_id = "%s|%s|%s:%s|%s:%s|3|%s" % (
                        self.read_number, self.chr, start, (start + self.read_length), right_start, (right_start + self.read_length), self.circle_id)
                else:
                    # right split read scenario
                    assert right_start <= self.chr_pos_end
                    assert (right_start + self.read_length) > self.chr_pos_end
                    left_read = fastafile.fetch(self.chr, start, (start + self.read_length))

                    # compute right dntps
                    left_dntps = self.chr_pos_end - right_start
                    right_dntps = self.read_length - left_dntps
                    left_split_read = fastafile.fetch(self.chr, right_start, self.chr_pos_end)
                    right_split_read = fastafile.fetch(self.chr, self.chr_pos_start, (self.chr_pos_start + right_dntps))
                    right_read = left_split_read + right_split_read
                    common_id = "%s|%s|%s:%s|%s:%s-%s:%s|2|%s" % (
                        self.read_number,self.chr, start, (start + self.read_length), right_start, self.chr_pos_end, self.chr_pos_start,
                        (self.chr_pos_start, right_dntps), self.circle_id)


        else:
            # non split read scenario
            left_read = fastafile.fetch(self.chr, start, (start + self.read_length))
            # correct right read start
            right_read = fastafile.fetch(self.chr, right_start, (right_start + self.read_length))
            common_id = "%s|%s|%s:%s|%s:%s|0|%s" % (
                self.read_number, self.chr, start, (start + self.read_length), right_start, (right_start + self.read_length), self.circle_id)

        return(right_read,left_read,common_id)



    def simulate_perfect_read(self,right_read,left_read,common_id):
        # put all together
        # unique identifiers for right and left reads
        right_read_id = "2:N:0:CGCTGTG"
        right_id = common_id + "  " + right_read_id
        left_read_id = "1:N:0:CGCTGTG"
        left_id = common_id + "  " + left_read_id
        quality = "I" * self.read_length
        # get the reverse complement of the right read
        right_read = Seq(right_read, generic_dna)
        right_read = right_read.reverse_complement()

        left_read = left_read.upper()

        right_read = right_read.upper()


        fastq_left = "@%s\n%s\n+\n%s\n" % (left_id, left_read, quality)
        fastq_right = "@%s\n%s\n+\n%s\n" % (right_id, right_read, quality)

        right_record = SeqIO.read(StringIO(fastq_right), "fastq")
        left_record = SeqIO.read(StringIO(fastq_left), "fastq")
        return (left_record, right_record)


    def simulate_read_with_errors(self,right_read, left_read, common_id,ins_rate1,ins_rate2,del_rate1,
                                    del_rate2,pid):
        # put all together
        # unique identifiers for right and left reads
        dir = os.getcwd()
        os.chdir("temp_files_%s" % pid)

        right_read_id = "2:N:0:CGCTGTG"
        right_id = common_id + "space" + right_read_id
        left_read_id = "1:N:0:CGCTGTG"
        left_id = common_id + "space" + left_read_id

        # attemp to use art to simulate the quality scores and the error rate
        #create a one read genome
        left_fasta = open("left_read_%s.fa" % (self.process), "w")
        left_fasta.write(">" + left_id + "\n" + str(left_read) + "\n")
        # sim the read with art
        left_fasta.close()

        sp.call("art_illumina -q -na -ss HS25  -ir %s -ir2 %s -dr %s -dr2 %s -nf 0 -i left_read_%s.fa -l %s -f 1 -o left%s" %
                (ins_rate1,ins_rate2,del_rate1,del_rate2,self.process,self.read_length,self.process),
                shell=True,stdout=sp.DEVNULL, stderr=sp.STDOUT)


        with open("left%s.fq" % (self.process), 'r') as left:
            left_read = left.read().replace('space', '   ').replace('1:N:0:CGCTGTG-1', '1:N:0:CGCTGTG')




        # get the reverse complement of the right read
        right_read = Seq(right_read, generic_dna)
        right_read = right_read.reverse_complement()

        right_fasta = open("right_read_%s.fa" % (self.process), "w")
        right_fasta.write(">" + right_id + "\n" + str(right_read) + "\n")
        right_fasta.close()
        # sim the read with art

        sp.call("art_illumina -na -q -ss HS25  -ir %s -ir2 %s -dr %s -dr2 %s -nf 0 -i right_read_%s.fa -l %s -f 1 -o right%s" %
                (ins_rate1,ins_rate2,del_rate1,del_rate2,self.process,self.read_length,self.process),
                shell=True,stdout=sp.DEVNULL, stderr=sp.STDOUT)

        with open("right%s.fq" % (self.process), 'r') as right:
            right_read = right.read().replace('space', '   ').replace('1:N:0:CGCTGTG-1', '2:N:0:CGCTGTG')

        #sometimes the reading fails. I introduce this to capture it
        try:

            right_record = SeqIO.read(StringIO(right_read), "fastq")
            left_record = SeqIO.read(StringIO(left_read), "fastq")
            os.chdir(dir)
            return (left_record, right_record)
        except ValueError as v:

            warnings.warn('Catched ValueError in a sampling round. Skipping')
            os.chdir(dir)
            return(None)











