from AOT_biomaps.AOT_Recon.AlgebraicRecon import AlgebraicRecon
from AOT_biomaps.AOT_Recon.ReconEnums import ReconType, ProcessType, NoiseType
from AOT_biomaps.AOT_Recon.AOT_Optimizers.PDHG import CP, CP_KL, CP_TV
from AOT_biomaps.AOT_Recon.ReconEnums import OptimizerType

import os
from datetime import datetime
import numpy as np

class PrimalDualRecon(AlgebraicRecon):
    """
    This class implements the convex reconstruction process.
    It currently does not perform any operations but serves as a template for future implementations.
    """
    def __init__(self, alpha, theta=1.0, L=None, **kwargs):
        super().__init__(**kwargs)
        self.reconType = ReconType.Convex
        self.alpha = alpha # regularization parameter
        self.theta = theta # relaxation parameter (between 1 and 2)
        self.L = L # norme spectrale de l'opérateur linéaire défini par les matrices P et P^T

    def run(self, processType=ProcessType.PYTHON, withTumor=True):
        """
        This method is a placeholder for the convex reconstruction process.
        It currently does not perform any operations but serves as a template for future implementations.
        """
        if(processType == ProcessType.CASToR):
            self._convexReconCASToR(withTumor)
        elif(processType == ProcessType.PYTHON):
            self._convexReconPython(withTumor)
        else:
            raise ValueError(f"Unknown convex reconstruction type: {processType}")

    def _convexReconCASToR(self, withTumor):
        raise NotImplementedError("CASToR convex reconstruction is not implemented yet.")


    def save(self, withTumor=True):
        """
        Save the reconstruction results to the specified directory, for Bayesian reconstruction.

        Parameters:
            withTumor (bool): If True, saves the reconstruction with tumor; otherwise, saves the reconstruction without tumor.
        """
        if self.saveDir is None:
            raise ValueError("Save directory is not specified. Please set saveDir before saving.")

        # Création du répertoire de résultats
        date = datetime.now().strftime("%d%m")
        results_dir = os.path.join(self.saveDir, f'results_{date}_{self.optimizer.value}_Alpha_{self.alpha}_Theta_{self.theta}_L_{self.L}')

        os.makedirs(results_dir, exist_ok=True)  # Crée le répertoire (évite les erreurs si déjà existant)

        # Sauvegarde des données
        if withTumor:
            if not hasattr(self, 'reconPhantom') or not self.reconPhantom or len(self.reconPhantom) == 0:
                raise ValueError("Reconstructed phantom is empty. Run reconstruction first.")
            np.save(os.path.join(results_dir, 'reconPhantom.npy'), np.array(self.reconPhantom))
        else:
            if not hasattr(self, 'reconLaser') or not self.reconLaser or len(self.reconLaser) == 0:
                raise ValueError("Reconstructed laser is empty. Run reconstruction first.")
            np.save(os.path.join(results_dir, 'reconLaser.npy'), np.array(self.reconLaser))

        # Sauvegarde des indices (commun aux deux cas)
        np.save(os.path.join(results_dir, 'reconIndices.npy'), np.array(self.indices))
        print(f"Reconstruction results saved to {results_dir}")


    def _convexReconPython(self, withTumor):
        if withTumor:
            y=self.experiment.AOsignal_withTumor
        else:
            y=self.experiment.AOsignal_withoutTumor

        if self.optimizer == OptimizerType.CP_TV:
            self.reconPhantom = CP_TV(
                self.SMatrix, 
                y=y, 
                alpha=self.alpha, 
                theta=self.theta, 
                numIterations=self.numIterations, 
                isSavingEachIteration=self.isSavingEachIteration,
                L=self.L, 
                withTumor=withTumor,
                device=None
                )
        elif self.optimizer == OptimizerType.CP_KL:
            self.reconPhantom = CP_KL(
                self.SMatrix, 
                y=y, 
                alpha=self.alpha, 
                theta=self.theta, 
                numIterations=self.numIterations, 
                isSavingEachIteration=self.isSavingEachIteration,
                L=self.L, 
                withTumor=withTumor,
                device=None
            )
        else:
            raise ValueError(f"Optimizer value must be CP_TV or CP_KL, got {self.optimizer}")

            



   
