"""
:Author: Pierre Barbier de Reuille <pierre.barbierdereuille@gmail.com>

This modules implement functions to test and plot parametric regression.
"""

from __future__ import division, print_function, absolute_import
from .curve_fitting import CurveFitting
from numpy import sort, iterable, argsort, std, abs, sqrt, arange, pi, c_
from pylab import figure, title, legend, plot, xlabel, ylabel, subplot, clf, ylim, hist, suptitle, gca
from .compat import izip
from . import bootstrap
from itertools import chain
from scipy.special import erfinv, gamma
from scipy import stats
#try:
#    from cy_kernel_smoothing import SpatialAverage
#except ImportError:
from .kernel_smoothing import SpatialAverage, LocalLinearKernel1D
import inspect
from .compat import unicode_csv_writer as csv_writer
from collections import namedtuple

smoothing = LocalLinearKernel1D

import sys
if sys.version_info >= (3,):
    CSV_WRITE_FLAGS = "wt"
else:
    CSV_WRITE_FLAGS = "wb"

def plot_dist_residuals(res):
    """
    Plot the distribution of the residuals.

    :returns: the handle toward the histogram and the plot of the fitted normal distribution
    """
    ph = hist(res,normed=True)
    xr = arange(res.min(), res.max(), (res.max()-res.min())/1024)
    yr = stats.norm(0, res.std()).pdf(xr)
    pn = plot(xr, yr, 'r--')
    xlabel('Residuals')
    ylabel('Frequency')
    title('Distributions of the residuals')
    return ph, pn

def plot_residuals(xname, xdata, res_desc, res):
    """
    Plot the residuals against the X axis

    :param str     xname:    Name of the X axis
    :param ndarray xdata:    1D array with the X data
    :param str     res_desc: Name of the Y axis
    :param ndarray res: 1D   array with the residuals

    The shapes of ``xdata`` and ``res`` must be the same

    :returns: The handles of the the plots of the residuals and of the smoothed residuals.
    """
    p_res = plot(xdata, res, '+', label='residuals')[0]
    plot([xdata.min(), xdata.max()], [0,0], 'r--')
    av = LocalLinearKernel1D(xdata, res)
    xr = arange(xdata.min(), xdata.max(), (xdata.max()-xdata.min())/1024)
    rr = av(xr)
    p_smooth = plot(xr, rr, 'g', label='smoothed residuals')
    xlabel(xname)
    ylabel("Residuals")
    ymin, ymax = ylim()
    ymax = max(ymax, -ymin)
    ylim(-ymax, ymax)
    title("Residuals (%s) vs. fitted" % (res_desc,))
    return p_res, p_smooth

def scaled_location_plot(yname, yopt, scaled_res):
    """
    Plot the scaled location, given the dependant values and scaled residuals.

    :param str     yname:      Name of the Y axis
    :param ndarray yopt:       Estimated values
    :param ndarray scaled_res: Scaled residuals

    :returns: the handles for the data and the smoothed curve
    """

    scr = sqrt(abs(scaled_res))
    p_scaled = plot(yopt, scr, '+')[0]
    av = LocalLinearKernel1D(yopt, scr)
    xr = arange(yopt.min(), yopt.max(), (yopt.max() - yopt.min())/1024)
    rr = av(xr)
    p_smooth = plot(xr, rr, 'g')[0]
    expected_mean = 2**(1/4)*gamma(3/4)/sqrt(pi)
    plot([yopt.min(), yopt.max()], [expected_mean, expected_mean], 'r--')
    title('Scale-location')
    xlabel(yname)
    ylabel('$|$Normalized residuals$|^{1/2}$')
    gca().set_yticks([0,1,2])
    return [p_scaled, p_smooth]

def qqplot(scaled_res, normq):
    """
    Draw a Q-Q Plot from the sorted, scaled residuals (i.e. residuals sorted
    and normalized by their standard deviation)

    :param ndarray scaled_res: Scaled residuals
    :param ndarray normq:      Expected value for each scaled residual, based on its quantile.

    :returns: handle to the data plot
    """
    qqp = []
    qqp += plot(normq, scaled_res, '+');
    qqp += plot(normq, normq, 'r--');
    xlabel('Theoretical quantiles');
    ylabel('Normalized residuals');
    title('Normal Q-Q plot');
    return qqp

ResultStruct = namedtuple('ResultStruct', """fct fct_desc param_names xdata ydata xname yname res_name residuals popt res
        yopts eval_points interpolation sorted_yopts scaled_res normq CI CIs CIresults""")

def fit_evaluation(fit, xdata, ydata, eval_points=None,
        CI=(), CIresults = None, xname="X", yname="Y", fct_desc = None, param_names = (), residuals=None, res_name = 'Standard'):
    """
    This function takes the output of a curve fitting experiment and store all the relevant information for evaluating
    its success in the result.

    :type  fit: fitting object
    :param fit: object configured for the fitting

    :type  xdata: ndarray of shape (N,) or (k,N) for function with k prefictors
    :param xdata: The independent variable where the data is measured

    :type  ydata: ndarray
    :param ydata: The dependant data

    :type  eval_points: ndarray or None
    :param eval_points: Contain the list of points on which the result must be expressed. It is
        used both for plotting and for the bootstrapping.

    :type  CI: tuple of int
    :param CI: List of confidence intervals to calculate. If empty, none are calculated.

    :type  xname: string
    :param xname: Name of the X axis

    :type  yname: string
    :param yname: Name of the Y axis

    :type  fct_desc: string
    :param fct_desc: Formula of the function

    :type  param_names: tuple of strings
    :param param_names: Name of the various parameters

    :type  residuals: callable
    :param residuals: Residual function

    :type  res_desc: string
    :param res_desc: Description of the residuals

    :rtype: :py:class:`ResultStruct`
    :returns: Data structure summarising the fitting and its evaluation
    """
    popt = fit.popt
    pcov = fit.pcov
    res = fit.res
    infodict = fit.infodict

    if CI:
        CIs = CIresults.CIs
    else:
        CIs = []


    yopts = fit(xdata)
    if eval_points is None:
        yvals = yopts
        eval_points = xdata
    else:
        yvals = fit(eval_points)

    scaled_res, res_IX, prob, normq = residual_measures(res)
    sorted_yopts = yopts[res_IX]

    result = {}
    result["fct"] = fit
    result["fct_desc"] = fct_desc
    result["param_names"] = param_names
    result["xdata"] = xdata
    result["ydata"] = ydata
    result["xname"] = xname
    result["yname"] = yname
    result["res_name"] = res_name
    result["residuals"] = residuals
    #result["args"] = fit.args
    result["popt"] = popt
    result["res"] = res
    result["yopts"] = yopts
    result["eval_points"] = eval_points
    result["interpolation"] = yvals
    result["sorted_yopts"] = sorted_yopts
    result["scaled_res"] = scaled_res
    result["normq"] = normq
    result["CI"] = CI
    result["CIs"] = CIs
    #result["CIparams"] = CIparams
    result["CIresults"] = CIresults
    #print("estimate jacobian = %s" % result["extra_output"][-1]["est_jacobian"])
    return ResultStruct(**result)

ResidualMeasures = namedtuple("ResidualMeasures", "scaled_res res_IX prob normq")

def residual_measures(res):
    """
    Compute quantities needed to evaluate the quality of the estimation, based solely on the residuals.

    :rtype: :py:class:`ResidualMeasures`
    :returns: the scaled residuals, their ordering, the theoretical quantile for each residuals, and the expected value
        for each quantile.
    """
    IX = argsort(res)
    scaled_res = res[IX]/std(res)

    prob = (arange(len(scaled_res))+0.5) / len(scaled_res)
    normq = sqrt(2)*erfinv(2*prob-1);

    return ResidualMeasures(scaled_res, IX, prob, normq)

_restestfields = "res_figure residuals scaled_residuals qqplot dist_residuals"
ResTestResult = namedtuple("ResTestResult", _restestfields)
Plot1dResult = namedtuple("Plot1dResult", "figure estimate data CIs " + _restestfields)

def plot1d(result, loc=0, fig = None, res_fig = None):
    """
    Use matplotlib to display the result of a fit, and return the list of plots used

    :rtype: :py:class:`Plot1dResult`
    :returns: hangles to the various figures and plots
    """
    if fig is None:
        fig = figure()
    else:
        try:
            figure(fig)
        except TypeError:
            figure(fig.number)

    p_est = plot(result.eval_points, result.interpolation, label='estimated')[0]
    p_data = plot(result.xdata, result.ydata, '+', label='data')[0]
    p_CIs = []
    if result.CI:
        for p, (low, high) in izip(result.CI,result.CIs[0]):
            l = plot(result.eval_points, low, '--', label='%g%% CI' % (p,))[0]
            h = plot(result.eval_points, high, l.get_color()+'--')[0]
            p_CIs += [l,h]
    if result.param_names:
        param_strs = ", ".join("%s=%g" % (n,v) for n,v in izip(result.param_names, result.popt))
    else:
        param_strs = ", ".join("%g" % v for v in result.popt)
    param_strs = "$%s$" % (param_strs,)

    title("Estimated function %s with params %s" % (result.fct_desc, param_strs))

    xlabel(result.xname)
    ylabel(result.yname)
    legend(loc=loc)

    plots = {"figure": fig, "estimate": p_est, "data": p_data, "CIs": p_CIs}

    prt = plot_residual_tests(
            result.xdata, result.yopts, result.res,
            "{0} with params {1}".format(result.fct_desc, param_strs),
            result.xname, result.yname, result.res_name, result.sorted_yopts, result.scaled_res,
            result.normq, res_fig)

    plots.update(prt._asdict())

    return Plot1dResult(**plots)

def plot_residual_tests(xdata, yopts, res, fct_name, xname = "X", yname = 'Y', res_name = "residuals",
        sorted_yopts = None, scaled_res = None, normq = None, fig = None):
    """
    Plot, in a single figure, all four residuals evaluation plots: :py:func:`plot_residuals`,
    :py:func:`plot_dist_residuals`, :py:func:`scaled_location_plot` and :py:func:`qqplot`.

    :param ndarray xdata:        Explaining variables
    :param ndarray yopt:         Optimized explained variables
    :param str     fct_name:     Name of the fitted function
    :param str     xname:        Name of the explaining variables
    :param str     yname:        Name of the dependant variables
    :param str     res_name:     Name of the residuals
    :param ndarray sorted_yopts: ``yopt``, sorted to match the scaled residuals
    :param ndarray scaled_res:   Scaled residuals
    :param ndarray normq:        Estimated value of the quantiles for a normal distribution

    :type  fig: handle or None
    :param fig: Handle of the figure to put the plots in, or None to create a new figure

    :rtype: :py:class:`ResTestResult`
    :returns: The handles to all the plots
    """
    if fig is None:
        fig = figure()
    else:
        try:
            figure(fig)
        except TypeError:
            figure(fig.number)

    plot1 = subplot(2,2,1)
# First subplot is the residuals
    if len(xdata.shape) == 1 or xdata.shape[1] == 1:
        p_res = plot_residuals(xname, xdata.squeeze(), res_name, res)
    else:
        p_res = plot_residuals(yname, yopts, res_name, res)

    if scaled_res is None or sorted_yopts is None or normq is None:
        scaled_res, res_IX, _, normq = residual_measures(res)
        sorted_yopts = yopts[res_IX]

    plot2 = subplot(2,2,2)
    p_scaled = scaled_location_plot(yname, sorted_yopts, scaled_res)

    subplot(2,2,3)
# Q-Q plot
    qqp = qqplot(scaled_res, normq)

    subplot(2,2,4)
# Distribution of residuals
    drp = plot_dist_residuals(res)

    suptitle("Residual Test for {}".format(fct_name))

    return ResTestResult(fig, p_res, p_scaled, qqp, drp)

def write1d(outfile, result, res_desc, CImethod):
    """
    Write the result of a fitting and its evaluation to a CSV file.

    :param str          outfile:  Name of the file to write to
    :param ResultStruct result:   Result of the fitting evaluation (e.g. output of :py:func:`fit_evaluation`)
    :param str          res_desc: Description of the residuals (in more details than just the name of the residuals)
    :param str          CImethod: Description of the confidence interval estimation method
    """
    with open(outfile, CSV_WRITE_FLAGS) as f:
        w = csv_writer(f)
        w.writerow(["Function",result.fct.fct.description])
        w.writerow(["Residuals",result.res_name,res_desc])
        w.writerow(["Parameter","Value"])
        for pn, pv in izip(result.param_names, result.popt):
            w.writerow([pn, "%.20g" % pv])
        #TODO w.writerow(["Regression Evaluation"])
        w.writerow([])
        w.writerow(["Data"])
        w.writerow([result.xname, result.yname, result.fct_desc, "Residuals: %s" % result.res_name])
        w.writerows(c_[result.xdata, result.ydata, result.yopts, result.res])
        w.writerow([])
        w.writerow(['Model validation'])
        w.writerow([result.yname, 'Normalized residuals', 'Theoretical quantiles'])
        w.writerows(c_[result.sorted_yopts, result.scaled_res, result.normq])
        if result.eval_points is not result.xdata:
            w.writerow([])
            w.writerow(["Interpolated data"])
            w.writerow([result.yname, result.yname])
            w.writerows(c_[result.eval_points, result.interpolation])
        if result.CI:
            w.writerow([])
            w.writerow(["Confidence interval"])
            w.writerow(["Method",CImethod])
            head = ["Parameters"] + list(chain(*[["%g%% - low" % v, "%g%% - high" % v] for v in result.CI]))
            w.writerow(head)
            #print(result.CIs[1])
            for cis in izip(result.param_names, *chain(*result.CIs[1])):
                cistr = [cis[0]] + ["%.20g" % v for v in cis[1:]]
                w.writerow(cistr)
            w.writerow([result.yname])
            head[0] = result.xname
            w.writerow(head)
            w.writerows(c_[tuple(chain([result.eval_points], *result.CIs[0]))])

def test():
    import residuals
    from numpy.random import rand, randn
    from pylab import plot, savefig, clf, legend, arange, figure, title, show
    from curve_fit import curve_fit

    def test(x,params):
        p0, p1, p2 = params
        return p0 + p1*x + p2*x**2

    init = (10,1,1)
    target = (10,4,1.2)
    print("Target parameters: %s" % (target,))
    x = 6*rand(200) - 3
    y = test(x, target)*(1+0.2*randn(x.shape[0]))
    xr = arange(-3, 3, 0.01)
    yr = test(xr,target)

    res = residuals.get('Log residual')

    result = plot_fit(test, x, y, init, eval_points=xr,
                      param_names=("p_0", "p_1", "p_2"), CI=(95,99), fct_desc="$y = p_0 + p_1 x + p_2 x^2$",
                      loc='upper left', fit_kwrds={"residuals":res}, shuffle_args={"add_residual":res.invert},
                      res_desc=res.name)

    result = plot_fit(test, x, y, init, eval_points=xr, shuffle_method=bootstrap.bootstrap_regression,
                      param_names=("p_0", "p_1", "p_2"), CI=(95,99), fct_desc="$y = p_0 + p_1 x + p_2 x^2$",
                      loc='upper left', fit_kwrds={"residuals":res},
                      res_desc=res.name)

    show()
    return locals()

if __name__ == "__main__":
    test()

# /home/barbier/prog/python/curve_fitting/test.csv
