import numpy as np
import itertools
import pandas as pd

from pathlib import Path
from PIL import Image
from scipy.optimize import linear_sum_assignment
from scipy.signal import correlate2d
from joblib import Memory

from plipy.ddl import DDLInpaintingConv


DATA_HOME = Path(__file__).resolve().parents[1] / "data/letters"
FIGURES = Path(__file__).resolve().parents[1] / "figures"
NUM_EXP = 10

mem = Memory(location='.', verbose=0)


def build_dictionary(images):
    D = np.zeros((len(images), 1, images[0].shape[0], images[0].shape[1]))
    for i in range(len(images)):
        D[i, 0] = images[i]
    return D / np.linalg.norm(D, axis=(2, 3))[:, :, None, None]


def cost_matrix(D, Dref):
    C = np.zeros((D.shape[0], Dref.shape[0]))
    for i in range(D.shape[0]):
        for j in range(Dref.shape[0]):
            C[i, j] = correlate2d(D[i, 0], Dref[j, 0]).max()
    return C


def compute_score(D, Dref):
    C = cost_matrix(D, Dref)
    row_ind, col_ind = linear_sum_assignment(C, maximize=True)
    score = C[row_ind, col_ind]
    return score.mean(), row_ind, col_ind


im_N = Image.open(str(DATA_HOME / "N.png")).convert("L")
im_e = Image.open(str(DATA_HOME / "e.png")).convert("L")
im_u = Image.open(str(DATA_HOME / "u.png")).convert("L")
im_r = Image.open(str(DATA_HOME / "r.png")).convert("L")
im_I = Image.open(str(DATA_HOME / "I.png")).convert("L")
im_P = Image.open(str(DATA_HOME / "P.png")).convert("L")
im_S = Image.open(str(DATA_HOME / "S.png")).convert("L")
im_text = Image.open(str(DATA_HOME / "text_neurips.png")).convert("L")

im_N_array = np.array(im_N) / 255.
im_e_array = np.array(im_e) / 255.
im_u_array = np.array(im_u) / 255.
im_r_array = np.array(im_r) / 255.
im_I_array = np.array(im_I) / 255.
im_P_array = np.array(im_P) / 255.
im_S_array = np.array(im_S) / 255.
im_text_array = np.array(im_text) / 255.

list_images = [im_N_array,
               im_e_array,
               im_u_array,
               im_r_array,
               im_I_array,
               im_P_array,
               im_S_array]


Dref = build_dictionary(list_images)


@mem.cache()
def run_test(params, num):

    lambd = params["lambd"]
    rho = params["rho"]
    snr = params["snr"]

    sigma_data = 1
    sigma_noise = np.sqrt(sigma_data**2 / snr)

    omega = np.random.random(im_text_array.shape)
    omega = (omega > rho).astype(float)

    noise = np.random.normal(scale=sigma_noise, size=im_text_array.shape)
    im_noisy = np.clip(im_text_array + noise, 0, 1)

    im_inpainting = omega * im_noisy

    # Dictionary parameters
    patch_dim = 12
    n_atoms = 30

    # Algo parameters
    n_iter = 20
    scores1 = np.zeros(num)
    scores2 = np.zeros(num)

    for i in range(num):

        # Synthesis patches
        ddl_inpainting = DDLInpaintingConv(n_atoms, n_iter, lambd=lambd,
                                           kernel_size=patch_dim)
        ddl_inpainting.training_process(im_inpainting[None, :, :],
                                        omega[None, :, :])
        D1 = ddl_inpainting.get_prior()

        # Synthesis random
        ddl_inpainting = DDLInpaintingConv(n_atoms, n_iter, lambd=lambd,
                                           kernel_size=patch_dim)
        ddl_inpainting.training_process(im_inpainting[None, :, :],
                                        omega[None, :, :],
                                        init="random")
        D2 = ddl_inpainting.get_prior()

        score1, _, _ = compute_score(D1, Dref)
        score2, _, _ = compute_score(D2, Dref)
        scores1[i] = score1
        scores2[i] = score2

    return {"patches": scores1.mean(), "random": scores2.mean()}


if __name__ == '__main__':

    hyperparams = {
        "lambd": [0.1, 1, 10],
        "rho": list(np.arange(0.1, 1, 0.1)),
        "snr": list(np.unique(np.logspace(0.7, 2, num=5, dtype=int)))
    }

    keys, values = zip(*hyperparams.items())
    permuts_params = [dict(zip(keys, v)) for v in itertools.product(*values)]

    dico_results = {}

    for params in permuts_params:

        # Experiment
        print(params)

        try:
            results = run_test(params, NUM_EXP)

            # Storing results
            for key in params.keys():
                if key not in dico_results:
                    dico_results[key] = [params[key]]
                else:
                    dico_results[key].append(params[key])

            for key in results:
                if key not in dico_results:
                    dico_results[key] = [results[key]]
                else:
                    dico_results[key].append(results[key])

        except (KeyboardInterrupt, SystemExit):
            raise

        except:
            pass

    # Data Frame
    path = Path(__file__).resolve().parents[1]
    results = pd.DataFrame(dico_results)
    results.to_csv(str(path / "results/results_cdl.csv"))
