import argparse, sys, os, pickle
from os.path import splitext, basename
import numpy as np
import pylab as pl
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

__location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
sys.path.append(os.path.split(__location__)[0])
from helpers import parse_output_dir, parse_input_dir

color_list = ['#66c2a5', '#fc8d62','#8da0cb']
color_dict = {'C':'#66c2a5',
              'K': '#fc8d62',
              'R': '#8da0cb'}

parser = argparse.ArgumentParser(description='Using tsvs generated by graph_fingerprints, plot a subset under eachother'
                                             'for publication.')
parser.add_argument('--in-dir', type=str, required=True,
                    help='Output directory of graph_fingerprints')
parser.add_argument('--min-p', type=float, default=1.0,
                    help='Minimum occurence for particular value to be included as peak [default: 1.0]')
parser.add_argument('--resolution', default=1.0)
parser.add_argument('--idx', nargs='+', type=str, required=True,
                    help='For each type, indicate idx of the fingerprint that should be included. Format as follows:'
                          'PROTEIN_IDX')
parser.add_argument('--out-svg', type=str, required=True)
args = parser.parse_args()

tsv_list = parse_input_dir(args.in_dir, pattern='*.tsv')
idx_tups = [idx.rsplit('_', 1) for idx in args.idx]
idx_dict = {idx[0]: int(idx[1]) for idx in idx_tups}
nb_segments = 100 / args.resolution


plot_dict = {}
for tsv_fn in tsv_list:
    prot_id = splitext(basename(tsv_fn))[0]
    cidx = idx_dict[prot_id]
    fps = np.loadtxt(tsv_fn)
    fps[fps < args.min_p] = 0.0
    fps[fps > 0] = 1.0
    nb_tags = round(fps.shape[1] / nb_segments)
    plot_dict[prot_id] = fps[cidx, :].reshape(nb_tags,-1)


# --- collect plotting elements ---
ls_list = []
colors = []

for i1, pid in enumerate(plot_dict):
    for i2, fp in enumerate(plot_dict[pid]):
        segments = [[(it, i1), (it, i1 + fs)] for it, fs in enumerate(fp) if fs != 0]
        ls_list.append(LineCollection(segments, colors=len(segments) * [color_list[i2]]))

fig, ax = plt.subplots(figsize=(6, 2.9375))
for ls in ls_list: ax.add_collection(ls)
plt.xlim(0, 100)
plt.ylim(0, len(plot_dict))

id_list = list(plot_dict)
plt.yticks([i + 0.5 for i in list(range(len(id_list)))], id_list)
plt.xticks([0, 25, 50, 75, 100])
plt.xlabel('$E_{FRET}$')

fig.savefig(args.out_svg)
