"""
This script performs a simulation to test the number of inversions required to
  separate the genes of a fusion event back to the original state, in which they
  don't overlap at all.
"""
configfile: "config.yaml"

import time
from random import randrange
import sys

def fusion_mixing_reversion_simulation(
        groupname, num_genes_on_chrom,
        num_genes_A, num_genes_B,
        num_days_simulation, outfile):
    """
    This performs the simulation given a few parameters.
     The 'groupname' to use for the filename
     How many genes are on the chromosome in question.
     How many genes are in group A.
     How many genes are in group B.

    First make a pseudochrom and fuse
     ......A...AA...A...A.AAA.......,,,,,B,,,,,B,B,,,,BBBB,,,B,,,,,,
    Then each generation, invert somewhere in the chromosome.
     ......A...AA...A...A.B,,,,,.......AAA,,,,,B,B,,,,BBBB,,,B,,,,,,
    After each inversion, check if all the As and Bs are on either side
      AAAAAAAABBBBBBBB or BBBBBBBBAAAAAAAA
    If the genes match that, track the number of generations.
    At the end of the program, print out a vector of how many generations
      it took in each case to return to the original phenotype of A and B
      being on non-overlapping intervals on the chromosome.

    Runs the simulation for N days
    """
    # how many days for simulation
    when_to_end_simulation = time.time() + (60*60) * 24 * num_days_simulation

    num_genes = num_genes_on_chrom
    num1 = num_genes_A
    num0 = num_genes_B

    mixingleft = "".join(["x"]*(int(num_genes/2) - int((num1 + num0)/2) ))
    for i in range(num1):
        thisdex = randrange(len(mixingleft)+1)
        mixingleft = mixingleft[:thisdex] + '1' + mixingleft[thisdex:]

    mixingright = "".join(["y"]*(int(num_genes/2) - int((num1 + num0)/2) ))
    for i in range(num0):
        thisdex = randrange(len(mixingright)+1)
        mixingright = mixingright[:thisdex] + '0' + mixingright[thisdex:]

    mixing = mixingleft + mixingright

    posA = "".join((["1"]*num1) + (["0"]*num0))
    posB = "".join((["0"]*num0) + (["1"]*num1))

    generations_counter = {}
    mixing2 = mixing
    counter = 1
    # go until we say to stop
    while True:
        ssdone = False
        start = -1
        stop = -1
        while not ssdone:
            start = randrange(len(mixing2)+1)
            stop  = randrange(len(mixing2)+1)
            if start != stop:
                ssdone = True
        start, stop = sorted([start, stop])
        #print(start, stop)
        ptA = mixing2[:start]
        #ptBpre = mixing2[start:stop]
        ptB = mixing2[start:stop][::-1]
        ptC = mixing2[stop:]
        mixing2 = ptA + ptB + ptC
        test = mixing2.replace("x","").replace("y","")
        #print(test, posA, posB)
        if (test == posA) or (test == posB):
            if counter not in generations_counter:
                generations_counter[counter] = 0
            generations_counter[counter] += 1
            #print(test, posA, posB)
            if counter > 1000:
                print("Found one solution at {} generations".format(counter))
            counter = 1

        if counter % 10000 == 0:
            print("  Done with {} iterations \r".format(counter), end = "\r")
        if time.time() > when_to_end_simulation:
            break
        counter = counter + 1

    print()
    with open(outfile, "w") as outhandle:
        print("{}\t{}\t{}\t{}\t{}\t{}".format(
            groupname,
            num_genes_on_chrom,
            num_genes_A,
            num_genes_B,
            num_days_simulation,
            generations_counter), file = outhandle)


rule all:
    input:
        expand("inversion_simulation/individual_results/{sample}.results",
               sample = config["samples"])

rule shuffle_simulation:
    output:
        results = "inversion_simulation/individual_results/{sample}.results"
    params:
        samplename = lambda wildcards: wildcards.sample,
        num_genes_on_chrom = lambda wildcards: config["samples"][wildcards.sample]["num_genes_on_chromosome"],
        num_genes_on_A = lambda wildcards: config["samples"][wildcards.sample]["num_genes_on_A"],
        num_genes_on_B = lambda wildcards: config["samples"][wildcards.sample]["num_genes_on_B"],
        time = config["num_days"]
    threads: 1
    run:
        fusion_mixing_reversion_simulation(params.samplename,
                params.num_genes_on_chrom, params.num_genes_on_A,
                params.num_genes_on_B, params.time, output.results)
