import os

# Use sans-serif fonts for plotting
import matplotlib
matplotlib.rcParams[u'font.family'] = [u'sans-serif']
matplotlib.rcParams[u'mathtext.default'] = u'regular'

import matplotlib.pyplot as plt
%matplotlib inline
# %matplotlib notebook # Alternative to `%matplotlib inline` that will make plots interactive

# Import mfoutparser

try:
    # If mfoutparser is in PYTHONPATH (or is installed), then use a direct import
    import mfoutparser as mf
    
except ImportError:
    # Attempt to demo mfoutparser without installing by importing from directory
    mfoutparser_path = '../../mfoutparser'
    
    print("Module mfoutparser was not found in PYTHONPATH. Looking for module in directory '{:s}'\n".format(mfoutparser_path))
    
    if os.path.exists(mfoutparser_path):
        import imp
        mf = imp.load_package('mf', mfoutparser_path)
        print("Module mfoutparser was found in the directory '{:s}' and imported.\n".format(mfoutparser_path))
    else:
        raise ImportError("Module mfoutparser could not be found in the directory '{:s}'.\n".format(mfoutparser_path) + \
        "This demonstration will not run until the module is located.")

help(mf)

"""
Explanation: ModelFree Parser Demo
Arthur G. Palmer, III and Michelle L. Gill
2015/12/06
This IPython notebook demonstartes how to parse various types of ModelFree STAR output files with the library mfoutparser. The output of each of the files is returned in two variables, one called loops, which is a dictionary containing all of the table  information. (The name is because tables are called loops in the ModelFree STAR file.) The second variable, tags, is a dictionary containing all other information from various output file sections that does not reside in a table.
The ModelFree output files are parsed with the parse_mfout command, which takes a path to the output file as input and returns the two dictionaries mentioned above.
The tag and loop data can all be written to a file using the function write_all_to_file, which takes the dictionaries containing the tag and loop data as inputs in addition to an optional file prefix. It will generate a tab-delimited file for each of the tables and the miscellaneous items in the tags variable.
General data selection can be accomplished with get_data_selection as well as with any Pandas-dataframe compatible method.
A correlation matrix of the correlation values can be created with make_correlation_matrices and written to a file using the function write_correlation_matrix_to_file. 
mfoutparser has been tested on python 2.7 and 3.4. It requires the Numpy (tested on version 1.10.1) and Pandas (tested on version 0.17.1) libraries. Matplotlib is also required if plotting of the data is desired.
End of explanation
"""


help(mf.parse_mfout)

help(mf.write_all_to_file)

# Create the output directory if necessary
output_directory = 'output'
if not os.path.exists(output_directory):
    os.mkdir(output_directory)

# Parse the ModelFree output file
input_directory = 'input_data'
mfoutfilename = os.sep.join([input_directory, 'mfout.singlefield'])
tags, loops = mf.parse_mfout(mfoutfilename)

# Write everything to a file
output_filename = os.sep.join([output_directory, 'mfout_singlefield'])
mf.write_all_to_file(tags, loops, output_filename)

"""
Explanation: Single Field Data
Input file: mfout.singlefield
End of explanation
"""


tags.keys()

tags['header']

tags['chi_square']

"""
Explanation: Overview of tags and loops data structure
Preview the returned tags.
End of explanation
"""


loops.keys()

loops['header_1']

loops['header_2']

loops['header_3']

loops['chi_square']

loops['relaxation']

"""
Explanation: Preview the returned loops (tables).
End of explanation
"""


help(mf.get_data_selection)

"""
Explanation: Selecting and plotting data
Extract data using a given parameter as an inequality using the get_data_selection function. The selection is based on selector_dict which is a dictionary whose key corresponds to the column name and value corresponds to the selected value. Multiple criteria can be used, as shown below. 
The function will accept a specific (in)equality parameter as a string ('==', '&lt;', '&gt;', '&lt;=', '&gt;=') and '==' is the default. To combine selections with different inequalities, simply call the command additional times on the previous command's output.
End of explanation
"""


selector_dict = {'residue':5}
residue_5_rates = mf.get_data_selection(loops['relaxation'], selector_dict, '<=')
residue_5_rates

"""
Explanation: This command will extract all relaxation data for residues whose numbers are <= 5:
End of explanation
"""


selector_dict = {'relaxation_rate_name':'R1', 'field':500.13}
r1_data = mf.get_data_selection(loops['relaxation'], selector_dict)
r1_data

"""
Explanation: This command will extract data for a given relaxation parameter at a specific field strength (R1 data at 500 MHz).
End of explanation
"""


fig = plt.gcf()
fig.set_size_inches(5,3)

ax = plt.axes()

# The extra ".values" property is sometimes required with certain Matplotlib and Pandas version combinations
ax.errorbar(r1_data.residue.values, r1_data.value.values, yerr=r1_data.uncertainty.values, 
             color='blue', marker='o', ls='', capthick=1.0)

# Using periods followed by the column name is a shortcut for this:
# ax.errorbar(r1_data['residue'].values, r1_data['value'].values, yerr=r1_data['uncertainty'].values, 
#              color='blue', marker='o', ls='', capthick=1.0)

ax.set_ylim(0,3)
ax.set_xlim(0, r1_data.residue.max()+1)
ax.set_xlabel('Residue')
ax.set_ylabel('$R_1$ $s^{-1}$')

# Set the fontsize for the label and tick labels
fontsize = 12.0
ax.xaxis.label.set_fontsize(fontsize)
ax.yaxis.label.set_fontsize(fontsize)
for tick in ax.xaxis.get_ticklabels():
    tick.set_fontsize(fontsize)
for tick in ax.yaxis.get_ticklabels():
    tick.set_fontsize(fontsize)

fig.tight_layout()

output_filename = os.sep.join([output_directory, 'R1_500MHz_plot.pdf'])
fig.savefig(output_filename)
output_filename = os.sep.join([output_directory, 'R1_500MHz_plot.png'])
fig.savefig(output_filename, dpi=300)

"""
Explanation: Plot the selected R1 data.
End of explanation
"""


loops['model_1']

selector_dict = {'model_free_name':'S2'}
s2_data = mf.get_data_selection(loops['model_1'], selector_dict)

fig = plt.gcf()
fig.set_size_inches(5,3)

ax = plt.axes()
ax.errorbar(s2_data.residue.values, s2_data.fit_value.values, yerr=s2_data.sim_error.values, 
             color='blue', marker='o', ls='', capthick=1.0)

ax.set_ylim(0,1)
ax.set_xlim(0, s2_data.residue.max()+1)
ax.set_xlabel('Residue')
ax.set_ylabel('$S^2$')

# Set the fontsize for the label and tick labels
fontsize = 12.0
ax.xaxis.label.set_fontsize(fontsize)
ax.yaxis.label.set_fontsize(fontsize)
for tick in ax.xaxis.get_ticklabels():
    tick.set_fontsize(fontsize)
for tick in ax.yaxis.get_ticklabels():
    tick.set_fontsize(fontsize)

fig.tight_layout()

output_filename = os.sep.join([output_directory, 'S2_plot.pdf'])
fig.savefig(output_filename)
output_filename = os.sep.join([output_directory, 'S2_plot.png'])
fig.savefig(output_filename, dpi=300)

"""
Explanation: Select a Model-free parameter and plot.
End of explanation
"""


help(mf.make_correlation_matrices)

correlation_matrix = mf.make_correlation_matrices(loops['correlation_matrix'])
correlation_matrix

"""
Explanation: Creation of a correlation matrix
The correlation table can also be manipulated so the values are converted to a true matrix representation. Values that are missing are undefined.
End of explanation
"""


selector_dict = {'residue':2}
residue2_correlation_matrix = mf.get_data_selection(correlation_matrix, selector_dict)
residue2_correlation_matrix

"""
Explanation: The matrix corresponding to a single residue can be selected as before.
End of explanation
"""


residue2_correlation_matrix.values

"""
Explanation: And just the values of this matrix can be extracted as a numpy two-dimensional array for use in statistical analysis.
End of explanation
"""


residue2_correlation_matrix.values[:, 2:]

"""
Explanation: Parts of the matrix can be extracted as normal with numpy arrays. Here are just the $S^2_s$, $\theta$, and $\tau_e$ elements.
End of explanation
"""


help(mf.write_correlation_matrix_to_file)

output_filename = os.sep.join([output_directory, 'mfout_singlefield_loop'])
mf.write_correlation_matrix_to_file(correlation_matrix, output_filename)

"""
Explanation: The correlation matrix can be written to a file as well. The filename will have the text correlation_matrix_pivot added to it.
End of explanation
"""


# Parse the ModelFree output file
mfoutfilename = os.sep.join([input_directory, 'mfout.compare'])
tags, loops = mf.parse_mfout(mfoutfilename)

# Write everything to a file
output_filename = os.sep.join([output_directory, 'mfout_compare'])
mf.write_all_to_file(tags, loops, output_filename)

loops['F_dist']

loops['model_2']

"""
Explanation: Multiple Model Data
Input file: mfout.compare
End of explanation
"""


# Parse the ModelFree output file
mfoutfilename = os.sep.join([input_directory, 'mfout.multifield'])
tags, loops = mf.parse_mfout(mfoutfilename)

# Write everything to a file
output_filename = os.sep.join([output_directory, 'mfout_multifield'])
mf.write_all_to_file(tags, loops, output_filename)

loops['relaxation']

"""
Explanation: Multiple Field Data
Input file: mfout.multifield
End of explanation
"""


