###############################################
##Dmitry Sutormin, 2020##
##N-to-C-terminus assymetry in protein domain secondary structure elements composition##

#Pareses DSSP output generated for representative structures with Run_DSSP.py
#Get statistics of secondary structure elements at N- and C-termini of protein domains.
###############################################

#######
#Packages to be imported.
#######

import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as st

#######
#Variables to be defined.
#######

#Path to DSSP data file.
DSSP_data_inpath="C:\\Users\sutor\OneDrive\ThinkPad_working\Sutor\Documents\Skoltech_PhD\Structural_bioinformatics\Project\DSSP\DSSP\DSSP_Representative_pdb_domains.txt"

#Path to output file.
DSSP_data_outpath="C:\\Users\sutor\OneDrive\ThinkPad_working\Sutor\Documents\Skoltech_PhD\Structural_bioinformatics\Project\DSSP\\"

#######
#Read DSSP output.
#######

def read_dssp_data(DSSP_inpath):

    #Read DSSP data keep in dictionary.
    DSSP_data_dict={}
    filein=open(DSSP_inpath, 'r')
    for line in filein:
        line=line.rstrip().split('\t')
        if line[0] not in ['PFAM_id']:
            PFAM_id=line[0]
            PDB_id=line[1]
            Chain_id=line[2]
            Start=int(line[3])
            End=int(line[4])
            Resolution=float(line[5])
            SSE_string=line[6]
            Phi_list=[float(x) for x in line[7].lstrip('[').rstrip(']').split(', ')]
            Psi_list=[float(x) for x in line[8].lstrip('[').rstrip(']').split(', ')]
            DSSP_data_dict[PFAM_id]=[PDB_id, Chain_id, Start, End, Resolution, SSE_string, Phi_list, Psi_list]

    filein.close()
    return DSSP_data_dict

#######
#Filter data, split data by domain length.
#######

def define_length_groups(DSSP_data_dict, min_len, thr_len):

    Domain_length_ar=[]
    Discarded_short_structures={}
    Short_structures={}
    Long_structures={}
    for pfam_id, dssp_data in DSSP_data_dict.items():
        domain_len=len(dssp_data[5])
        Domain_length_ar.append(domain_len)

        #Classify domains by length.
        if domain_len<min_len:
            Discarded_short_structures[pfam_id]=dssp_data
        elif min_len<=domain_len<thr_len:
            Short_structures[pfam_id]=dssp_data
        elif domain_len>=thr_len:
            Long_structures[pfam_id]=dssp_data

    #Plot domain length distribution.
    fig, plot=plt.subplots(1,1,figsize=(3,3), dpi=100)
    plot.hist(Domain_length_ar, bins=50, rwidth=0.85, color='#94fff1', edgecolor='black', linewidth=0.1)
    plot.axvline(x=min_len, ls='--', linewidth=0.7, color='black')
    plot.axvline(x=thr_len, ls='--', linewidth=0.7)
    plot.set_xlabel('Domain length, aa')
    plot.set_ylabel('Number of structures')
    plot.set_yscale('log')
    plot.set_title('Representative structures\nafter DSSP')
    plot.annotate(f'Number of\nstructures\n{len(Domain_length_ar)}', (0.5,0.7), xycoords='axes fraction')
    plot.annotate(f'x<{min_len} : {len(Discarded_short_structures)}', (0.5,0.6), xycoords='axes fraction', size=7)
    plot.annotate(f'{min_len}<=x<{thr_len} : {len(Short_structures)}', (0.5,0.5), xycoords='axes fraction', size=7)
    plot.annotate(f'x>={thr_len} : {len(Long_structures)}', (0.5,0.4), xycoords='axes fraction', size=7)

    plt.tight_layout()
    plt.show()

    return Short_structures, Long_structures

#######
#Take phi, psi angles for N- and C-termini.
#######

def phi_psi_N_to_C(structures_dict, window_width):

    phi_N=[]
    phi_C=[]
    psi_N=[]
    psi_C=[]
    phi=[]
    psi=[]
    for pfam_id, dssp_data in structures_dict.items():
        phi_list=dssp_data[6]
        psi_list=dssp_data[7]
        phi_N+=phi_list[:window_width]
        phi_C+=phi_list[-window_width:]
        psi_N+=psi_list[:window_width]
        psi_C+=psi_list[-window_width:]
        phi+=phi_list
        psi+=psi_list

    return phi_N, phi_C, psi_N, psi_C, phi, psi

#######
#Create secondary structure element frequency matrix.
#######

def ss_element_frequency_matrix(structures_dict, window_width):

    #Create positioned matrix of secondary structure elements.
    ss_matrix_N=[]
    ss_matrix_C=[]
    for i in range(window_width):
        column_N=[]
        column_C=[]
        for pfam_id, dssp_data in structures_dict.items():
            SSE_string=dssp_data[5]
            column_N.append(SSE_string[i])
            column_C.append(SSE_string[-window_width+i])
        ss_matrix_N.append(column_N)
        ss_matrix_C.append(column_C)

    #Create position frequency matrix of secondary structure elements.
    DSSP_alphabet=['H', 'B', 'E', 'G', 'I', 'T', 'S', '-']
    ss_pfm_N={}
    ss_pfm_C={}
    ss_pfm_conf_N={}
    ss_pfm_conf_C={}
    for letter_code in DSSP_alphabet:
        #Keep frequences of ss elements.
        frequency_row_N=[]
        frequency_row_C=[]
        #Keep boundaries of confident interval.
        conf_upper_N=[]
        conf_upper_C=[]
        conf_lower_N=[]
        conf_lower_C=[]
        for i in range(len(ss_matrix_N)):
            column_letter_freq_N=ss_matrix_N[i].count(letter_code)/float(len(ss_matrix_N[i]))
            confident_interval_N=st.binom.interval(0.95, len(ss_matrix_N[i]), column_letter_freq_N, loc=0)
            lower_N=confident_interval_N[0]/len(ss_matrix_N[i])
            upper_N=confident_interval_N[1]/len(ss_matrix_N[i])
            conf_lower_N.append(lower_N)
            conf_upper_N.append(upper_N)
            frequency_row_N.append(column_letter_freq_N)

            column_letter_freq_C=ss_matrix_C[i].count(letter_code)/float(len(ss_matrix_C[i]))
            confident_interval_C=st.binom.interval(0.95, len(ss_matrix_C[i]), column_letter_freq_C, loc=0)
            lower_C=confident_interval_C[0]/len(ss_matrix_C[i])
            upper_C=confident_interval_C[1]/len(ss_matrix_C[i])
            conf_lower_C.append(lower_C)
            conf_upper_C.append(upper_C)
            frequency_row_C.append(column_letter_freq_C)

        ss_pfm_N[letter_code]=frequency_row_N
        ss_pfm_C[letter_code]=frequency_row_C

        ss_pfm_conf_N[letter_code+'_upper']=conf_upper_N
        ss_pfm_conf_N[letter_code+'_lower']=conf_lower_N
        ss_pfm_conf_C[letter_code+'_upper']=conf_upper_C
        ss_pfm_conf_C[letter_code+'_lower']=conf_lower_C

    print(ss_pfm_N)
    print(ss_pfm_C)
    return ss_matrix_N, ss_matrix_C, ss_pfm_N, ss_pfm_C, ss_pfm_conf_N, ss_pfm_conf_C

#######
#Enrichment of secondary structure elements N- over C-terminus.
#######

def ss_ele_enrichment(ss_pfm_N, ss_pfm_C):

    enrichment_N_to_C_dict={}
    for letter_code, frequency_row_N in ss_pfm_N.items():
        frequency_N=np.array(frequency_row_N)
        frequency_C=np.array(ss_pfm_C[letter_code][::-1])
        enrichment_N_to_C=frequency_N/frequency_C
        enrichment_N_to_C_dict[letter_code]=enrichment_N_to_C

    return enrichment_N_to_C_dict

#######
#Terminal beta-strands relations: simultaneous occurence or independant?.
#######

def termini_dependance(ss_matrix_N, ss_matrix_C, ss_pfm_N, ss_pfm_C, local_window_width):

    DSSP_alphabet=['H', 'B', 'E', 'G', 'I', 'T', 'S', '-']

    #Calculate observed frequences of elements co-occurence coordinate-wise.
    Observed_matrices_dict={}
    for letter_code_1 in DSSP_alphabet:
        for letter_code_2 in DSSP_alphabet:
            Ni_Cj_freq_2d_ar=[]
            for Ni in range(local_window_width):
                Cj_freq_1d_ar=[]
                for Cj in range(local_window_width):
                    N_column_i=ss_matrix_N[Ni]
                    C_column_j=ss_matrix_C[-Cj-1]
                    Ni_Cj_counts=0
                    for structure_k in range(len(N_column_i)):
                        if (N_column_i[structure_k]==letter_code_1) and (C_column_j[structure_k]==letter_code_2):
                            Ni_Cj_counts+=1
                    Ni_Cj_frequency=Ni_Cj_counts/float(len(N_column_i))
                    Cj_freq_1d_ar.append(Ni_Cj_frequency)

                Ni_Cj_freq_2d_ar.append(Cj_freq_1d_ar)

            Ni_Cj_freq_2d_ar_np=np.array(Ni_Cj_freq_2d_ar)
            Observed_matrices_dict[letter_code_1+letter_code_2]=Ni_Cj_freq_2d_ar_np

    #Calculate expected frequences of elements co-occurence coordinate-wise.
    Expected_matrices_dict={}
    for letter_code_1 in DSSP_alphabet:
        for letter_code_2 in DSSP_alphabet:
            Expected_Ni_Cj_freq_2d_ar=[]
            for Ni in range(local_window_width):
                Expected_Cj_freq_1d_ar=[]
                for Cj in range(local_window_width):
                    Ni_frequency=ss_pfm_N[letter_code_1][Ni]
                    Cj_frequency=ss_pfm_C[letter_code_2][-Cj-1]
                    Expected_Ni_Cj_frequency=Ni_frequency*Cj_frequency
                    Expected_Cj_freq_1d_ar.append(Expected_Ni_Cj_frequency)
                Expected_Ni_Cj_freq_2d_ar.append(Expected_Cj_freq_1d_ar)

            Expected_Ni_Cj_freq_2d_ar_np=np.array(Expected_Ni_Cj_freq_2d_ar)
            Expected_matrices_dict[letter_code_1+letter_code_2]=Expected_Ni_Cj_freq_2d_ar_np

    #Calculate observed over expected ratio of frequences of elements co-occurence coordinate-wise.
    Obs_over_exp_matrices_dict={}
    for letter_code_1 in DSSP_alphabet:
        for letter_code_2 in DSSP_alphabet:
            Obs_over_exp_freq_matrix=np.divide(Observed_matrices_dict[letter_code_1+letter_code_2], Expected_matrices_dict[letter_code_1+letter_code_2])
            Obs_over_exp_matrices_dict[letter_code_1+letter_code_2]=Obs_over_exp_freq_matrix

            print(letter_code_1, letter_code_2, Obs_over_exp_freq_matrix)

    return Obs_over_exp_matrices_dict

#######
#Analyse N-to-C asymmetry in secondary structure elements frequency.
#######

def distribution_of_frequences(ss_pfm_N, ss_pfm_C):
    ss_pfm=ss_pfm_N+ss_pfm_C
    print(len(ss_pfm))

    fig, plot=plt.subplots(1,1,figsize=(4,4), dpi=100)
    weights=np.ones_like(ss_pfm)/(len(ss_pfm)) #Taken from https://stackoverflow.com/questions/42481698/probability-density-histogram-with-matplotlib-doesnt-make-sense
    plot.hist(ss_pfm, bins=10, weights=weights)
    plot.set_xlabel('Frequency of ss element')
    plot.set_ylabel('Fraction of positions')
    plt.show()

    return

#######
#Analyse N-to-C asymmetry in secondary structure elements frequency.
#######

def N_to_C_asymmetry(ss_pfm_N_short, ss_pfm_C_short, ss_pfm_conf_N_short, ss_pfm_conf_C_short, ss_pfm_N_long, ss_pfm_C_long, ss_pfm_conf_N_long, ss_pfm_conf_C_long, window_width):

    ##Plot distribution of frequences, get confidential interval.
    ##Distribution are far from being normal, estimation is not good -> depricated.
    #distribution_of_frequences(ss_pfm_N_short['H'], ss_pfm_C_short['H'])
    #distribution_of_frequences(ss_pfm_N_long['H'], ss_pfm_C_long['H'])

    #Plot frequency of ss elements as a function of distance from N- and C-termini.
    X_N=range(window_width)
    X_C=range(60, 60+window_width)
    xticks_ar=list(range(0,111,10))
    xticklabels_ar=[0,10,20,30,40,50,-50,-40,-30,-20,-10,0]
    xticks_spec=[-4, 114]
    xticklabels_spec=['N-', '-C']
    fig, plot=plt.subplots(2,2,figsize=(12,6), dpi=100)
    plot[0,0].plot(X_N, ss_pfm_N_short['H'], color='#94fff1', linewidth=2, label=r'$\alpha$-helix short domains')
    plot[0,0].fill_between(X_N, ss_pfm_conf_N_short['H_lower'], ss_pfm_conf_N_short['H_upper'], color='#94fff1', alpha=0.3, linewidth=0)
    plot[0,0].plot(X_C, ss_pfm_C_short['H'], color='#94fff1', linewidth=2)
    plot[0,0].fill_between(X_C, ss_pfm_conf_C_short['H_lower'], ss_pfm_conf_C_short['H_upper'], color='#94fff1', alpha=0.3, linewidth=0)

    plot[0,0].plot(X_N, ss_pfm_N_long['H'], color='#ab658c', linewidth=2, label=r'$\alpha$-helix long domains')
    plot[0,0].fill_between(X_N, ss_pfm_conf_N_long['H_lower'], ss_pfm_conf_N_long['H_upper'], color='#ab658c', alpha=0.3, linewidth=0)
    plot[0,0].plot(X_C, ss_pfm_C_long['H'], color='#ab658c', linewidth=2)
    plot[0,0].fill_between(X_C, ss_pfm_conf_C_long['H_lower'], ss_pfm_conf_C_long['H_upper'], color='#ab658c', alpha=0.3, linewidth=0)

    plot[0,0].plot(X_N, ss_pfm_N_short['E'], color='#59acff', linewidth=2, label=r'$\beta$-strand short domains')
    plot[0,0].fill_between(X_N, ss_pfm_conf_N_short['E_lower'], ss_pfm_conf_N_short['E_upper'], color='#59acff', alpha=0.3, linewidth=0)
    plot[0,0].plot(X_C, ss_pfm_C_short['E'], color='#59acff', linewidth=2)
    plot[0,0].fill_between(X_C, ss_pfm_conf_C_short['E_lower'], ss_pfm_conf_C_short['E_upper'], color='#59acff', alpha=0.3, linewidth=0)

    plot[0,0].plot(X_N, ss_pfm_N_long['E'], color='#ffec59', linewidth=2, label=r'$\beta$-strand long domains')
    plot[0,0].fill_between(X_N, ss_pfm_conf_N_long['E_lower'], ss_pfm_conf_N_long['E_upper'], color='#ffec59', alpha=0.3, linewidth=0)
    plot[0,0].plot(X_C, ss_pfm_C_long['E'], color='#ffec59', linewidth=2)
    plot[0,0].fill_between(X_C, ss_pfm_conf_C_long['E_lower'], ss_pfm_conf_C_long['E_upper'], color='#ffec59', alpha=0.3, linewidth=0)

    plot[0,0].set_xticks(xticks_ar)
    plot[0,0].set_xticklabels(xticklabels_ar)
    plot[0,0].set_xticks(xticks_spec, minor=True)
    plot[0,0].set_xticklabels(xticklabels_spec, minor=True)
    plot[0,0].tick_params(axis='x', which='minor', length=0)
    plot[0,0].set_xlabel('Distance, aa')
    plot[0,0].set_ylabel('Frequency')
    plot[0,0].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='upper center')

    plot[0,1].plot(X_N, ss_pfm_N_short['B'], color='#94fff1', linewidth=2, label=r'$\beta$-bridge short domains')
    plot[0,1].fill_between(X_N, ss_pfm_conf_N_short['B_lower'], ss_pfm_conf_N_short['B_upper'], color='#94fff1', alpha=0.3, linewidth=0)
    plot[0,1].plot(X_C, ss_pfm_C_short['B'], color='#94fff1', linewidth=2)
    plot[0,1].fill_between(X_C, ss_pfm_conf_C_short['B_lower'], ss_pfm_conf_C_short['B_upper'], color='#94fff1', alpha=0.3, linewidth=0)

    plot[0,1].plot(X_N, ss_pfm_N_long['B'], color='#ab658c', linewidth=2, label=r'$\beta$-bridge long domains')
    plot[0,1].fill_between(X_N, ss_pfm_conf_N_long['B_lower'], ss_pfm_conf_N_long['B_upper'], color='#ab658c', alpha=0.3, linewidth=0)
    plot[0,1].plot(X_C, ss_pfm_C_long['B'], color='#ab658c', linewidth=2)
    plot[0,1].fill_between(X_C, ss_pfm_conf_C_long['B_lower'], ss_pfm_conf_C_long['B_upper'], color='#ab658c', alpha=0.3, linewidth=0)

    plot[0,1].plot(X_N, ss_pfm_N_short['G'], color='#59acff', linewidth=2, label=r'3-10-helix short domains')
    plot[0,1].fill_between(X_N, ss_pfm_conf_N_short['G_lower'], ss_pfm_conf_N_short['G_upper'], color='#59acff', alpha=0.3, linewidth=0)
    plot[0,1].plot(X_C, ss_pfm_C_short['G'], color='#59acff', linewidth=2)
    plot[0,1].fill_between(X_C, ss_pfm_conf_C_short['G_lower'], ss_pfm_conf_C_short['G_upper'], color='#59acff', alpha=0.3, linewidth=0)

    plot[0,1].plot(X_N, ss_pfm_N_long['G'], color='#ffec59', linewidth=2, label=r'3-10-helix long domains')
    plot[0,1].fill_between(X_N, ss_pfm_conf_N_long['G_lower'], ss_pfm_conf_N_long['G_upper'], color='#ffec59', alpha=0.3, linewidth=0)
    plot[0,1].plot(X_C, ss_pfm_C_long['G'], color='#ffec59', linewidth=2)
    plot[0,1].fill_between(X_C, ss_pfm_conf_C_long['G_lower'], ss_pfm_conf_C_long['G_upper'], color='#ffec59', alpha=0.3, linewidth=0)

    plot[0,1].set_xticks(xticks_ar)
    plot[0,1].set_xticklabels(xticklabels_ar)
    plot[0,1].set_xticks(xticks_spec, minor=True)
    plot[0,1].set_xticklabels(xticklabels_spec, minor=True)
    plot[0,1].tick_params(axis='x', which='minor', length=0)
    plot[0,1].set_xlabel('Distance, aa')
    plot[0,1].set_ylabel('Frequency')
    plot[0,1].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='center')

    plot[1,0].plot(X_N, ss_pfm_N_short['I'], color='#94fff1', linewidth=2, label=r'$\pi$-helix short domains')
    plot[1,0].fill_between(X_N, ss_pfm_conf_N_short['I_lower'], ss_pfm_conf_N_short['I_upper'], color='#94fff1', alpha=0.3, linewidth=0)
    plot[1,0].plot(X_C, ss_pfm_C_short['I'], color='#94fff1', linewidth=2)
    plot[1,0].fill_between(X_C, ss_pfm_conf_C_short['I_lower'], ss_pfm_conf_C_short['I_upper'], color='#94fff1', alpha=0.3, linewidth=0)

    plot[1,0].plot(X_N, ss_pfm_N_long['I'], color='#ab658c', linewidth=2, label=r'$\pi$-helix long domains')
    plot[1,0].fill_between(X_N, ss_pfm_conf_N_long['I_lower'], ss_pfm_conf_N_long['I_upper'], color='#ab658c', alpha=0.3, linewidth=0)
    plot[1,0].plot(X_C, ss_pfm_C_long['I'], color='#ab658c', linewidth=2)
    plot[1,0].fill_between(X_C, ss_pfm_conf_C_long['I_lower'], ss_pfm_conf_C_long['I_upper'], color='#ab658c', alpha=0.3, linewidth=0)

    plot[1,0].plot(X_N, ss_pfm_N_short['T'], color='#59acff', linewidth=2, label=r'turn short domains')
    plot[1,0].fill_between(X_N, ss_pfm_conf_N_short['T_lower'], ss_pfm_conf_N_short['T_upper'], color='#59acff', alpha=0.3, linewidth=0)
    plot[1,0].plot(X_C, ss_pfm_C_short['T'], color='#59acff', linewidth=2)
    plot[1,0].fill_between(X_C, ss_pfm_conf_C_short['T_lower'], ss_pfm_conf_C_short['T_upper'], color='#59acff', alpha=0.3, linewidth=0)

    plot[1,0].plot(X_N, ss_pfm_N_long['T'], color='#ffec59', linewidth=2, label=r'turn long domains')
    plot[1,0].fill_between(X_N, ss_pfm_conf_N_long['T_lower'], ss_pfm_conf_N_long['T_upper'], color='#ffec59', alpha=0.3, linewidth=0)
    plot[1,0].plot(X_C, ss_pfm_C_long['T'], color='#ffec59', linewidth=2)
    plot[1,0].fill_between(X_C, ss_pfm_conf_C_long['T_lower'], ss_pfm_conf_C_long['T_upper'], color='#ffec59', alpha=0.3, linewidth=0)

    plot[1,0].set_xticks(xticks_ar)
    plot[1,0].set_xticklabels(xticklabels_ar)
    plot[1,0].set_xticks(xticks_spec, minor=True)
    plot[1,0].set_xticklabels(xticklabels_spec, minor=True)
    plot[1,0].tick_params(axis='x', which='minor', length=0)
    plot[1,0].set_xlabel('Distance, aa')
    plot[1,0].set_ylabel('Frequency')
    plot[1,0].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='center')

    plot[1,1].plot(X_N, ss_pfm_N_short['S'], color='#94fff1', linewidth=2, label=r'bend short domains')
    plot[1,1].fill_between(X_N, ss_pfm_conf_N_short['S_lower'], ss_pfm_conf_N_short['S_upper'], color='#94fff1', alpha=0.3, linewidth=0)
    plot[1,1].plot(X_C, ss_pfm_C_short['S'], color='#94fff1', linewidth=2)
    plot[1,1].fill_between(X_C, ss_pfm_conf_C_short['S_lower'], ss_pfm_conf_C_short['S_upper'], color='#94fff1', alpha=0.3, linewidth=0)

    plot[1,1].plot(X_N, ss_pfm_N_long['S'], color='#ab658c', linewidth=2, label=r'bend long domains')
    plot[1,1].fill_between(X_N, ss_pfm_conf_N_long['S_lower'], ss_pfm_conf_N_long['S_upper'], color='#ab658c', alpha=0.3, linewidth=0)
    plot[1,1].plot(X_C, ss_pfm_C_long['S'], color='#ab658c', linewidth=2)
    plot[1,1].fill_between(X_C, ss_pfm_conf_C_long['S_lower'], ss_pfm_conf_C_long['S_upper'], color='#ab658c', alpha=0.3, linewidth=0)

    plot[1,1].plot(X_N, ss_pfm_N_short['-'], color='#59acff', linewidth=2, label=r'unstructured short domains')
    plot[1,1].fill_between(X_N, ss_pfm_conf_N_short['-_lower'], ss_pfm_conf_N_short['-_upper'], color='#59acff', alpha=0.3, linewidth=0)
    plot[1,1].plot(X_C, ss_pfm_C_short['-'], color='#59acff', linewidth=2)
    plot[1,1].fill_between(X_C, ss_pfm_conf_C_short['-_lower'], ss_pfm_conf_C_short['-_upper'], color='#59acff', alpha=0.3, linewidth=0)

    plot[1,1].plot(X_N, ss_pfm_N_long['-'], color='#ffec59', linewidth=2, label=r'unstructured long domains')
    plot[1,1].fill_between(X_N, ss_pfm_conf_N_long['-_lower'], ss_pfm_conf_N_long['-_upper'], color='#ffec59', alpha=0.3, linewidth=0)
    plot[1,1].plot(X_C, ss_pfm_C_long['-'], color='#ffec59', linewidth=2)
    plot[1,1].fill_between(X_C, ss_pfm_conf_C_long['-_lower'], ss_pfm_conf_C_long['-_upper'], color='#ffec59', alpha=0.3, linewidth=0)

    plot[1,1].set_xticks(xticks_ar)
    plot[1,1].set_xticklabels(xticklabels_ar)
    plot[1,1].set_xticks(xticks_spec, minor=True)
    plot[1,1].set_xticklabels(xticklabels_spec, minor=True)
    plot[1,1].tick_params(axis='x', which='minor', length=0)
    plot[1,1].set_xlabel('Distance, aa')
    plot[1,1].set_ylabel('Frequency')
    plot[1,1].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='center')

    plt.tight_layout()
    plt.show()

    return

#######
#Enrichment of ss elements at N-terminus over the elements at C-terminus.
#######

def N_to_C_enrichment(enrichment_N_to_C_dict_short, enrichment_N_to_C_dict_long, window_width):

    #Plot enrichment of ss elements as a function of distance from termini.
    X=range(window_width)
    xticks_ar=list(range(0,51,10))
    xticklabels_ar=[0,10,20,30,40,50]
    xticks_spec=[-4]
    xticklabels_spec=['N-\nC-', ]
    fig, plot=plt.subplots(2,2,figsize=(12,6), dpi=100)
    plot[0,0].plot(X, enrichment_N_to_C_dict_short['H'], color='#94fff1', linewidth=2, label=r'$\alpha$-helix short domains')
    plot[0,0].plot(X, enrichment_N_to_C_dict_long['H'], color='#ab658c', linewidth=2, label=r'$\alpha$-helix long domains')
    plot[0,0].plot(X, enrichment_N_to_C_dict_short['E'], color='#59acff', linewidth=2, label=r'$\beta$-strand short domains')
    plot[0,0].plot(X, enrichment_N_to_C_dict_long['E'], color='#ffec59', linewidth=2, label=r'$\beta$-strand long domains')
    plot[0,0].set_xticks(xticks_ar)
    plot[0,0].set_xticklabels(xticklabels_ar)
    plot[0,0].set_xticks(xticks_spec, minor=True)
    plot[0,0].set_xticklabels(xticklabels_spec, minor=True)
    plot[0,0].tick_params(axis='x', which='minor', length=0)
    plot[0,0].set_xlabel('Distance, aa')
    plot[0,0].set_ylabel('Enrichment p(N)/p(C)')
    plot[0,0].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='upper center')

    plot[0,1].plot(X, enrichment_N_to_C_dict_short['B'], color='#94fff1', linewidth=2, label=r'$\beta$-bridge short domains')
    plot[0,1].plot(X, enrichment_N_to_C_dict_long['B'], color='#ab658c', linewidth=2, label=r'$\beta$-bridge long domains')
    plot[0,1].plot(X, enrichment_N_to_C_dict_short['G'], color='#59acff', linewidth=2, label=r'3-10-helix short domains')
    plot[0,1].plot(X, enrichment_N_to_C_dict_long['G'], color='#ffec59', linewidth=2, label=r'3-10-helix long domains')
    plot[0,1].set_xticks(xticks_ar)
    plot[0,1].set_xticklabels(xticklabels_ar)
    plot[0,1].set_xticks(xticks_spec, minor=True)
    plot[0,1].set_xticklabels(xticklabels_spec, minor=True)
    plot[0,1].tick_params(axis='x', which='minor', length=0)
    plot[0,1].set_xlabel('Distance, aa')
    plot[0,1].set_ylabel('Enrichment p(N)/p(C)')
    plot[0,1].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='upper center')

    plot[1,0].plot(X, enrichment_N_to_C_dict_short['I'], color='#94fff1', linewidth=2, label=r'$\pi$-helix short domains')
    plot[1,0].plot(X, enrichment_N_to_C_dict_long['I'], color='#ab658c', linewidth=2, label=r'$\pi$-helix long domains')
    plot[1,0].plot(X, enrichment_N_to_C_dict_short['T'], color='#59acff', linewidth=2, label=r'turn short domains')
    plot[1,0].plot(X, enrichment_N_to_C_dict_long['T'], color='#ffec59', linewidth=2, label=r'turn long domains')
    plot[1,0].set_xticks(xticks_ar)
    plot[1,0].set_xticklabels(xticklabels_ar)
    plot[1,0].set_xticks(xticks_spec, minor=True)
    plot[1,0].set_xticklabels(xticklabels_spec, minor=True)
    plot[1,0].tick_params(axis='x', which='minor', length=0)
    plot[1,0].set_xlabel('Distance, aa')
    plot[1,0].set_ylabel('Enrichment p(N)/p(C)')
    plot[1,0].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='upper center')

    plot[1,1].plot(X, enrichment_N_to_C_dict_short['S'], color='#94fff1', linewidth=2, label=r'bend short domains')
    plot[1,1].plot(X, enrichment_N_to_C_dict_long['S'], color='#ab658c', linewidth=2, label=r'bend long domains')
    plot[1,1].plot(X, enrichment_N_to_C_dict_short['-'], color='#59acff', linewidth=2, label=r'unstructured short domains')
    plot[1,1].plot(X, enrichment_N_to_C_dict_long['-'], color='#ffec59', linewidth=2, label=r'unstructured long domains')
    plot[1,1].set_xticks(xticks_ar)
    plot[1,1].set_xticklabels(xticklabels_ar)
    plot[1,1].set_xticks(xticks_spec, minor=True)
    plot[1,1].set_xticklabels(xticklabels_spec, minor=True)
    plot[1,1].tick_params(axis='x', which='minor', length=0)
    plot[1,1].set_xlabel('Distance, aa')
    plot[1,1].set_ylabel('Enrichment p(N)/p(C)')
    plot[1,1].legend(fontsize=8.5, ncol=2, handlelength=0.7, frameon=False, columnspacing=0.7, loc='upper center')

    plt.tight_layout()
    plt.show()

    return

#######
#Plot co-occurence of ss elements at termini.
#######

def plot_co_occurence(Obs_over_exp_matrices_short, Obs_over_exp_matrices_long, local_window_width):

    fig, plot=plt.subplots(2,2,figsize=(7,7), dpi=100)
    ticks_ar=list(range(-1, local_window_width, int(local_window_width/10)))
    ticks_ar[0]=0
    print(list(ticks_ar))
    ticklabels_ar=np.array(list(ticks_ar))+1
    plot00=plot[0,0].imshow(Obs_over_exp_matrices_short['EH'], cmap='gnuplot', vmin=0.3, vmax=2.1, interpolation='nearest')
    plot[0,0].set_title(r'$\beta-\alpha$ short domains')
    plo00_cbar=plot[0,0].figure.colorbar(plot00, ax=plot[0,0], shrink=0.7)
    plot[0,0].set_xticks(ticks_ar)
    plot[0,0].set_xticklabels(ticklabels_ar)
    plot[0,0].set_yticks(ticks_ar)
    plot[0,0].set_yticklabels(ticklabels_ar)
    plot[0,0].set_xlabel(r'Distance from C-terminus, aa ($\alpha$)')
    plot[0,0].set_ylabel(r'Distance from N-terminus, aa ($\beta$)')
    plo00_cbar.ax.set_ylabel('p(Obs)/p(Exp)', rotation=-90, va="bottom")

    plot01=plot[0,1].imshow(Obs_over_exp_matrices_long['EH'], cmap='gnuplot', vmin=0.3, vmax=2.1, interpolation='nearest')
    plot[0,1].set_title(r'$\beta-\alpha$-helix long domains')
    plo01_cbar=plot[0,1].figure.colorbar(plot01, ax=plot[0,1], shrink=0.7)
    plot[0,1].set_xticks(ticks_ar)
    plot[0,1].set_xticklabels(ticklabels_ar)
    plot[0,1].set_yticks(ticks_ar)
    plot[0,1].set_yticklabels(ticklabels_ar)
    plot[0,1].set_xlabel(r'Distance from C-terminus, aa ($\alpha$)')
    plot[0,1].set_ylabel(r'Distance from N-terminus, aa ($\beta$)')
    plo01_cbar.ax.set_ylabel('p(Obs)/p(Exp)', rotation=-90, va="bottom")

    plot10=plot[1,0].imshow(Obs_over_exp_matrices_short['HE'], cmap='gnuplot', vmin=0.3, vmax=2.1, interpolation='nearest')
    plot[1,0].set_title(r'$\alpha-\beta$-strand short domains')
    plo10_cbar=plot[0,0].figure.colorbar(plot10, ax=plot[1,0], shrink=0.7)
    plot[1,0].set_xticks(ticks_ar)
    plot[1,0].set_xticklabels(ticklabels_ar)
    plot[1,0].set_yticks(ticks_ar)
    plot[1,0].set_yticklabels(ticklabels_ar)
    plot[1,0].set_xlabel(r'Distance from C-terminus, aa ($\beta$)')
    plot[1,0].set_ylabel(r'Distance from N-terminus, aa ($\alpha$)')
    plo10_cbar.ax.set_ylabel('p(Obs)/p(Exp)', rotation=-90, va="bottom")

    plot11=plot[1,1].imshow(Obs_over_exp_matrices_long['HE'], cmap='gnuplot', vmin=0.3, vmax=2.1, interpolation='nearest')
    plot[1,1].set_title(r'$\alpha-\beta$-strand long domains')
    plo11_cbar=plot[1,1].figure.colorbar(plot11, ax=plot[1,1], shrink=0.7)
    plot[1,1].set_xticks(ticks_ar)
    plot[1,1].set_xticklabels(ticklabels_ar)
    plot[1,1].set_yticks(ticks_ar)
    plot[1,1].set_yticklabels(ticklabels_ar)
    plot[1,1].set_xlabel(r'Distance from C-terminus, aa ($\beta$)')
    plot[1,1].set_ylabel(r'Distance from N-terminus, aa ($\alpha$)')
    plo11_cbar.ax.set_ylabel('p(Obs)/p(Exp)', rotation=-90, va="bottom")

    plt.tight_layout()
    plt.show()
    return

#######
#Wrapper function.
#######

def wrapper(DSSP_inpath):
    #Define domain length trhresholds.
    min_len=50
    thr_len=130

    #Define distance from termini to analyse.
    window_width=50
    local_window_width=50

    #Read DSSP data.
    DSSP_data_dict=read_dssp_data(DSSP_inpath)
    #Classify domains by length.
    Short_structures, Long_structures=define_length_groups(DSSP_data_dict, min_len, thr_len)

    #Get phi, psi angles for N- and C-termini.
    sphi_N, sphi_C, spsi_N, spsi_C, sphi, spsi=phi_psi_N_to_C(Short_structures, window_width)
    lphi_N, lphi_C, lpsi_N, lpsi_C, lphi, lpsi=phi_psi_N_to_C(Long_structures, window_width)

    #Compute position frequency matrices.
    ss_matrix_N_short, ss_matrix_C_short, ss_pfm_N_short, ss_pfm_C_short, ss_pfm_conf_N_short, ss_pfm_conf_C_short=ss_element_frequency_matrix(Short_structures, window_width)
    ss_matrix_N_long, ss_matrix_C_long, ss_pfm_N_long, ss_pfm_C_long, ss_pfm_conf_N_long, ss_pfm_conf_C_long=ss_element_frequency_matrix(Long_structures, window_width)

    #Plot frequency of ss elements as a function of a distance from termini.
    N_to_C_asymmetry(ss_pfm_N_short, ss_pfm_C_short, ss_pfm_conf_N_short, ss_pfm_conf_C_short, ss_pfm_N_long, ss_pfm_C_long, ss_pfm_conf_N_long, ss_pfm_conf_C_long, window_width)

    #Enrichment of ss elements N- over C-terminus.
    enrichment_N_to_C_dict_short=ss_ele_enrichment(ss_pfm_N_short, ss_pfm_C_short)
    enrichment_N_to_C_dict_long=ss_ele_enrichment(ss_pfm_N_long, ss_pfm_C_long)

    #Plot enrichment of frequency of ss elements at N-terminus over C-terminus.
    N_to_C_enrichment(enrichment_N_to_C_dict_short, enrichment_N_to_C_dict_long, window_width)

    #Analyse co-occurence of secondary structure elements at protein termini.
    Obs_over_exp_matrices_short=termini_dependance(ss_matrix_N_short, ss_matrix_C_short, ss_pfm_N_short, ss_pfm_C_short, local_window_width)
    Obs_over_exp_matrices_long=termini_dependance(ss_matrix_N_long, ss_matrix_C_long, ss_pfm_N_long, ss_pfm_C_long, local_window_width)

    #Plot co-occurence of secondary structure elements at protein termini.
    plot_co_occurence(Obs_over_exp_matrices_short, Obs_over_exp_matrices_long, local_window_width)

    return

wrapper(DSSP_data_inpath)


#!/usr/bin/python
# -*- coding: utf-8 -*-
# EFOSC2 polarisation photometry script v1.0
# Created by Adam Higgins
# Email: abh13@le.ac.uk

__doc__ = """ Script runs photometry for the polarisation images from EFOSC2
and outputs flux information for ordinary and extraordinary beams.
For usage please go to https://github.com/abh13/EFOSC2_Scripts.

File names for each half-wave plate angle should be:
0ang.fits,
225ang.fits,
45ang.fits,
675ang.fits
"""

from astropy.io import fits
from astropy.table import vstack
from astropy.stats import sigma_clip
from astropy.stats import sigma_clipped_stats
from astropy.visualization import SqrtStretch
from astropy.visualization import ZScaleInterval
from astropy.visualization.mpl_normalize import ImageNormalize
from photutils import DAOStarFinder
from photutils import CircularAperture
from photutils import aperture_photometry
from photutils import RectangularAnnulus
from photutils.utils import calc_total_error
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os
import sys
import argparse

def get_args():
	""" Parse command line arguments """
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("Directory",metavar="DIR",type=str,action="store",
		help="Required directory")
	parser.add_argument("--ap",type=float,default=2.0,dest='aperture',
		help="Source aperture diameter X*FWHM (default = 2.0)")
	parser.add_argument("--fwhm",type=float,default=0,dest='fwhm',
		help="Manually set the FWHM (default = AUTO)")

	args = parser.parse_args()
	folder_path = args.__dict__['Directory']
	apermul = args.aperture
	fwhm = args.fwhm
	return folder_path,apermul,fwhm

def efosc2_pol_phot(folder_path,apermul,fwhm):
	""" Script runs photometry for EFOSC2 polarimetry images """

	# Read in four wave plate angle files and set up arrays for later
	file_ang0 = os.path.join(folder_path,'0ang.fits')
	file_ang225 = os.path.join(folder_path,'225ang.fits')
	file_ang45 = os.path.join(folder_path,'45ang.fits')
	file_ang675 = os.path.join(folder_path,'675ang.fits')

	files = [file_ang0,file_ang225,file_ang45,file_ang675]
	angle = ['0','225','45','675']
	ang_dec = ['0','22.5','45','67.5']
	label = ['$0^{\circ}$ image','$22.5^{\circ}$ image',
		'$45^{\circ}$ image','$67.5^{\circ}$ image']

	# Set up array to store the number of sources per half-wave plate
	numsource = []

	# Loop over files for the four wave plate files
	for k in range(0,len(angle),1):

		# Open fits file, extract pixel flux data and remove saturated pixels
		try:
			hdulist = fits.open(files[k])
			image_data = hdulist[0].data

		except FileNotFoundError:
			raise FileNotFoundError("Cannot find the input fits file(s).")

		# Remove bad pixels and mask edges
		image_data[image_data > 60000] = 0
		image_data[image_data < 0] = 0
		rows = len(image_data[:,0])
		cols = len(image_data[0,:])
		hdulist.close()

		# Calculate estimate of background using sigma-clipping and calculate
		# number of pixels used in the background region that were not
		# clipped! This is done in a small area near the optical axis.
		go_bmean, go_bmedian, go_bstd = sigma_clipped_stats(image_data
			[510:568,520:580],sigma=3.0,maxiters=5)
		ge_bmean, ge_bmedian, ge_bstd = sigma_clipped_stats(image_data
			[446:504,520:580],sigma=3.0,maxiters=5)
		mask_o = sigma_clip(image_data[510:568,520:580],sigma=3.0,maxiters=5,
			masked=True)
		mask_e = sigma_clip(image_data[446:504,520:580],sigma=3.0,maxiters=5,
			masked=True)
		ann_area_o = np.ma.MaskedArray.count(mask_o)
		ann_area_e = np.ma.MaskedArray.count(mask_e)

		# Detect sources using DAO star finder
		daofind_o = DAOStarFinder(fwhm=5,threshold=5*go_bstd,
			exclude_border=True)
		daofind_e = DAOStarFinder(fwhm=5,threshold=5*ge_bstd,
			exclude_border=True)
		sources_o = daofind_o(image_data[522:552,535:565])
		sources_e = daofind_e(image_data[462:492,535:565])

		if (sources_o is None or sources_e is None):
			raise ValueError("No source detected in image")

		if len(sources_o) != len(sources_e):
			raise ValueError("Unequal number of sources in o and e images!")

		glob_bgm = [go_bmean,ge_bmean]
		glob_bgerr = [go_bstd,ge_bstd]

		# Convert the source centroids back into detector pixels
		sources_o['xcentroid'] = sources_o['xcentroid'] + 535
		sources_o['ycentroid'] = sources_o['ycentroid'] + 522
		sources_e['xcentroid'] = sources_e['xcentroid'] + 535
		sources_e['ycentroid'] = sources_e['ycentroid'] + 462

		# Estimate the FWHM of the source by simulating a 2D Gaussian
		# This is only done on the 0 angle image ensuring aperture sizes
		# are equal for all half-wave plate angles. If a user specified
		# FWHM is given, then the estimation is not used.
		if fwhm == 0.0:
			xpeaks_o = []
			xpeaks_e = []
			ypeaks_o = []
			ypeaks_e = []
			fwhm = []

			for i in range(0,len(sources_o),1):
				data_o = image_data[525:550,535:565]
				xpeaks_o.append(int(sources_o[i]['xcentroid']) - 535)
				ypeaks_o.append(int(sources_o[i]['ycentroid']) - 525)

				data_e = image_data[465:490,535:560]
				xpeaks_e.append(int(sources_e[i]['xcentroid']) - 535)
				ypeaks_e.append(int(sources_e[i]['ycentroid']) - 465)

				min_count_o = np.min(data_o)
				min_count_e = np.min(data_e)
				max_count_o = data_o[ypeaks_o[i],xpeaks_e[i]]
				max_count_e = data_e[ypeaks_o[i],xpeaks_e[i]]
				half_max_o = (max_count_o + min_count_o)/2
				half_max_e = (max_count_e + min_count_e)/2

				# Crude calculation for each source
				nearest_above_x_o = ((np.abs(data_o[ypeaks_o[i],
					xpeaks_o[i]:-1] - half_max_o)).argmin())
				nearest_below_x_o = ((np.abs(data_o[ypeaks_o[i],0:
					xpeaks_o[i]] - half_max_o)).argmin())
				nearest_above_x_e = ((np.abs(data_e[ypeaks_e[i],
					xpeaks_e[i]:-1] - half_max_e)).argmin())
				nearest_below_x_e = ((np.abs(data_e[ypeaks_e[i],0:
					xpeaks_e[i]] - half_max_e)).argmin())
				nearest_above_y_o = ((np.abs(data_o[ypeaks_o[i]:-1,
					xpeaks_o[i]] - half_max_o)).argmin())
				nearest_below_y_o = ((np.abs(data_o[0:ypeaks_o[i],
					xpeaks_o[i]] - half_max_o)).argmin())
				nearest_above_y_e = ((np.abs(data_e[ypeaks_e[i]:-1,
					xpeaks_e[i]] - half_max_e)).argmin())
				nearest_below_y_e = ((np.abs(data_e[0:ypeaks_e[i],
					xpeaks_e[i]] - half_max_e)).argmin())
				fwhm.append((nearest_above_x_o + (xpeaks_o[i] -
					nearest_below_x_o)))
				fwhm.append((nearest_above_y_o + (ypeaks_o[i] -
					nearest_below_y_o)))
				fwhm.append((nearest_above_x_e + (xpeaks_e[i] -
					nearest_below_x_e)))
				fwhm.append((nearest_above_y_e + (ypeaks_e[i] -
					nearest_below_y_e)))

			fwhm = np.mean(fwhm)

		# Stack both ord and exord sources together
		tot_sources = vstack([sources_o,sources_e])

		# Store the ordinary and extraordinary beam source images and
		# create apertures for aperture photometry
		positions = np.swapaxes(np.array((tot_sources['xcentroid'],
			tot_sources['ycentroid']),dtype='float'),0,1)
		aperture = CircularAperture(positions, r=0.5*apermul*fwhm)
		phot_table = aperture_photometry(image_data,aperture)

		# Set up arrays of ord and exord source parameters
		s_id = np.zeros([len(np.array(phot_table['id']))])
		xp = np.zeros([len(s_id)])
		yp = np.zeros([len(s_id)])
		fluxbgs = np.zeros([len(s_id)])
		mean_bg = np.zeros([len(s_id)])
		bg_err = np.zeros([len(s_id)])
		s_area = []

		for i in range(0,len(np.array(phot_table['id'])),1):
			s_id[i] = np.array(phot_table['id'][i])
			xpos = np.array(phot_table['xcenter'][i])
			ypos = np.array(phot_table['ycenter'][i])
			xp[i] = xpos
			yp[i] = ypos
			s_area.append(np.pi*(0.5*apermul*fwhm)**2)
			j = i%2
			fluxbgs[i] = (phot_table['aperture_sum'][i] -
				aperture.area*glob_bgm[j])
			mean_bg[i] = glob_bgm[j]
			bg_err[i] = glob_bgerr[j]

		# Create and save the image in z scale and overplot the ordinary and
		# extraordinary apertures and local background annuli if applicable
		fig = plt.figure()
		zscale = ZScaleInterval(image_data)
		norm = ImageNormalize(stretch=SqrtStretch(),interval=zscale)
		image = plt.imshow(image_data,cmap='gray',origin='lower',norm=norm)
		bg_annulus_o = RectangularAnnulus((550,539),w_in=0.1,w_out=60,h_out=58,
			theta=0)
		bg_annulus_e = RectangularAnnulus((550,475),w_in=0.1,w_out=60,h_out=58,
			theta=0)
		bg_annulus_o.plot(color='skyblue',lw=1.5,alpha=0.5)
		bg_annulus_e.plot(color='lightgreen',lw=1.5,alpha=0.5)

		for i in range(0,len(np.array(phot_table['id'])),1):
			aperture = CircularAperture((xp[i],yp[i]),r=0.5*apermul*fwhm)

			if i < int(len(np.array(phot_table['id']))/2):
				aperture.plot(color='blue',lw=1.5,alpha=0.5)

			else:
				aperture.plot(color='green',lw=1.5,alpha=0.5)

		plt.xlim(500,600)
		plt.ylim(425,575)
		plt.title(label[k])
		image_fn = folder_path + angle[k] + '_image.png'
		fig.savefig(image_fn)

		# Create dataframes for photometry results
		cols = ['xpix','ypix','fluxbgs','sourcearea','meanbg','bgerr',
			'bgarea']
		df_o = pd.DataFrame(columns=cols)
		df_e = pd.DataFrame(columns=cols)

		for i in range(0,len(np.array(phot_table['id'])),1):
			if 0 <= i < int(len(np.array(phot_table['id']))/2):
				df_o = df_o.append({cols[0]:xp[i],cols[1]:yp[i],
					cols[2]:fluxbgs[i],cols[3]:s_area[i],cols[4]:mean_bg[i],
					cols[5]:bg_err[i],cols[6]:ann_area_o},ignore_index=True)

			else:
				df_e = df_e.append({cols[0]:xp[i],cols[1]:yp[i],
					cols[2]:fluxbgs[i],cols[3]:s_area[i],cols[4]:mean_bg[i],
					cols[5]:bg_err[i],cols[6]:ann_area_e},ignore_index=True)

		# Save dataframes to text files
		df_o.to_string(folder_path+'angle'+angle[k]+'_ord.txt',
			index=False,justify='left')

		df_e.to_string(folder_path+'angle'+angle[k]+'_exord.txt',
			index=False,justify='left')

		# Save the number of sources in each beam to a list
		numsource.append(int(len(np.array(phot_table['id']))/2))

	# Print number of sources per half-wave plate image and FWHM
	print("FWHM =",fwhm,"pixels")
	for i in range(0,len(numsource),1):
		print("No of sources detected at",ang_dec[i],"degrees:",numsource[i])

	return 0

def main():
	""" Run script from command line """
	folder_path,apermul,fwhm = get_args()
	return efosc2_pol_phot(folder_path,apermul,fwhm)

if __name__ == '__main__':
    sys.exit(main())


import csv, json, argparse, copy, re, os, urllib2

import numpy as np
from scipy.spatial import distance
from sklearn import manifold, metrics, decomposition, preprocessing

import igraph

import jsmin

import rdkit
from rdkit import Chem, DataStructs, Geometry
from rdkit.DataStructs import cDataStructs
from rdkit.Chem import Draw, AllChem, Scaffolds, Lipinski, Crippen, rdMolDescriptors, TemplateAlign
from rdkit.Chem.Scaffolds import MurckoScaffold

PROPS_ORDER = ["mw", "hba", "hbd", "rb", "rc", "arc", "logp", "tpsa"]

PROP2FNC = {
    "mw": rdMolDescriptors.CalcExactMolWt,
    "hba": Lipinski.NumHAcceptors,
    "hbd": Lipinski.NumHDonors,
    "rb": Lipinski.NumRotatableBonds,
    "rc": Lipinski.RingCount,
    "arc": Lipinski.NumAromaticRings,
    "logp": Crippen.MolLogP,
    "tpsa": rdMolDescriptors.CalcTPSA,
}

PROP2LABEL = {
    "mw": "Molecular weight",
    "hba": "H-bond acceptors",
    "hbd": "H-bond donors",
    "rb": "Rotatable bonds",
    "rc": "Rings",
    "arc": "Aromatic rings",
    "logp": "cLogP",
    "tpsa": "TPSA"
}

FP2FNC = {
    "ecfp4": lambda rdmol: AllChem.GetMorganFingerprintAsBitVect(rdmol, radius=2, nBits=1024),
    "ecfp6": lambda rdmol: AllChem.GetMorganFingerprintAsBitVect(rdmol, radius=3, nBits=1024),
    "apfp": lambda rdmol: AllChem.GetHashedAtomPairFingerprintAsBitVect(rdmol, nBits=1024),
    "ttfp": lambda rdmol: AllChem.GetHashedTopologicalTorsionFingerprintAsBitVect(rdmol, nBits=1024),
    "maccs": lambda rdmol: AllChem.GetMACCSKeysFingerprint(rdmol),
}

AVAILABLE_METRICS = ["Tanimoto", "Dice", "Cosine", "Sokal", "Russel", "RogotGoldberg", "AllBit", "Kulczynski", "McConnaughey", "Asymmetric", "BraunBlanquet"]

class ChemSpace():

    def __init__(self):
        self.category_field = False
        self.category_field_delimiter = False
        self.label_field = False
        self.compound_structure_field = False
        self.sdf = False
        self.write_structures = True
        self.fp = "ecfp4"
        self.fingerprint_field = False
        self.metric = "Tanimoto"

        if self.metric not in AVAILABLE_METRICS:
            raise Exception("Metric '{}' not found in available similarity metrics: {}".format(self.metric, AVAILABLE_METRICS))

        self.index2rdmol = {}
        self.index2fpobj = {}

    def read_csv(self, filename, delimiter=",", header=False, missing_value=False, remove_columns=False):
        """Reads data from the CSV file"""
        print("Reading file: {}".format(filename))

        self.filename = filename
        with open(self.filename, "r") as input_file:
            reader = csv.reader(input_file, delimiter=delimiter)
            rows = [row for row in reader]

        self.read_data(rows, header, missing_value, remove_columns)

    def read_sdf(self, filename):
        """Reads data from a sdf file"""
        print("Reading file: {}".format(filename))
        self.sdf = True
        self.header = False
        self.data = []
        self.filename = filename
        self.index2rdmol = {}
        self.index2fpobj = {}
        self.index2props = {}
        self.index2category = {}
        self.index2label = {}
        self.index2id = {}

        molsupplier = Chem.SDMolSupplier(str(filename))
        not_parsed = []

        for index, m in enumerate(molsupplier):
            try:
                Chem.SanitizeMol(m)
                self.index2rdmol[index] = m
                self.index2fpobj[index] = FP2FNC[self.fp](m)
                self.index2props[index] = m.GetPropsAsDict()
                self.index2id[index] = index

            except Exception, e:
                print(e)
                not_parsed.append(index)

        self.index_order = self.index2rdmol.keys()
        self.index_order.sort()
        self.index2row = {i: [] for i in self.index_order}
        self.data = self.index2row.values()

        if self.label_field is not False and self.label_field in self.index2props[self.index_order[0]]:
            self.index2label = {i: self.index2props[i].get(self.label_field) for i in self.index_order}

        if self.category_field is not False and self.category_field in self.index2props[self.index_order[0]]:
            self.index2category = {i: self.index2props[i].get(self.category_field) for i in self.index_order}

        self.__create_chemspace_format__()

    def add_compounds_from_file(self, filename, delimiter=","):
        print("Reading compounds: {}".format(filename))
        self.filename = filename

        with open(self.filename, "r") as input_file:
            reader = csv.reader(input_file, delimiter=delimiter)
            rows = [row for row in reader]

        self.add_compounds(rows)

    def add_category(self, category):
        if not "categories" in self.chemical_space:
            self.chemical_space["categories"] = []

        self.chemical_space["categories"].append(category)

    def add_compounds(self, rows):
        """Reads data in a form of list of lists (tuples)"""
        self.compounds = {r[0]: r[1] for r in rows}
        self.chemical_space["compounds"] = {}
        self.__parse_compounds__()

        for key in self.chemical_space["points"]:
            if key in self.id2rdmol:
                self.chemical_space["compounds"][key] = {"structure": self.__get_compound__(key)}

    def read_data(self, rows, header=False, missing_value=False, remove_columns=False):
        """Reads data in a form of list of lists (tuples)"""
        self.header = header
        self.missing_value = missing_value
        data_start = 0

        self.data = rows
        self.index2id = {}
        self.index2row = {}
        self.index2compound = {}
        self.index2label = {}
        self.index2category = {}

        if self.header:
            self.header = self.data[0]
            self.data = self.data[1:]

        if self.header:
            if remove_columns is not False and len(remove_columns) > 0:
                for col in remove_columns:
                    self.__remove_field__(col)

            if self.compound_structure_field and self.compound_structure_field in self.header:
                self.index2compound = self.__extract_field__(self.compound_structure_field)
                self.__read_compounds__()

            if self.label_field and self.label_field in self.header:
                self.index2label = self.__extract_field__(self.label_field)

            if self.category_field and self.category_field in self.header:
                self.index2category = self.__extract_field__(self.category_field)

            if self.fingerprint_field and self.fingerprint_field in self.header:
                self.index2fp = self.__extract_field__(self.fingerprint_field)
                self.index2fpobj = {}

                for index, fp in self.index2fp.items():
                    self.index2fpobj[index] = self.__get_bitvect_for_fp__(fp)

            # remove ID field
            self.header.pop(0)

        self.index2id = {i: row[0] for i, row in enumerate(self.data)}
        self.index2row = {i: [round(float(v), 2) if v not in ["", None, "None", self.missing_value] else None for v in row[1:]] for i, row in enumerate(self.data)}
        self.index_order = [i for i, row in enumerate(self.data)]
        self.data = [self.index2row[i] for i in self.index_order]

        if self.missing_value is not False:
            self.data, self.missing_values_indexes = self.__impute_missing_values__(self.data)
            # self.original_data = self.__return_missing_values__(copy.deepcopy(self.data), self.missing_values_indexes)

        # self.original_data = copy.deepcopy(self.index2row)
        self.__create_chemspace_format__()

    def __read_compounds__(self):
        for i, smi in self.index2compound.items():
            try:
                self.index2rdmol[i] = Chem.MolFromSmiles(smi)
                self.index2fpobj[i] = FP2FNC[self.fp](self.index2rdmol[i])

            except Exception, e:
                print(e)
                self.index2rdmol[i] = None
                self.index2fpobj[i] = None

    def __remove_field__(self, field):
        if field in self.header:
            index = self.header.index(field)

            if index is not False:
                self.header.pop(index)

                for i, row in enumerate(self.data):
                    self.data[i].pop(index)

    def __extract_field__(self, field):
        index2value = {}

        if field in self.header:
            index = self.header.index(field)

            if index is not False:
                self.header.pop(index)

                for i, row in enumerate(self.data):
                    index2value[i] = row[index]
                    self.data[i].pop(index)

        return index2value

    def __impute_missing_values__(self, data):
        datatype2impute = {"numeric": {"strategy":"mean",
                                        "value": lambda x: round(float(value), 3)},
                           "binary": {"strategy":"most_frequent",
                                      "value": lambda x: int(value)}
                           }

        missing_values_indexes = []

        for i, row in enumerate(self.data):
            missing_values_indexes.append([j for j, v in enumerate(row) if v == self.missing_value])

            for j, value in enumerate(row):
                if value == self.missing_value:
                    data[i][j] = np.nan
        imputer = preprocessing.Imputer(missing_values="NaN", strategy=datatype2impute["numeric"]["strategy"])
        #error when using median strategy - minus one dimension in imputed data... omg
        imputed_data = [list(row) for row in imputer.fit_transform(self.data)]
        imputed_data = [[datatype2impute["numeric"]["value"](value) for value in row] for row in imputed_data]
        return imputed_data, missing_values_indexes

    def __return_missing_values__(self, data, missing_values_indexes):
        for i, indexes in enumerate(missing_values_indexes):
            if indexes:
                for index in indexes:
                    data[i][index] = None
        return data

    def __create_chemspace_format__(self):
        self.chemical_space = {"points": {}}

        for index in self.index_order:
            self.chemical_space["points"][index] = {"object_ids": [self.index2id[index]]}

        if len(self.index2category):
            self.__parse_categories__()

        if len(self.index2label):
            for index, label in self.index2label.items():
                self.chemical_space["points"][index]["label"] = label

        for index, row in self.index2row.items():
            self.chemical_space["points"][index]["features"] = copy.copy(row)

        if self.header:
            current_header = self.chemical_space.get("feature_names", [])
            current_header.extend(self.header)
            self.chemical_space["feature_names"] = current_header

        if len(self.index2rdmol) and self.write_structures:
            self.chemical_space["compounds"] = {}

            for index, rdmol in self.index2rdmol.items():
                # self.chemical_space["compounds"][index] = {"structure": self.__get_compound__(rdmol), "smiles": Chem.MolToSmiles(rdmol, True)}

                self.chemical_space["compounds"][index] = {"smiles": Chem.MolToSmiles(rdmol, True)}

    def __parse_categories__(self):
        category2ids = {}

        for index, category in self.index2category.items():
            categories = [category] if self.category_field_delimiter is False else [c.strip() for c in category.split(self.category_field_delimiter)]

            for c in categories:
                if c in category2ids:
                    category2ids[c].add(index)
                else:
                    category2ids[c] = {index}

        if not "categories" in self.chemical_space:
            self.chemical_space["categories"] = []

        for c, ids in category2ids.items():
            self.chemical_space["categories"].append({"label": c, "points": list(ids)})

    def add_paths(self, paths):
        if not self.chemical_space.get("paths", False):
            self.chemical_space["paths"] = []

        self.chemical_space["paths"].extend(paths)

    def add_paths_from_file(self):
        pass

    def add_physico_chemical_properties(self):
        print("Calculating physico-chemical properties: {} compounds".format(len(self.index2rdmol)))
        self.pcp = True
        if len(self.index2rdmol):
            count = len(self.index2rdmol)
            i = 0

            id2pcp = {}
            for index, rdmol in self.index2rdmol.items():
                if i%100 == 0 or i == count:
                    print("{}/{}".format(i, count))

                id2pcp[index] = self.__get_pcp_for_rdmol__(rdmol)
                i+=1

            empty = [None for x in PROP2LABEL]
            for i, index in enumerate(self.index_order):

                if id2pcp.get(index, False):
                    pcps = id2pcp[index]
                else:
                    pcps = empty

                self.chemical_space["points"][index]["features"].extend(pcps)
                self.data[i].extend(pcps)

            current_header = self.chemical_space.get("feature_names", [])
            current_header.extend([PROP2LABEL[prop] for prop in PROPS_ORDER])
            self.chemical_space["feature_names"] = current_header
            self.original_data = copy.deepcopy(self.data)

    def __get_pcp_for_rdmol__(self, rdmol):
        return [round(PROP2FNC[prop](rdmol), 2) for prop in PROPS_ORDER]

    def __get_compound__(self, rdmol):
        if rdmol is not None:
            Chem.Kekulize(rdmol)
            AllChem.Compute2DCoords(rdmol)
            compound = {"atoms": {}}
            atoms = [a for a in rdmol.GetAtoms()]
            bond_types = []
            for i, a in enumerate(atoms, 1):
                number = a.GetIdx()
                position = rdmol.GetConformer().GetAtomPosition(number)

                compound["atoms"][number] = {
                    "bonds": {b.GetEndAtomIdx():b.GetBondTypeAsDouble() for b in a.GetBonds() if b.GetEndAtomIdx() != number},
                    "symbol": a.GetSymbol(),
                    "charge": a.GetFormalCharge(),
                    "coordinates": [round(position.x, 3), round(position.y, 3)]
                }

                bond_types.extend(compound["atoms"][number]["bonds"].values())
        else:
            compound = None

        return compound

    def normalize_data(self, feature_range=(0,1)):
        """Normalizes data to a scale from 0 to 1."""
        print("Data normalization (scale): {}".format(feature_range))

        min_max_scaler = preprocessing.MinMaxScaler(feature_range)
        self.data = min_max_scaler.fit_transform(self.data)
        self.data = [[round(v, 3) for v in row] for row in self.data]

    def __calculate_distance_matrix__(self, similarity_threshold):
        print("\nCalculating distance matrix: {} compounds".format(len(self.index2fpobj)))

        self.dist_matrix = {x:[] for x in self.index_order}
        self.edges = []
        self.index2edges = {}

        fps_count = len(self.index_order)

        for i, index_1 in enumerate(self.index_order):
            self.index2edges[index_1] = []

            if i%100 == 0 or i == fps_count:
                print("{}/{}".format(i, fps_count))

            for j, index_2 in enumerate(self.index_order[i:], i):
                sim = DataStructs.FingerprintSimilarity(self.index2fpobj[index_1], self.index2fpobj[index_2], metric=getattr(DataStructs, "{}Similarity".format(self.metric)))
                self.dist_matrix[index_1].append(1-sim)

                if index_1 != index_2:
                    self.dist_matrix[index_2].append(1-sim)

                    if sim >= similarity_threshold:
                        self.edges.append((index_1, index_2))
                        self.index2edges[index_1].append([index_2])

    def __get_edges__(self, similarity_threshold=0.7, k=2):
        print("\nCalculating edges [similarity threshold={}]: {} compounds".format(similarity_threshold, len(self.index2fpobj)))
        self.edges = []
        self.index2edges = {}
        count = len(self.index_order)

        for i, index in enumerate(self.index_order):
            if (i+1)%100 == 0:
                print("{}/{}".format(i, count))

            values = [[idx, v] for idx, v in zip(self.index_order, self.dist_matrix[index]) if idx != index]
            values.sort(key=lambda x: x[1])

            if 1-values[1][1] >= similarity_threshold:
                self.index2edges[index] = []

                for v in values:
                    if 1-v[1] >= similarity_threshold:
                        self.edges.append((index, v[0]))
                        self.index2edges[index].append([v[0]])

                        if len(self.index2edges[index]) == k:
                            break
                    else:
                        break

        print("EDGES: {}".format(len(self.edges)))

    def __convert_fps_to_bitvects__(self, fps):
        converted = []

        for fp in fps:
            row = [fp[0], self.__get_bitvect_for_fp__(fp[1:])]
            converted.append(row)

        return converted

    def __get_bitvect_for_fp__(self, fp):
        if type(fp) is list and len(fp) == 1:
            fp = fp[0]
        bitvect = cDataStructs.ExplicitBitVect(len(fp))
        on_indexes = [i for i, b in enumerate(fp) if int(b) == 1]
        bitvect.SetBitsFromList(on_indexes)
        return bitvect

    def arrange(self, by="fps", fps=[], method="pca", similarity_threshold=0.7, add_edges=False, k=None):
        self.index2edges = False
        self.edges = False
        self.dist_matrix = False
        bitvects = False

        if type(method) is not list:
            methods = [method]
        else:
            methods = method

        for method in methods:
            if by == "scaffolds":
                self.__arrange_by_scaffolds__()

                g = igraph.Graph(len(self.index_order))
                print("\nCalculating Chemical Space Network...")
                feature_names = ["CSN1", "CSN2"]

                g.add_edges(self.edges)
                layout = g.layout_fruchterman_reingold()
                coords = layout.coords

            elif by == "dm" or method == "sas":
                if len(fps) == 0:
                    for index in self.index_order:
                        fps.append(self.index2fpobj[index])
                    bitvects = True

                elif type(fps[0][1]) in [unicode, int] and not bitvects:
                    fps = self.__convert_fps_to_bitvects__(fps)
                    bitvects = True

                if not self.dist_matrix:
                    self.__calculate_distance_matrix__(similarity_threshold)

                dist_matrix = np.matrix([np.array(self.dist_matrix[index]) for index in self.index_order])
                g = igraph.Graph(len(self.index_order))

                if method == "csn":
                    print("\nCalculating Chemical Space Network...")
                    feature_names = ["CSN1", "CSN2"]
                    if k is not None:
                        # k = len(self.index_order)
                        self.__get_edges__(similarity_threshold=similarity_threshold, k=k)

                    print("Fruchterman-Reingold Layout calculation...")
                    g.add_edges(self.edges)
                    layout = g.layout_fruchterman_reingold()
                    coords = layout.coords

                elif method == "mds":
                    print("\nCalculating MDS...")
                    feature_names = ["MDS1", "MDS2"]
                    # sklearn implementation
                    mds = manifold.MDS(n_components=2, dissimilarity='precomputed')
                    coords = mds.fit_transform(dist_matrix)

                    # igraph implementation
                    # layout = g.layout_mds(dist_matrix, 2, arpack_options=igraph.ARPACKOptions(iter=1000))
                    # coords = layout.coords

                elif method == "pca":
                    print("\nCalculating PCA...")
                    feature_names = ["PC1", "PC2"]
                    pca = decomposition.PCA(n_components=2)
                    coords = pca.fit_transform(dist_matrix)

                elif method == "fa":
                    print("\nCalculating Factor Analysis...")
                    feature_names = ["FA1", "FA2"]
                    fa = decomposition.FactorAnalysis(n_components=2)
                    coords = fa.fit_transform(dist_matrix)

                elif method == "isomap":
                    print("\nCalculating Isomap...")
                    feature_names = ["Isomap1", "Isomap2"]
                    isomap = manifold.Isomap(n_neighbors=200, n_components=2)
                    coords = isomap.fit_transform(dist_matrix)

                elif method == "tsne":
                    print("\nCalculating t-SNE...")
                    feature_names = ["t-SNE1", "t-SNE2"]
                    tsne = manifold.TSNE(n_components=2, metric='precomputed')
                    coords = tsne.fit_transform(dist_matrix)

                elif method == "sas":
                    print("\nCalculating SAS...")
                    feature_names = ["Similarity", "Activity difference"]
                    self.chemical_space = {"points": {}, "feature_names": ["SALI"]}
                    ai = self.header.index(self.activity_field)
                    ids = []
                    coords = []

                    for i, index_1 in enumerate(self.index_order[:-1]):
                        for j, index_2 in enumerate(self.index_order[i:], i):
                            if i != j:
                                activity_diff = round(abs(float(self.data[i][ai]) - float(self.data[j][ai])), 2)
                                distance = self.dist_matrix[i][j]
                                distance = distance if distance > 0 else 0.01

                                sali = round(activity_diff/distance, 2)
                                coord = [round(1 - self.dist_matrix[i][j], 2), activity_diff]
                                self.chemical_space["points"]["{}_{}".format(index_1, index_2)] = {"features": [sali]}
                                ids.append("{}_{}".format(index_1, index_2))
                                coords.append(coord)

                    self.index_order = ids

            elif by in ["data", "fps"]:
                if by == "fps":
                    if len(fps) == 0 and len(self.index2fpobj):
                        for index in self.index_order:
                            fps.append(self.index2fpobj[index])

                    data = [[int(b) for b in fp] for fp in fps]
                else:
                    data = self.data

                if method == "pca":
                    print("\nCalculating PCA...")
                    feature_names = ["PC1", "PC2"]
                    pca = decomposition.PCA(n_components=2)
                    coords = pca.fit_transform(data)

                elif method == "fa":
                    print("\nCalculating Factor Analysis...")
                    feature_names = ["FA1", "FA2"]
                    fa = decomposition.FactorAnalysis(n_components=2)
                    coords = fa.fit_transform(data)

            if method in ["csn", "nn"] or add_edges or by == "scaffolds":
                if self.dist_matrix is False and self.edges is False:
                    self.__calculate_distance_matrix__()

                if self.edges is False:
                    if k is None:
                        k = len(self.index_order)
                    self.__get_edges__(similarity_threshold=similarity_threshold, k=k)

                if by != "scaffolds":
                    for cid, es in self.index2edges.items():
                        if not self.chemical_space["points"][cid].get("links", False):
                            self.chemical_space["points"][cid]["links"] = []

                        for e in es:
                            self.chemical_space["points"][cid]["links"].extend(e)

            index2coords = {index:coords[i] for i, index in enumerate(self.index_order)}

            for index, values in self.chemical_space["points"].items():
                if index in index2coords:
                    point_features = self.chemical_space["points"][index]["features"]
                    features = [round(index2coords[index][0], 3), round(index2coords[index][1], 3)]
                    features.extend(point_features)
                    self.chemical_space["points"][index]["features"] = features

                else:
                    self.chemical_space["points"].pop(index, None)

            feature_names.extend(self.chemical_space.get("feature_names", []))
            self.chemical_space["feature_names"] = feature_names

    def __arrange_by_scaffolds__(self, align_by_scaffold=True):
        self.scaffold2indexes = {}
        self.scaffold2rdmol = {}
        self.index2scaffold = {}

        for index, rdmol in self.index2rdmol.items():
            AllChem.Compute2DCoords(rdmol)

            if align_by_scaffold:
                scaffold = Scaffolds.MurckoScaffold.GetScaffoldForMol(rdmol)
                AllChem.Compute2DCoords(scaffold)
                scaffold_smiles = Chem.MolToSmiles(scaffold)
                self.scaffold2rdmol[scaffold_smiles] = scaffold

                matched = rdmol.GetSubstructMatch(scaffold)
                coords = [rdmol.GetConformer().GetAtomPosition(x) for x in matched]
                coords2D = [Geometry.Point2D(pt.x,pt.y) for pt in coords]

                coordDict = {}
                for i,coord in enumerate(coords2D):
                    coordDict[matched[i]] = coord

                AllChem.Compute2DCoords(rdmol, coordMap=coordDict)

                if scaffold_smiles in self.scaffold2indexes:
                    self.scaffold2indexes[scaffold_smiles].append(index)
                else:
                    self.scaffold2indexes[scaffold_smiles] = [index]

        self.scaffold2indexes = {scaffold: indexes for scaffold, indexes in self.scaffold2indexes.items() if len(indexes) > 1}

        for index, scaffold in enumerate(self.scaffold2indexes.keys(), len(self.index_order)):
                self.index2scaffold[index] = scaffold
                self.index2rdmol[index] = self.scaffold2rdmol[scaffold]
                self.index_order.append(index)

        self.edges = []
        self.index2edges = {}

        for index_1, scaffold in self.index2scaffold.items():
            self.index2edges[index_1] = []

            for index_2 in self.scaffold2indexes[scaffold]:
                self.edges.append((index_1, index_2))
                self.index2edges[index_1].append(index_2)

        self.__add_scaffolds_to_chemical_space__()

    def __add_scaffolds_to_chemical_space__(self):
        for index, scaffold in self.index2scaffold.items():
            self.chemical_space["points"][index] = {
                # "features": self.__get_pcp_for_rdmol__(self.scaffold2rdmol[scaffold]),
                "object_ids": self.scaffold2indexes[scaffold],
                "links": self.scaffold2indexes[scaffold],
            }
            if self.pcp:
                self.chemical_space["points"][index]["features"] = [None for i in range(len(self.chemical_space["feature_names"]) - len(PROPS_ORDER))]
                self.chemical_space["points"][index]["features"].extend(self.__get_pcp_for_rdmol__(self.scaffold2rdmol[scaffold]))
            self.chemical_space["compounds"][index] = {"smiles": scaffold, "color": "red"}

    def get_chemspace_compound_from_smiles(self, smi):
        mol_obj = Chem.MolFromSmiles(smi)
        if mol_obj is not None:
            AllChem.Compute2DCoords(mol_obj)
            mol = self.__get_compound__(mol=mol_obj)
        else:
            mol = False
        return mol

    def export_chemical_space_as_html(self, htmldir=".", ):
        """Export a simple HTML page with embedded chemical space and dependencies into a given directory."""
        if not os.path.exists(htmldir):
            os.makedirs(htmldir)

        chemspace_json = self.export_chemical_space_as_json(minify=True, dump=True)

        libs = [
            ("chemspace-0.2.0.min.js", "https://openscreen.cz/software/chemspace/static/js/chemspace-0.2.0.min.js"),
            ("jquery-3.3.1.min.js", "https://code.jquery.com/jquery-3.3.1.min.js"),
            ("konva.min.js", "https://cdn.rawgit.com/konvajs/konva/1.7.6/konva.min.js")
        ]

        js_html = []
        for l in libs:
            js_html.append("<script src='{}'></script>".format(l[0]))

        settings = {
            "target": "chemspace"
        }

        template = """<html>
        <head>
            {}
            <script>
                $(document).ready(function() {{
                    var data = {};
                    var chemspace = new ChemSpace({});
                    chemspace.read_data(data);
                    chemspace.draw();
                }});
            </script>
        </head>

        <body>
            <div id="chemspace"></div>
        </body>
        </html>""".format('\n'.join(js_html), chemspace_json, json.dumps(settings))

        for l in libs:
            lib, url = l
            try:
                source = urllib2.urlopen(url)
                source_html = source.read()

                with open(os.path.join(htmldir, lib), "w") as output:
                    output.write(source_html)
            except urllib2.URLError, e:
                raise Exception("""
                        \nCan't download file {}.\nPlease check your internet connection and try again.\nIf the error persists there can be something wrong with the InCHlib server.\n""".format(url)
                    )

        with open(os.path.join(htmldir, "chemspace.html"), "w") as output:
            output.write(template)

    def export_chemical_space_as_json(self, filename=None, minify=False, dump=True):
        """Returns space in a JSON format or exports it to the file specified by the filename parameter."""
        space_json = self.chemical_space

        if minify:
            space_json = json.dumps(space_json)
            space_json = self.__minify_data(space_json)
        elif dump:
            space_json = json.dumps(space_json, indent=4)

        if filename:
            output = open(filename, "w")
            output.write(space_json)

        return space_json

    def __minify_data(self, data):
        return jsmin.jsmin(data)

def _process_(arguments):
    s = ChemSpace()
    s.sdf_file = False
    s.write_structures = False if arguments.dont_write_structures else True
    s.fp = arguments.fingerprint
    s.category_field = arguments.category_field
    s.category_field_delimiter = arguments.category_field_delimiter
    s.label_field = arguments.label_field
    s.activity_field = arguments.activity_field
    s.compound_structure_field = arguments.compound_structure_field
    s.fingerprint_field = arguments.fingerprint_field
    s.metric = arguments.similarity_metric

    if arguments.data_file.split(".")[-1].lower() == "sdf":
        s.read_sdf(arguments.data_file)
    else:
        s.read_csv(arguments.data_file, arguments.data_delimiter, arguments.data_header, arguments.missing_values, arguments.remove_columns)

    if s.compound_structure_field is not False or s.sdf == True:
        if arguments.physico_chemical_properties:
            s.add_physico_chemical_properties()

    if arguments.normalize:
        s.normalize_data()

    if arguments.arrange_by:
        s.arrange(
            method=arguments.dimensional_reduction_method,
            similarity_threshold=float(arguments.compound_similarity_threshold),
            add_edges=arguments.add_edges,
            by=arguments.arrange_by,
            k=arguments.knn
        )

    if arguments.html_dir:
        s.export_chemical_space_as_html(arguments.html_dir)
    elif arguments.output_file:
        s.export_chemical_space_as_json(arguments.output_file, minify=arguments.minify_output)
    else:
        print(s.export_chemical_space_as_json(minify=arguments.minify_output))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("data_file", type=str, help="csv(text) data file with delimited values or a sdf file")
    parser.add_argument("-dh", "--data_header", default=False, help="whether the first row of data file is a header", action="store_true")
    parser.add_argument("-dd", "--data_delimiter", type=str, default=",", help="delimiter of values in data file")
    parser.add_argument("-o", "--output_file", type=str, help="the name of output file")
    parser.add_argument("-fpf", "--fingerprint_field", type=str, default=False, help="set a fingerprint field name in case it is in the data file")
    parser.add_argument("-cf", "--category_field", type=str, default=False, help="set a category field name in case it is in the data file")
    parser.add_argument("-cfd", "--category_field_delimiter", type=str, default=False, help="a category field delimiter")
    parser.add_argument("-lf", "--label_field", type=str, default=False, help="set a label field name in case it is in the data file")
    parser.add_argument("-af", "--activity_field", type=str, default=False, help="set an activity field name in case it is in the data file")
    parser.add_argument("-csf", "--compound_structure_field", type=str, default=False, help="the name of a column with a compound structure")
    parser.add_argument("-fp", "--fingerprint", type=str, default="ecfp4", help="fingerprint used for a compound representation (ecfp4, ecfp6, maccs, topological, atom_pairs)")
    parser.add_argument("-arr", "--arrange_by", default=False, help="arrange data by compound structures (distance matrix) or by input data (data/fps)", type=str)
    parser.add_argument("-cst", "--compound_similarity_threshold", default=0.7, help="compound similarity threshold")
    parser.add_argument("-drm", "--dimensional_reduction_method", nargs='+', type=str, default="pca", help="which method use for dimensional reduction (pca/isomap/csn)")
    parser.add_argument("-dws", "--dont_write_structures", default=False, help="dont write structures to output file", action="store_true")
    parser.add_argument("-min", "--minify_output", default=False, help="minify the JSON output format", action="store_true")
    parser.add_argument("-html", "--html_dir", type=str, default=False, help="the directory to store HTML page with dependencies")
    parser.add_argument("-pcp", "--physico_chemical_properties", default=False, help="calculate basic phyisico-chemical properties and add them ass features", action='store_true')
    parser.add_argument("-edges", "--add_edges", default=False, help="add edges based on compound similarity to the graph", action='store_true')
    parser.add_argument("-n", "--normalize", default=False, help="normalize data to [0, 1] range", action="store_true")
    parser.add_argument("-mv", "--missing_values", type=str, default=False, help="define the string representating missing values in the data")
    parser.add_argument("-k", "--knn", type=int, default=None, help="the number of neighbours (k) used for the construction of csn using the nn method")
    parser.add_argument("-sm", "--similarity_metric", type=str, default="Tanimoto", help="similarity metric")
    parser.add_argument('-rmc','--remove_columns', nargs='+', default=False, help='columns in data that should not be used')

    args = parser.parse_args()
    _process_(args)


#!/usr/bin/env python
# coding: utf-8

# # Python Crash Course

# Hello! This is a quick intro to programming in Python to help you hit the ground running with the _12 Steps to Navier–Stokes_.
#
# There are two ways to enjoy these lessons with Python:
#
# 1. You can download and install a Python distribution on your computer. One option is the free [Anaconda Scientific Python](https://store.continuum.io/cshop/anaconda/) distribution. Another is [Canopy](https://www.enthought.com/products/canopy/academic/), which is free for academic use.  Our recommendation is Anaconda.
#
# 2. You can run Python in the cloud using [Wakari](https://wakari.io/) web-based data analysis, for which you need to create a free account. (No software installation required!)
#
# In either case, you will probably want to download a copy of this notebook, or the whole AeroPython collection. We recommend that you then follow along each lesson, experimenting with the code in the notebooks, or typing the code into a separate Python interactive session.
#
# If you decided to work on your local Python installation, you will have to navigate in the terminal to the folder that contains the .ipynb files. Then, to launch the notebook server, just type:
# ipython notebook
#
# You will get a new browser window or tab with a list of the notebooks available in that folder. Click on one and start working!

# ## Libraries

# Python is a high-level open-source language.  But the _Python world_ is inhabited by many packages or libraries that provide useful things like array operations, plotting functions, and much more. We can import libraries of functions to expand the capabilities of Python in our programs.
#
# OK! We'll start by importing a few libraries to help us out. First: our favorite library is **NumPy**, providing a bunch of useful array operations (similar to MATLAB). We will use it a lot! The second library we need is **Matplotlib**, a 2D plotting library which we will use to plot our results.
# The following code will be at the top of most of your programs, so execute this cell first:

# In[1]:

# <-- comments in python are denoted by the pound sign, like this one

import numpy as np                # we import the array library
from matplotlib import pyplot    # import plotting library

# We are importing one library named `numpy` and we are importing a module called `pyplot` of a big library called `matplotlib`.
# To use a function belonging to one of these libraries, we have to tell Python where to look for it. For that, each function name is written following the library name, with a dot in between.
# So if we want to use the NumPy function [linspace()](http://docs.scipy.org/doc/numpy/reference/generated/numpy.linspace.html), which creates an array with equally spaced numbers between a start and end, we call it by writing:

# In[2]:

myarray = np.linspace(0, 5, 10)
myarray

# If we don't preface the `linspace()` function with `numpy`, Python will throw an error.

# In[4]:

myarray = np.linspace(0, 5, 10)

# The function `linspace()` is very useful. Try it changing the input parameters!
#
# **Import style:**
#
# You will often see code snippets that use the following lines
# ```Python
# import numpy as np
# import matplotlib.pyplot as plt
# ```
# What's all of this import-as business? It's a way of creating a 'shortcut' to the NumPy library and the pyplot module. You will see it frequently as it is in common usage, but we prefer to keep out imports explicit. We think it helps with code readability.
#
# **Pro tip:**
#
# Sometimes, you'll see people importing a whole library without assigning a shortcut for it (like `from numpy import *`). This saves typing but is sloppy and can get you in trouble. Best to get into good habits from the beginning!
#
#
# To learn new functions available to you, visit the [NumPy Reference](http://docs.scipy.org/doc/numpy/reference/) page. If you are a proficient `Matlab` user, there is a wiki page that should prove helpful to you: [NumPy for Matlab Users](http://wiki.scipy.org/NumPy_for_Matlab_Users)

# ## Variables

# Python doesn't require explicitly declared variable types like C and other languages.

# In[5]:

a = 5        #a is an integer 5
b = 'five'   #b is a string of the word 'five'
c = 5.0      #c is a floating point 5

# In[6]:

type(a)

# In[7]:

type(b)

# In[8]:

type(c)

# Note that if you divide an integer by an integer that yields a remainder, the result will be converted to a float.  (This is *different* from the behavior in Python 2.7, beware!)

# ## Whitespace in Python

# Python uses indents and whitespace to group statements together.  To write a short loop in C, you might use:
#
#     for (i = 0, i < 5, i++){
#        printf("Hi! \n");
#     }

# Python does not use curly braces like C, so the same program as above is written in Python as follows:

# In[9]:

for i in range(5):
    print("Hi \n")

# If you have nested for-loops, there is a further indent for the inner loop.

# In[10]:

for i in range(3):
    for j in range(3):
        print(i, j)

    print("This statement is within the i-loop, but not the j-loop")

# ## Slicing Arrays

# In NumPy, you can look at portions of arrays in the same way as in `Matlab`, with a few extra tricks thrown in.  Let's take an array of values from 1 to 5.

# In[11]:

myvals = np.array([1, 2, 3, 4, 5])
print(myvals)

# Python uses a **zero-based index**, so let's look at the first and last element in the array `myvals`

# In[12]:

print(myvals[0], myvals[4])

# There are 5 elements in the array `myvals`, but if we try to look at `myvals[5]`, Python will be unhappy, as `myvals[5]` is actually calling the non-existant 6th element of that array.

# In[13]:

print(myvals[5])

# Arrays can also be 'sliced', grabbing a range of values.  Let's look at the first three elements

# In[14]:

print(myvals[0:3])

# Note here, the slice is inclusive on the front end and exclusive on the back, so the above command gives us the values of `myvals[0]`, `myvals[1]` and `myvals[2]`, but not `myvals[3]`.

# ## Assigning Array Variables

# One of the strange little quirks/features in Python that often confuses people comes up when assigning and comparing arrays of values.  Here is a quick example.  Let's start by defining a 1-D array called $a$:

# In[15]:

a = np.linspace(1,5,5)

# In[16]:

a

# OK, so we have an array $a$, with the values 1 through 5.  I want to make a copy of that array, called $b$, so I'll try the following:

# In[17]:

b = a

# In[18]:

b

# Great.  So $a$ has the values 1 through 5 and now so does $b$.  Now that I have a backup of $a$, I can change its values without worrying about losing data (or so I may think!).

# In[19]:

a[2] = 17

# In[20]:

a

# Here, the 3rd element of $a$ has been changed to 17.  Now let's check on $b$.

# In[21]:

b

# And that's how things go wrong!  When you use a statement like $a = b$, rather than copying all the values of $a$ into a new array called $b$, Python just creates an alias (or a pointer) called $b$ and tells it to route us to $a$.  So if we change a value in $a$ then $b$ will reflect that change (technically, this is called *assignment by reference*).  If you want to make a true copy of the array, you have to tell Python to copy every element of $a$ into a new array.  Let's call it $c$.

# In[22]:

c = a.copy()

# Now, we can try again to change a value in $a$ and see if the changes are also seen in $c$.

# In[23]:

a[2] = 3

# In[24]:

a

# In[25]:

c

# OK, it worked!  If the difference between `a = b` and `a = b.copy()` is unclear, you should read through this again.  This issue will come back to haunt you otherwise.

# ## Learn More

# There are a lot of resources online to learn more about using NumPy and other libraries. Just for kicks, here we use Jupyter's feature for embedding videos to point you to a short video on YouTube on using NumPy arrays.

# In[26]:

from IPython.display import YouTubeVideo
# a short video about using NumPy arrays, from Enthought
YouTubeVideo('vWkb7VahaXQ')

# In[27]:

from IPython.core.display import HTML
def css_styling():
    styles = open("../styles/custom.css", "r").read()
    return HTML(styles)
css_styling()


from __future__ import print_function
import numpy as np
import random
import datetime
import os
from scipy.ndimage.interpolation import zoom
import matplotlib.pyplot as plt

def write_time_and_date():
    os.environ['TZ'] = 'CST6CDT'
    return "Time" + datetime.datetime.now().strftime("_%H_%M_%S_") + \
           "Date" + datetime.datetime.now().strftime("_%Y_%m_%d")

def results2(res, number_isotopes_displayed):

    index = [i[0] for i in sorted(enumerate(res), key=lambda x:x[1])]
    index = list(reversed(index))
    for i in range(number_isotopes_displayed):
        print((isotopes[index[i]], round(res[index[i]], 3)))

def load_template_spectra_from_folder(parent_folder,
                                      spectrum_identifier,
                                      LLD=10):
    '''
    inputs: partent_folder, spectrum_identifier
    output: dictionary containing all template spectra from a folder.

    Load template spectrum data into a dictionary. This allows templates from
    different folders to be loaded into different dictionaries.

    '''

    temp_dict = {}

    def normalize_spectrum(ID):
        temp_spectrum = read_spectrum(parent_folder + ID + spectrum_identifier)
        temp_spectrum[0:LLD] = 0
        return temp_spectrum / np.max(temp_spectrum)

    for i in range(len(isotopes)):
        # Fixes background spectra name issue
        if i >= len(isotopes) - 3:
            spectrum_identifier = ''

        temp_dict[isotopes[i]] = normalize_spectrum(isotopes_GADRAS_ID[i])

    return temp_dict

def zoom_spectrum(spectrum, zoom_strength):
    spectrum = np.abs(zoom(spectrum, zoom_strength))
    if zoom_strength < 1.0:
        spectrum = np.lib.pad(spectrum,
                              (0, 1024 - spectrum.shape[0]),
                              'constant',
                              constant_values=0)
    if zoom_strength > 1.0:
        spectrum = spectrum[0:1024]

    return spectrum

def read_spectrum(filename):
    '''
    Reads spectrum from .spe files.
    Works with silver detector and GADRAS formatted spectra.
    '''
    spectrum = np.empty(1024)

    with open(filename, 'rb') as f:

        content = f.readlines()

        if RepresentsInt(content[8]):
            for i in range(1024):
                # spectra begins at index 8
                spectrum[i] = int(content[8 + i])
        else:
            for i in range(1024):
                # spectra begins at index 12
                spectrum[i] = int(content[12 + i])

    return spectrum

def create_simplex(number_samples, number_categories):
    # make an empty array
    k = np.zeros([number_samples, number_categories + 1])
    # Make a sorted array of random variables
    a = np.sort(np.random.uniform(0,
                                  1,
                                  [number_samples, number_categories - 1]),
                axis=1)
    # Zero pad left side
    k[:, 0] = 0
    # Put sorted array in new array
    k[:, 1:number_categories] = a
    # One pad right side
    k[:, number_categories] = 1
    # Take the difference of adjacent elements
    temp_simplex = np.diff(k)
    return temp_simplex

def shuffle_simplex(simplex):
    # last term from simpex is always background, by convention
    temp = random.sample(simplex[:-1], len(simplex) - 1)

    # 29 here because 29 isotopes plus one background super-isotope
    shuffled_array = np.pad(temp, [0, 29 - len(temp)], 'constant')

    np.random.shuffle(shuffled_array)

    # Add background
    shuffled_array = np.append(shuffled_array, simplex[-1])

    return shuffled_array

def visualize_simplex(key, index1, index2):
    new_list = []

    for i in range(len(key)):
        if key[i][index1] > 0 and key[i][index2] > 0:
            new_list.append([key[i][index1], key[i][index2]])

    plt.scatter(np.array(new_list)[:, 0], np.array(new_list)[:, 1])
    return new_list

def RepresentsInt(s):
    '''
    Helper funtion to see if a string represents an integer
    '''
    try:
        int(s)
        return True
    except ValueError:
        return False

isotopes = [
    'Am241',
    'Ba133',
    'Co57',
    'Co60',
    'Cs137',
    'Cr51',
    'Eu152',
    'Ga67',
    'I123',
    'I125',
    'I131',
    'In111',
    'Ir192',
    'U238',
    'Lu177m',
    'Mo99',
    'Np237',
    'Pd103',
    'Pu239',
    'Pu240',
    'Ra226',
    'Se75',
    'Sm153',
    'Tc99m',
    'Xe133',
    'Tl201',
    'Tl204',
    'U233',
    'U235',
    'Back_Th',
    'Back_U',
    'Back_K',
]

isotopes_GADRAS_ID = [
    '241AM',
    '133BA',
    '57CO',
    '60CO',
    '137CS',
    '51CR',
    '152EU',
    '67GA',
    '123I',
    '125I',
    '131I',
    '111IN',
    '192IR',
    '238U',
    '177MLU',
    '99MO',
    '237NP',
    '103PD',
    '239PU',
    '240PU',
    '226RA',
    '75SE',
    '153SM',
    '99TCM',
    '133XE',
    '201TL',
    '204TL',
    '233U',
    '235U',
    'ThoriumInSoil.spe',
    'UraniumInSoil.spe',
    'PotassiumInSoil.spe',
]

isotopes_sources_GADRAS_ID = [
    '241AM',
    '133BA',
    '57CO',
    '60CO',
    '137CS',
    '51CR',
    '152EU',
    '67GA',
    '123I',
    '125I',
    '131I',
    '111IN',
    '192IR',
    '238U',
    '177MLU',
    '99MO',
    '237NP',
    '103PD',
    '239PU',
    '240PU',
    '226RA',
    '75SE',
    '153SM',
    '99TCM',
    '133XE',
    '201TL',
    '204TL',
    '233U',
    '235U'
]

def sample_spectrum(iso_DRF, ncounts):
    '''
    Input:
    isoDRF: the 1024x1 vector containing the spectrum to be sampled.
            Does not need to be normalized.
    Output:
    ncounts: the 1024x1 vector containing the sampled spectrum.

    Method:
    Normalize isoDRF, and it is effectively a probability density function
    Calculate the cumulative distribution function
    Generate uniform random numbers to sample the cdf
    '''

    pdf = iso_DRF / sum(iso_DRF)
    cdf = np.cumsum(pdf)

    # take random samples and generate spectrum
    t_all = np.random.rand(np.int(ncounts))
    spec = pdf * 0
    for t in t_all:
        pos = np.argmax(cdf > t)
        spec[pos] = spec[pos] + 1
    return spec