###############################################################################
#                                                                             #
# Copyright (C) 2004-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/>.       #
#                                                                             #
###############################################################################

# Module docstring.
"""The R1 and R2 exponential relaxation curve fitting optimisation functions."""

# Python module imports.
from numpy import average
from re import search

# relax module imports.
from dep_check import C_module_exp_fn
from lib.errors import RelaxError, RelaxLenError
from specific_analyses.relax_fit.parameters import assemble_param_vector, assemble_scaling_matrix

# C modules.
if C_module_exp_fn:
    from target_functions.relax_fit import setup, func, back_calc_I


def back_calc(spin=None, relax_time_id=None):
    """Back-calculation of peak intensity for the given relaxation time.

    @keyword spin:              The spin container.
    @type spin:                 SpinContainer instance
    @keyword relax_time_id:     The ID string for the desired relaxation time.
    @type relax_time_id:        str
    @return:                    The peak intensity for the desired relaxation time.
    @rtype:                     float
    """

    # Create the initial parameter vector.
    param_vector = assemble_param_vector(spin=spin)

    # Create a scaling matrix.
    scaling_matrix = assemble_scaling_matrix(spin=spin, scaling=False)

    # The keys.
    keys = list(spin.peak_intensity.keys())

    # The peak intensities and times.
    values = []
    errors = []
    times = []
    for key in keys:
        values.append(spin.peak_intensity[key])
        errors.append(spin.peak_intensity_err[key])
        times.append(cdp.relax_times[key])

    # The scaling matrix in a diagonalised list form.
    scaling_list = []
    for i in range(len(scaling_matrix)):
        scaling_list.append(scaling_matrix[i, i])

    # Initialise the relaxation fit functions.
    setup(num_params=len(spin.params), num_times=len(cdp.relax_times), values=values, sd=errors, relax_times=times, scaling_matrix=scaling_list)

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

    # Get the data back.
    results = back_calc_I()

    # Return the correct peak height.
    return results[keys.index(relax_time_id)]


def func_wrapper(params):
    """Wrapper function for the C module, for converting numpy arrays.

    @param params:  The parameter array from the minimisation code.
    @type params:   numpy array
    @return:        The function value generated by the C module.
    @rtype:         float
    """

    # Call the C code.
    chi2 = func(params.tolist())

    # Return the chi2 value.
    return chi2


def dfunc_wrapper(params):
    """Wrapper function for the C module, for converting numpy arrays.

    The currently does nothing.
    """


def d2func_wrapper(params):
    """Wrapper function for the C module, for converting numpy arrays.

    The currently does nothing.
    """


def grid_search_setup(spin=None, param_vector=None, lower=None, upper=None, inc=None, scaling_matrix=None):
    """The grid search setup function.

    @keyword spin:              The spin data container.
    @type spin:                 SpinContainer instance
    @keyword param_vector:      The parameter vector.
    @type param_vector:         numpy array
    @keyword lower:             The lower bounds of the grid search which must be equal to the
                                number of parameters in the model.  This optional argument is
                                only used when doing a grid search.
    @type lower:                array of numbers
    @keyword upper:             The upper bounds of the grid search which must be equal to the
                                number of parameters in the model.  This optional argument is
                                only used when doing a grid search.
    @type upper:                array of numbers
    @keyword inc:               The increments for each dimension of the space for the grid
                                search.  The number of elements in the array must equal to the
                                number of parameters in the model.  This argument is only used
                                when doing a grid search.
    @type inc:                  array of int
    @keyword scaling_matrix:    The scaling matrix.
    @type scaling_matrix:       numpy diagonal matrix
    @return:                    A tuple of the grid size and the minimisation options.  For the
                                minimisation options, the first dimension corresponds to the
                                model parameter.  The second dimension is a list of the number
                                of increments, the lower bound, and upper bound.
    @rtype:                     (int, list of lists [int, float, float])
    """

    # The length of the parameter array.
    n = len(param_vector)

    # Make sure that the length of the parameter array is > 0.
    if n == 0:
        raise RelaxError("Cannot run a grid search on a model with zero parameters.")

    # Lower bounds.
    if lower != None and len(lower) != n:
        raise RelaxLenError('lower bounds', n)

    # Upper bounds.
    if upper != None and len(upper) != n:
        raise RelaxLenError('upper bounds', n)

    # Increments.
    if isinstance(inc, list) and len(inc) != n:
        raise RelaxLenError('increment', n)
    elif isinstance(inc, int):
        inc = [inc]*n

    # Set up the default bounds.
    if not lower:
        # Init.
        lower = []
        upper = []

        # Loop over the parameters.
        for i in range(n):
            # Relaxation rate (from 0 to 20 s^-1).
            if spin.params[i] == 'rx':
                lower.append(0.0)
                upper.append(20.0)

            # Intensity
            elif search('^i', spin.params[i]):
                # Find the ID of the first time point.
                min_time = min(cdp.relax_times.values())
                for key in list(cdp.relax_times.keys()):
                    if cdp.relax_times[key] == min_time:
                        id = key
                        break

                # Defaults.
                lower.append(0.0)
                upper.append(average(spin.peak_intensity[id]))

    # Diagonal scaling of minimisation options.
    lower_new = []
    upper_new = []
    for i in range(n):
        lower_new.append(lower[i] / scaling_matrix[i, i])
        upper_new.append(upper[i] / scaling_matrix[i, i])

    # Return the minimisation options.
    return inc, lower_new, upper_new
