from matplotlib import pyplot as plt
from multiprocessing import Pool
from multiprocessing import cpu_count
import numpy as np
from scipy.misc import imread, imsave
import scipy.ndimage
import sklearn.linear_model as lm
import sys
import os

import utils


def sense(img, k=1000, basis="wvt", wvt_level=4, alpha=None):
    """ Senses an image with k measurements

    :param img: 2d array
        image
    :param k: int
        number of measurements
    :param basis: str
        either 'wvt', 'dct'
    :param wvt_level: int
        level of wavelet transform
    :param alpha: float or None
        regularization parameter
        if None: alpha is first found with CV (takes long)
    :return: 2d array
        sensed image
    """
    print "Image size:", img.shape
    img_f = img.flatten()
    print "Build sensing matrix"

    A = np.random.normal(0, 1, len(img_f)*k).astype(np.float32).\
        reshape(k, img.shape[0], img.shape[1])

    print "Measurement"
    b = np.dot(A.reshape(k, len(img_f)), img_f)

    if basis == "wvt":
        print "Wavelets"
        trans_A = [utils.dwt2(A[i].reshape(img.shape), level=wvt_level).
                     astype(np.float16).flatten() for i in range(k)]
    elif basis == "dct":
        print "DCT"
        trans_A = [utils.dct2(A[i].reshape(img.shape)).
                       astype(np.float16).flatten() for i in range(k)]
    else:
        raise Exception("Unknown basis")

    A = None

    if alpha:
        lasso = lm.Lasso(alpha=alpha, max_iter=100000, normalize=True)
        print "Fit"
        lasso.fit(trans_A, b)
    else:
        lasso_cv = lm.LassoCV(n_jobs=cpu_count(), max_iter=100000,
                              normalize=True)
        print "Fit"
        lasso_cv.fit(trans_A, b)
        print "Alpha: %.6f" % lasso_cv.alpha_
        lasso = lm.Lasso(alpha=lasso_cv.alpha_, max_iter=100000,
                         normalize=True)
        print "Fit"
        lasso.fit(trans_A, b)

    if basis == "wvt":
        return utils.idwt2(lasso.coef_.reshape(img.shape), level=wvt_level)
    elif basis == "dct":
        return utils.idct2(lasso.coef_.reshape(img.shape))
    else:
        raise Exception("Unknown basis")


def _sense_block_thread(args):
    rec = sense(args[1], args[2], wvt_level=args[3], basis=args[4],
                alpha=args[5])
    return args[0], rec


def sense_blocks(img, ratio, blocksize=64, wvt_level=3, alpha=None,
                 basis="wvt", n_processes=None):
    """ Recovers image parallel with blocking

    :param img: 2d array
        image
    :param mask: 2d bool array
        mask
    :param blocksize: int
        defines edgelength of blocking squares
    :param wvt_level: int
        level of wavelet transform
    :param alpha: float or None
        regularization parameter
        if None: alpha is first found with CV (takes long)
    :param n_processes: int or None
        number of processes to be used
        if None: all cores are used
    :return: 2d array
        recovered image
    """
    if not n_processes:
        n_processes = cpu_count()

    my_img = np.pad(img, [[blocksize/2, blocksize/2],
                          [blocksize/2, blocksize/2]], mode="reflect")

    k = int((blocksize/2)**2 * ratio)

    params = []

    for x_pos in range(0, my_img.shape[0]-blocksize/2, blocksize/2):
        for y_pos in range(0, my_img.shape[1]-blocksize/2, blocksize/2):
            params.append([[x_pos, y_pos],
                           my_img[x_pos: x_pos + blocksize,
                                  y_pos: y_pos + blocksize],
                           k, wvt_level, basis, alpha])

    print("N jobs: %d" % (len(params)))
    if n_processes > 1:
        pool = Pool(n_processes)
        results = pool.map(_sense_block_thread, params)
        pool.close()
        pool.join()
    else:
        results = map(_sense_block_thread, params)

    r_img = np.zeros(my_img.shape, dtype=np.float32)