#!/usr/bin/python
# bodenstein.py
# by Walter Dal'Maz Silva
# 18th October 2016 (version 0.1.0)

__version__ = '0.1.0'

DEBUG_MODULE = False

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------

import os
import numpy as np
import pandas as pd
from datetime import date
import matplotlib.pyplot as plt
import xutils.xpython as _xu
from .aux._distributions import bodenstein, johnson_SB

# -----------------------------------------------------------------------------
# Auxiliary functions
# -----------------------------------------------------------------------------

def _plot_distribution(filename, xdata, ydata, yfit, **kwargs):
    save   = _xu._get_kwarg('save',   kwargs, False)
    folder = _xu._get_kwarg('folder', kwargs, '.')

    plt.clf()
    plt.figure(figsize=(6,3))
    plt.xlabel('$\\theta$')
    plt.ylabel('E($\\theta$)')
    #plt.xlim(np.min(xdata), np.max(xdata))
    plt.ylim(np.min(ydata), 1.05*np.max(ydata))
    plt.plot(xdata, ydata, 'b-')
    plt.plot(xdata, yfit,  'r-')
    plt.legend(['data', 'fit'])
    if save:
        figname = os.path.basename(filename).split('.')[0]
        figname = os.path.join(folder, figname)
        plt.savefig(figname, dpi=400)
    else:
        plt.show()

# -----------------------------------------------------------------------------
# Fit functions
# -----------------------------------------------------------------------------

def fit_johnson(filename, **kwargs):
    """ Fit a Johnson_SB distribution function to data.

        filename : csv or text file containing data

        returns : time axis, fit function and fit parameters
    """
    if DEBUG_MODULE: _xu._print_head(fit_johnson.__name__, kwargs)

    # Parse arguments from kwargs
    delimiter = _xu._get_kwarg('delimiter', kwargs, ',')
    guess     = _xu._get_kwarg('guess',     kwargs, (1.4,991,5.4,2.6))
    maxfev    = _xu._get_kwarg('maxfev',    kwargs, 10000)
    tlabel    = _xu._get_kwarg('tlabel',    kwargs, 't(s)')
    flabel    = _xu._get_kwarg('flabel',    kwargs, 'E(s)')
    tzero     = _xu._get_kwarg('tzero',     kwargs, 0.0)
    plotfit   = _xu._get_kwarg('plotfit',   kwargs, False)


    # Read data file 
    # TODO test here
    data = pd.read_csv(filename, header=0, delimiter=delimiter)


    # Get axis
    ctime = data[tlabel]
    cfunc = data[flabel]


    # Reposition axis
    rcond = (data[tlabel] >= tzero)
    ctime = data.loc[rcond, (tlabel)]
    cfunc = data.loc[rcond, (flabel)]
    ctime -= min(ctime)
    cfunc -= min(cfunc)


    # Optimize function
    try:
        src = filename
        mod = johnson_SB
        fit = _xu._fit_data(src, mod, ctime, cfunc, guess, maxfev=maxfev)
        popt, pcov, res = fit
    except RuntimeError:
        return None


    fitfmt = 'p0=[{:.16e}, {:.16e},\n\t{:.16e}, {:.16e}]'.format(*popt)


    if DEBUG_MODULE and plotfit:
        _plot_distribution(filename, ctime, cfunc, res)


    return ctime, res, fitfmt



def fit_bodenstein(filename, **kwargs):
    """ Fit a Bodenstein distribution function to data.

        filename : csv or text file containing data
    """
    if DEBUG_MODULE: _xu._print_head(fit_bodenstein.__name__, kwargs)

    # Parse arguments from kwargs
    delimiter = _xu._get_kwarg('delimiter', kwargs, ',')
    guess     = _xu._get_kwarg('guess',     kwargs, 10.0)
    maxfev    = _xu._get_kwarg('maxfev',    kwargs, 100000)
    tlabel    = _xu._get_kwarg('tlabel',    kwargs, 't(s)')
    flabel    = _xu._get_kwarg('flabel',    kwargs, 'E(s)')
    tzero     = _xu._get_kwarg('tzero',     kwargs, 0.0)
    plotfit   = _xu._get_kwarg('plotfit',   kwargs, False)

    print('\n  Dealing with', filename)

    # Find a Johnson fit for the peak
    conf_dict = {
        'delimiter' : delimiter,
        'tlabel'    : tlabel,
        'flabel'    : flabel,
        'tzero'     : tzero,
        'plotfit'   : plotfit
    }

    fit = fit_johnson(filename, **conf_dict)
    if fit is None:
        print('  Failed to fit Johnson_SB', filename)
        return

    # Unpack fit return
    ctime, cfunc, fitpars = fit

    # Transform into reduced time
    ctheta = np.array(ctime/np.trapz(ctime*cfunc, ctime))
    ftheta = np.array(ctime*cfunc)


    # Crop zero for fitting
    ctheta = ctheta[np.where(ctheta > 0)]
    ftheta = ftheta[np.where(ctheta > 0)]


    # Optimize function
    try:
        src = filename
        mod = bodenstein
        fit = _xu._fit_data(src, mod, ctheta, ftheta, guess, maxfev=maxfev)
        popt, pcov, res = fit
    except RuntimeError:
        print('  Error: %s, %s' % (fit_bodenstein.__name__,filename))
        return


    # Open report
    try:
        funcname= fit_bodenstein.__name__
        outfolder = _xu._get_outfolder('results-' + funcname)
        namefmt   = funcname  + '-%Y-%m-%d.txt'
        outfile = str(date.today().strftime(namefmt))
        outfile = os.path.join(outfolder, outfile)

        report = open(outfile, 'a')
        report.write(filename)
        report.write('\nBodenstein number = %.3f' % popt[0])
        report.write('\nIntegral          = %.3f' % (np.trapz(res, ctheta)))
        report.write('\nJohnson parameters: \n\t' + fitpars)
        report.write('\n\n\n')
    except IOError:
        print('  Cannot write report')
        pass
    finally:
        report.close()

    if plotfit:
        plot = {
            'save'   : True,
            'folder' : outfolder
        }
        _plot_distribution(filename, ctheta, ftheta, res, **plot)

# -----------------------------------------------------------------------------
# EOF
# -----------------------------------------------------------------------------
