from AOT_biomaps.Config import config
from AOT_biomaps.AOT_Experiment.Tomography import Tomography, hex_to_binary_profile
from .ReconEnums import ReconType
from .ReconTools import mse, ssim, get_phase_deterministic

import os
import numpy as np
import matplotlib.pyplot as plt
from abc import ABC, abstractmethod
from tqdm import trange


class Recon(ABC):
    def __init__(self, experiment, saveDir = None, isGPU = config.get_process() == 'gpu', isMultiCPU = True):
        self.reconPhantom = None
        self.reconLaser = None
        self.experiment = experiment
        self.reconType = None
        self.saveDir = saveDir
        self.MSE = None
        self.SSIM = None
        self.CRC = None

        self.isGPU = isGPU
        self.isMultiCPU = isMultiCPU

        if str(type(self.experiment)) != str(Tomography):
            raise TypeError(f"Experiment must be of type {Tomography}")

    @abstractmethod
    def run(self,withTumor = True):
        pass

    def save(self, withTumor=True, overwrite=False, date=None, show_logs=True):
        """
        Save the reconstruction results (reconPhantom is with tumor, reconLaser is without tumor) and indices of the saved recon results, in numpy format.

        Args:
            withTumor (bool): If True, saves reconPhantom. If False, saves reconLaser. Default is True.
            overwrite (bool): If False, does not save if the file already exists. Default is False.

        Warnings:
            reconPhantom and reconLaser are lists of 2D numpy arrays, each array corresponding to one iteration.
        """
        isExisting, filepath = self.checkExistingFile(date=date)
        if isExisting and not overwrite:
            return
        
        filename = 'reconPhantom.npy' if withTumor else 'reconLaser.npy'
        filepathRecon = os.path.join(filepath, filename)

        if withTumor:
            if not self.reconPhantom or len(self.reconPhantom) == 0:
                raise ValueError("Reconstructed phantom is empty. Run reconstruction first.")
            np.save(filepathRecon, np.array(self.reconPhantom))
        else:
            if not self.reconLaser or len(self.reconLaser) == 0:
                raise ValueError("Reconstructed laser is empty. Run reconstruction first.")
            np.save(filepathRecon, np.array(self.reconLaser))

        if self.indices is not None and len(self.indices) > 0:
            filepathIndices = os.path.join(filepath, "indices.npy")
            np.save(filepathIndices, np.array(self.indices))

        if show_logs:
            print(f"Reconstruction results saved to {os.path.dirname(filepath)}")

    @abstractmethod
    def checkExistingFile(self, date = None):
        pass

    def calculateCRC(self, use_ROI=True):
        """
        Computes the Contrast Recovery Coefficient (CRC) for all ROIs combined or globally.
        For analytic reconstruction: returns a single CRC value.
        For iterative reconstruction: returns a list of CRC values (one per iteration).
        If iteration is specified, returns CRC for that specific iteration only.

        :param iteration: Specific iteration index (optional). If None, computes for all iterations.
        :param use_ROI: If True, computes CRC for all ROIs combined. If False, computes global CRC.
        :return: CRC value or list of CRC values.
        """
        if self.reconType is None:
            raise ValueError("Run reconstruction first")

        if self.reconLaser is None or self.reconLaser == []:
            raise ValueError("Reconstructed laser is empty. Run reconstruction first.")
        if self.reconPhantom is None or self.reconPhantom == []:
            raise ValueError("Reconstructed phantom is empty. Run reconstruction first.")

        # Handle empty reconstructions
        if self.reconLaser is None or self.reconLaser == []:
            print("Reconstructed laser is empty. Running reconstruction without tumor...")
            self.run(withTumor=False, isSavingEachIteration=True)

        # Get the ROI mask(s) from the phantom if needed
        if use_ROI:
            self.experiment.OpticImage.find_ROI()
            global_mask = np.logical_or.reduce(self.experiment.OpticImage.maskList)

        # Analytic reconstruction case
        if self.reconType is ReconType.Analytic:
            if use_ROI:
                recon_ratio = np.mean(self.reconPhantom[global_mask]) / np.mean(self.reconLaser[global_mask])
                lambda_ratio = np.mean(self.experiment.OpticImage.phantom[global_mask]) / np.mean(self.experiment.OpticImage.laser.intensity[global_mask])
            else:
                recon_ratio = np.mean(self.reconPhantom) / np.mean(self.reconLaser)
                lambda_ratio = np.mean(self.experiment.OpticImage.phantom) / np.mean(self.experiment.OpticImage.laser.intensity)

            self.CRC =(recon_ratio - 1) / (lambda_ratio - 1)

        # Iterative reconstruction case
        else:
            iterations = range(len(self.reconPhantom))

            crc_list = []
            for it in iterations:
                if use_ROI:
                    recon_ratio = np.mean(self.reconPhantom[it][global_mask]) / np.mean(self.reconLaser[it][global_mask])
                    lambda_ratio = np.mean(self.experiment.OpticImage.phantom[global_mask]) / np.mean(self.experiment.OpticImage.laser.intensity[global_mask])
                else:
                    recon_ratio = np.mean(self.reconPhantom[it]) / np.mean(self.reconLaser[it])
                    lambda_ratio = np.mean(self.experiment.OpticImage.phantom) / np.mean(self.experiment.OpticImage.laser.intensity)

                crc_list.append((recon_ratio - 1) / (lambda_ratio - 1))

            self.CRC = crc_list

    def calculateMSE(self):
        """
        Calculate the Mean Squared Error (MSE) of the reconstruction.

        Returns:
            mse: float or list of floats, Mean Squared Error of the reconstruction
        """
                
        if self.reconPhantom is None or self.reconPhantom == []:
            raise ValueError("Reconstructed phantom is empty. Run reconstruction first.")

        if self.reconType in (ReconType.Analytic, ReconType.DeepLearning):
            self.MSE = mse(self.experiment.OpticImage.phantom, self.reconPhantom)

        elif self.reconType in (ReconType.Algebraic, ReconType.Bayesian, ReconType.Convex):
            self.MSE = []
            for theta in self.reconPhantom:
                self.MSE.append(mse(self.experiment.OpticImage.phantom, theta))
  
    def calculateSSIM(self):
        """
        Calculate the Structural Similarity Index (SSIM) of the reconstruction.

        Returns:
            ssim: float or list of floats, Structural Similarity Index of the reconstruction
        """

        if self.reconPhantom is None or self.reconPhantom == []:
            raise ValueError("Reconstructed phantom is empty. Run reconstruction first.")
    
        if self.reconType in (ReconType.Analytic, ReconType.DeepLearning):
            data_range = self.reconPhantom.max() - self.reconPhantom.min()
            self.SSIM = ssim(self.experiment.OpticImage.phantom, self.reconPhantom, data_range=data_range)

        elif self.reconType in (ReconType.Algebraic, ReconType.Bayesian):
            self.SSIM = []
            for theta in self.reconPhantom:
                data_range = theta.max() - theta.min()
                ssim_value = ssim(self.experiment.OpticImage.phantom, theta, data_range=data_range)
                self.SSIM.append(ssim_value)
    
    def show(self, withTumor=True, savePath=None, scale='same'):
        """
        Display the reconstructed images.
        Args:
            withTumor (bool): If True, displays reconPhantom. If False, displays reconLaser. Default is True.
            savePath (str): Path to save the figure. If None, the figure is not saved. Default is None.
            scale (str): Scale for the aspect ratio of the plots. Default is 'same'. Options are 'same' or 'auto'.
        """
        if withTumor:
            if self.reconPhantom is None:
                raise ValueError("Reconstructed phantom with tumor is empty. Run reconstruction first.")
            if isinstance(self.reconPhantom, (list, tuple)) and len(self.reconPhantom) == 0:
                raise ValueError("Reconstructed phantom with tumor is empty. Run reconstruction first.")
            if isinstance(self.reconPhantom, list):
                image = self.reconPhantom[-1]
            else:
                image = self.reconPhantom
            if self.experiment.OpticImage is None:
                fig, axs = plt.subplots(1, 1, figsize=(10, 10))
            else:
                fig, axs = plt.subplots(1, 2, figsize=(20, 10))
                if scale == 'same':
                    vmin = 0
                    vmax = 1
                elif scale == 'auto':
                    vmin = np.min(self.experiment.OpticImage.phantom)
                    vmax = np.max(self.experiment.OpticImage.phantom)
                    
                # Phantom original
                im1 = axs[1].imshow(
                    self.experiment.OpticImage.phantom,
                    cmap='hot',
                    vmin=vmin,
                    vmax=vmax,
                    extent=(
                        self.experiment.params.general['Xrange'][0],
                        self.experiment.params.general['Xrange'][1],
                        self.experiment.params.general['Zrange'][1],
                        self.experiment.params.general['Zrange'][0]
                    ),
                    aspect='equal'  
                )
                axs[1].set_title("Phantom with tumor")
                axs[1].set_xlabel("x (mm)", fontsize=12)
                axs[1].set_ylabel("z (mm)", fontsize=12)
                axs[1].tick_params(axis='both', which='major', labelsize=8)
            if scale == 'same':
                vmin = 0
                vmax = 1
            elif scale == 'auto':
                vmin = np.min(image)
                vmax = np.max(image)
            # Phantom reconstruit
            im0 = axs[0].imshow(
                image,
                cmap='hot',
                vmin=vmin,
                vmax=vmax,
                extent=(
                    self.experiment.params.general['Xrange'][0],
                    self.experiment.params.general['Xrange'][1],
                    self.experiment.params.general['Zrange'][1],
                    self.experiment.params.general['Zrange'][0]
                ),
                aspect='equal'  
            )
            axs[0].set_title("Reconstructed phantom with tumor")
            axs[0].set_xlabel("x (mm)", fontsize=12)
            axs[0].set_ylabel("z (mm)", fontsize=12)
            axs[0].tick_params(axis='both', which='major', labelsize=8)
            axs[0].tick_params(axis='y', which='both', left=False, right=False, labelleft=False)
        else:
            if self.reconLaser is None:
                raise ValueError("Reconstructed laser without tumor is empty. Run reconstruction first.")
            if isinstance(self.reconLaser, (list, tuple)) and len(self.reconLaser) == 0:
                raise ValueError("Reconstructed laser without tumor is empty. Run reconstruction first.")
            if isinstance(self.reconLaser, list):
                image = self.reconLaser[-1]
            else:
                image = self.reconLaser
            if self.experiment.OpticImage is None:
                fig, axs = plt.subplots(1, 1, figsize=(10, 10))
            else:
                fig, axs = plt.subplots(1, 2, figsize=(20, 10))
                # Laser original
                im1 = axs[1].imshow(
                    self.experiment.OpticImage.laser.intensity,
                    cmap='hot',
                    vmin=0,
                    vmax=np.max(self.experiment.OpticImage.laser.intensity),
                    extent=(
                        self.experiment.params.general['Xrange'][0],
                        self.experiment.params.general['Xrange'][1],
                        self.experiment.params.general['Zrange'][1],
                        self.experiment.params.general['Zrange'][0]
                    ),
                    aspect='equal'  
                )
                axs[1].set_title("Laser without tumor")
                axs[1].set_xlabel("x (mm)", fontsize=12)
                axs[1].set_ylabel("z (mm)", fontsize=12)
                axs[1].tick_params(axis='both', which='major', labelsize=8)
            # Laser reconstruit
            im0 = axs[0].imshow(
                image,
                cmap='hot',
                vmin=0,
                vmax=np.max(self.experiment.OpticImage.laser.intensity),
                extent=(
                    self.experiment.params.general['Xrange'][0],
                    self.experiment.params.general['Xrange'][1],
                    self.experiment.params.general['Zrange'][1],
                    self.experiment.params.general['Zrange'][0]
                ),
                aspect='equal'
            )
            axs[0].set_title("Reconstructed laser without tumor")
            axs[0].set_xlabel("x (mm)", fontsize=12)
            axs[0].set_ylabel("z (mm)", fontsize=12)
            axs[0].tick_params(axis='both', which='major', labelsize=8)
            axs[0].tick_params(axis='y', which='both', left=False, right=False, labelleft=False)

        # Colorbar commune
        fig.subplots_adjust(bottom=0.2)
        cbar_ax = fig.add_axes([0.25, 0.08, 0.5, 0.03])
        cbar = fig.colorbar(im0, cax=cbar_ax, orientation='horizontal')
        cbar.set_label('Normalized Intensity', fontsize=12)
        cbar.ax.tick_params(labelsize=8)

        plt.subplots_adjust(wspace=0.3)

        if savePath is not None:
            if not os.path.exists(savePath):
                os.makedirs(savePath)
            if withTumor:
                plt.savefig(os.path.join(savePath, 'recon_with_tumor.png'), dpi=300, bbox_inches='tight')
            else:
                plt.savefig(os.path.join(savePath, 'recon_without_tumor.png'), dpi=300, bbox_inches='tight')

        plt.show()

    def parse_and_demodulate(self, withTumor=True):

        if withTumor:
            AOsignal = self.experiment.AOsignal_withTumor
        else:
            AOsignal = self.experiment.AOsignal_withoutTumor
        delta_x = self.experiment.params.general['dx']  # en m 
        n_piezos = self.experiment.params.acoustic['probe']['num_elements']
        demodulated_data = {}
        structured_buffer = {} 

        for i in trange(len(self.experiment.AcousticFields), desc="Demodulating AO signals"):   
            label = self.experiment.AcousticFields[i].getName_field()
            
            parts = label.split("_")
            hex_pattern = parts[0]
            angle_code = parts[-1]
            
            # Angle
            if angle_code.startswith("1"):
                angle_deg = -int(angle_code[1:])
            else:
                angle_deg = int(angle_code)
            angle_rad = np.deg2rad(angle_deg)
            
            # Onde Plane (f_s = 0)
            if set(hex_pattern.lower().replace(" ", "")) == {'f'}:
                fs_key = 0.0 # fs_key est en mm^-1 (0.0 mm^-1)
                demodulated_data[(fs_key, angle_rad)] = np.array(AOsignal[:,i])
                continue
                
            # Onde Structurée
            profile = hex_to_binary_profile(hex_pattern, n_piezos)
            
            # Calcul FS (Fréquence de Structuration)
            ft_prof = np.fft.fft(profile)
            # On regarde uniquement la partie positive non DC
            idx_max = np.argmax(np.abs(ft_prof[1:len(profile)//2])) + 1
            freqs = np.fft.fftfreq(len(profile), d=delta_x)
            
            # freqs est en m^-1 car delta_x est en mètres.
            fs_m_inv = abs(freqs[idx_max]) 
            
            # *** CORRECTION 1: Conversion de f_s en mm^-1 (mm^-1 est utilisé dans iRadon) ***
            fs_key = fs_m_inv / 1000.0 # Fréquence spatiale en mm^-1

            
            if fs_key == 0: continue

            # Calcul de la Phase (Shift)
            phase = get_phase_deterministic(profile)
            
            # Stockage par (fs, theta) et phase
            key = (fs_key, angle_rad)
            if key not in structured_buffer:
                structured_buffer[key] = {}
            
            # La moyenne est nécessaire si plusieurs acquisitions ont la même phase (pour le SNR) 
            if phase in structured_buffer[key]:
                structured_buffer[key][phase] = (structured_buffer[key][phase] + np.array(AOsignal[:,i])) / 2
            else:
                structured_buffer[key][phase] = np.array(AOsignal[:,i])


        
        for (fs, theta), phases in structured_buffer.items():
            s0 = phases.get(0.0, 0)
            s_pi_2 = phases.get(np.pi/2, 0)
            s_pi = phases.get(np.pi, 0)
            s_3pi_2 = phases.get(3*np.pi/2, 0)

            # Assurer que les zéros sont des vecteurs de la bonne taille
            example = next(val for val in phases.values() if not isinstance(val, int))
            if isinstance(s0, int): s0 = np.zeros_like(example)
            if isinstance(s_pi, int): s_pi = np.zeros_like(example)
            if isinstance(s_pi_2, int): s_pi_2 = np.zeros_like(example)
            if isinstance(s_3pi_2, int): s_3pi_2 = np.zeros_like(example)

            real = s0 - s_pi
            imag = s_pi_2 - s_3pi_2 
              
            demodulated_data[(fs, theta)] = (real - 1j * imag) / (2/np.pi)
            
        return demodulated_data
