from AOT_biomaps.AOT_Recon.ReconTools import power_method, gradient_cpu, gradient_gpu, div_cpu, div_gpu, proj_l2, gradient_KL
from AOT_biomaps.Config import config

import torch
from tqdm import trange

def chambolle_pock_TV_cpu(SMatrix, y, alpha, theta, numIterations, isSavingEachIteration, L, withTumor):
    device = torch.device("cpu")
    A = torch.tensor(SMatrix, dtype=torch.float32, device=device)
    y = torch.tensor(y, dtype=torch.float32, device=device)

    T, Z, X, N = SMatrix.shape
    A_flat = A.permute(0, 3, 1, 2).reshape(T * N, Z * X)
    y_flat = y.reshape(-1)

    P = lambda x: torch.matmul(A_flat, x.ravel())
    PT = lambda y: torch.matmul(A_flat.T, y)

    if L is None:
        L = power_method(P, PT, y_flat, Z, X, isGPU=True)

    sigma = 1.0 / L
    tau = 1.0 / L

    x = torch.zeros(Z * X, device=device)
    p = torch.zeros((2, Z, X), device=device)
    q = torch.zeros_like(y_flat)
    x_tilde = x.clone()

    I_reconMatrix = [x.reshape(Z, X).cpu().numpy()]

    if withTumor:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (TV : Gaussian Noise) α:{alpha:.4f} L: {L:.4f} ---- WITH TUMOR ---- processing on single CPU ----"
    else:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (TV : Gaussian Noise) α:{alpha:.4f} L: {L:.4f} ---- WITHOUT TUMOR ---- processing on single CPU ----"

    for iteration in trange(numIterations, desc=description):
        p = proj_l2(p + sigma * gradient_cpu(x_tilde.reshape(Z, X)), alpha)
        q = (q + sigma * P(x_tilde) - sigma * y_flat) / (1.0 + sigma)

        x_old = x
        x = x + tau * div_cpu(p).ravel() - tau * PT(q)
        x_tilde = x + theta * (x - x_old)

        if iteration % 1 == 0:
            I_reconMatrix.append(x.reshape(Z, X).cpu().numpy())

    return I_reconMatrix if isSavingEachIteration else I_reconMatrix[-1]


def chambolle_pock_TV_gpu(SMatrix, y, alpha, theta, numIterations, isSavingEachIteration, L, withTumor):

    device = torch.device(f"cuda:{config.select_best_gpu()}")
    A = torch.tensor(SMatrix, dtype=torch.float32, device=device)
    y = torch.tensor(y, dtype=torch.float32, device=device)
    T, Z, X, N = SMatrix.shape
    A_flat = A.permute(0, 3, 1, 2).reshape(T * N, Z * X)
    y_flat = y.reshape(-1)
    P = lambda x: torch.matmul(A_flat, x.ravel())
    PT = lambda y: torch.matmul(A_flat.T, y)

    if L is None:
        L = power_method(P, PT, y_flat, Z, X, isGPU=True)

    sigma = 1.0 / L
    tau = 1.0 / L
    x = torch.zeros(Z * X, device=device)
    p = torch.zeros((2, Z, X), device=device)
    q = torch.zeros_like(y_flat)
    x_tilde = x.clone()
    I_reconMatrix = [x.reshape(Z, X).cpu().numpy()]

    if withTumor:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (TV : Gaussian Noise) α:{alpha:.4f} L: {L:.4f} ---- WITH TUMOR ---- processing on GPU no.{torch.cuda.current_device()} ----"
    else:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (TV: Gaussian Noise) α:{alpha:.4f} L: {L:.4f} ---- WITHOUT TUMOR ---- processing on GPU no.{torch.cuda.current_device()} ----"

    for iteration in trange(numIterations, desc=description):
        p = proj_l2(p + sigma * gradient_gpu(x_tilde.reshape(Z, X)), alpha)
        q = (q + sigma * P(x_tilde) - sigma * y_flat) / (1.0 + sigma)
        x_old = x
        x = x + tau * div_gpu(p).ravel() - tau * PT(q)
        x_tilde = x + theta * (x - x_old)
        if iteration % 1 == 0:
            I_reconMatrix.append(x.reshape(Z, X).cpu().numpy())
    return I_reconMatrix if isSavingEachIteration else I_reconMatrix[-1]

def chambolle_pock_KL_cpu(SMatrix, y, alpha, theta, numIterations, isSavingEachIteration, L, withTumor):

    device = torch.device("cpu")
    A = torch.tensor(SMatrix, dtype=torch.float32, device=device)
    y = torch.tensor(y, dtype=torch.float32, device=device)
    T, Z, X, N = SMatrix.shape
    A_flat = A.permute(0, 3, 1, 2).reshape(T * N, Z * X)
    y_flat = y.reshape(-1)

    P = lambda x: torch.matmul(A_flat, x.ravel())
    PT = lambda y: torch.matmul(A_flat.T, y)

    # Initialisation
    x = torch.ones(Z * X, device=device)  # Initialisation positive pour KL
    p = torch.zeros((2, Z, X), device=device)
    q = torch.zeros_like(y_flat)
    x_tilde = x.clone()

    # Paramètres de l'algorithme
    sigma = 0.1
    tau = 0.1

    I_reconMatrix = [x.reshape(Z, X).cpu().numpy()]

    if L is None:
        L = power_method(P, PT, y_flat, Z, X, isGPU=True)

    if withTumor:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (KL : Poisson noise) α:{alpha:.4f} L: {L:.4f} ---- WITH TUMOR ---- processing on single CPU ----"
    else:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (KL : Poisson noise) α:{alpha:.4f} L: {L:.4f} ---- WITHOUT TUMOR ---- processing on single CPU ----"
    
    for iteration in trange(numIterations, desc= description):
        # Mise à jour de p (gradient de la variation totale)
        p = proj_l2(p + sigma * gradient_cpu(x_tilde.reshape(Z, X)), alpha)

        # Mise à jour de q (gradient de la divergence de Kullback-Leibler)
        Ax = P(x_tilde)
        q = q + sigma * gradient_KL(Ax, y_flat)

        # Mise à jour de x
        x_old = x
        x = x - tau * (div_cpu(p).ravel() + PT(q))
        x = torch.clamp(x, min=1e-10)  # Assurez-vous que x reste positif
        x_tilde = x + theta * (x - x_old)

        if iteration % 1 == 0:
            I_reconMatrix.append(x.reshape(Z, X).cpu().numpy())

    return I_reconMatrix if isSavingEachIteration else I_reconMatrix[-1]

def chambolle_pock_KL_gpu(SMatrix, y, alpha, theta, numIterations, isSavingEachIteration, L, withTumor):
    device = torch.device(f"cuda:{config.select_best_gpu()}")
    A = torch.tensor(SMatrix, dtype=torch.float32, device=device)
    y = torch.tensor(y, dtype=torch.float32, device=device)
    T, Z, X, N = SMatrix.shape
    A_flat = A.permute(0, 3, 1, 2).reshape(T * N, Z * X)
    y_flat = y.reshape(-1)

    P = lambda x: torch.matmul(A_flat, x.ravel())
    PT = lambda y: torch.matmul(A_flat.T, y)

    if L is None:
        L = power_method(P, PT, y_flat, Z, X, isGPU=True)

    # Initialisation
    x = torch.ones(Z * X, device=device)  # Initialisation positive pour KL
    p = torch.zeros((2, Z, X), device=device)
    q = torch.zeros_like(y_flat)
    x_tilde = x.clone()


    # Paramètres de l'algorithme
    sigma = 0.1
    tau = 0.1

    I_reconMatrix = [x.reshape(Z, X).cpu().numpy()]
    
    if withTumor:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (KL : Poisson noise) α:{alpha:.4f} L: {L:.4f} ---- WITH TUMOR  ---- processing on GPU no.{torch.cuda.current_device()} ----"
    else:
        description = f"AOT-BioMaps -- Primal/Dual Recontruction Tomography : Chambolle-Pock (KL : Poisson noise) α:{alpha:.4f} L: {L:.4f} ---- WITHOUT TUMOR  ---- processing on GPU no.{torch.cuda.current_device()} ----"

    for iteration in trange(numIterations, desc=description):
        # Mise à jour de p (gradient de la variation totale)
        p = proj_l2(p + sigma * gradient_gpu(x_tilde.reshape(Z, X)), alpha)

        # Mise à jour de q (gradient de la divergence de Kullback-Leibler)
        Ax = P(x_tilde)
        q = q + sigma * gradient_KL(Ax, y_flat)

        # Mise à jour de x
        x_old = x
        x = x - tau * (div_gpu(p).ravel() + PT(q))
        x = torch.clamp(x, min=1e-10)  # Assurez-vous que x reste positif
        x_tilde = x + theta * (x - x_old)

        if iteration % 1 == 0:
            I_reconMatrix.append(x.reshape(Z, X).cpu().numpy())

    return I_reconMatrix if isSavingEachIteration else I_reconMatrix[-1]
