import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from tqdm import tqdm
from pathlib import Path
from PIL import Image

from utils import create_patches_overlap
from plipy.ddl import DDLInpainting


def iter_algo(y, D, A, n_iter, L, algo="fista", lambd=0.1):
    out = np.zeros((D.shape[1], y.shape[1]))

    if algo == "fista":
        out_old = out.copy()
        t_old = 1.

    step = 1. / L

    for i in range(n_iter):

        # Gradient descent
        out = out - step * D.T @ (A * (A * (D @ out) - y))

        # Thresholding
        thresh = np.abs(out) - step * lambd
        out = np.sign(out) * np.maximum(0, thresh)

        if algo == "fista":
            t = 0.5 * (1 + np.sqrt(1 + 4 * t_old * t_old))
            z = out + ((t_old-1) / t) * (out - out_old)
            out_old = out.copy()
            t_old = t
            out = z

    return out


class DeepISTA(nn.Module):
    def __init__(self, n_iter, D, L, algo="fista", lambd=0.1,
                 device=None, n_iter_backprop=None):
        super().__init__()

        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Regularization parameter
        self.lambd = lambd

        # Algorithm unrolled and parameters
        self.algo = algo
        self.n_iter = n_iter

        if n_iter_backprop is None:
            self.backprop_iter = n_iter
        else:
            self.backprop_iter = n_iter_backprop

        # Shape
        self.dim_x = D.shape[1]

        # Initial dictionary
        self.init_D = D
        self.lipschitz = L

    def forward(self, y):
        """
        (F)ISTA-like forward pass
        """
        out = torch.zeros((self.dim_x, y.shape[1]),
                          dtype=torch.float,
                          device=self.device)

        if self.algo == "fista":
            out_old = out.clone()
            t_old = 1

        step = 1. / self.lipschitz

        torch.set_grad_enabled(False)

        for i in range(self.n_iter):
            if i >= self.n_iter - self.backprop_iter:
                torch.set_grad_enabled(True)
            # Gradient descent
            out = out - step * torch.matmul(
                self.dictionary.t(),
                self.A * (self.A * torch.matmul(self.dictionary, out) - y)
                )
            thresh = torch.abs(out) - step * self.lambd
            out = torch.sign(out) * F.relu(thresh)

            if self.algo == "fista":
                t = 0.5 * (1 + np.sqrt(1 + 4 * t_old * t_old))
                z = out + ((t_old-1) / t) * (out - out_old)
                out_old = out.clone()
                t_old = t
                out = z

        return out

    def cost(self, y, x):
        """
        LASSO cost function
        """
        signal = torch.matmul(self.dictionary, x)
        product = self.A * signal
        res = product - y
        l2 = (res * res).sum()
        l1 = torch.abs(x).sum()

        return 0.5 * l2 + self.lambd * l1

    def compute_gradient(self, data_y, A):
        """
        Training procedure
        """
        # Dimension
        self.dim_y = data_y.shape[0]

        # Operator
        self.A = torch.from_numpy(A).float().to(self.device)

        # Dictionary
        if self.init_D is None:
            choice = np.random.choice(data_y.shape[1], self.n_components)
            dico = data_y[:, choice]
            self.dictionary = nn.Parameter(
                torch.tensor(dico, device=self.device, dtype=torch.float)
                )
        else:
            dico_tensor = torch.from_numpy(self.init_D).float().to(self.device)
            self.dictionary = nn.Parameter(dico_tensor)

        # Data
        Y_tensor = torch.from_numpy(data_y).float().to(self.device)

        if self.backprop_iter == 0:
            with torch.no_grad():
                out = self.forward(Y_tensor)
        else:
            out = self.forward(Y_tensor)
        loss = self.cost(Y_tensor, out)
        loss.backward()

        grad = self.dictionary.grad.to("cpu").detach().numpy()

        self.dictionary.grad.zero_()

        return grad


def compute_lipschitz(D):
    iterations = 50
    u = np.random.random(D.shape[1])
    for i in range(iterations):
        u = D.T @ D @ u
        norme = np.linalg.norm(u)
        u /= norme
    return norme


def compute_angles_dico(data, mask, D, n_iter, lambd,
                        iterations, iterations_backprop):

    lipschitz = compute_lipschitz(D)

    N_data = 1000
    n_iter = 10000
    x_result = iter_algo(data[:, :N_data], D, mask[:, :N_data],
                         n_iter, lipschitz, algo="fista", lambd=lambd)
    dista = DeepISTA(n_iter=20, D=D, L=lipschitz, lambd=lambd)
    g_star = mask[:, :N_data] * (D @ x_result - data[:, :N_data]) @ x_result.T

    angles = np.zeros((len(iterations), len(iterations_backprop)))

    for i in range(len(iterations)):
        for j in range(len(iterations_backprop)):
            dista.n_iter = iterations[i]
            dista.backprop_iter = min(iterations_backprop[j], 100)
            g2 = dista.compute_gradient(data[:, :N_data], mask[:, :N_data])
            den = np.sqrt(np.trace(g2.T @ g2) * np.trace(g_star.T @ g_star))
            angle2 = np.trace(g2.T @ g_star) / den
            angles[i, j] = angle2

    return angles


N_JOBS = 4
RHO = 0.5
M_PATCH = 10
COMPONENTS = 128

path = Path(__file__).resolve().parents[1]
im = Image.open(str(path / "data/flowers.png"))
im_gray = im.convert("L")
im_gray_resized = im_gray.resize((128, 128), Image.ANTIALIAS)
im_to_process = np.array(im_gray_resized) / 255.

omega = np.ones_like(im_to_process)
noise = np.random.normal(scale=0.1, size=im_to_process.shape)
im_inpainting = omega * im_to_process + noise

data, mask = create_patches_overlap(im_inpainting, M_PATCH, A=omega)

lambd = 0.1
ddl = DDLInpainting(COMPONENTS, 20, lambd)
ddl.keep_dico = True
ddl.fit(data, mask)

path_dico = ddl.path_optim

n_iter = 10000
iterations = np.unique(np.logspace(0, 2, dtype=int, num=30))
iterations_backprop = np.unique(np.logspace(0, 2, dtype=int, num=30))
iterations_backprop = np.concatenate([np.array([0]), iterations_backprop])

indexes = np.unique(
    np.logspace(0, np.log10(len(path_dico)-1), dtype=int, num=20)
    )
angles = np.zeros((len(indexes), len(iterations), len(iterations_backprop)))

for i in tqdm(range(len(indexes))):
    angles[i] = compute_angles_dico(data, mask, path_dico[indexes[i]],
                                    n_iter, lambd, iterations,
                                    iterations_backprop)

np.save(str(path / "results/angles_iterations.npy"), iterations)
np.save(
    str(path / "results/angles_iterations_backprop.npy"),
    iterations_backprop
    )
np.save(str(path / "results/angles_g1_g2.npy"), angles)

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from tqdm import tqdm
from pathlib import Path
from PIL import Image

from utils import create_patches_overlap
from plipy.ddl import DDLInpainting


def iter_algo(y, D, A, n_iter, L, algo="fista", lambd=0.1):
    out = np.zeros((D.shape[1], y.shape[1]))

    if algo == "fista":
        out_old = out.copy()
        t_old = 1.

    step = 1. / L

    for i in range(n_iter):

        # Gradient descent
        out = out - step * D.T @ (A * (A * (D @ out) - y))

        # Thresholding
        thresh = np.abs(out) - step * lambd
        out = np.sign(out) * np.maximum(0, thresh)

        if algo == "fista":
            t = 0.5 * (1 + np.sqrt(1 + 4 * t_old * t_old))
            z = out + ((t_old-1) / t) * (out - out_old)
            out_old = out.copy()
            t_old = t
            out = z

    return out


class DeepISTA(nn.Module):
    def __init__(self, n_iter, D, L, algo="fista", lambd=0.1,
                 device=None, n_iter_backprop=None):
        super().__init__()

        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Regularization parameter
        self.lambd = lambd

        # Algorithm unrolled and parameters
        self.algo = algo
        self.n_iter = n_iter

        if n_iter_backprop is None:
            self.backprop_iter = n_iter
        else:
            self.backprop_iter = n_iter_backprop

        # Shape
        self.dim_x = D.shape[1]

        # Initial dictionary
        self.init_D = D
        self.lipschitz = L

    def forward(self, y):
        """
        (F)ISTA-like forward pass
        """
        out = torch.zeros((self.dim_x, y.shape[1]),
                          dtype=torch.float,
                          device=self.device)

        if self.algo == "fista":
            out_old = out.clone()
            t_old = 1

        step = 1. / self.lipschitz

        torch.set_grad_enabled(False)

        for i in range(self.n_iter):
            if i >= self.n_iter - self.backprop_iter:
                torch.set_grad_enabled(True)
            # Gradient descent
            out = out - step * torch.matmul(
                self.dictionary.t(),
                self.A * (self.A * torch.matmul(self.dictionary, out) - y)
                )
            thresh = torch.abs(out) - step * self.lambd
            out = torch.sign(out) * F.relu(thresh)

            if self.algo == "fista":
                t = 0.5 * (1 + np.sqrt(1 + 4 * t_old * t_old))
                z = out + ((t_old-1) / t) * (out - out_old)
                out_old = out.clone()
                t_old = t
                out = z

        return out

    def cost(self, y, x):
        """
        LASSO cost function
        """
        signal = torch.matmul(self.dictionary, x)
        product = self.A * signal
        res = product - y
        l2 = (res * res).sum()
        l1 = torch.abs(x).sum()

        return 0.5 * l2 + self.lambd * l1

    def compute_gradient(self, data_y, A):
        """
        Training procedure
        """
        # Dimension
        self.dim_y = data_y.shape[0]

        # Operator
        self.A = torch.from_numpy(A).float().to(self.device)

        # Dictionary
        if self.init_D is None:
            choice = np.random.choice(data_y.shape[1], self.n_components)
            dico = data_y[:, choice]
            self.dictionary = nn.Parameter(
                torch.tensor(dico, device=self.device, dtype=torch.float)
                )
        else:
            dico_tensor = torch.from_numpy(self.init_D).float().to(self.device)
            self.dictionary = nn.Parameter(dico_tensor)

        # Data
        Y_tensor = torch.from_numpy(data_y).float().to(self.device)

        if self.backprop_iter == 0:
            with torch.no_grad():
                out = self.forward(Y_tensor)
        else:
            out = self.forward(Y_tensor)
        loss = self.cost(Y_tensor, out)
        loss.backward()

        grad = self.dictionary.grad.to("cpu").detach().numpy()

        self.dictionary.grad.zero_()

        return grad


def compute_lipschitz(D):
    iterations = 50
    u = np.random.random(D.shape[1])
    for i in range(iterations):
        u = D.T @ D @ u
        norme = np.linalg.norm(u)
        u /= norme
    return norme


def compute_angles_dico(data, mask, D, n_iter, lambd,
                        iterations, iterations_backprop):

    lipschitz = compute_lipschitz(D)

    N_data = 1000
    n_iter = 10000
    x_result = iter_algo(data[:, :N_data], D, mask[:, :N_data],
                         n_iter, lipschitz, algo="fista", lambd=lambd)
    dista = DeepISTA(n_iter=20, D=D, L=lipschitz, lambd=lambd)
    g_star = mask[:, :N_data] * (D @ x_result - data[:, :N_data]) @ x_result.T

    angles = np.zeros((len(iterations), len(iterations_backprop)))

    for i in range(len(iterations)):
        for j in range(len(iterations_backprop)):
            dista.n_iter = iterations[i]
            dista.backprop_iter = min(iterations_backprop[j], 100)
            g2 = dista.compute_gradient(data[:, :N_data], mask[:, :N_data])
            den = np.sqrt(np.trace(g2.T @ g2) * np.trace(g_star.T @ g_star))
            angle2 = np.trace(g2.T @ g_star) / den
            angles[i, j] = angle2

    return angles


N_JOBS = 4
RHO = 0.5
M_PATCH = 10
COMPONENTS = 128

path = Path(__file__).resolve().parents[1]
im = Image.open(str(path / "data/flowers.png"))
im_gray = im.convert("L")
im_gray_resized = im_gray.resize((128, 128), Image.ANTIALIAS)
im_to_process = np.array(im_gray_resized) / 255.

omega = np.ones_like(im_to_process)
noise = np.random.normal(scale=0.1, size=im_to_process.shape)
im_inpainting = omega * im_to_process + noise

data, mask = create_patches_overlap(im_inpainting, M_PATCH, A=omega)

lambd = 0.1
ddl = DDLInpainting(COMPONENTS, 20, lambd)
ddl.keep_dico = True
ddl.fit(data, mask)

path_dico = ddl.path_optim

n_iter = 10000
iterations = np.unique(np.logspace(0, 2, dtype=int, num=30))
iterations_backprop = np.unique(np.logspace(0, 2, dtype=int, num=30))
iterations_backprop = np.concatenate([np.array([0]), iterations_backprop])

indexes = np.unique(
    np.logspace(0, np.log10(len(path_dico)-1), dtype=int, num=20)
    )
angles = np.zeros((len(indexes), len(iterations), len(iterations_backprop)))

for i in tqdm(range(len(indexes))):
    angles[i] = compute_angles_dico(data, mask, path_dico[indexes[i]],
                                    n_iter, lambd, iterations,
                                    iterations_backprop)

np.save(str(path / "results/angles_iterations.npy"), iterations)
np.save(
    str(path / "results/angles_iterations_backprop.npy"),
    iterations_backprop
    )
np.save(str(path / "results/angles_g1_g2.npy"), angles)

# Gradient descent # Thresholding # Regularization parameter # Algorithm unrolled and parameters # Shape # Initial dictionary (F)ISTA-like forward pass # Gradient descent LASSO cost function Training procedure # Dimension # Operator # Dictionary # Data