##############################################################################
#                                                                             #
# Copyright (C) 2007-2014 Edward d'Auvergne                                   #
#                                                                             #
# This file is part of the program relax (http://www.nmr-relax.com).          #
#                                                                             #
# This program is free software: you can redistribute it and/or modify        #
# it under the terms of the GNU General Public License as published by        #
# the Free Software Foundation, either version 3 of the License, or           #
# (at your option) any later version.                                         #
#                                                                             #
# This program is distributed in the hope that it will be useful,             #
# but WITHOUT ANY WARRANTY; without even the implied warranty of              #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
# GNU General Public License for more details.                                #
#                                                                             #
# You should have received a copy of the GNU General Public License           #
# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#                                                                             #
###############################################################################

"""Script for creating a Grace plot of peak intensity differences.

This script creates a Grace plot of Ix - Ix(theta), the difference between the measured peak and the back calculated peak intensity for each spin system x.  Ix(theta) is back using the parameter vector theta = [Rx, I0], where Rx is either the R1 or R2 relaxation and I0 is the initial peak intensity.  The plot consists of distributions of intensity for each residue at each measured relaxation period.  The average and standard of these distributions are also plotted.

The resultant plot is useful for finding bad points or bad spectra when fitting exponential curves determine the R1 and R2 relaxation rates.  If the averages deviate systematically from zero, bias in the spectra or fitting will be clearly revealed.

To use this script, R1 or R2 exponential curve fitting must have previously have been carried out the program state saved to the file 'rx.save' (either with or without the .gz or .bz2 ).  The file name of the saved state can be changed at the top of this script.  It is to note that the same version of relax should be used for creating the saved state as the program state, these files are neither backwards nor forwards compatible.  Only the parameter exponential fit is currently supported.
"""

# Python module imports.
from numpy import float64, array, identity, sqrt, zeros

# relax module imports.
from pipe_control.mol_res_spin import spin_loop
from target_functions.relax_fit import back_calc_I, func, setup


# The name of the save file generated by 'relax_fit.py'.
SAVE_FILE = 'rx.save'


def back_calc():
    """Back calculate the peak intensities.

    The simple two parameter exponential curve (Rx, I0) is assumed.
    """

    # Loop over the spins.
    for spin in spin_loop():
        # Skip deselected spins.
        if not spin.select:
            continue

        # The parameter vector.
        param_vector = array([spin.rx, spin.i0], float64)

        # Initialise the relaxation fit functions.
        setup(num_params=len(spin.params), num_times=len(cdp.relax_times), values=spin.ave_intensities, sd=cdp.sd, relax_times=cdp.relax_times, scaling_matrix=identity(2, float64))

        # Make a single function call.  This will cause back calculation and the data will be stored in the C module.
        func(param_vector)

        # Get the data and store it in the spin specific data structure.
        spin.fit_int = back_calc_I()


def calc_ave_sd():
    """Calculate the average difference and sd between the measured and fitted intensities.

    The standard deviation formula is:
                  ___________________________
                 /   1
        sd =    /  ----- * sum({Xi - Xav}^2)],
              \/   n - 1

    where n is the total number of selected spins, Xi is the peak intensity difference for spin i,
    and Xav is the peak intensity difference averaged across all spins.
    """

    # Diff array, std deviation array, and number of spins.
    diff_array = zeros(sum(cdp.num_spectra), float64)
    sd_array = zeros(sum(cdp.num_spectra), float64)
    num_spins = 0


    # Calculate the average difference.
    ###################################

    # Loop over the spins.
    for spin in spin_loop():
        # Skip deselected spins.
        if not spin.select:
            continue

        # Loop over the intensities.
        index = 0
        for i in range(len(spin.peak_intensity)):
            for j in range(len(spin.peak_intensity[i])):
                # Add the difference between the measured and fitted intensity to the diff array.
                diff_array[index] = diff_array[index] + (spin.peak_intensity[i][j] - spin.fit_int[i])

                # Increment the index.
                index = index + 1

        # The number of selected spins.
        num_spins = num_spins + 1

    # Average difference.
    diff_array = diff_array / num_spins


    # Calculate the standard deviations.
    ####################################

    # Loop over the spins.
    for spin in spin_loop():
        # Skip deselected spins.
        if not spin.select:
            continue

        # Loop over the intensities.
        index = 0
        for i in range(len(spin.peak_intensity)):
            for j in range(len(spin.peak_intensity[i])):
                # Calculate the sum of squares.
                sd_array[index] = sd_array[index] + ((spin.peak_intensity[i][j] - spin.fit_int[i]) - diff_array[index])**2

                # Increment the index.
                index = index + 1

    # The standard deviations.
    sd_array = sqrt(sd_array / (num_spins - 1))


    # Return the values.
    ####################

    return diff_array, sd_array


def grace_header(file, xmin, xmax, ymin, ymax):
    """Write the formatted Grace header."""

    # Grace version!
    file.write("@version 50118\n")

    # Graph G0.
    file.write("@with g0\n")

    # The graph zoom level.
    file.write("@    world %.1f, %.1f, %.1f, %.1f\n" % (xmin, ymin, xmax, ymax))

    # X-axis label.
    file.write("@    xaxis  label \"\qRelaxation time period (s)\Q\"\n")

    # X-axis specific settings.
    file.write("@    xaxis  label char size 1.48\n")
    file.write("@    xaxis  tick major size 0.75\n")
    file.write("@    xaxis  tick major linewidth 0.5\n")
    file.write("@    xaxis  tick minor linewidth 0.5\n")
    file.write("@    xaxis  tick minor size 0.45\n")
    file.write("@    xaxis  ticklabel char size 1.00\n")

    # Y-axis label.
    file.write("@    yaxis  label \"\qPeak intensity differences (I\\sx\\N - I\\sx\\N(\\xq\\f{}))\Q\"\n")

    # Y-axis specific settings.
    file.write("@    yaxis  label char size 1.48\n")
    file.write("@    yaxis  tick major size 0.75\n")
    file.write("@    yaxis  tick major linewidth 0.5\n")
    file.write("@    yaxis  tick minor linewidth 0.5\n")
    file.write("@    yaxis  tick minor size 0.45\n")
    file.write("@    yaxis  ticklabel char size 1.00\n")

    # Frame.
    file.write("@    frame linewidth 0.5\n")

    # Symbols.
    file.write("@    s0 symbol 9\n")
    file.write("@    s0 symbol size 0.45\n")
    file.write("@    s0 symbol linewidth 0.5\n")
    file.write("@    s0 line linestyle 0\n")
    file.write("@    s1 symbol 8\n")
    file.write("@    s1 symbol size 0.45\n")
    file.write("@    s1 symbol linewidth 0.5\n")
    file.write("@    s1 line linestyle 0\n")
    file.write("@    s2 symbol 1\n")
    file.write("@    s2 symbol size 1.00\n")
    file.write("@    s2 symbol fill pattern 1\n")
    file.write("@    s2 symbol linewidth 0.5\n")
    file.write("@    s2 line linestyle 3\n")


def grace_plot(ave, sd):
    """Grace plot of the intensity differences."""

    # Open the file.
    file = open('differences.agr', 'w')

    # Y-axis min and max.
    ymin = 2.5*min(ave - sd)
    ymax = 2.5*max(ave + sd)

    # Grace header.
    grace_header(file, xmin=0, xmax=cdp.relax_times[-1], ymin=ymin, ymax=ymax)


    # First time point difference distributions.
    ############################################

    # First graph, first data set.
    file.write("@target G0.S0\n")
    file.write("@type xy\n")

    # Loop over the individual time points.
    for i in range(len(cdp.num_spectra)):
        # Loop over the spins.
        for spin in spin_loop():
            # Skip deselected spins.
            if not spin.select:
                continue

            # Grace data point.
            file.write("%-30s%-30s\n" % (repr(cdp.relax_times[i]), repr(spin.peak_intensity[i][0] - spin.fit_int[i])))

    # End the graph.
    file.write("&\n")


    # Second time point difference distributions.
    #############################################

    # First graph, second data set.
    file.write("@target G0.S1\n")
    file.write("@type xy\n")

    # Loop over the individual time points.
    for i in range(len(cdp.num_spectra)):
        # Loop over the spins.
        for spin in spin_loop():
            # Skip deselected spins.
            if not spin.select:
                continue

            # Grace data point.
            if len(spin.peak_intensity[i]) == 2:
                file.write("%-30s%-30s\n" % (repr(cdp.relax_times[i]), repr(spin.peak_intensity[i][1] - spin.fit_int[i])))

    # End the graph.
    file.write("&\n")


    # Average and sd.
    #################

    # First graph, third data set.
    file.write("@target G0.S2\n")
    file.write("@type xydy\n")

    # Loop over the data.
    index = 0
    for i in range(len(cdp.num_spectra)):
        for j in range(cdp.num_spectra[i]):
            # Grace data point.
            file.write("%-30s%-30s%-30s\n" % (repr(cdp.relax_times[i]), repr(ave[index]), repr(sd[index])))

            # Increment the index.
            index = index + 1

    # End the graph.
    file.write("&\n")

    # Close the file.
    file.close()


# Load the program state containing saved by the 'relax_fit.py' sample script.
state.load(SAVE_FILE)

# Back calculate the peak intensities from the fitted parameters.
back_calc()

# Calculate the average difference and standard deviations for each time point.
ave, sd = calc_ave_sd()

# Create a Grace plot of the differences.
grace_plot(ave, sd)

# View the graph.
grace.view(file='differences.agr', dir=None)
