import numpy as np 
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed
import mne
import scipy
from scipy.signal import *
from scipy.stats import gamma
import scipy.special as sc

def reorder(ori, shuffled):
    return np.where(shuffled.reshape(-1, 1) == ori.reshape(1, -1))[1]


def parallel_process(array, function, n_jobs=16, use_kwargs=False, front_num=3):
    """
        A parallel version of the map function with a progress bar. 

        Args:
            array (array-like): An array to iterate over.
            function (function): A python function to apply to the elements of array
            n_jobs (int, default=16): The number of cores to use
            use_kwargs (boolean, default=False): Whether to consider the elements of array as dictionaries of 
                keyword arguments to function 
            front_num (int, default=3): The number of iterations to run serially before kicking off the parallel job. 
                Useful for catching bugs
        Returns:
            [function(array[0]), function(array[1]), ...]
    """
    #We run the first few iterations serially to catch bugs
    if front_num > 0:
        front = [function(**a) if use_kwargs else function(a) for a in array[:front_num]]
    #If we set n_jobs to 1, just run a list comprehension. This is useful for benchmarking and debugging.
    if n_jobs==1:
        return front + [function(**a) if use_kwargs else function(a) for a in tqdm(array[front_num:])]
    #Assemble the workers
    with ProcessPoolExecutor(max_workers=n_jobs) as pool:
        #Pass the elements of array into function
        if use_kwargs:
            futures = [pool.submit(function, **a) for a in array[front_num:]]
        else:
            futures = [pool.submit(function, a) for a in array[front_num:]]
        kwargs = {
            'total': len(futures),
            'unit': 'it',
            'unit_scale': True,
            'leave': True
        }
        #Print out the progress as tasks complete
        for f in tqdm(as_completed(futures), **kwargs):
            pass
    out = []
    #Get the results from the futures. 
    for i, future in enumerate(futures):
        try:
            out.append(future.result())
        except Exception as e:
            out.append(e)
    return front + out


def read_raw(raw_path):
    """Read raw data from an EDF file
    
    Parameters
    ----------
    raw_path : str
        Path to the EDF file
    
    Returns
    -------
    data : numpy array
    channels : numpy array"""
    raw_path = validate_type(raw_path, 'raw_path', str)
    raw = mne.io.read_raw_edf(raw_path)
    raw_channels = raw.info['ch_names']
    data = []

    for raw_ch in raw_channels:
        ch_data = raw.get_data(raw_ch) * 1E6
        data.append(ch_data)

    data = np.squeeze(data)
    channels = np.array(raw_channels)
    return data, channels


def preprocess(raw, sample_freq, filter_freq, rp=0.5, rs=93, space=0.5):
    """Preprocess the raw data
    Parameters
    ----------
    raw : numpy array like of float, shape (n_samples,)
        raw data
    sample_freq : float
        sample frequency
    filter_freq : list of float, shape (2,)
        filter frequency
    rp : float | 0.5
        max loss in passband (dB)
    rs : float | 93
        min attenuation in stopband (dB)
    space : float

    Returns
    -------
    filtered : array
        filtered data
    """
    # check raw data
    raw = validate_type(raw, 'raw', np.ndarray)
    if raw.ndim != 1:
        raise ValueError('raw must be a 1D array')
    # check sample frequency
    sample_freq = validate_param(sample_freq, 'sample_freq', (int, float))
    # check filter frequency
    filter_freq = validate_filter_freq(filter_freq)

    nyq = sample_freq / 2
    # MGNM
    if filter_freq[1] >= .99 * nyq:
        filter_freq[1] = nyq * .99
    low, high = filter_freq

    scale = 0
    while 0 < low < 1:
        low *= 10
        scale += 1
    low = filter_freq[0] - (space * 10 ** (-1 * scale))

    scale = 0
    while high < 1:
        high *= 10
        scale += 1
    high = filter_freq[1] + (space * 10 ** (-1 * scale))  
    stop_freq = [low, high]
    ps_order, wst = cheb2ord([filter_freq[0] / nyq, filter_freq[1] / nyq], [stop_freq[0] / nyq, stop_freq[1] / nyq], rp,
                                rs)
    z, p, k = cheby2(ps_order, rs, wst, btype='bandpass', analog=0, output='zpk')
    sos = zpk2sos(z, p, k)    

    filtered = sosfilt(sos, raw)
    filtered = sosfilt(sos, np.flipud(filtered))
    filtered = np.flipud(filtered)

    return filtered


def validate_type(var, var_name, var_type):
    """Validate the type of a variable"""
    if not isinstance(var, var_type):
        raise TypeError('{} must be a {}'.format(var_name, var_type))
    return var


def validate_param(var, var_name, var_type):
    """Validate the type and value of a variable"""
    var = validate_type(var, var_name, var_type)
    if var < 0:
        raise ValueError('{} must not be a negative {}'.format(var_name, var_type))
    return var


def validate_filter_freq(filter_freq):
    """Validate filter frequency
    Parameters
    ----------
    filter_freq : array
        filter frequency
    Returns
    -------
    filter_freq : array
        filter frequency
    """
    filter_freq = validate_type(filter_freq, 'filter_freq', (list, tuple, np.ndarray))
    if len(filter_freq) != 2:
        raise ValueError('filter_freq must have two elements')
    if not all(isinstance(f, (int, float)) for f in filter_freq):
        raise TypeError('filter_freq must be a list of int or float')   
    if filter_freq[0] >= filter_freq[1]:
        raise ValueError('filter_freq[0] must be smaller than filter_freq[1]')
    if filter_freq[0] <= 0:
        raise ValueError('filter_freq[0] must be positive')
    return filter_freq


def compute_rms(filtered, sample_freq, window_size=6*1e-3, detector=None):
    """Compute RMS
    Parameters
    ----------
    filtered : numpy array like of float, shape (n_samples,)
        filtered data
    sample_freq : float | int
        sample frequency
    window_size : float | int
        window size
    detector : str
        detector type
    Returns
    -------
    rms : array of float, shape (n_samples,)
        RMS
    """
    # check filtered data
    filtered = validate_type(filtered, 'filtered', np.ndarray)
    if filtered.ndim != 1:
        raise ValueError('filtered must be a 1D array')
    # check sample frequency
    sample_freq = validate_param(sample_freq, 'sample_freq', (int, float))
    # check window size 
    window_size = validate_param(window_size, 'window_size', (int, float))
    # check detector
    if detector is not None:
        detector = validate_type(detector, 'detector', str)
        if detector not in ['RMS', 'MNI']:
            raise ValueError('detector must be RMS or MNI')

    # compute RMS
    if detector == 'RMS':
        rms_window = round(window_size * sample_freq)
    elif detector == 'MNI':
        rms_window = round(0.002 * sample_freq)
    
    if rms_window % 2 == 0:
        rms_window += 1
    
    temp = np.square(filtered)
    temp = lfilter(np.ones(rms_window), 1, temp, axis=0) / rms_window
    rms = np.zeros(len(temp))
    rms[:int(-np.floor(rms_window / 2))] = temp[int(np.floor(rms_window / 2)):]
    rms = np.sqrt(rms)

    return rms


def compute_wavelet_entropy(sample_freq, epoch_samples, low, high, seed, dev_cycles=3):
    """Compute wavelet entropy
    Parameters
    ----------
    sample_freq : int | float
        sample frequency
    epoch_samples : int | float
        epoch samples
    low : int | float
        low frequency
    high : int | float
        high frequency
    dev_cycles : int | float | 3 | optional
        number of cycles
    seed: int | None
        seed for random number generator
    Returns
    -------
    wavelet_entropy : int | float
        wavelet entropy"""
    #check inputs
    sample_freq = validate_param(sample_freq, 'sample_freq', (int, float))
    epoch_samples = validate_param(epoch_samples, 'epoch_samples', (int, float))
    low = validate_param(low, 'low', (int, float))
    high = validate_param(high, 'high', (int, float))
    dev_cycles = validate_param(dev_cycles, 'dev_cycles', (int, float))
    
    np.random.seed(seed)
    wavelet_entropy_max = np.zeros((100,1))
    for i in range(0, 100):
        segment = np.random.rand(epoch_samples, 1)
        auto_corr = correlate(segment, segment, mode='full') / np.sum(segment ** 2)
        w_coef = gabor_transform_wait(auto_corr, sample_freq, low, high, dev_cycles)
        prob_energy = np.mean(w_coef ** 2, 1) / np.sum(np.mean(w_coef ** 2, 1))
        wavelet_entropy_max[i] = -np.sum(prob_energy * np.log(prob_energy))

    del segment, prob_energy, w_coef

    return np.median(wavelet_entropy_max)


def gabor_transform_wait(pv_signal, sample_rate, min_freq, max_freq, dev_cycles=3):
    """This function calculates the Wavelet Transform using a Gaussian modulated window (Gabor Wavelet).
    Parameters
    ----------
    pv_signal : array
        Signal to process
    sample_rate : float | int
        sample rate
    min_freq : float | int
        Min frequency (in Hz) to process from
    max_freq : float | int
        Max frequency (in Hz) to process to
    dev_cycles : float | int
         the number of cycles to include for the wavelet transform at every frequency.
   
    Returns
    -------
    w_coef : array
        wavelet transform output matrix (n_freqs x n_samples) Time in rows, Frequency in columns
    """
    # check inputs
    pv_signal = validate_type(pv_signal, 'pv_signal', np.ndarray)
    sample_rate = validate_param(sample_rate, 'sample_rate', (int, float))
    min_freq = validate_param(min_freq, 'min_freq', (int, float))
    max_freq = validate_param(max_freq, 'max_freq', (int, float))
    dev_cycles = validate_param(dev_cycles, 'dev_cycles', (int, float))
    pv_signal = np.reshape(pv_signal[:], (-1))
    freq_seg = len(list(range(min_freq, max_freq, 5)))
    freq_step = int((max_freq - min_freq) / (freq_seg - 1))
    freq_axis = np.arange(min_freq, max_freq + freq_step, freq_step)
    freq_axis = np.flip(freq_axis)
    if len(pv_signal) % 2 == 0:
        pv_signal = pv_signal[0:len(pv_signal) - 1]
    
    time_axis = (np.arange(0, len(pv_signal))) / sample_rate
    s_len = len(time_axis)
    s_half_len = (s_len // 2) + 1

    w_axis = (2 * np.pi / s_len) * (np.arange(0, s_len))
    w_axis = w_axis * sample_rate
    w_axis_half = w_axis[:s_half_len]

    dev_sec = ((1 / freq_axis) * dev_cycles).reshape((-1, 1))
    tmp1 = (np.repeat(w_axis_half[None,:], len(freq_axis), 0) - 2 * np.pi * freq_axis[:,None]) ** 2
    tmp = np.exp(-0.5 * tmp1 * (dev_sec ** 2))
    win_fft = np.zeros((len(freq_axis), s_len))
    win_fft[:, :s_half_len] = tmp
    win_fft = win_fft * np.sqrt(s_len) / np.linalg.norm(win_fft, axis=1)[:, None] #21 *999
    signal_fft = scipy.fft.fft(pv_signal, len(pv_signal), axis=0)
    gabor_wt = scipy.fft.ifft(signal_fft * win_fft, axis=1) / np.sqrt(dev_sec)
    gabor_wt = abs(gabor_wt)

    return gabor_wt


def pos_baseline(baseline_window, rms, epoch_time, sample_freq, thr_perc):
    """Calculate baseline for positive peaks.

    Parameters
    ----------
    baseline_window : array
        baseline window
    rms : array
        rms
    epoch_time : int|float
        epoch time
    sample_freq : int|float
        sample frequency
    thr_perc : int|float
        threshold percentage
    Returns
    ------- 
    res : array
        baseline
    """
    # check inputs
    baseline_window = validate_type(baseline_window, 'baseline_window', np.ndarray)
    if baseline_window.ndim != 2 and baseline_window.shape[1] != 1:
        raise ValueError('baseline_window must be a 2D array with shape (n_samples, 1)')
    if not np.equal(baseline_window, 0).any() or not np.equal(baseline_window, 1).any():
        raise ValueError('baseline_window contains only 0 and 1')
    rms = validate_type(rms, 'rms', np.ndarray)
    epoch_time = validate_param(epoch_time, 'epoch_time', (int, float))
    sample_freq = validate_param(sample_freq, 'sample_freq', (int, float))
    thr_perc = validate_param(thr_perc, 'thr_perc', (int, float))

    # calculate baseline
    window_thrd = round(epoch_time * sample_freq)
    bw_tmp = np.zeros((len(baseline_window)+2, 1))
    bw_tmp[1:-1] = baseline_window
    win_baseline = np.diff(bw_tmp, axis=0)
    win_baseline_up = np.where(win_baseline == 1)[0] + 1
    win_baseline_down = np.where(win_baseline == -1)[0]
    
    idx_ini = []
    for i in range(len(win_baseline_up)):
        idx_add = list(range(win_baseline_up[i], win_baseline_down[i], window_thrd))
        idx_ini.extend(idx_add)
    idx_ini = np.array(idx_ini)
    idx_end = idx_ini + window_thrd - 1
    idx_rem = idx_end <= len(rms)
    idx_ini = idx_ini[idx_rem]
    idx_end = idx_end[idx_rem]       

    res = np.zeros(rms.shape)
    for i in range(len(idx_ini)):
        sect = np.sort(rms[int(idx_ini[i]) - 1:int(idx_end[i])])
        gam_params = fb_gamfit(sect)
        percent = fb_gamcdf(sect, gam_params[0], gam_params[1])
        p_idxs = np.where(percent <= thr_perc)
        pi = p_idxs[0][-1]
        res[int(idx_ini[i]) - 1:int(idx_end[i])] = sect[pi]

    return res


def neg_baseline(rms, min_win, sample_freq, epo_CHF, per_CHF):
    """Calculate baseline for negative peaks.

    Parameters
    ----------
    rms : array
    min_win : int
        minimum HFO Time window
    sample_freq : int|float 
        sample frequency
    epo_CHF : int|float
        Continuous High Frequency Epoch
    per_CHF : int|float
        Continuous High Frequency Percentile Threshold
    
    Returns
    -------
    res : array
        baseline
    """
    # check inputs
    rms = validate_type(rms, 'rms', np.ndarray)
    min_win = validate_param(min_win, 'min_win', (int, float))
    sample_freq = validate_param(sample_freq, 'sample_freq', (int, float))
    epo_CHF = validate_param(epo_CHF, 'epo_CHF', (int, float))
    per_CHF = validate_param(per_CHF, 'per_CHF', (int, float))

    # calculate baseline
    min_win = min_win * sample_freq
    window_chf = round(epo_CHF * sample_freq)
    idx_ini = np.arange(0, len(rms), window_chf)
    idx_end = idx_ini + window_chf - 1
    idx_rem = idx_end <= len(rms)
    idx_ini = idx_ini[idx_rem]
    idx_end = idx_end[idx_rem]

    res = np.zeros(rms.shape)
    for i in range(len(idx_ini)):
        sect = np.sort(rms[int(idx_ini[i]):int(idx_end[i]) + 1], axis=0)
        thres_last = np.max(sect)
        while True:
            if sum(abs(sect)) == 0:
                break
            # MLE of params of gamma distribution fit to the data(sect), return gam_params = (shape, scale)
            if np.all(sect != 0):
                gam_params = gamma.fit(sect, floc=0)
            else:
                gam_params = gamma_fit(sect)
            percent = gamma.cdf(sect, gam_params[0], gam_params[1], gam_params[2])
            p_idxs = np.where(percent <= per_CHF)
            thres_new = sect[p_idxs[0][-1]]
            energy_over = np.zeros(len(sect)+2)
            energy_over[1:-1] = sect >= thres_new
            jumps = np.diff(energy_over, axis=0)
            jump_up = np.where(jumps == 1)[0]
            jump_down = np.where(jumps == -1)[0] - 1
            dist = jump_down - jump_up

            select = np.where(dist > min_win)[0]
            if select.size == 0:
                break
            ini = jump_up[select]
            end = jump_down[select]

            for ii in range(len(ini)):
                sect[ini[ii]: end[ii] + 1] = 0
                sect = np.sort(sect,axis=0)

            thres_last = thres_new

        res[int(idx_ini[i]):int(idx_end[i]) + 1] = thres_last

    return res


def gampdf(x, a, b):
    """Gamma probability density function.

    Parameters
    ----------
    x : array
        input values
    a : int|float
        parameter of the gamma distribution
    b : int|float
        parameter of the gamma distribution

    Returns
    -------
    res : array
        the probability density function evaluated at x
    """
    # check inputs
    x = validate_type(x, 'x', np.ndarray)
    a = validate_type(a, 'a', (int, float))
    b = validate_type(b, 'b', (int, float))
    # calculate
    sz = len(x)
    pdf = np.zeros(sz)

    k = np.where(np.logical_or(a <= 0, b <= 0, np.isnan(x)))[0]
    if k.size != 0:
        pdf[k] = np.nan

    k = np.where(np.logical_and(x > 0, 0 < a <= 1 and b > 0))[0]
    if k.size != 0:
        if np.isscalar(a) and np.isscalar(b):
            pdf[k] = (x[k] ** (a-1)) * np.exp(-x[k] / b) / sc.gamma(a) / (b ** a)
        else:
            pdf[k] = (x[k] ** (a[k]-1)) * np.exp(-x[k] / b[k]) / sc.gamma(a[k]) / (b[k] ** a[k])

    k = np.where(np.logical_and(x > 0, a > 1 and b > 0))[0]
    if k.size != 0:
        if np.isscalar(a) and np.isscalar(b):
            pdf[k] = np.exp(-a * np.log(b) + (a-1) * np.log(x[k]) - x[k] / b - sc.gammaln(a))
        else:
            pdf[k] = np.exp(-a[k] * np.log(b[k]) + (a[k]-1) * np.log(x[k]) - x[k] / b[k] - sc.gammaln(a[k]))

    return pdf

            
def gamlike(p, x):
    """Negative log-likelihood function for the Gamma distribution over vector x, with the given parameters p[0] and p[1].

    Parameters
    ----------
    p : array
        parameters of the gamma distribution (shape, scale)
    x : array
        input values

    Returns
    -------
    res : array
        the negative log-likelihood function evaluated at x with the given parameters
    """
    # check inputs
    p = validate_type(p, 'p', (list, np.ndarray))
    if len(p) != 2:
        raise ValueError('p must be a list or array with 2 elements')
    p[0] = validate_type(p[0], 'p[0]', (int, float))
    p[1] = validate_type(p[1], 'p[1]', (int, float))
    x = validate_type(x, 'x', np.ndarray)

    with np.errstate(divide='ignore'): 
        return -np.sum(np.log(gampdf(x, p[0], p[1])))


def gamfit_search(a, avg, x):
    """Search for the best fit of gamma distribution.

    Parameters
    ----------
    a : int|float
    avg : int|float
    x : array

    Returns
    -------
    res : array
    """
    # check inputs
    a = validate_type(a, 'a', (int, float))
    avg = validate_type(avg, 'avg', (int, float))
    x = validate_type(x, 'x', np.ndarray)  
    # calculate
    return -gamlike([a, avg / a], x)


def nmsmax(func, x, stopit=[], savit=[], *args):
    """Nelder-Mead simplex method for direct search optimization.

    Parameters
    ----------
    func : function
        function to be optimized
    x : list|np.ndarray 
        initial guess
    stopit : list|np.ndarray 
        stopping criteria
    saveit : list|np.ndarray 
        saving criteria
    args : array|tuple|list
        additional arguments for the function

    Returns
    -------
    x[0] : array
        vector yielding largest function value found
    fmax : float   
        function value at x
    nf : int
        number of function evaluations
    """
    x = validate_type(x, 'x', (list, np.ndarray))
    stopit = validate_type(stopit, 'stopit', (list, np.ndarray))
    savit = validate_type(savit, 'savit', (list, np.ndarray))

    x0 = x[:]
    n = len(x0)

    stopit = [1e-3, np.inf, np.inf, 0, 0]
    tol, trace = stopit[0], stopit[4]

    V = np.hstack((np.zeros((n, 1)), np.eye(n)))
    f = np.zeros((n + 1, 1))
    V[:, 0] = x0
    f[0] = func(1, *args)
    fmax_old = f[0]

    k, m = 0, 0

    scale = np.max((np.linalg.norm(np.array(x), np.inf), 1))
    alpha = scale / (n * np.sqrt(2)) * np.array([np.sqrt(n + 1) - 1 + n, np.sqrt(n + 1) - 1])
    V[:, 1:n + 1] = x0 + alpha[1] * np.ones((n, 1)) * np.ones((1, n))
    for i in range(1, n + 1):
        V[i - 1, i] = x0[i - 1] + alpha[0]
        x[:] = V[:, i]
        f[i] = func(x[0], *args)
    nf = n + 1
    temp = np.sort(f)
    idx = np.where(temp)[0]
    idx = idx[::-1]
    f = f[idx]
    V = V[:, idx]

    alpha, beta, gamma = 1, 1 / 2, 2

    while True:
        k += 1
        fmax = f[0]
        if fmax > fmax_old:
            if len(savit) != 0:
                x[:] = V[:, 0]
        fmax_old = fmax
        # three stopping tests
        # Stopping Test 1 - f reached target value
        if fmax >= stopit[2]:
            break
        # Stopping Test 2 - too many f-evals
        if nf >= stopit[1]:
            break
        # Stopping Test 3 - converged
        v1 = V[:, 0]
        size_simplex = np.linalg.norm(V[:, 1:n + 1] - np.tile(v1, n), 1) / np.max((1, np.linalg.norm(v1, 1)))
        if size_simplex <= tol:
            break

        # One step of the Nelder-Mead simplex algorithm
        vbar = np.transpose(np.sum(np.transpose(V[:, 0:n]) / n))
        vr = (1 + alpha) * vbar - alpha * V[:, n]
        x[:] = vr
        fr = func(x[0], *args)
        nf += 1
        vk, fk = vr, fr
        if fr > f[n - 1]:
            if fr > f[0]:
                ve = gamma * vr + (1 - gamma) * vbar
                x[:] = ve
                fe = func(x[0], *args)
                nf += 1
                if fe > f[0]:
                    vk, fk = ve, fe
        else:
            vt = V[:, n]
            ft = f[n]
            if fr > ft:
                vt, ft = vr, fr
            vc = beta * vt + (1 - beta) * vbar
            x[:] = vc
            fc = func(x[0], *args)
            nf += 1
            if fc > f[n - 1]:
                vk, fk = vc, fc
            else:
                for j in range(1, n):
                    V[:, j] = (V[:, 0] + V[:, j]) / 2
                    x[:] = V[:, j]
                    f[j] = func(x[0], *args)
                nf += (n - 1)
                vk = (V[:, 0] + V[:, n]) / 2
                x[:] = vk
                fk = func(x[0], *args)
                nf += 1
        V[:, n] = vk
        f[n] = fk
        temp = np.sort(f)
        idx = np.where(temp)[0]
        idx = idx[::-1]
        f = f[idx]
        V = V[:, idx]

    x[:] = V[:, 0]
    return x[0], fmax[0], nf


def fb_gamfit(x):
    """the maximumlikelihood estimator for the Gamma distribution for x

    Parameters
    ----------
    x : array
        data

    Returns
    ------- 
    res : list of float, length 2
        [a, b] where a is the shape and b is the scale
    """
    # check inputs
    x = validate_type(x, 'x', np.ndarray)
    # fit gamma distribution
    avg = np.mean(x)
    a, _, _ = nmsmax(gamfit_search, [1], [], [], avg, x)
    b = a / avg
    res = [a, 1 / b]
    return res


def fb_gamcdf(x, a, b):
    """the cumulative distribution function for the Gamma distribution for x given a and b

    Parameters
    ----------
    x : array
        data
    a : float | int
        shape
    b : float | int
        scale

    Returns
    -------
    cdf : array
        cumulative distribution function values for x 
    """

    # check inputs
    x = validate_type(x, 'x', np.ndarray)
    a = validate_type(a, 'a', (int, float))
    b = validate_type(b, 'b', (int, float))

    # calculate cdf
    sz = len(x)
    cdf = np.zeros(sz)

    k = np.where(np.logical_or(a <= 0, b <= 0, np.isnan(x)))[0]
    if k.size != 0:
        cdf[k] = np.nan

    k = np.where(np.logical_and(x > 0, a > 0 and b > 0))[0]
    if k.size != 0:
        if np.isscalar(a) and np.isscalar(b):
            cdf[k] = sc.gammainc(a, x[k]/b)
        else:
            cdf[k] = sc.gammainc(a[k], x[k]/b[k])
    return cdf


def gamma_fit(x):
    n = len(x)
    freq = 1
    scalex = np.sum(freq*x) / n
    x = x / scalex
    xbar = 1
    s2 = np.sum(freq*(x-xbar)**2) / n
    s2 = s2 * n / (n - 1)
    ahat = xbar ** 2 / s2
    bhat = s2 / xbar
    parmhat = [ahat, 0, bhat*scalex]
    # parmhat[2] = parmhat[2] * scalex
    return parmhat





    
