import os, sys
import numpy as np
import scipy.constants
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as colors
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.font_manager import FontProperties
from matplotlib.patches import Rectangle
import matplotlib.patches as patches
from subprocess import Popen, PIPE
import subprocess, shlex
from time import gmtime, strftime
import math
from matplotlib import path
import scipy.integrate as integrate
import multiprocessing as mp
import time
# import argparse
import re
from scipy import stats
from numba import jit
from scipy.spatial import cKDTree

# 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'])]

def get_values(file_name, in_list):
    with open(file_name, 'r') as xvg_input:
        for line_nr, line in enumerate(xvg_input.readlines()):
            if len(line) > 0:
                if line[0] not in ['@','#']:
                    if len(in_list) < int(line.split()[0]):
                        in_list.append([])
                    in_list[int(line.split()[0])-1].append(float(line.split()[1]))
    return in_list

def analysis(file_name):
    in_list=[]
    for time in range(5, 105, 5):
        in_list = get_values(file_name+str(time)+'.xvg', in_list)
    od_values = np.array(in_list)

    od_mean, od_std = np.array([]), np.array([])
    for timepoint in range(len(od_values[0,:])):
        od_mean = np.append(od_mean, np.mean(np.sqrt((od_values[:,timepoint]-od_values[:,-1])**2)))
        od_std= np.append(od_std, np.std(np.sqrt((od_values[:,timepoint]-od_values[:,-1])**2)))
    return od_mean, od_std

files = ['POPC_pal_od_', 'POPC_ol_od_','POPS_pal_od_', 'POPS_ol_od_']
titles = ['POPC Palmitoyl', 'POPC Oleoyl', 'POPS Palmitoyl', 'POPS Oleoyl']
loc_start = '/media/owen/WD_drive/cg2at_testing/cg2at_quality_control/POPC_7-POPS_1_CHOL_2/'
col = ['orange', 'blue']

plt.figure(1, figsize=(15,20))
for sup_val, sup_file in enumerate(files, start=1):
    
    plt.subplot(4,1,sup_val)
    plt.title(titles[sup_val-1] ,  fontproperties=font1, fontsize=30,y=0.75)#+' gaussian height'
    for loc_val, loc in enumerate(['charmmgui','CG2AT']):
        od_mean, od_std = analysis(loc_start+loc+'/analysis/order_param/time_course/'+sup_file)
        plt.fill_between(np.arange(5, 105, 5), od_mean-od_std, od_mean+od_std, alpha=0.5, facecolor=col[loc_val])
        plt.plot(np.arange(5, 105, 5), od_mean,linewidth=4, color=col[loc_val], label=loc)
    plt.xticks(np.arange(0, 110,10), fontproperties=font1,  fontsize=20)#
    plt.yticks(np.arange(0, 0.061,0.02), fontproperties=font1,  fontsize=20)#
    plt.ylim(0,0.06)
    plt.ylabel('', fontproperties=font2,fontsize=30)

    plt.legend(shadow=False, fancybox=False, loc='upper right', prop={'size': 20})

    if sup_val<4:
        plt.tick_params(axis='both', which='major', width=3, length=5, labelsize=35, direction='in', pad=10, right=False, top=False,labelbottom=False)
    else:
        plt.tick_params(axis='both', which='major', width=3, length=5, labelsize=35, direction='in', pad=10, right=False, top=False)
        plt.xlabel('Time (ns)', fontproperties=font2,fontsize=30)
plt.subplots_adjust(top=0.967, bottom=0.1,left=0.087,right=0.99, hspace=0.2, wspace=0.16)

plt.savefig('order_param_time.png', dpi=500)

plt.show()