#!/usr/bin/python
#
# Simulates an MDP-Strategy

import math
import os
import sys, code
import resource
import copy
import itertools
import random
from PIL import Image
import os, pygame, pygame.locals
from pybrain3.rl.environments import Environment
from pybrain3.rl.environments import Task
from pybrain3.rl.agents import LearningAgent
from pybrain3.rl.experiments import Experiment
from my_pybrain.my_explorer import MyUCBExplorer
from my_pybrain.my_explorer import MyGreedyExplorer
from my_pybrain.my_table import MyActionValueTable
from my_pybrain.my_learner import MyQ, SARSA
from pybrain3.utilities import abstractMethod

import numpy as np
from itertools import product
from copy import deepcopy
from scipy import argmax
from scipy import where
from random import choice
import importlib

# from scenario_9x9_shield_multi3 import Shield
# from cycling_enemy_shield_incl_enemy_multi3 import Shield

np.set_printoptions(threshold=np.inf)

import argparse

parser = argparse.ArgumentParser(description='Simulator')
parser.add_argument(dest="png_file_base")
parser.add_argument("-c", "--collect-data", dest="collect_data_file", help="Provide a file for collecting convergence data")
parser.add_argument('-g', "--gen-spec", dest='gen_spec', help='Generate shield files', action='store_true', default=False)
parser.add_argument('-l', "--load", dest='load_file', help='Load Q-Table from file')
parser.add_argument('-s', "--save", dest='save_file', help='Save Q-Table to file')
parser.add_argument('-t', "--train", dest='train', help='Training activated', type=float, default=.2)
parser.add_argument('-o', "--shield_options", dest='shield_options', help='Number of actions the shield can choose of. 0 disables the shield', type=int, default=1)
parser.add_argument('-n', "--negative-reward", dest='neg_reward', help='Indicated whether negative reward should be used for unsafe actions', action='store_true', default=False)
parser.add_argument('-p', "--huge-negative-reward", dest='huge_neg_reward', help='Indicated whether a huge negative reward should be used for unsafe actions', action='store_true', default=False)
parser.add_argument('-r', "--sarsa", dest='sarsa', help='Indicated whether to use SARSA or default Q-learning', action='store_true', default=False)
parser.add_argument("--num-steps", dest='num_steps', help='Number of interactions', type=int, default=1000000)

args = parser.parse_args()

collect_data_file = args.collect_data_file
gen_spec = args.gen_spec
specFile = args.png_file_base
shield_options = args.shield_options
load_file = args.load_file
save_file = args.save_file
exploration = args.train
neg_reward = args.neg_reward
huge_neg_reward = args.huge_neg_reward
MAX_STEPS = args.num_steps

pngfile = Image.open(specFile)
pngFileBasis = specFile[0:specFile.rfind(".png")]
path = pngFileBasis[:pngFileBasis.rfind(os.path.sep)]

# ==================================
# Settings
# ==================================
MAGNIFY = 64

# ==================================
# Read parameter file
# ==================================
parameterFileName = pngFileBasis+".params"
allParams = {}
for a in open(parameterFileName,"r").readlines():
    a = a.strip()
    if len(a)>0 and a[0]!='#':
        posEqual = a.index("=")
        allParams[a[0:posEqual].strip()] = a[posEqual+1:].strip()

# ==================================
# Parse parameter file
# ==================================
initX = int(allParams["initX"])
initY = int(allParams["initY"])
positionUpdateNoise = float(allParams["positionUpdateNoise"])
WALL = int(allParams["wall"])
NORMAL_FIELD = int(allParams["normalField"])
NUMBER_OF_COLORS = int(allParams["numberOfColors"])

enemies_enabled = False
if "enemies" in allParams:
    try:
        mode_name = allParams["enemies"]
        mode_name = path + "." + mode_name[:mode_name.rfind(".py")]
        enemy_handler = importlib.import_module(mode_name.replace(os.path.sep, ".")).EnemyHandler()
        enemies_enabled = True
    except ImportError as e:
        print ("Could not find file " + enemy_handler_file)
        print (e)
        exit()

bombs = []
if "bombs" in allParams:
    # careful with evil evals
    bombs = eval(allParams["bombs"])
else:
    bombs = []

# ==================================
# Read input image
# ==================================

if shield_options > 0:
    try:
        mod_name = pngFileBasis + "_" + str(shield_options)
        Shield = importlib.import_module(mod_name.replace(os.path.sep, ".")).Shield
    except ImportError as e:
        print ("Could not find file " + pngFileBasis + "_" + str(shield_options) + ".py")
        print (e)
        exit()

else:
    from no_shield import Shield

xsize = pngfile.size[0]
ysize = pngfile.size[1]
imageData = pngfile.getdata()
palette = pngfile.getpalette()

# for i in range(len(imageData)):
#     print (imageData[i])

if "colorOrder" in allParams:
    colors = eval(allParams["colorOrder"])
else:
    assert(max(imageData) == NUMBER_OF_COLORS + 1)
    colors = range(max(imageData) + 1)
    colors.remove(WALL)
    colors.remove(NORMAL_FIELD)

# ==================================
# Construct MDP --> States
# ==================================
stateMapper = {}
for xA in range(0,xsize):
    for yA in range(0,ysize):
        for (csf,payoff) in [(x, 0) for x in range(NUMBER_OF_COLORS)] + [(0,1)]:
            if (imageData[xA+yA*xsize]!=WALL):
                stateNum = len(stateMapper)
                stateMapper[(xA,yA,csf,payoff)] = stateNum

# print (stateMapper)
BAD_STATE = len(stateMapper)
# Add error state
errorState = len(stateMapper)
errorStateKey = (-1,-1,0,0)
stateMapper[errorStateKey] = errorState

# ==================================
# Construct MDP --> Transition file
# ==================================

# First, a function that computes the possible/likely
# transitions when going from a (x,y)-cell into some
# direction. It computes the image of the complete cell
# and then performs probability-weighting according to
# the areas of overlap
def computeSuccs(xpos,ypos,direction):

    # If direction is "4", this means no move
    if (direction==4):
        return [(xpos,ypos,1.0)]

    finalSuccs = []
    errorProb = 0.0
    if (direction==0):
        succs = [(xpos+1,ypos),(xpos+1,ypos+1)]
    elif (direction==1):
        succs = [(xpos,ypos+1),(xpos-1,ypos+1)]
    elif (direction==2):
        succs = [(xpos-1,ypos),(xpos-1,ypos-1)]
    elif (direction==3):
        succs = [(xpos,ypos-1),(xpos+1,ypos-1)]

    if succs[0][0]<0:
        errorProb += 1-positionUpdateNoise
    elif succs[0][0]>=xsize:
        errorProb += 1-positionUpdateNoise
    elif succs[0][1]<0:
        errorProb += 1-positionUpdateNoise
    elif succs[0][1]>=ysize:
        errorProb += 1-positionUpdateNoise
    else:
        finalSuccs.append((succs[0][0],succs[0][1],1-positionUpdateNoise))

    if succs[1][0]<0:
        errorProb += positionUpdateNoise
    elif succs[1][0]>=xsize:
        errorProb += positionUpdateNoise
    elif succs[1][1]<0:
        errorProb += positionUpdateNoise
    elif succs[1][1]>=ysize:
        errorProb += positionUpdateNoise
    else:
        finalSuccs.append((succs[1][0],succs[1][1],positionUpdateNoise))

    if errorProb>0.0:
        finalSuccs.append((-1,-1,errorProb))

    return finalSuccs

# Iterate over all cells and compute transition probabilities
transitionLines = []
overallNofTransitions = 0
for xA in range(0,xsize):
    for yA in range(0,ysize):
        for (csf,payoff) in [(x,0) for x in range(NUMBER_OF_COLORS)] + [(0,1)]:
            if (imageData[xA+yA*xsize]!=WALL):
                sourceState = stateMapper[(xA,yA,csf,payoff)]
                overallNofTransitions += 5
                for dirA in [0,1,2,3,4]: # Action 4 is standing still
                    errorProb = 0
                    succA = computeSuccs(xA,yA,dirA)
                    for (destXA,destYA,probA) in succA:
                        if destXA==-1:
                            errorProb += probA
                        elif (imageData[destXA+destYA*xsize]==WALL):
                            errorProb += probA
                        else:
                            if imageData[destXA+destYA*xsize]==colors[csf]:
                                csfPrime = csf + 1
                                payoffPrime = 1
                            else:
                                csfPrime = csf
                                payoffPrime = 0
                            if csfPrime==NUMBER_OF_COLORS:
                                csfPrime = 0
                            else:
                                payoffPrime = 0

                            # transitionLines.append([sourceState,dirA,stateMapper[(destXA,destYA,csfPrime,payoffPrime)],probA*0.99999])
                            transitionLines.append([sourceState,dirA,stateMapper[(destXA,destYA,csfPrime,payoffPrime)],probA])
                    # errorProb += 0.00001*(1-errorProb)
                    if errorProb>0:
                        transitionLines.append([sourceState,dirA,errorState,errorProb])

# ==================================
# Prepare reverse state mapper and
# Searchable transition list
# ==================================
reverseStateMapper = {}
for (a,b) in stateMapper.items():
    reverseStateMapper[b] = a
transitionLists = {}
for (a,b,c,d) in transitionLines:
    if not (a,b) in transitionLists:
        transitionLists[(a,b)] = [(c,d)]
    else:
        transitionLists[(a,b)].append((c,d))

NUMBER_OF_BITS = int(math.ceil(math.log((len(reverseStateMapper) - 1) / (NUMBER_OF_COLORS + 1), 2)))
# print ("Number of bits for states:", NUMBER_OF_BITS)
danger_zone = [(7, 6), (7, 7), (7, 8), (8, 6), (8, 7), (8, 8), (9, 6), (9, 7), (9, 8)]
max_steps_in_zone = 3

num_steps_on_bomb = 3

# recharging_zone = [(5, 10)]
# max_steps_in_zone = 20
# danger_zone = []
# for state in xrange(0, len(reverseStateMapper) - 1, 5): #exclude error state
#     (x, y, _, _) = reverseStateMapper[state]
#     danger_zone.append((x,y))

# danger_zone = [(7, 6), (7, 7), (8, 6), (8, 7)]
# max_steps_in_zone = 3

if gen_spec:
    with open("avoid_walls_shield.dfa", "w") as file:
        directions = [0, 1, 2, 3]
        transitions = []

        for combination in sum([list(map(list, itertools.combinations(directions, i))) for i in range(5)], []):
            sensors_enc = [str(x + 1 if x in combination else -(x + 1)) for x in directions]
            for action in range(4):
                action_enc = [str(-(idx + 5) if x == '0' else (idx + 5)) for idx, x in enumerate(list(bin(action)[2:].rjust(3, '0')))]
                target_state = 1 if action not in combination else 2
                transitions.append("1 {0} {1} {2}\n".format(target_state, " ".join(sensors_enc), " ".join(action_enc)))
        action_enc = [str(-(idx + 5) if x == '0' else (idx + 5)) for idx, x in enumerate(list(bin(4)[2:].rjust(3, '0')))]
        transitions.append("1 1 {0}\n".format(" ".join(action_enc)))

        #print unused action transitions
        for action in range(5, 8):
            action_enc = [str(-(idx + 5) if x == '0' else (idx + 5)) for idx, x in enumerate(list(bin(action)[2:].rjust(3, '0')))]
            transitions.append("1 2 {0}\n".format(" ".join(action_enc)))

        # print 'bad' state loop
        transitions.append("2 2\n")

        #print header & start/end states
        file.write("dfa 2 4 3 1 1 {0}\n1\n2\n".format(len(transitions)))

        #print transitions
        file.write("".join(transitions))
        file.write("1 sensor_right\n")
        file.write("2 sensor_down\n")
        file.write("3 sensor_left\n")
        file.write("4 sensor_up\n")

        for bit in range(1, 4):
            file.write("{0} o{1}\n".format(4 + bit, 4 - bit))

    #shield preventing collision with second robot
    with open("enemy_shield.dfa", "w") as file:
        #
        # x x x x x  1  6 11 16 21
        # x x x x x  2  7 12 17 22
        # x x o x x  3  8 13 18 23
        # x x x x x  4  9 14 19 24
        # x x x x x  5 10 15 20 25
        #
        # state 0 means no enemy in range

        transitions = []
        unused_states = [0, 13]
        for (enemy_x, enemy_y) in list(product(range(1, 6), repeat=2)):
            if enemy_x == 3 and enemy_y == 3:
                continue
            enemy_state = 5 * (enemy_x - 1) + enemy_y
            if abs(enemy_x - 3) + abs(enemy_y - 3) > 2:
                unused_states.append(enemy_state)
                continue
            num_state_bits = 5
            for action in range(4):
                action_allowed = True
                for enemy_action in range(5):
                    enemy_next = filter(lambda t: t[0] != -1, list(map(lambda t: (t[0], t[1]) if t[2] > 0 else (-1, -1), computeSuccs(enemy_x, enemy_y, enemy_action))))
                    next = filter(lambda t: t[0] != -1, list(map(lambda t: (t[0], t[1]) if t[2] > 0 else (-1, -1), computeSuccs(3, 3, action))))

                    intersection = set(next).intersection(set(enemy_next))
                    if len(intersection) > 0:
                        action_allowed = False
                        break
                state_enc = [str(-(idx + 1) if x == '0' else (idx + 1)) for idx, x in enumerate(list(bin(enemy_state)[2:].rjust(num_state_bits, '0')))]
                action_enc = [str(-(idx + 1 + num_state_bits) if x == '0' else (idx + 1 + num_state_bits)) for idx, x in enumerate(list(bin(action)[2:].rjust(3, '0')))]

                transitions.append("1 {0} {1} {2}\n".format(1 if action_allowed else 2, " ".join(state_enc), " ".join(action_enc)))

        action_enc = [str(-(idx + 1 + num_state_bits) if x == '0' else (idx + 1 + num_state_bits)) for idx, x in enumerate(list(bin(4)[2:].rjust(3, '0')))]
        transitions.append("1 1 " + " ".join(action_enc) + "\n")
  #
        #print unused action transitions
        for action in range(5, 8):
            action_enc = [str(-(idx + 1 + num_state_bits) if x == '0' else (idx + 1 + num_state_bits)) for idx, x in enumerate(list(bin(action)[2:].rjust(3, '0')))]
            transitions.append("1 2 " + " ".join(action_enc) + "\n")

        for state in unused_states:
            state_enc = [str(-(idx + 1) if x == '0' else (idx + 1)) for idx, x in enumerate(list(bin(state)[2:].rjust(num_state_bits, '0')))]
            transitions.append("1 1 " + " ".join(state_enc) + "\n")

        #print ununsed state transitions
        for state in range(26, int(math.pow(2, num_state_bits))):
            state_enc = [str(-(idx + 1) if x == '0' else (idx + 1)) for idx, x in enumerate(list(bin(state)[2:].rjust(num_state_bits, '0')))]
            transitions.append("1 1 " + " ".join(state_enc) + "\n")

        #print final state transition
        transitions.append("2 2\n")

        # print header
        file.write("dfa 2 {0} 3 1 1 {1}\n1\n2\n".format(num_state_bits, len(transitions)))

        #print transitions
        for transition in transitions:
            file.write(transition)

        # print labels
        for bit in range(1, num_state_bits + 1):
            file.write("{0} e{1}\n".format(bit, num_state_bits + 1 - bit))
        for bit in range(1, 4):
            file.write("{0} o{1}\n".format(num_state_bits + bit, 4 - bit))

    with open("bomb_shield.dfa", "w") as file:

        transitions = []
        for state in range(1, num_steps_on_bomb + 1):
            transitions.append("{0} 1 -1".format(state))

            actions = range(8)
            actions.remove(4) # remove stay
            for action in actions:
                action_enc = [str(-(idx + 2) if x == '0' else (idx + 2)) for idx, x in enumerate(list(bin(action)[2:].rjust(3, '0')))]
                transitions.append("{0} 1 1 {1}".format(state, " ".join(action_enc)))

            action_enc = [str(-(idx + 2) if x == '0' else (idx + 2)) for idx, x in enumerate(list(bin(4)[2:].rjust(3, '0')))]
            transitions.append("{0} {1} 1 {2}".format(state, state + 1, " ".join(action_enc)))

        transitions.append("{0} {0}".format(num_steps_on_bomb + 1))

        file.write("dfa {0} 1 3 1 1 {1}\n1\n{0}\n".format(num_steps_on_bomb + 1, len(transitions)))
        file.write("\n".join(transitions))
        file.write("\n1 b\n")
        for bit in range(1, 4):
            file.write("{0} o{1}\n".format(1 + bit, 4 - bit))

    #shield for danger zones
    with open("danger_zone_shield.dfa", "w") as file:
        zone = set(danger_zone)
        zones = {}
        #compute zones:
        current_zone = 1
        while len(zone) > 0:
            zones[current_zone] = set()
            for (x, y) in zone:
                at_boundary = False
                for action in range(4):
                    # we are looking for an action which leads for sure out of the danger zone
                    succs = computeSuccs(x, y, action)
                    at_boundary = True
                    for (new_x, new_y, prob) in succs:
                        if prob > 0 and ((new_x,new_y) in zone or (new_x,new_y,0,0) not in stateMapper):
                            # in the danger zone
                            at_boundary = False
                            break
                    if at_boundary:
                        break
                if at_boundary:
                   zones[current_zone].add((x,y))

            zone -= zones[current_zone]
            current_zone += 1

        # print (zones)

        end_state = max_steps_in_zone + 1
        transitions = []
        num_state_bits = 7
        for num_steps_in_zone in range(1,end_state):
            print ("state " + str(num_steps_in_zone))
            for state in xrange(0, len(reverseStateMapper) - 1, 5): #exclude error state
                (x,y,_,_) = reverseStateMapper[state]
                state_enc = [str(-(idx + 1) if bit == '0' else (idx + 1)) for idx, bit in enumerate(list(bin(state / 5)[2:].rjust(num_state_bits, '0')))]
                zone_idx = 0
                for idx, zone in zones.iteritems():
                    if (x,y) in zone:
                        zone_idx = idx
                        break

                if zone_idx == 0:
                    transitions.append("{0} 1 {1}\n".format(num_steps_in_zone, " ".join(state_enc)))
                    continue
                max_acceptable_zone = max_steps_in_zone - num_steps_in_zone # maximal acceptable zone as target
                if zone_idx < max_acceptable_zone or max(zones.keys()) <= max_acceptable_zone:
                    transitions.append("{0} {1} {2}\n".format(num_steps_in_zone, num_steps_in_zone + 1, " ".join(state_enc)))
                    continue

                print( "max_zone: " + str(max_acceptable_zone))
                if zone_idx <= max_acceptable_zone + 1:
                    for action in range(5):
                        succs = computeSuccs(x, y, action)
                        next_zone_idx = 0
                        for (next_x,next_y,prob) in succs:
                            if prob == 0: continue
                            for idx, zone in zones.iteritems():
                                if (next_x,next_y) in zone:
                                    next_zone_idx = max(next_zone_idx, idx)
                        print( "action " + str(action) + " leads to zone: " + str(next_zone_idx))
                        next_state = end_state if next_zone_idx > max_acceptable_zone else (num_steps_in_zone + 1 if next_zone_idx > 0 else 1)
                        action_enc = [str(-(idx + 1 + num_state_bits) if bit == '0' else (idx + 1 + num_state_bits)) for idx, bit in enumerate(list(bin(action)[2:].rjust(3, '0')))]
                        transitions.append("{0} {1} {2} {3}\n".format(num_steps_in_zone, next_state, " ".join(state_enc), " ".join(action_enc)))
                    for action in range(5,8):
                        action_enc = [str(-(idx + 1 + num_state_bits) if bit == '0' else (idx + 1 + num_state_bits)) for idx, bit in enumerate(list(bin(action)[2:].rjust(3, '0')))]
                        transitions.append("{0} {1} {2} {3}\n".format(num_steps_in_zone, end_state, " ".join(state_enc), " ".join(action_enc)))

                    continue

                # this should never happen .. do whatever we want to
                transitions.append("{0} 1 {1}\n".format(num_steps_in_zone, " ".join(state_enc)))

            #print ununsed state transitions
            for state in xrange((len(reverseStateMapper) - 1) / 5, int(math.pow(2, num_state_bits))):
                state_enc = [str(-(idx + 1) if x == '0' else (idx + 1)) for idx, x in enumerate(list(bin(state)[2:].rjust(num_state_bits, '0')))]
                transitions.append("{0} 1 {1}\n".format(num_steps_in_zone, " ".join(state_enc)))

        transitions.append("{0} {0}\n".format(end_state))
        # print header
        file.write("dfa {0} {1} 3 1 1 {2}\n1\n{0}\n".format(num_steps_in_zone + 1, num_state_bits, len(transitions)))

        #print transitions
        for transition in transitions:
            file.write(transition)

        # print labels
        for bit in range(1, num_state_bits + 1):
            file.write("{0} i{1}\n".format(bit, num_state_bits + 1 - bit))
        for bit in range(1, 4):
            file.write("{0} o{1}\n".format(num_state_bits + bit, 4 - bit))

    exit()

# =========================================
# Initialize interactive display
# =========================================
pygame.init()
displayInfo = pygame.display.Info()
MAGNIFY = min(MAGNIFY,displayInfo.current_w*3/4/xsize)
MAGNIFY = min(MAGNIFY,displayInfo.current_h*3/4/ysize)

class Map(Environment):
    def __init__(self):
        self.reset()

    def reset(self):
        # print "reset called"
        self.state = 0
        self.penalty = 0

    def performAction(self, action):
        error = len(reverseStateMapper) - 1
        # action = int(action[0])

        actions = action[action != -1]
        actions = list(map(int, actions))

        # print action
        # state_enc = map(int, list(bin(self.state / (NUMBER_OF_COLORS + 1))[2:].rjust(NUMBER_OF_BITS, '0')))

        encoded_actions = []
        for a in actions:
            encoded_actions.append(list(map(int, list(bin(a)[2:].rjust(3, '0')))))

        (robotXA, robotYA, csf, payoff) = reverseStateMapper[self.state]

        # simulate sensors
        state_enc = []
        for a in range(4):
            # print computeSuccs(robotXA, robotYA, a)
            succs = filter(lambda t: t[2] > 0, computeSuccs(robotXA, robotYA, a))
            valid = True
            for succ in succs:
                if succ[0] == -1 or not (succ[0], succ[1], 0, 0) in stateMapper:
                    valid = False
                    break
            state_enc.append(0 if valid else 1)

        # print state_enc
 # print "action" + str(encoded_actions[0])

        if enemies_enabled:
            enemy_state = 0
            (robotXA,robotYA,csf,payoff) = reverseStateMapper[level.state]
            for enemy in enemy_handler.getEnemyPositions():
                x_diff = abs(enemy[0] - robotXA)
                y_diff = abs(enemy[1] - robotYA)
                if x_diff + y_diff <= 2:
                    enemy_state = (enemy[0] - robotXA + 2) * 5 + (enemy[1] - robotYA + 3)
                    break
            enemy_state_enc = list(map(int, list(bin(enemy_state)[2:].rjust(5, '0'))))

            state_enc.extend(enemy_state_enc)

        # print "Colors seen so far:", csf
        if len(bombs) > 0:
            state_enc.append(1 if (robotXA + 1, robotYA + 1) in bombs else 0)
        for enc_action in encoded_actions:
            state_enc.extend(enc_action)

        # print state_enc
        corr_action = shield.tick(state_enc)

        # print corr_action

        corr_action = int("".join(list(map(str, corr_action[:len(corr_action) -1]))), 2)

        if (actions[0] != corr_action) and huge_neg_reward:
            self.penalty += 1.

        if (actions[0] != corr_action) and neg_reward and args.sarsa:
            self.penalty += 0.1
            # qvalue = self.module.getValue(self.laststate, action)
            # self.module.updateValue(self.laststate, action, qvalue + self.alpha * ((-1 if self.neg_reward else self.lastreward) - qvalue))
            #experiment.acc_reward -= .3
        #     print False
        used_actions = []
        for a in actions:
            if a == corr_action: break
            used_actions.append(a)
            # learner.explorer.n_values.params.reshape(learner.explorer.n_values.numRows,learner.explorer.n_values.numColumns)[self.state, a] += 1
        if huge_neg_reward:
            action = actions[0]
        else:
            action = corr_action

        used_actions.append(action)

        while len(used_actions) < 5:
            used_actions.append(-1)

        agent.lastaction = used_actions

        transitionList = transitionLists[(self.state, action)]

        dest = None
        randomNumber = random.random()
        for (a,b) in transitionList:
            if randomNumber<=b:
                dest = a
                randomNumber = 123.0
            else:
                randomNumber -= b
        # Rounding error?
        if (dest==None):
            dest = transitionList[0][0]

        if dest == len(reverseStateMapper) - 1:
            experiment.acc_reward -= 1
            self.penalty += 1
            # self.reset()
            if shield_options > 0 and not args.huge_neg_reward:
                print ("Shields are not allowed to make errors!")
                exit()
            transitionList = transitionLists[(self.state, 4)]
            dest = None
            randomNumber = random.random()
            for (a,b) in transitionList:
                if randomNumber<=b:
                    dest = a
                    randomNumber = 123.0
                else:
                    randomNumber -= b
            # Rounding error?
            if (dest==None):
                dest = transitionList[0][0]

        # learner.explorer.n_values.params.reshape(learner.explorer.n_values.numRows,learner.explorer.n_values.numColumns)[self.state, action] += 1

        self.state = dest

    def getSensors(self):
        return [self.state]

class VisitAllColors(Task):
    def __init__(self, env):
        Task.__init__(self, env)
        self.last_reward = 0

    def getReward(self):
        # if (reverseStateMapper[self.env.state][3] != 0):
        # print "all colors visited"
        ret = self.last_reward
        self.last_reward = reverseStateMapper[self.env.state][3] - self.env.penalty
        self.env.penalty = 0
        return self.last_reward

class MyExperiment(Experiment):
    def __init__(self, task, agent):
        Experiment.__init__(self, task, agent)

        agent.learner.explorer.experiment = self
        # agent.learner.module.getValue()

        self.screen = pygame.display.set_mode(((xsize+2)*MAGNIFY,(ysize+2)*MAGNIFY))
        pygame.display.set_caption('Policy Visualizer')
        self.clock = pygame.time.Clock()

        self.screenBuffer = pygame.Surface(self.screen.get_size())
        self.screenBuffer = self.screenBuffer.convert()
        self.screenBuffer.fill((64, 64, 64)) # Dark Gray

        self.bombImage = pygame.image.load("bomb_image.png")
        self.bombImage = pygame.transform.scale(self.bombImage, (MAGNIFY - 2, MAGNIFY - 2))

        self.isPaused = False
        self.isCrashed = False
        self.speed = 10
        self.num = 0
        self.robotXA = -1
        self.robotYA = -1
        self.bomb_counter = 0

        self.count = 0
        self.acc_reward = 0
        self.collect_data = False
        if collect_data_file != None:
            self.collect_data = True
            self.collect_episode_data_file = open(collect_data_file + "_episodelen.data", "w")
            self.collect_reward_data_file = open(collect_data_file + "_avg_reward.data", "w")

    def _oneInteraction(self):
        global draw

        resetInThisRound = False

        # Process events
        for event in pygame.event.get():
            if event.type == pygame.locals.QUIT or (event.type == pygame.locals.KEYDOWN and event.key in [pygame.locals.K_ESCAPE,pygame.locals.K_q]):
                return
            if (event.type == pygame.locals.KEYDOWN and event.key == pygame.locals.K_SPACE):
                controller.params.reshape(controller.numRows, controller.numColumns).tofile("test.table")
                self.isPaused = not self.isPaused
            if (event.type == pygame.locals.KEYDOWN and event.key == pygame.locals.K_r):
                resetInThisRound = True
            if (event.type == pygame.locals.KEYDOWN and event.key == pygame.locals.K_PLUS):
                self.speed += 1
            if (event.type == pygame.locals.KEYDOWN and event.key == pygame.locals.K_MINUS):
                self.speed = max(self.speed-1,1)
            if (event.type == pygame.locals.KEYDOWN and event.key == pygame.locals.K_d):
                draw = not draw

        # if self.isCrashed:
  #           self.isCrashed = False
  #           # level.reset()
  #
        # Update
        if resetInThisRound:
            print ("reset")
            level.reset()

        old = (self.robotXA, self.robotYA)
        (self.robotXA,self.robotYA,csf,payoff) = reverseStateMapper[level.state]

        if not self.isCrashed and enemies_enabled:
            enemy_handler.update(old)
            for e in enemy_handler.getEnemyPositions():
                if (self.robotXA, self.robotYA) == e:
                    self.isCrashed = True
                    level.penalty += 1
                    self.acc_reward -= 1
                    if shield_options > 0 and not args.huge_neg_reward:
                        print ("Shields are not allowed to make errors!")
                        exit()
                    break

        if (self.robotXA + 1, self.robotYA + 1) in bombs:
            self.bomb_counter += 1
            if self.bomb_counter == 4:
                self.isCrashed = True
                level.penalty += 1
                self.acc_reward -= 1
                if shield_options > 0 and not args.huge_neg_reward:
                    print ("Shields are not allowed to make errors!")
                    exit()
        else:
            self.bomb_counter = 0

        if draw:
            q_max = 0
            for state in range(len(reverseStateMapper) - 1):
                q_max = max(q_max, max(controller.getActionValues(state)))

            # Draw Field
            for x in xrange(0,xsize):
                for y in xrange(0,ysize):
                    paletteColor = imageData[y*xsize+x]
                    color = palette[paletteColor*3:paletteColor*3+3]
                    pygame.draw.rect(self.screenBuffer,color,((x+1)*MAGNIFY,(y+1)*MAGNIFY,MAGNIFY,MAGNIFY),0)

            # Draw boundary
            if self.robotXA==-1 or self.isCrashed:
                boundaryColor = (255,0,0)
            else:
                boundaryColor = (64,64,64)
            pygame.draw.rect(self.screenBuffer,boundaryColor,(0,0,MAGNIFY*(xsize+2),MAGNIFY),0)
            pygame.draw.rect(self.screenBuffer,boundaryColor,(0,MAGNIFY,MAGNIFY,MAGNIFY*(ysize+1)),0)
            pygame.draw.rect(self.screenBuffer,boundaryColor,(MAGNIFY*(xsize+1),MAGNIFY,MAGNIFY,MAGNIFY*(ysize+1)),0)
            pygame.draw.rect(self.screenBuffer,boundaryColor,(MAGNIFY,MAGNIFY*(ysize+1),MAGNIFY*xsize,MAGNIFY),0)
            # pygame.draw.rect(screenBuffer,boundaryColor,(0,0,MAGNIFY*(xsize+2),MAGNIFY),0)

            # Draw cell frames
            for x in xrange(0,xsize):
                for y in xrange(0,ysize):
                    pygame.draw.rect(self.screenBuffer,(0,0,0),((x+1)*MAGNIFY,(y+1)*MAGNIFY,MAGNIFY,MAGNIFY),1)
                    if (x+1,y+1) in bombs:
                        self.screenBuffer.blit(self.bombImage, ((x+1)*MAGNIFY+1,(y+1)*MAGNIFY+1))
            pygame.draw.rect(self.screenBuffer,(0,0,0),(MAGNIFY-1,MAGNIFY-1,MAGNIFY*xsize+2,MAGNIFY*ysize+2),1)

            # Draw "Good" Robot
            if self.robotXA!=-1:
                pygame.draw.circle(self.screenBuffer, (192,32,32), ((self.robotXA+1)*MAGNIFY+MAGNIFY/2,(self.robotYA+1)*MAGNIFY+MAGNIFY/2) , MAGNIFY/3-2, 0)
                pygame.draw.circle(self.screenBuffer, (255,255,255), ((self.robotXA+1)*MAGNIFY+MAGNIFY/2,(self.robotYA+1)*MAGNIFY+MAGNIFY/2) , MAGNIFY/3-1, 1)
                pygame.draw.circle(self.screenBuffer, (0,0,0), ((self.robotXA+1)*MAGNIFY+MAGNIFY/2,(self.robotYA+1)*MAGNIFY+MAGNIFY/2) , MAGNIFY/3, 1)

            # Draw "Bad" Robots
            if enemies_enabled:
                for (e_x, e_y) in enemy_handler.getEnemyPositions():
                    pygame.draw.circle(self.screenBuffer, (32,32,192), ((e_x+1)*MAGNIFY+MAGNIFY/2,(e_y+1)*MAGNIFY+MAGNIFY/2) , MAGNIFY/3-2, 0)
                    pygame.draw.circle(self.screenBuffer, (255,255,255), ((e_x+1)*MAGNIFY+MAGNIFY/2,(e_y+1)*MAGNIFY+MAGNIFY/2) , MAGNIFY/3-1, 1)
                    pygame.draw.circle(self.screenBuffer, (0,0,0), ((e_x+1)*MAGNIFY+MAGNIFY/2,(e_y+1)*MAGNIFY+MAGNIFY/2) , MAGNIFY/3, 1)

            # zone_width = danger_zone[-1][0] - danger_zone[0][0] + 1
     #        zone_height = danger_zone[-1][1] - danger_zone[0][1] + 1
     # pygame.draw.rect(screenBuffer,(200,200,0),(MAGNIFY*(danger_zone[0][0]+1),MAGNIFY*(danger_zone[0][1]+1),MAGNIFY*zone_width,MAGNIFY*zone_height),5)

            # Flip!
            self.screen.blit(self.screenBuffer, (0, 0))
            pygame.display.flip()

            # Make the transition
            if not self.isPaused:
                # Done
                self.clock.tick(self.speed)
            else:
                self.clock.tick(3)

        self.acc_reward += payoff * 10
        if self.collect_data:
            self.count += 1
            if payoff > 0:
                self.collect_episode_data_file.write(str(self.count) + "\n")
                self.count = 0
            if self.stepid % 100 == 0:
                self.collect_reward_data_file.write(str(self.acc_reward / 100.) + "\n")
                self.acc_reward = 0
            if self.stepid % 100000 == 0:
                pass

        if self.stepid % 100 == 0:
            sys.stdout.write("\033[K")
            sys.stdout.write("[{2}{3}] ({0}/{1}) | alpha = {4} | epsilon = {5}\n".format(self.stepid, MAX_STEPS, '#'*int(math.floor(self.stepid/float(MAX_STEPS)*20)), ' '*int((20 - math.floor(self.stepid/float(MAX_STEPS)*20))), learner.alpha, learner.explorer.exploration))
            sys.stdout.write("\033[F")

        if self.stepid >= MAX_STEPS:
            print ("\nSimulation done!")

            sys.exit()

        if payoff > 0:
            # episode done
            if save_file != None:
                controller.params.reshape(controller.numRows, controller.numColumns).tofile(save_file)
            learner.alpha *= 1.#0.999
            learner.explorer.exploration *= 1.#0.999

        self.isCrashed = False
        if not self.isPaused:
            return Experiment._oneInteraction(self)
        else: return self.stepid

# ==================================
# Call main program
# ==================================

#
# def enemy_random(enemy, good):
#     possible_next_positions = set([])
#     for action in range(5):
#         next = computeSuccs(enemy[0], enemy[1], action)
#         for t in next:
#             if t[0] != -1 and t[2] > 0 and (t[0], t[1], 0, 0) in stateMapper:
#                 possible_next_positions.add((t[0], t[1]))
#
#     next_position_invalid = True
#     while next_position_invalid:
#         idx = random.randint(0, len(possible_next_positions) - 1)
#         enemy = list(possible_next_positions)[idx]
#         # we do not allow to drive at the old position of the good robot
#         next_position_invalid = enemy == good
#
#     return enemy
#
#
#
# enemies = []

shield = Shield()
level = Map()
task = VisitAllColors(level)
controller = MyActionValueTable(len(reverseStateMapper) - 1, 5)
if load_file != None:
    controller.initialize(np.fromfile(load_file))
else:
    controller.initialize(0.)
alpha = .2
gamma = .95
if not args.sarsa:
    learner = MyQ(alpha, gamma, neg_reward)
    learner.explorer = MyGreedyExplorer(shield_options, exploration)
elif args.sarsa:
    learner = SARSA(alpha, gamma)
learner.explorer = MyGreedyExplorer(shield_options, exploration)
learner.explorer._setModule(controller)
agent = LearningAgent(controller, learner)

draw = False
EXPLORATION_FACTOR = exploration

experiment = MyExperiment(task, agent)
while 1:
    experiment.doInteractions(100)

    agent.learn()
    agent.reset()


import itertools
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import seaborn as sns
import pandas as pd
from typing import Union
import numpy as np
import re
import sys

from ..utils import compl, sbs_annotation_converter
from ..context import context96, context78, context83, context1536, context_composite, signature_composite, signature_cosmic, signature_DBS, signature_ID, context_polymerase96

def stacked_bar(H: pd.DataFrame, ref_type: str, figsize: tuple = (8,8)):
    """
    Plot stacked barchart & normalized stacked barchart.
    --------------------------------------
    Args:
        * H: matrix output from NMF
        * figsize: size of figure (int,int)

    Returns:
        * figure

    Example usage:
        plot_bar(H)
    """
    H = H.iloc[:,:-3].copy()
    # Map signature etiology
    if ref_type in ['pcawg_COMPOSITE', 'pcawg_COMPOSITE96', 'pcawg_SBS', 'pcawg_SBS96_ID', 'pcawg_SBS_ID']:
        H.columns = H.columns.map(lambda x: x[x.index('SBS') : x.index('_')]).map(signature_composite)
    elif ref_type in ['cosmic3', 'cosmic3_exome']:
        H.columns = H.columns.map(lambda x: x[x.index('SBS'):]).map(signature_cosmic)
    elif ref_type == 'cosmic3_DBS':
        H.columns = H.columns.map(lambda x: x[x.index('DBS'):]).map(signature_DBS)
    elif ref_type == 'cosmic3_ID':
        H.columns = H.columns.map(lambda x: x[x.index('ID'):]).map(signature_ID)

    # Sort H matrix by mutation burden for relevant mutation type
    H['sum'] = H.sum(1)
    H = H.sort_values('sum', ascending=False)

    fig,axes = plt.subplots(2,1,figsize=figsize, sharex=True)

    H.iloc[:,:-1].plot(
        kind='bar',
        stacked=True,
        ax=axes[0],
        width=1.0,
        rasterized=True
    )

    axes[0].set_xticklabels([])
    axes[0].set_xticks([])
    axes[0].set_ylabel('Counts', fontsize=20)

    H_norm = H.iloc[:,:-1].div(H['sum'].values,axis=0)
    H_norm.plot(
        kind='bar',
        stacked=True,
        ax=axes[1],
        width=1.0,
        rasterized=True
    )

    axes[1].set_xticklabels([])
    axes[1].set_xticks([])
    axes[1].set_xlabel('Samples', fontsize=16)
    axes[1].set_ylabel('Fractions', fontsize=20)
    axes[1].get_legend().remove()
    axes[1].set_ylim([0,1])

    return fig

def _map_sbs_sigs_back(df: pd.DataFrame) -> pd.Series:
    """
    Map Back Single-Base Substitution Signatures.
    -----------------------
    Args:
        * df: pandas.core.frame.DataFrame with index to be mapped

    Returns:
        * pandas.core.series.Series with matching indices to context96
    """
    def _check_to_flip(x, ref):
        if x in ref:
            return x
        else:
            return compl(x)

    if df.index.name is None: df.index.name = 'index'
    df_idx = df.index.name

    if ">" in df.index[0]:
        # Already in arrow format
        context_s = df.reset_index()[df_idx].apply(sbs_annotation_converter)
    else:
        # Already in word format
        context_s = df.reset_index()[df_idx]

    return context_s.apply(lambda x: _check_to_flip(x, context96.keys()))

def _map_id_sigs_back(df: pd.DataFrame) -> pd.Series:
    """
        Map Back Insertion-Deletion Signatures.
        -----------------------
        Args:
            * df: pandas.core.frame.DataFrame with index to be mapped

        Returns:
            * pandas.core.series.Series with matching indices to context83
        """
    if df.index.name is None: df.index.name = 'index'
    df_idx = df.index.name

    context_s = df.reset_index()[df_idx]

    def _convert_from_cosmic(x):
        if x in context83:
            return x
        i1, i2, i3, i4 = x.split('_')
        pre = i2 if i3 == '1' else i3
        main = i1.lower() + ('m' if i2 == 'MH' else '')
        if main == 'del':
            post = str(int(i4[0]) + 1) + i4[1:]
        else:
            post = i4
        return pre + main + post

    return context_s.apply(_convert_from_cosmic)

def signature_barplot(W: pd.DataFrame, contributions: Union[int, pd.Series] = 1):
    """
    Plots signatures from W-matrix for Single-Base Substitutions
    --------------------------------------
    Args:
        * W: W-matrix
        * contributions: Series of total contributions, np.sum(H), from each
            signature if W is normalized; else, 1

    Returns:
        * fig

    Example usage:
        signature_barplot(W, np.sum(H))
    """
    W = W.copy()
    W.index = _map_sbs_sigs_back(W)

    # Fill in any missing contexts
    for c in context96:
        if c not in W.index:
            W.loc[c] = 0

    # Sort contexts
    W.sort_index(inplace=True)

    # Extract columns corresponding to signatures
    sig_columns = [c for c in W if c.startswith('S')]

    # Calculate total number of mutations at each context for every signature
    if isinstance(contributions, pd.Series):
        W = W[sig_columns] * contributions[sig_columns]
    else:
        W = W[sig_columns] * contributions

    # Determine number of signatures
    n_sigs = len(sig_columns)

    # Initialize SBS C>N and T>N mutations and their contexts
    # For each context, iterate through C>N and T>N mutations, and take reverse complement
    # of context for A>N mutations
    context_label = []
    change_map = {'CA': [], 'CG': [], 'CT': [], 'TA': [], 'TC': [], 'TG': []}
    for p in itertools.product('ACGT', 'ACGT'):
        context = ''.join(p)
        # Reverse complement of context
        compl_context = compl(context, reverse=True)
        context_label.append('-'.join(context))
        for key in change_map:
            if key.startswith('C'):
                change_map[key].append(key + context)
            else:
                # Complement of mutation + reverse complement of context
                change_map[key].append(compl(key) + compl_context)
    color_map = {'CA': 'cyan', 'CG': 'red', 'CT': 'yellow', 'TA': 'purple', 'TC': 'green', 'TG': 'blue'}

    # Plot contributions
    x_coords = range(16)
    fig, axes = plt.subplots(nrows=n_sigs, ncols=6, figsize=(20, 2.5 * n_sigs), sharex='col', sharey='row')
    for row, sig in enumerate(sig_columns):
        for col, chg in enumerate(['CA', 'CG', 'CT', 'TA', 'TC', 'TG']):
            if n_sigs == 1:
                ax = axes[col]
            else:
                ax = axes[row, col]
            bar_heights = W[sig].loc[change_map[chg]]
            ax.bar(x_coords, bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map[chg], rasterized=True)
            ax.set_xlim(-.55, 15.55)
            if row == 0:
                ax.set_title('>'.join(chg), fontsize=18)
                if col == 0:
                    ax.text(51.2 / 16, 1.3, 'Mutational Signatures', transform=ax.transAxes,
                            horizontalalignment='center', fontsize=24)
            if row < n_sigs - 1:
                ax.tick_params(axis='x', length=0)
            else:
                ax.set_xticks(x_coords)
                ax.set_xticklabels(context_label, fontfamily='monospace', rotation='vertical')
                if col == 0:
                    ax.text(51.2 / 16, -.4, 'Motifs', transform=ax.transAxes, horizontalalignment='center', fontsize=20,
                            fontweight='bold')
            if col > 0:
                ax.tick_params(axis='y', length=0)
            if col == 5:
                ax.text(1.05, .5, sig, fontsize=14, rotation=270, transform=ax.transAxes, verticalalignment='center')

    plt.subplots_adjust(wspace=.08, hspace=.15)
    fig.text(.08, .5, 'Contributions', rotation='vertical', verticalalignment='center', fontsize=20, fontweight='bold')

    return fig

def signature_barplot_DBS(W, contributions):
    """
    Plots signatures from W-matrix for Doublet-Base Substitutions
    --------------------------------------
    Args:
        * W: W-matrix
        * contributions: Series of total contributions, np.sum(H), from each
            signature if W is normalized; else, 1

    Returns:
        * fig

    Example usage:
        signature_barplot_DBS(W, np.sum(H))
    """
    W = W.copy()
    for c in context78:
        if c not in W.index:
            W.loc[c] = 0
    W.sort_index(inplace=True)
    sig_columns = [c for c in W if c.startswith('S')]
    if isinstance(contributions, pd.Series):
        W = W[sig_columns] * contributions[sig_columns]
    else:
        W = W[sig_columns] * contributions

    n_sigs = len(sig_columns)

    ref_map = {'AC': [], 'AT': [], 'CC': [], 'CG': [], 'CT': [], 'GC': [], 'TA': [], 'TC': [], 'TG': [], 'TT': []}
    for x in W.index:
        ref_map[x[:2]].append(x)
    x_coords = {ref: range(len(sigs)) for ref, sigs in ref_map.items()}

    color_map = {'AC': '#99CCFF', 'AT': '#0000FF', 'CC': '#CCFF99', 'CG': '#00FF00', 'CT': '#FF99CC',
                 'GC': '#FF0000', 'TA': '#FFCC99', 'TC': '#FF8000', 'TG': '#CC99FF', 'TT': '#8000FF'}
    fig, axes = plt.subplots(nrows=n_sigs, ncols=10, figsize=(20, 2.5 * n_sigs), sharex='col',
                             sharey='row', gridspec_kw={'width_ratios': (3, 2, 3, 2, 3, 2, 2, 3, 3, 3)})
    for row, sig in enumerate(sig_columns):
        for col, ref in enumerate(ref_map):
            if n_sigs == 1:
                ax = axes[col]
            else:
                ax = axes[row, col]
            bar_heights = W[sig].loc[ref_map[ref]]
            ax.bar(x_coords[ref], bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map[ref],
                   rasterized=True)
            ax.set_xlim(-.55, x_coords[ref][-1] + .55)
            if row == 0:
                ax.set_title(ref)
                if col == 0:
                    ax.text(44.5 / 6, 1.2, 'Mutational Signatures', transform=ax.transAxes,
                            horizontalalignment='center', fontsize=24)
            if row < n_sigs - 1:
                ax.tick_params(axis='x', length=0)
            else:
                xlabels = [x[3:] for x in ref_map[ref]]
                ax.set_xticks(x_coords[ref])
                ax.set_xticklabels(xlabels, fontfamily='monospace', rotation='vertical')
                if col == 0:
                    ax.text(44.5 / 6, -.3, 'Motifs', transform=ax.transAxes, horizontalalignment='center', fontsize=20,
                            fontweight='bold')
            if col > 0:
                ax.tick_params(axis='y', length=0)
            if col == 9:
                ax.text(1.05, .5, sig, fontsize=14, rotation=270, transform=ax.transAxes, verticalalignment='center')

    plt.subplots_adjust(wspace=.08, hspace=.15)
    fig.text(.08, .5, 'Contributions', rotation='vertical', verticalalignment='center', fontsize=20, fontweight='bold')

    return fig

def signature_barplot_ID(W, contributions):
    """
    Plots signatures from W-matrix for Insertions-Deletions
    --------------------------------------
    Args:
        * W: W-matrix
        * contributions: Series of total contributions, np.sum(H), from each
            signature if W is normalized; else, 1

    Returns:
        * fig

    Example usage:
        signature_barplot_ID(W, np.sum(H))
    """
    W = W.copy()
    W.index = _map_id_sigs_back(W)
    for c in context83:
        if c not in W.index:
            W.loc[c] = 0
    W = W.loc[context83]
    sig_columns = [c for c in W if c.startswith('S')]
    if isinstance(contributions, pd.Series):
        W = W[sig_columns] * contributions[sig_columns]
    else:
        W = W[sig_columns] * contributions

    n_sigs = len(sig_columns)
    group_map = {'Cdel': [], 'Tdel': [], 'Cins': [], 'Tins': [],
                 '2del': [], '3del': [], '4del': [], '5+del': [],
                 '2ins': [], '3ins': [], '4ins': [], '5+ins': [],
                 '2delm': [], '3delm': [], '4delm': [], '5+delm': []}
    for x in W.index:
        group = re.search('.+?(?=[\d])', x).group(0)
        group_map[group].append(x)
    x_coords = {group: range(len(sigs)) for group, sigs in group_map.items()}

    color_map = {'Cdel': '#FFCC99', 'Tdel': '#FF8000', 'Cins': '#00FF00', 'Tins': '#00BB00',
                 '2del': '#FF99CC', '3del': '#FF3377', '4del': '#FF0000', '5+del': '#880000',
                 '2ins': '#99CCFF', '3ins': '#3377FF', '4ins': '#0000FF', '5+ins': '#000088',
                 '2delm': '#CC99FF', '3delm': '#9966FF', '4delm': '#8000FF', '5+delm': '#6000AA'}

    fig, axes = plt.subplots(nrows=n_sigs, ncols=16, figsize=(20, 2.5 * n_sigs), sharex='col',
                             sharey='row', gridspec_kw={'width_ratios': (6,) * 12 + (1, 2, 3, 5)})
    for row, sig in enumerate(sig_columns):
        for col, group in enumerate(group_map):
            if n_sigs == 1:
                ax = axes[col]
            else:
                ax = axes[row, col]
            bar_heights = W[sig].loc[group_map[group]]
            ax.bar(x_coords[group], bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map[group],
                   rasterized=True)
            ax.set_xlim(-.55, x_coords[group][-1] + .55)
            if row == 0:
                ax.set_title(re.search('[\d+CT]+', group).group(0), color=color_map[group])
                if col == 0:
                    ax.text(44.5 / 6, 1.3, 'Mutational Signatures', transform=ax.transAxes,
                            horizontalalignment='center', fontsize=24)
                if group == 'Tdel':
                    ax.text(-.02, 1.16, '1bp deletions at repeats', fontsize=10, transform=ax.transAxes,
                            horizontalalignment='center', color=color_map[group])
                if group == 'Tins':
                    ax.text(-.02, 1.16, '1bp insertions at repeats', fontsize=10, transform=ax.transAxes,
                            horizontalalignment='center', color=color_map[group])
                if group == '4del':
                    ax.text(-.02, 1.16, '>1bp deletions at repeats', fontsize=10, transform=ax.transAxes,
                            horizontalalignment='center', color=color_map[group])
                if group == '4ins':
                    ax.text(-.02, 1.16, '>1bp insertions at repeats', fontsize=10, transform=ax.transAxes,
                            horizontalalignment='center', color=color_map[group])
                if group == '4delm':
                    ax.text(.8, 1.16, '>1bp deletions with microhomology', fontsize=10, transform=ax.transAxes,
                            horizontalalignment='center', color=color_map[group])
            if row < n_sigs - 1:
                ax.tick_params(axis='x', length=0)
            else:
                xlabels = [re.search('[\d+]+$', x).group(0) for x in group_map[group]]
                ax.set_xticks(x_coords[group])
                ax.set_xticklabels(xlabels, fontfamily='monospace')
                if col == 0:
                    ax.text(44.5 / 6, -.3, 'Motifs', transform=ax.transAxes, horizontalalignment='center', fontsize=20,
                            fontweight='bold')
            if col > 0:
                ax.tick_params(axis='y', length=0)
            if col == 15:
                ax.text(1.05, .5, sig, fontsize=14, rotation=270, transform=ax.transAxes, verticalalignment='center')

    plt.subplots_adjust(wspace=.08, hspace=.15)
    fig.text(.08, .5, 'Contributions', rotation='vertical', verticalalignment='center', fontsize=20, fontweight='bold')

    return fig

def signature_barplot_composite(W: pd.DataFrame, contributions: Union[int, pd.Series] = 1):
    """
    Plot signatures from W-matrix for SBS, DBS, and IDs from composite W matrix
    --------------------------------------
    Args:
        * W: W-matrix
        * contributions: Series of total contributions, np.sum(H), from each
            signature if W is normalized; else, 1
    Returns:
       * fig
    Example usage:
        signature_barplot(W, np.sum(H))
    """

    W = W.copy()
    # Fill in missing features
    composite_index = list(context96)+list(context78)+list(context83)
    for c in composite_index:
        if c not in list(W.index):
            W.loc[c] = 0
    W = W.reindex(composite_index)

    # Get signature labels
    sig_columns = [c for c in W if c.startswith('S')]
    n_sigs = len(sig_columns)

    # Evaluate contributions
    if isinstance(contributions, pd.Series):
        W = W[sig_columns] * contributions[sig_columns]
    else:
        W = W[sig_columns] * contributions

    #### x coordinates for SBS contributions
    context_label = []
    change_map = {'CA': [], 'CG': [], 'CT': [], 'TA': [], 'TC': [], 'TG': []}
    for p in itertools.product('ACGT', 'ACGT'):
        context = ''.join(p)
        compl_context = compl(context, reverse=True)
        context_label.append('-'.join(context))
        for key in change_map:
            if key.startswith('C'):
                change_map[key].append(key + context)
            else:
                change_map[key].append(compl(key) + compl_context)
    color_map_sbs = {'CA': 'cyan', 'CG': 'red', 'CT': 'yellow', 'TA': 'purple', 'TC': 'green', 'TG': 'blue'}
    x_coords_sbs = range(16)

    ##### x coordinates for DBS contributions
    ref_map = {'AC': [], 'AT': [], 'CC': [], 'CG': [], 'CT': [], 'GC': [], 'TA': [], 'TC': [], 'TG': [], 'TT': []}
    for x in context78:
        ref_map[x[:2]].append(x)
    x_coords_dbs = {ref: range(len(sigs)) for ref, sigs in ref_map.items()}
    color_map_dbs = {'AC': '#99CCFF', 'AT': '#0000FF', 'CC': '#CCFF99', 'CG': '#00FF00', 'CT': '#FF99CC',
                 'GC': '#FF0000', 'TA': '#FFCC99', 'TC': '#FF8000', 'TG': '#CC99FF', 'TT': '#8000FF'}

    ##### x coordinates for ID contributions
    group_map = {'Cdel': [], 'Tdel': [], 'Cins': [], 'Tins': [],
                 '2del': [], '3del': [], '4del': [], '5+del': [],
                 '2ins': [], '3ins': [], '4ins': [], '5+ins': [],
                 '2delm': [], '3delm': [], '4delm': [], '5+delm': []}
    for x in context83:
        group = re.search('.+?(?=[\d])', x).group(0)
        group_map[group].append(x)
    x_coords_id = {group: range(len(sigs)) for group, sigs in group_map.items()}

    color_map_id = {'Cdel': '#FFCC99', 'Tdel': '#FF8000', 'Cins': '#00FF00', 'Tins': '#00BB00',
                 '2del': '#FF99CC', '3del': '#FF3377', '4del': '#FF0000', '5+del': '#880000',
                 '2ins': '#99CCFF', '3ins': '#3377FF', '4ins': '#0000FF', '5+ins': '#000088',
                 '2delm': '#CC99FF', '3delm': '#9966FF', '4delm': '#8000FF', '5+delm': '#6000AA'}

    # Include spaces to separate feature types
    all_columns = ['CA', 'CG', 'CT', 'TA', 'TC', 'TG'] + ['space'] + list(ref_map) + ['space'] + list(group_map)
    fig, axes = plt.subplots(nrows=n_sigs, ncols=34, figsize=(60,2.5*n_sigs), sharex='col',
                             gridspec_kw={'width_ratios': (16,)*6 + (1,)+ (9,6,9,6,9,6,6,9,9,9) + (1,) + (6,)*12+(1,2,3,5)})
    max_height = 0  # Maximum height for scaling y-axis per feature type per signature
    # Iterate through signatures, such that each row plots mutational landscape for a signature
    for row, sig in enumerate(sig_columns):
        for col, ref in enumerate(all_columns):
            if n_sigs == 1:
                ax = axes[col]
            else:
                ax = axes[row,col]
            if col in [6,17]:  # Space between feature types...Remove ax and move to next feature (column)
                ax.remove()
                continue
            # For SBS portion, iterate through 6 SNV types (C>A, C>T, C>G, T>A...)
            if col < 6:
                bar_heights = W[sig].loc[change_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_sbs, bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map_sbs[ref], rasterized=True)
                ax.set_xlim(-.55, 15.55)
                if row == 0:
                    ax.set_title('>'.join(ref), fontsize=18)
                    if col == 0:
                        ax.text(8.1, 1.3, 'Mutational Signatures', transform=ax.transAxes,
                                horizontalalignment='center', fontsize=24)
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0, labelbottom=False)
                else:
                    ax.set_xticks(x_coords_sbs)
                    ax.set_xticklabels(context_label, fontfamily='monospace', rotation='vertical')
                    if col == 0:
                        ax.text(8.1, -.4, 'Motifs', transform = ax.transAxes, horizontalalignment='center', fontsize=20)
                if col == 5:
                    if n_sigs == 1:
                        for axis in axes[:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    else:
                        for axis in axes[row,:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    max_height = 0

            # For DBS portion
            elif col < 17:
                bar_heights = W[sig].loc[ref_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_dbs[ref], bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map_dbs[ref],
                       rasterized=True)
                ax.set_xlim(-.55, x_coords_dbs[ref][-1] + .55)
                if row == 0:
                    ax.set_title(ref)
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0)
                else:
                    xlabels = [x[3:] for x in ref_map[ref]]
                    ax.set_xticks(x_coords_dbs[ref])
                    ax.set_xticklabels(xlabels, fontfamily='monospace', rotation='vertical')
                if col == 15:
                    if n_sigs == 1:
                        for axis in axes[6:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    else:
                        for axis in axes[row,6:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    max_height = 0

            # For ID portion
            else:
                bar_heights = W[sig].loc[group_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_id[ref], bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map_id[ref],
                       rasterized=True)
                ax.set_xlim(-.55, x_coords_id[ref][-1] + .55)
                if row == 0:
                    ax.set_title(re.search('[\d+CT]+', ref).group(0), color=color_map_id[ref])
                    if ref == 'Tdel':
                        ax.text(-.02, 1.16, '1bp deletions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == 'Tins':
                        ax.text(-.02, 1.16, '1bp insertions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == '4del':
                        ax.text(-.02, 1.16, '>1bp deletions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == '4ins':
                        ax.text(-.02, 1.16, '>1bp insertions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == '4delm':
                        ax.text(.8, 1.16, '>1bp deletions with microhomology', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0)
                else:
                    xlabels = [re.search('[\d+]+$', x).group(0) for x in group_map[ref]]
                    ax.set_xticks(x_coords_id[ref])
                    ax.set_xticklabels(xlabels, fontfamily='monospace')
                if col == 33:
                    if n_sigs == 1:
                        for axis in axes[16:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    else:
                        for axis in axes[row,16:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    max_height = 0

            if col not in [0,7,18]:
                ax.tick_params(axis='y', which='both',length=0, labelleft=False)
            if col == 33:
                ax.text(1.05, .5, sig, fontsize=14, rotation=270, transform=ax.transAxes, verticalalignment='center')

    # Set titles and organize plot
    plt.subplots_adjust(wspace=.12, hspace=.15)
    fig.text(.105, .5, 'Contributions', rotation='vertical', verticalalignment='center', fontsize=20, fontweight='bold')
    return fig

def signature_barplot_sbs_id(W: pd.DataFrame, contributions: Union[int, pd.Series] = 1):
    """
    Plot signatures from W-matrix for SBS, DBS, and IDs from composite W matrix
    --------------------------------------
    Args:
        * W: W-matrix
        * contributions: Series of total contributions, np.sum(H), from each
            signature if W is normalized; else, 1

    Returns:
       * fig

    Example usage:
        signature_barplot(W, np.sum(H))
    """
    W = W.copy()

    # Fill in missing features and sort
    composite_index = list(context96)+list(context83)
    for c in composite_index:
        if c not in list(W.index):
            W.loc[c] = 0
    W = W.reindex(composite_index)

    # Get signature labels
    sig_columns = [c for c in W if c.startswith('S')]
    n_sigs = len(sig_columns)

    # Evaluate contributions
    if isinstance(contributions, pd.Series):
        W = W[sig_columns] * contributions[sig_columns]
    else:
        W = W[sig_columns] * contributions

    #### x coordinates for SBS contributions
    context_label = []
    change_map = {'CA': [], 'CG': [], 'CT': [], 'TA': [], 'TC': [], 'TG': []}
    for p in itertools.product('ACGT', 'ACGT'):
        context = ''.join(p)
        compl_context = compl(context, reverse=True)
        context_label.append('-'.join(context))
        for key in change_map:
            if key.startswith('C'):
                change_map[key].append(key + context)
            else:
                change_map[key].append(compl(key) + compl_context)
    color_map_sbs = {'CA': 'cyan', 'CG': 'red', 'CT': 'yellow', 'TA': 'purple', 'TC': 'green', 'TG': 'blue'}
    x_coords_sbs = range(16)

    ##### x coordinates for ID contributions
    group_map = {'Cdel': [], 'Tdel': [], 'Cins': [], 'Tins': [],
                 '2del': [], '3del': [], '4del': [], '5+del': [],
                 '2ins': [], '3ins': [], '4ins': [], '5+ins': [],
                 '2delm': [], '3delm': [], '4delm': [], '5+delm': []}
    for x in context83:
        group = re.search('.+?(?=[\d])', x).group(0)
        group_map[group].append(x)
    x_coords_id = {group: range(len(sigs)) for group, sigs in group_map.items()}

    color_map_id = {'Cdel': '#FFCC99', 'Tdel': '#FF8000', 'Cins': '#00FF00', 'Tins': '#00BB00',
                 '2del': '#FF99CC', '3del': '#FF3377', '4del': '#FF0000', '5+del': '#880000',
                 '2ins': '#99CCFF', '3ins': '#3377FF', '4ins': '#0000FF', '5+ins': '#000088',
                 '2delm': '#CC99FF', '3delm': '#9966FF', '4delm': '#8000FF', '5+delm': '#6000AA'}

    all_columns = ['CA', 'CG', 'CT', 'TA', 'TC', 'TG'] + ['space'] + list(group_map)

    fig, axes = plt.subplots(nrows=n_sigs, ncols=23, figsize=(60,2.5*n_sigs), sharex='col',
                             gridspec_kw={'width_ratios': (16,)*6 + (1,) + (6,)*12+(1,2,3,5)})
    max_height = 0
    for row, sig in enumerate(sig_columns):
        for col, ref in enumerate(all_columns):
            if n_sigs == 1:
                ax = axes[col]
            else:
                ax = axes[row,col]
            if col == 6:
                ax.remove()
                continue
            # For SBS portion
            if col < 6:
                bar_heights = W[sig].loc[change_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_sbs, bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map_sbs[ref], rasterized=True)
                ax.set_xlim(-.55, 15.55)
                if row == 0:
                    ax.set_title('>'.join(ref), fontsize=18)
                    if col == 0:
                        ax.text(5.5, 1.3, 'Mutational Signatures', transform=ax.transAxes,
                                horizontalalignment='center', fontsize=24)
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0, labelbottom=False)
                else:
                    ax.set_xticks(x_coords_sbs)
                    ax.set_xticklabels(context_label, fontfamily='monospace', rotation='vertical')
                    if col == 0:
                        ax.text(5.5, -.4, 'Motifs', transform = ax.transAxes, horizontalalignment='center', fontsize=20)
                if col == 5:
                    if n_sigs == 1:
                        for axis in axes[:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    else:
                        for axis in axes[row,:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    max_height = 0

            # For ID portion
            else:
                bar_heights = W[sig].loc[group_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_id[ref], bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map_id[ref],
                       rasterized=True)
                ax.set_xlim(-.55, x_coords_id[ref][-1] + .55)
                if row == 0:
                    ax.set_title(re.search('[\d+CT]+', ref).group(0), color=color_map_id[ref])
                    if ref == 'Tdel':
                        ax.text(-.02, 1.16, '1bp deletions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == 'Tins':
                        ax.text(-.02, 1.16, '1bp insertions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == '4del':
                        ax.text(-.02, 1.16, '>1bp deletions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == '4ins':
                        ax.text(-.02, 1.16, '>1bp insertions at repeats', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                    if ref == '4delm':
                        ax.text(.8, 1.16, '>1bp deletions with microhomology', fontsize=10, transform=ax.transAxes,
                                horizontalalignment='center', color=color_map_id[ref])
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0)
                else:
                    xlabels = [re.search('[\d+]+$', x).group(0) for x in group_map[ref]]
                    ax.set_xticks(x_coords_id[ref])
                    ax.set_xticklabels(xlabels, fontfamily='monospace')
                if col == 22:
                    if n_sigs == 1:
                        for axis in axes[6:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    else:
                        for axis in axes[row,6:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    max_height = 0

            if col not in [0,7]:
                ax.tick_params(axis='y', which='both',length=0, labelleft=False)
            if col == 22:
                ax.text(1.05, .5, sig, fontsize=14, rotation=270, transform=ax.transAxes, verticalalignment='center')

    # Set titles and organize plot
    plt.subplots_adjust(wspace=.12, hspace=.15)
    fig.text(.105, .5, 'Contributions', rotation='vertical', verticalalignment='center', fontsize=20, fontweight='bold')
    return fig

def signature_barplot_polymerase(W: pd.DataFrame, contributions: Union[int, pd.Series] = 1):
    W = W.copy()
    # Fill in missing features
    for c in context_polymerase96:
        if c not in list(W.index):
            W.loc[c] = 0
    W = W.reindex(context_polymerase96)

    # Get signature labels
    sig_columns = [c for c in W if c.startswith('S')]
    n_sigs = len(sig_columns)

    # Evaluate contributions
    if isinstance(contributions, pd.Series):
        W = W[sig_columns] * contributions[sig_columns]
    else:
        W = W[sig_columns] * contributions

    #### X coordinates for SBS contributions
    context_label = []
    change_map = {'CA': [], 'CG': [], 'CT': [], 'TA': [], 'TC': [], 'TG': []}
    for p in itertools.product('ACGT','ACGT'):
        context = ''.join(p)
        compl_context = compl(context, reverse=True)
        context_label.append('-'.join(context))
        for key in change_map:
            if key.startswith('C'):
                change_map[key].append(key + context)
            else:
                change_map[key].append(compl(key) + compl_context)
    color_map_sbs = {'CA': 'cyan', 'CG': 'red', 'CT': 'yellow', 'TA': 'purple', 'TC': 'green', 'TG': 'blue'}
    x_coords_sbs = range(16)

    #### X coordinates for ID contributions
    group_map = {'INS': ['INS' + str(i+1) for i in range(4)], 'DEL': ['DEL' + str(i+1) for i in range(4)]}
    color_map_id = {'INS':'#FFCC99', 'DEL':'#FF8000'}
    x_coords_id = {'INS':range(0,4), 'DEL':range(0,4)}
    all_columns = [x for x in change_map.keys()] + ['space', 'INS', 'DEL']

    fig, axes = plt.subplots(nrows=n_sigs, ncols=9, figsize=(60,2.5*n_sigs), sharex='col',
                             gridspec_kw={'width_ratios': (16,)*6 + (1,) + (4,)*2})
    max_height = 0
    # Iterate through signatures
    for row, sig in enumerate(sig_columns):
        # iterate through columns
        for col, ref in enumerate(all_columns):
            if n_sigs == 1:
                ax = axes[col]
            else:
                ax  = axes[row,col]
                if col == 6:
                    ax.remove()
                    continue
            # For SBS portion
            if col < 6:
                bar_heights = W[sig].loc[change_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_sbs, bar_heights, width=.95, linewidth=1.5, edgecolor='gray', color=color_map_sbs[ref], rasterized=True)
                ax.set_xlim(-.55, 15.55)
                if row == 0:
                    ax.set_title('>'.join(ref), fontsize=18)
                    if col == 0:
                        ax.text(4, 1.3, 'Mutational Signatures', transform=ax.transAxes,
                                horizontalalignment='center', fontsize=24)
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0, labelbottom=False)
                else:
                    ax.set_xticks(x_coords_sbs)
                    ax.set_xticklabels(context_label, fontfamily='monospace', rotation='vertical')
                    if col == 0:
                        ax.text(4, -.4, 'Motifs', transform = ax.transAxes, horizontalalignment='center', fontsize=20)
                if col == 5:
                    if n_sigs == 1:
                        for axis in axes[:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    else:
                        for axis in axes[row,:col+1]:
                            axis.set_ylim(0,max_height + 0.1*max_height+1)
                    max_height = 0
            # For ID portion
            else:
                bar_heights = W[sig].loc[group_map[ref]]
                for height in bar_heights:
                    if height > max_height: max_height = height
                ax.bar(x_coords_id[ref], bar_heights, width=0.95, linewidth=1.5, edgecolor='gray', color=color_map_id[ref],
                       rasterized=True)
                ax.set_xlim(-.55, x_coords_id[ref][-1] + 0.55)
                # Set column titles
                if row == 0:
                    ax.set_title(ref, color=color_map_id[ref])
                if row < n_sigs - 1:
                    ax.tick_params(axis='x', length=0)
                else:
                    xlabels = ['1','2','3','4+']
                    ax.set_xticks(x_coords_id[ref])
                    ax.set_xticklabels(xlabels, fontfamily='monospace')
                if col == 8:
                    ax.text(1.05, .5, sig, fontsize=14, rotation=270, transform=ax.transAxes, verticalalignment="center")
    # Set titles and organize plot
    plt.subplots_adjust(wspace=.12, hspace=0.15)
    fig.text(0.105, 0.5, 'Contributions', rotation='vertical', verticalalignment='center', fontsize=20, fontweight='bold')
    return fig


#MIT License
#
#Copyright (c) 2020 Pierre Michel Joubert
#
#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 csv
import re
import random
from scipy.stats import gaussian_kde
import numpy as np
import scipy.interpolate
import sys

## USAGE ##
# this python script uses an input distribution of eccDNAs called using uniquely mapped reads to probabilistically call putative eccDNAs with multi-mapped reads
# this script specifically works on split reads with one side that is uniquely mapped
# options
# "known_srs_length" - input list of lengths of eccDNAs called using uniquely mapped reads
# "split_reads" - input list of multi-mapped split reads with secondary alignments
# "column_cutoff" - maximum length of output eccDNAs, usually determined by the maximum size of a molecule that can come out of a filtration column during circleseq library prep

known_srs_length = str(sys.argv[1])
split_reads = str(sys.argv[2])
column_cutoff = int(sys.argv[3])

# open the lengths of uniquely mapped eccDNAs and calculate a kernel density estimate
with open(known_srs_length, newline = '') as file: ## this is for each individual technical replicate
    file_reader = csv.reader(file, delimiter = '\t')
    dsn_list = [int(row[0]) for row in file_reader]
density = gaussian_kde(dsn_list)
kde_max = 0
while True:
    if density(kde_max) == 0:
        break
    else:
        kde_max += 100
xs = np.linspace(0,kde_max,1000)
densities = density(xs)
y_interp = scipy.interpolate.interp1d(xs,densities)

## define regex for processing split reads
start_pattern = "^([0-9]+)M.*[HS]$"
end_pattern = ".*[HS]([0-9]+)M$"

# use regex to grab the matches and nonmatches to the genome and count them
# side one the uniquely mapped side, side two is the multimapped side
# do this for each split read, to calculate all potential combinations of side one and side two later
def process_split_read(read):
    matches = re.findall(r'(\d+)([A-Z]{1})', read[6])
    matches_sums = {'M': 0, 'other': 0}
    for i in range(len(matches)):
        if matches[i][1] == 'M':
            matches_sums['M'] += int(matches[i][0])
        else:
            matches_sums['other'] += int(matches[i][0])
    if str(read[5]) == '-':
        sense = 'reverse'
    elif str(read[5]) == '+':
        sense = 'forward'
    string = read[6]
    if re.match(start_pattern, string):
        loc = 'start'
    elif re.match(end_pattern, string):
        loc = 'end'
    else:
        return False ## if regexp doesn't work then just drop the split
    if read[4] != 0:
        split_read_side_one.append([read[0], int(read[1]), int(read[2]), sense, read[6], loc, matches_sums['M'], matches_sums['other']])
    else:
        split_read_side_two.append([read[0], int(read[1]), int(read[2]), sense, read[6], loc, matches_sums['M'], matches_sums['other']])
    return True

# go through all possible split read combinations and filter down to a subset that represent eccDNAs
# then, of all of the possible combinations, randomly choose one
# random choice is weighted based off KDE from uniquely mapped eccDNAs above
def choose_split_reads(split_read_side_one, split_read_side_two_pre):
    if len(split_read_side_two_pre) == 0:
        return False
    split_read_side_two = np.array(split_read_side_two_pre, dtype=object)
    combos = []
    for i in range(len(split_read_side_one)):
        split_read_first_side = split_read_side_one[i]
        ## reduce to
        # same chromosome
        # same orientation
        # distance smaller than column cutoff
        # start and end pairs
        # properly oriented (eccDNA vs intron)
        # total of read lengths is around 150
        if ((split_read_first_side[3] == 'forward' and split_read_first_side[5] == 'start') or
            (split_read_first_side[3] == 'reverse' and split_read_first_side[5] == 'end')):
                reduced = split_read_side_two[np.logical_and.reduce((split_read_side_two[:, 0] == split_read_first_side[0],
                                       split_read_side_two[:, 3] == split_read_first_side[3],
                                       abs(split_read_side_two[:,1] - split_read_first_side[1]) < column_cutoff,
                                        split_read_side_two[:, 5] != split_read_first_side[5],
                                        split_read_side_two[:,1] < split_read_first_side[1],
                                        np.isclose((split_read_side_two[:,6] + split_read_first_side[6]).astype(int), (split_read_side_two[:,7] + split_read_first_side[7]).astype(int), atol=10)))]
        if ((split_read_first_side[3] == 'forward' and split_read_first_side[5] == 'end') or
            (split_read_first_side[3] == 'reverse' and split_read_first_side[5] == 'start')):
                reduced = split_read_side_two[np.logical_and.reduce((split_read_side_two[:, 0] == split_read_first_side[0],
                                       split_read_side_two[:, 3] == split_read_first_side[3],
                                       abs(split_read_side_two[:,1] - split_read_first_side[1]) < column_cutoff,
                                        split_read_side_two[:, 5] != split_read_first_side[5],
                                        split_read_side_two[:,1] > split_read_first_side[1],
                                        np.isclose((split_read_side_two[:,6] + split_read_first_side[6]).astype(int), (split_read_side_two[:,7] + split_read_first_side[7]).astype(int), atol=10)))]
        if len(reduced) != 0:
            for ii in range(len(reduced)):
                split_read_second_side = reduced[ii]
                start = min([split_read_first_side[1], split_read_first_side[2], split_read_second_side[1], split_read_second_side[2]])
                end = max([split_read_first_side[1], split_read_first_side[2], split_read_second_side[1], split_read_second_side[2]])
                combos.append([split_read_first_side[0], start, end])
    if combos == []: # any of the options left?
        return False
    densities = []
    for ii in range(len(combos)):
        try:
            densities.append(y_interp(abs(combos[ii][2]-combos[ii][1])))
        except ValueError:
            densities.append(0)
    densities_sum = sum(densities)
    if densities_sum == 0: ## any of the options possible?
        return False
    densities_array = np.array(densities)
    densities_ratio = densities_array/densities_sum
    choice = random.choices(combos, densities_ratio, k=1)[0]
    return choice

# run code without opening whole files into memory
# write chosen split reads
with open(split_reads, newline = '') as file:
    file_reader = csv.reader(file, delimiter = '\t')
    with open('singleunique_choices', 'w', newline = '') as confirmed:
        w = csv.writer(confirmed, delimiter = '\t')
        for row in file_reader:
            current_line = row
            current_read = current_line[3]
            split_read_side_one = []
            split_read_side_two = []
            while current_line[3] == current_read:
                process_split_read(current_line)
                try:
                    current_line = next(file_reader)
                except StopIteration:
                    break
            choice = choose_split_reads(split_read_side_one, split_read_side_two)
            if choice:
                w.writerow([choice[0], choice[1], choice[2]])