import os, sys 
import numpy as np
from pylab import *
from matplotlib.font_manager import FontProperties
import argparse
from subprocess import Popen, PIPE
import subprocess, shlex
from time import gmtime, strftime
from shutil import copyfile
from pathlib import Path
from collections import Counter  
import re
import time


# Fonts

alignment = {'horizontalalignment': 'center', 'verticalalignment': 'baseline'}
families = ['serif', 'sans-serif', 'cursive', 'fantasy', 'monospace']
styles = ['normal', 'italic', 'oblique']
font = FontProperties()

#  tick fonts
font1 = font.copy()
font1.set_family('sans-serif')
font1.set_style('normal')

# label fonts
font2 = font.copy()
font2.set_family('sans-serif')
font2.set_style('normal')
font2.set_weight('normal')
plt.rcParams['mathtext.default'] = 'regular'
plt.rcParams['axes.linewidth'] = 6
bbox = dict(boxstyle="round", fc="0.8")
cmaps = [('Perceptually Uniform Sequential', [
            'viridis', 'plasma', 'inferno', 'magma']),
         ('Sequential', [
            'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
            'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
            'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn']),
         ('Sequential (2)', [
            'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone', 'pink',
            'spring', 'summer', 'autumn', 'winter', 'cool', 'Wistia',
            'hot', 'afmhot', 'gist_heat', 'copper']),
         ('Diverging', [
            'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
            'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']),
         ('Qualitative', [
            'Pastel1', 'Pastel2', 'Paired', 'Accent',
            'Dark2', 'Set1', 'Set2', 'Set3',
            'tab10', 'tab20', 'tab20b', 'tab20c']),
         ('Miscellaneous', [
            'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
            'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg', 'hsv',
            'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]





### Note trajectory name cannot be the same as folder name

parser = argparse.ArgumentParser(description='checks for atom overlap', prog='CG2AT', epilog='Enjoy the program and best of luck!\n', allow_abbrev=True)

parser.add_argument('-f', help='xtc file',metavar='xtc',type=str)
args = parser.parse_args()
options = vars(args)


def get_ss():
    in1={}
    reset = []
    counter=0
    time_prev = 'x'
    with open(args.f, 'r') as xvg_input:
        for line_nr, line in enumerate(xvg_input.readlines()):
            if len(line) > 0:
                if not line.startswith('#'):
                    if line.split()[3] != time_prev:
                        counter = 0
                    if counter not in reset:
                        in1[counter]=[]
                        reset.append(counter)

                    in1[counter].append(line.split()[4])
                    counter+=1
                    time_prev = line.split()[3]
    return in1

SS = ['T', 'E', 'B', 'H', 'G', 'I', 'C']
SS_col = ['blue', 'orange', 'green', 'red', 'purple', 'brown', 'pink']
aa_ss = get_ss()
ss_all = []
for aa in aa_ss:
    ss_all.append(aa_ss[aa])
ss_all = np.array(ss_all)

time_line = [[],[],[],[],[],[],[]]
if len(ss_all[0]) > 1: 
    for time_stamp in range(len(ss_all[0])):
        d = {'T':0, 'E':0, 'B':0, 'H':0, 'G':0, 'I':0, 'C':0}
        unique, counts = np.unique(ss_all[:,time_stamp], return_counts=True)
        d.update(dict(zip(unique, counts)))
        for i_val, i in enumerate(SS):
            time_line[i_val].append((d[i]/np.sum(counts))*100)
else:
    d = {'T':0, 'E':0, 'B':0, 'H':0, 'G':0, 'I':0, 'C':0}
    ss_all = np.concatenate(ss_all)

    unique, counts = np.unique(ss_all, return_counts=True)
    # print(unique, counts)
    d.update(dict(zip(unique, counts))) 
    for i in d:
        print(i, np.round((d[i]/len(ss_all))*100,1), np.round(np.std(d[i]),1))
    sys.exit()

# for i_val, i in enumerate(SS):
#     print(i, np.round(np.mean(time_line[i_val]),1), np.round(np.std(time_line[i_val]),1))

plt.figure(1, figsize=(20,5))
for i_val, i in enumerate(SS):
    plt.axhline(np.mean(time_line[i_val]),xmin=-10, xmax=2000, linewidth=1, linestyle='--', color=SS_col[i_val]) # LHUHBU

    plt.plot(np.arange(0, 101,1), time_line[i_val], linewidth=3, label=i, color=SS_col[i_val])

    print(i, np.round(np.mean(time_line[i_val]),1), np.round(np.std(time_line[i_val]),1))
plt.ylim(0, 100)
# plt.xlabel('Time (10ps)', fontproperties=font2,fontsize=40)
# plt.ylabel('Percentage', fontproperties=font2,fontsize=40)
plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=7, mode="expand", borderaxespad=0.,  fontsize=20, markerscale=10)
plt.tick_params(axis='both', which='major', width=3, length=5, labelsize=40, direction='in', pad=20, right=False, top=False)
# plt.show()
plt.savefig('SS_time.png', dpi=600)
