# import itertools
# from functools import partial
# from scipy.stats import norm
# from scipy.sparse import csc_matrix, linalg as sla
# from scipy import sparse
# from scipy.optimize import minimize, minimize_scalar
# from collections import deque, namedtuple
import numpy as np
from networkx import Graph
from pygfl.solver import TrailSolver
from pygfl.trails import decompose_graph, save_chains
from pygfl.utils import chains_to_trails, calc_plateaus, hypercube_edges
from smoothfdr.smoothed_fdr import GaussianKnown
from smoothfdr.normix import *
from smoothfdr.utils import calc_fdr

def smooth_fdr(data, fdr_level, edges=None, initial_values=None, verbose=0, null_dist=None, signal_dist=None, num_sweeps=10, missing_val=None):
    flat_data = data.flatten()
    nonmissing_flat_data = flat_data

    if edges is None:
        if verbose:
            print('Using default edge set of a grid of same shape as the data: {0}'.format(data.shape))
        edges = hypercube_edges(data.shape)
        if missing_val is not None:
            if verbose:
                print('Removing all data points whose data value is {0}'.format(missing_val))
            edges = [(e1,e2) for (e1,e2) in edges if flat_data[e1] != missing_val and flat_data[e2] != missing_val]
            nonmissing_flat_data = flat_data[flat_data != missing_val]

    # Decompose the graph into trails
    g = Graph()
    g.add_edges_from(edges)
    chains = decompose_graph(g, heuristic='greedy')
    ntrails, trails, breakpoints, edges = chains_to_trails(chains)

    if null_dist is None:
        # empirical null estimation
        mu0, sigma0 = empirical_null(nonmissing_flat_data, verbose=max(0,verbose-1))
    elif isinstance(null_dist,GaussianKnown):
        mu0, sigma0 = null_dist.mean, null_dist.stdev
    else:
        mu0, sigma0 = null_dist
    null_dist = GaussianKnown(mu0, sigma0)

    if verbose:
        print('Empirical null: {0}'.format(null_dist))

    # signal distribution estimation
    if verbose:
        print('Running predictive recursion for {0} sweeps'.format(num_sweeps))
    if signal_dist is None:
        grid_x = np.linspace(max(-20, nonmissing_flat_data.min() - 1), min(nonmissing_flat_data.max() + 1, 20), 220)
        pr_results = predictive_recursion(nonmissing_flat_data, num_sweeps, grid_x, mu0=mu0, sig0=sigma0)
        signal_dist = GridDistribution(pr_results['grid_x'], pr_results['y_signal'])

    if verbose:
        print('Smoothing priors via solution path algorithm')

    solver = TrailSolver()
    solver.set_data(flat_data, edges, ntrails, trails, breakpoints)

    results = solution_path_smooth_fdr(flat_data, solver, null_dist, signal_dist, verbose=max(0, verbose-1))

    results['discoveries'] = calc_fdr(results['posteriors'], fdr_level)
    results['null_dist'] = null_dist
    results['signal_dist'] = signal_dist

    # Reshape everything back to the original data shape
    results['betas'] = results['betas'].reshape(data.shape)
    results['priors'] = results['priors'].reshape(data.shape)
    results['posteriors'] = results['posteriors'].reshape(data.shape)
    results['discoveries'] = results['discoveries'].reshape(data.shape)
    results['beta_iters'] = np.array([x.reshape(data.shape) for x in results['beta_iters']])
    results['prior_iters'] = np.array([x.reshape(data.shape) for x in results['prior_iters']])
    results['posterior_iters'] = np.array([x.reshape(data.shape) for x in results['posterior_iters']])

    return results

def smooth_fdr_known_dists(data, fdr_level, null_dist, signal_dist, edges=None, initial_values=None, verbose=0, missing_val=None):
    '''FDR smoothing where the null and alternative distributions are known
    (and not necessarily Gaussian). Both must define the function pdf.'''
    flat_data = data.flatten()
    nonmissing_flat_data = flat_data

    if edges is None:
        if verbose:
            print('Using default edge set of a grid of same shape as the data: {0}'.format(data.shape))
        edges = hypercube_edges(data.shape)
        if missing_val is not None:
            if verbose:
                print('Removing all data points whose data value is {0}'.format(missing_val))
            edges = [(e1,e2) for (e1,e2) in edges if flat_data[e1] != missing_val and flat_data[e2] != missing_val]
            nonmissing_flat_data = flat_data[flat_data != missing_val]

    # Decompose the graph into trails
    g = Graph()
    g.add_edges_from(edges)
    chains = decompose_graph(g, heuristic='greedy')
    ntrails, trails, breakpoints, edges = chains_to_trails(chains)

    if verbose:
        print('Smoothing priors via solution path algorithm')

    solver = TrailSolver()
    solver.set_data(flat_data, edges, ntrails, trails, breakpoints)

    results = solution_path_smooth_fdr(flat_data, solver, null_dist, signal_dist, verbose=max(0, verbose-1))

    results['discoveries'] = calc_fdr(results['posteriors'], fdr_level)
    results['null_dist'] = null_dist
    results['signal_dist'] = signal_dist

    # Reshape everything back to the original data shape
    results['betas'] = results['betas'].reshape(data.shape)
    results['priors'] = results['priors'].reshape(data.shape)
    results['posteriors'] = results['posteriors'].reshape(data.shape)
    results['discoveries'] = results['discoveries'].reshape(data.shape)
    results['beta_iters'] = np.array([x.reshape(data.shape) for x in results['beta_iters']])
    results['prior_iters'] = np.array([x.reshape(data.shape) for x in results['prior_iters']])
    results['posterior_iters'] = np.array([x.reshape(data.shape) for x in results['posterior_iters']])

    return results

def solution_path_smooth_fdr(data, solver, null_dist, signal_dist, min_lambda=0.20, max_lambda=1.5, lambda_bins=30, verbose=0, initial_values=None):
        '''Follows the solution path of the generalized lasso to find the best lambda value.'''
        lambda_grid = np.exp(np.linspace(np.log(max_l
ambda), np.log(min_lambda), lambda_bins))
        aic_trace = np.ze
ros(lambda_grid.shape) # The AIC score for each lambda value
        aicc_trace = np.zeros(lambda_grid.shape) # The AICc score for each lambda value (correcting for finite sample size)
        bic_trace = np.zeros(lambda_grid.shape) # The BIC score for each lambda value
        dof_trace = np.zeros(lambda_grid.shape) # The degrees of freedom of each final solution
        log_likelihood_trace = np.zeros(lambda_grid.shape)
        beta_trace = []
        u_trace = []
        w_trace = []
        c_trace = []
        results_trace = []
        best_idx = None
        best_plateaus = None
        for i, _lambda in enumerate(lambda_grid):
            if verbose:
                print('#{0} Lambda = {1}'.format(i, _lambda))

            # Fit to the final values
            results = fixed_penalty_smooth_fdr(data, solver, _lambda, null_dist, signal_dist,
                                               verbose=max(0,verbose - 1),
                                               initial_values=initial_values)

            if verbose:
                print('Calculating degrees of freedom')

            plateaus = calc_plateaus(results['beta'], solver.edges)
            dof_trace[i] = len(plateaus)

            if verbose:
                print('Calculating AIC')

            # Get the negative log-likelihood
            log_likelihood_trace[i] = -_data_negative_log_likelihood(data, results['c'], null_dist, signal_dist)

            # Calculate AIC = 2k - 2ln(L)
            aic_trace[i] = 2. * dof_trace[i] - 2. * log_likelihood_trace[i]
            
            # Calculate AICc = AIC + 2k * (k+1) / (n - k - 1)
            aicc_trace[i] = aic_trace[i] + 2 * dof_trace[i] * (dof_trace[i]+1) / (data.shape[0] - dof_trace[i] - 1.)

            # Calculate BIC = -2ln(L) + k * (ln(n) - ln(2pi))
            bic_trace[i] = -2 * log_likelihood_trace[i] + dof_trace[i] * (np.log(len(data)) - np.log(2 * np.pi))

            # Track the best model thus far
            if best_idx is None or bic_trace[i] < bic_trace[best_idx]:
                best_idx = i
                best_plateaus = plateaus

            # Save the final run parameters to use for warm-starting the nex