from __future__ import division
from __future__ import print_function

import threading
import multiprocessing as mp
from math import ceil
import logging

import numpy as np


# logging.basicConfig(filename="./log/thread_log.txt", filemode="w",
#                     level=logging.DEBUG, format='(%(threadName)-10s) %(message)s',)


def residual_reduction(uj):
    uj_plus = max(uj, 0)
    return 2*uj*uj_plus - uj_plus*uj_plus


def NOMP(x, H, gamma, epsilon, tau):
    if H.shape[1] != x.shape[0]:
        raise ValueError("H and x dimensions don't agree for matrix multiplication")
    if x.shape[1] != 1:
        raise ValueError("x should be a vector with x.shape[1] == 1")
    if gamma > H.shape[1]:
        raise ValueError("Sparsity constraint gamma can't be greater than dimension of w vector")
    s = H.dot(x)
    S = H.dot(H.T)
    k = s.shape[0]
    u = s
    R = 1.
    A = set()
    Ac = set(np.arange(k)) - A
    w = np.zeros(k)

    iteration_num = 0
    while len(A) < gamma:
        assert (len(A & Ac) == 0), "Error, A and Ac are not disjoint."
        iteration_num += 1
        Ac_list = list(Ac)
        reduction_values = [residual_reduction(u[j_candidate]) for j_candidate in Ac_list]
        max_idx = np.argmax(reduction_values)
        # noinspection PyTypeChecker
        j_star = Ac_list[max_idx]
        # noinspection PyTypeChecker
        reduction_value = reduction_values[max_idx]
        logging.log(0, "iteration_num=%d, max reduction value=%f, max reduction index=%d"
                    % (iteration_num, reduction_value, j_star))
        if reduction_value <= epsilon:
            logging.log(0, "iteration_num=%d, R=%f, reduction value=%f reached below eps=%f, returning w with "
                           "%d of %d elements" % (iteration_num, R, reduction_value, epsilon, len(A), gamma))
            return w
        assert (j_star not in A), "Overlap Error in A and Ac sets"
        A.add(j_star)
        Ac.remove(j_star)
        subiteration_num = 0
        while True:
            subiteration_num += 1
            for j in A:
                delta_j = w[j]
                w[j] = max(u[j] + w[j], 0)
                delta_j = w[j] - delta_j
                Sj = S[:, j].reshape(-1, 1)
                assert (u.shape == Sj.shape), "u and Sj don't match in shape"
                u = u - Sj*delta_j
            reduction = 0.
            for j in A:
                reduction += w[j]*(s[j] + u[j])
            R_old = R
            R = 1 - reduction
            if R_old != 0:
                tau_star = 1 - R/R_old
            else:
                tau_star = 0.
            logging.log(0, "iteration_num=%d, subiteration_num=%d, reduction=%f, tau_star=%f"
                        % (iteration_num, subiteration_num, reduction, tau_star))
            if R_old - R < tau*R_old or R == 0:
                logging.log(0, "\nBreaking from cyclic coordinate descent for iteration number=%d after "
                            "subiteration number=%d\n" % (iteration_num, subiteration_num))
                break
    assert (iteration_num == gamma), "gamma and num_iterations don't match"
    logging.log(0, "\nCompleted gamma=%d iterations, returning with full elements" % iteration_num)
    return w


def NOMP_chunk(X, W, H, N, chunk_size, gamma, thread_num):
    start_idx = chunk_size*thread_num
    for i in range(int(start_idx), int(min(start_idx + chunk_size, N))):
        W[i, :] = NOMP(X[i, :].reshape((-1, 1)), H, gamma=gamma, epsilon=0.001, tau=0.01)
    logging.log(0, "Thread/Process %d finished processing chunk %d : %d"
                % (thread_num, chunk_size*thread_num, chunk_size*(thread_num + 1)))
    return


def NOMP_chunk_mp(X, H, N, chunk_size, gamma, thread_num):
    start_idx = chunk_size*thread_num
    result = []
    for i in range(start_idx, min(start_idx + chunk_size, N)):
        result.append(NOMP(X[i, :].reshape((-1, 1)), H, gamma=gamma, epsilon=0.001, tau=0.01))
    # logging.log(0, "Thread/Process %d finished processing chunk %d : %d"
    #             % (thread_num, chunk_size*thread_num, chunk_size*(thread_num + 1)))
    result.append([start_idx, min(start_idx + chunk_size, N)])
    return result


def learn(X, W, H, gamma, num_threads=20, num_processes=3, parallelization=False):
    N = W.shape[0]
    if parallelization:
        if N < 3000:
            use_threading, use_multiprocessing = True, False
            logging.debug("Using threading with %d threads" % num_threads)
        else:
            use_multiprocessing, use_threading = True, False
            logging.debug("Using multiprocessing with %d processes" % num_processes)
    else:
        use_threading, use_multiprocessing = False, False
        logging.debug("Using serial computation")
    if not use_threading and not use_multiprocessing:
        for i in range(N):
            W[i, :] = NOMP(X[i, :].reshape((-1, 1)), H, gamma=gamma, epsilon=0.001, tau=0.001)
        return W

    # def NOMP_chunk(thread_num):
    #     start_idx = int(chunk_size * thread_num)
    #     for i in range(start_idx, min(int(start_idx + chunk_size), N)):
    #         W[i, :] = NOMP(X[i, :].reshape((-1, 1)), H, gamma=gamma, epsilon=0.001, tau=0.001)
    #     logging.debug("Thread %d finished processing chunk %d : %d"
    #                   % (thread_num, chunk_size * thread_num, chunk_size * (thread_num + 1)))
    #     return

    elif use_threading:
        chunk_size = ceil(N / num_threads)
        logging.log(0, "Chunk size = %f" % chunk_size)
        threads = []
        for thread_num in range(num_threads):
            thread = threading.Thread(target=NOMP_chunk, args=(X, W, H, N, chunk_size, gamma, thread_num, ))
            threads.append(thread)
            thread.start()
        for thread in threads:
            thread.join()
        return W
    elif use_multiprocessing:
        chunk_size = int(ceil(N / num_processes))
        logging.debug("Chunk size = %f" % chunk_size)
        p = mp.Pool(processes=num_processes)

        def collect_results(result):
            start_idx, end_idx = result[-1]
            W[start_idx:end_idx, :] = np.array(result[:-1])

        for i in range(num_processes):
            p.apply_async(NOMP_chunk_mp, args=(X, H, N, chunk_size, gamma, i), callback=collect_results)
        p.close()
        p.join()
        # processes = []
        # for process_num in range(num_processes):
        #     process = multiprocessing.Process(target=NOMP_chunk, args=(X, w, H, N, chunk_size, gamma, process_num))
        #     processes.append(process)
        # for process in processes:
        #     process.start()
        # for process in processes:
        #     process.join()
        # W = w
        # for i, row in enumerate(results):
        #     W[i, :] = row
        return W
