#!/usr/bin/env python
# -*- coding: utf-8 -*-
from os import path
import numpy as np
import pandas as pd
from scipy import stats
from joblib import Parallel, delayed
from tqdm import tqdm
from . import utils, afni, io


def estimate_acf(errts_file, mask_file, X_file=None, parts=None, censors=None):
    '''
    Estimate spatial auto-correlation function of fMRI volumes from GLM errts.
    Use AFNI's 3dFWHMx and ACF model: 
    ACF(r) = a * exp(-r*r/(2*b*b)) + (1-a)*exp(-r/c)
    
    The errts file can be too large to fit into memory all at once. It is 
    recommended to split the computation into several pieces, e.g., by parts=4, 
    or by runs inferred from an X_file (design matrix file), and average the 
    ACF parameters afterwards (geometric mean by default following 3dFWHMx).
    
    Sometimes you may only want to consider certain volumes in the errts file, 
    e.g., those actually used in 3dDeconvolve. Specify which TRs to keep using 
    `censors` argument, or it can also be inferred from an X_file.
    
    Parameters
    ----------
    errts_file : str
        errts file (generated by 3dDeconvolve, e.g., errts.prefix_REML.nii)
    mask_file : str
        Mask file within which ACF is estimated. Must not be too small.
    X_file : str
        Design matrix file (generated by 3dDeconvolve, e.g., X.prefix.1D) 
        from which runs and censors info can be extracted.
    parts : int, or list-of-int
        Specify how many parts would you like to split the computation into, e.g., 4.
        Or you can also specify the first volume of each part, e.g., [0, 100, 200, 300].
        This will overwrite runs info from X_file.
    censors : array-like
        Boolean array, True for good volumes.
        This will overwrite censors info from X_file.
        
    Returns
    -------
    acf : array of 3 floats
        The a, b, c parameters of 3dFWHMx's ACF model.
    fwhm : float
        The effective fwhm based on the ACF model.
    '''
    def parse_params(output):
        for k in range(len(output)):
            if output[k].startswith('# ACF model parameters'):
                params = np.float_(output[k+1].split())
        return params
    def run_3dFWHMx(sel):
        sel_cmd = '' if sel is None else f"'[{afni.get_subbrick_selector(sel)}]'"
        # res = utils.run(f"3dFWHMx -ACF 3dFWHMx.1D -mask {mask_file} {errts_file}{sel_cmd} -overwrite")
        mask_cmd = '' if mask_file=='none' else f"-mask {mask_file}"
        res = utils.run(f"3dFWHMx -ACF NULL {mask_cmd} {errts_file}{sel_cmd} -overwrite")
        return parse_params(res['output'])
    if censors is not None:
        sel = censors
    elif X_file is not None:
        sel = afni.get_censor_from_X(X_file)
    else:
        sel = None
    if parts is None and X_file is None:
        params = run_3dFWHMx(sel)
    else:
        if sel is None:
            sel = True
        if parts is None: # X_file will not be None
            groups = afni.get_runs_from_X(X_file, fmt='groups')
        else:
            N = afni.get_dims(errts_file)[-1]
            if np.iterable(parts): # parts is run_starts
                starts = parts
            elif parts < 0: # parts is -batch_size
                starts = np.arange(0, N, int(np.ceil(N/np.floor(N/-parts))))
            else: # parts is n_parts
                starts = np.arange(0, N, int(np.ceil(N/parts)))
            groups = np.concatenate([k*np.ones(int(d), dtype=int) for k, d in enumerate(np.diff(np.r_[starts, N]))])
        weights = np.array([np.sum(groups==g) for g in np.unique(groups)])[:,np.newaxis]
        params = np.array([run_3dFWHMx(sel&(groups==g)) for g in np.unique(groups)])
        # params = np.sum(weights*params, axis=0)/weights.sum()
        # 3dFWHMx uses geometric mean by default, for no good reason
        params = np.exp(np.sum(weights*np.log(params), axis=0)/weights.sum())
    acf, fwhm = params[:3], params[3]
    return acf, fwhm


def generate_simulated_data(acf, mask_file, out_dir, n=10000):
    '''
    '''
    mask_cmd = '' if mask_file=='none' else f"-mask {mask_file}"
    utils.run(f"3dClustSim -acf {' '.join([f'{p:.6f}' for p in acf])} \
        -iter {n} {mask_cmd} \
        -cmd {out_dir}/3dClustSim.ACF.cmd \
        -both -prefix {out_dir}/3dClustSim.ACF \
        -ssave:masked {out_dir}/ACFSim.nii.gz")


def voxel_corr_MonteCarlo_test(x, y, mask, errts_x, mask_sim=None, n=10000, r=10, 
    loader=None, estimator=None, n_jobs=1, visualize=False):
    '''
    Parameters
    ----------
    mask : str
        Filename for voxel pattern mask, used to extract x and y.
    errts_x : str
        Filename for errts corresponding to x dataset.
    mask_sim : str
        Filename for simulation mask, used to generate simulated noise dataset.
    stat : callable(x, y)
        The statistics function, corrcoef by default.
    '''
    # Compute observed statistics
    if estimator is None: # Pearson's corrcoef
        estimator = lambda x, y: np.corrcoef(x, y)[0,1]
    obs = estimator(x, y)
    # Generate simulated data
    if path.isdir(errts_x):
        # The simulated data is precomputed, as "{errts_x}/ACFSim_000000.nii.gz"
        sim_dir = errts_x
    else:
        # Estimate ACF (spatial smoothness of the real fMRI data)
        acf = estimate_acf(errts_x, mask_sim, parts=-100)[0]
        # Generate smoothed Gaussian noise data
        sim_dir = utils.temp_folder()
        generate_simulated_data(acf, mask_sim, sim_dir, n=n)
    # Compute simulated statistics
    if loader is None:
        mask = io.Mask(mask)
        loader = lambda sim_dir, k: mask.dump(f"{sim_dir}/ACFSim_{k:06d}.nii.gz")
    func = lambda sim_dir, k, y: estimator(loader(sim_dir, k), y)
    if n_jobs == 1:
        sim = [func(sim_dir, k, y) for k in tqdm(range(n), position=0)]
    else:
        sim = Parallel(n_jobs=n_jobs)(delayed(func)(sim_dir, k, y) for k in tqdm(range(n), position=0))
    # Compute p-value
    p_sim = stats.percentileofscore(sim, obs)
    p_sim = min(p_sim, 100-p_sim)/100*2 # Two-tailed p-value
    # Visualize
    if visualize:
        import matplotlib.pyplot as plt
        import seaborn as sns
        sns.kdeplot(sim, fill=True)
        plt.axvline(obs)
        plt.title(f"$p_{{sim}}=${p_sim:.3f}" if p_sim>=1e-3 else f"$p_{{sim}}<$0.001")
    return p_sim, sim



if __name__ == '__main__':
    pass