from ._mainRecon import Recon
from .ReconEnums import ReconType, AnalyticType, ProcessType
from AOT_biomaps.AOT_Experiment.Tomography import hex_to_binary_profile
from .ReconTools import fourierz_gpu, get_phase_deterministic, add_sincos_cpu, EvalDelayLawOS_center, ifourierx_gpu, rotate_theta_gpu, filter_radon_gpu, ifourierz_gpu  

import numpy as np
from tqdm import trange
import cupy as cp


class AnalyticRecon(Recon):
    def __init__(self, analyticType, Lc = None,**kwargs):
        super().__init__(**kwargs)
        self.reconType = ReconType.Analytic
        self.analyticType = analyticType
        if self.analyticType == AnalyticType.iRADON and Lc is None:
            raise ValueError("Lc parameter must be provided for iRADON analytic reconstruction.")
        self.Lc = Lc # in meters
        self.AOsignal_demoldulated = None

    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['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

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

    def _analyticReconPython(self,withTumor):
        """
        This method is a placeholder for the analytic reconstruction process in Python.
        It currently does not perform any operations but serves as a template for future implementations.
        
        Parameters:
            analyticType: The type of analytic reconstruction to perform (default is iFOURIER).
        """
        d_t = 1 / float(self.experiment.params.acoustic['f_saving'])
        t_array = np.arange(0, self.experiment.AOsignal_withTumor.shape[0])*d_t
        Z = t_array * self.experiment.params.acoustic['c0']
        X_m = np.arange(0, self.experiment.params.acoustic['num_elements'])* self.experiment.params.general['dx']
        dfX = 1 / (X_m[1] - X_m[0]) / len(X_m)
        if withTumor:
            self.AOsignal_demoldulated = self.parse_and_demodulate(withTumor=True)
            if self.analyticType == AnalyticType.iFOURIER:
                self.reconPhantom = self._iFourierRecon(
                    R = self.experiment.AOsignal_withTumor,
                    z = Z,    
                    X_m=X_m,
                    theta=self.experiment.theta,
                    decimation=self.experiment.decimations,
                    c=self.experiment.params.acoustic['c0'],
                    DelayLAWS=self.experiment.DelayLaw,
                    ActiveLIST=self.experiment.ActiveList,
                    withTumor=True,
                )
                    
            elif self.analyticType == AnalyticType.iRADON:
                self.reconPhantom = self._iRadonRecon(
                    R=self.experiment.AOsignal_withTumor,
                    z=Z,
                    X_m=X_m,
                    theta=self.experiment.theta,
                    decimation=self.experiment.decimations,
                    df0x=dfX,
                    Lc =self.Lc,
                    c=self.experiment.params.acoustic['c0'],
                    DelayLAWS=self.experiment.DelayLaw,
                    ActiveLIST=self.experiment.ActiveList,
                    withTumor=True)
            else:            
                raise ValueError(f"Unknown analytic type: {self.analyticType}")
        else:
            self.AOsignal_demoldulated = self.parse_and_demodulate(withTumor=False)
            if self.analyticType == AnalyticType.iFOURIER:
                self.reconLaser = self._iFourierRecon(
                    R = self.experiment.AOsignal_withoutTumor,
                    z = Z,    
                    X_m=X_m,
                    theta=self.experiment.theta,
                    decimation=self.experiment.decimations,
                    c=self.experiment.params.acoustic['c0'],
                    DelayLAWS=self.experiment.DelayLaw,
                    ActiveLIST=self.experiment.ActiveList,
                    withTumor=False,
                )
            elif self.analyticType == AnalyticType.iRADON:
                self.reconLaser = self._iRadonRecon(
                    R=self.experiment.AOsignal_withoutTumor,
                    z=Z,
                    X_m=X_m,
                    theta=self.experiment.theta,
                    decimation=self.experiment.decimations,
                    df0x=dfX,
                    Lc = self.Lc,
                    c=self.experiment.params.acoustic['c0'],
                    DelayLAWS=self.experiment.DelayLaw,
                    ActiveLIST=self.experiment.ActiveList,
                    withTumor=False)
            else:            
                raise ValueError(f"Unknown analytic type: {self.analyticType}")
    
    def _iFourierRecon(
        self,
        R,
        z, 
        X_m, 
        theta, 
        decimation,  
        c, 
        DelayLAWS, 
        ActiveLIST,
        withTumor,
    ):
        # --- 1. Préparation des données ---
        R = cp.asarray(R)
        z = cp.asarray(z)
        X_m = cp.asarray(X_m)
        theta = cp.asarray(theta)
        decimation = cp.asarray(decimation)
        DelayLAWS = cp.asarray(DelayLAWS)
        ActiveLIST = cp.asarray(ActiveLIST)
        # Normalisation DelayLAWS (ms → s si besoin)
        DelayLAWS_s = cp.where(cp.max(DelayLAWS) > 1e-3, DelayLAWS / 1000.0, DelayLAWS)
        # Regroupement tirs (décimation, angle)
        ScanParam_cpu = cp.asnumpy(
            cp.stack([decimation, cp.round(theta, 4)], axis=1)
        )

        _, ia_cpu, ib_cpu = np.unique(
            ScanParam_cpu, axis=0, return_index=True, return_inverse=True
        )

        ia = cp.asarray(ia_cpu)
        ib = cp.asarray(ib_cpu)

 
        # --- 2. Structuration complexe ---
        # add_sincos_gpu doit être l'équivalent GPU de add_sincos_cpu
        F_complex_cpu, theta_u_cpu, decim_u_cpu = add_sincos_cpu(
        cp.asnumpy(R),
        cp.asnumpy(decimation),
        np.radians(cp.asnumpy(theta))
        )
        M0 = EvalDelayLawOS_center(
            X_m,
            theta_u_cpu,
            DelayLAWS_s.T[:, ia],
            ActiveLIST.T[:, ia],
            c
        )

        # Transfert GPU (UNE FOIS)
        F_complex = cp.asarray(F_complex_cpu)
        theta_u = cp.asarray(theta_u_cpu)
        decim_u = cp.asarray(decim_u_cpu)

        Nz = z.size
        Nx = X_m.size
        dx = X_m[1] - X_m[0]

        X_grid, Z_grid = cp.meshgrid(X_m, z)
        idx0_x = Nx // 2
        # Angles uniques après compression
        angles_group, ia_u, ib_u = cp.unique(
            theta_u, return_index=True, return_inverse=True
        )
        Ntheta = angles_group.size

        I_final = cp.zeros((Nz, Nx), dtype=cp.complex64)

        # --- 3. Boucle InverseFourierX ---
        for i_ang in trange(Ntheta, desc=f"AOT-BioMaps -- Analytic Reconstruction Tomography: iFourier ({'with tumor' if withTumor else 'without tumor'})  ---- processing on single GPU ----", unit="angle"):

            # Grille (z, fx)
            F_fx_z = cp.zeros((Nz, Nx), dtype=cp.complex64)

            # Indices correspondant à cet angle
            indices = cp.where(ib_u == i_ang)[0]

            for idx in indices:
                n = int(decim_u[idx])
                trace_z = F_complex[:, idx]

                # Mapping positif
                ip = idx0_x + n
                if 0 <= ip < Nx:
                    F_fx_z[:, ip] = trace_z

                # Mapping négatif (symétrie hermitienne Matlab)
                if n != 0:
                    im = idx0_x - n
                    if 0 <= im < Nx:
                        col_conj = cp.zeros(Nz, dtype=cp.complex64)
                        col_conj[1:] = cp.conj(trace_z[:-1])
                        F_fx_z[:, im] = col_conj

            # Correction DC
            F_fx_z[:, idx0_x] *= 0.5

            # --- Inverse Fourier X (GPU) ---
            I_spatial = ifourierx_gpu(F_fx_z, dx)

            # --- Rotation spatiale GPU ---
            I_rot = rotate_theta_gpu(
                X_grid,
                Z_grid,
                I_spatial,
                -angles_group[i_ang],
                M0[i_ang, :]
            )

            # Somme incohérente (OriginIm = sum)
            I_final += I_rot

        return cp.real(I_final).get(), z, X_m

    def _iRadonRecon(
        self,
        R, 
        z, 
        X_m,
        theta, 
        decimation,
        df0x,
        Lc,
        c,
        DelayLAWS,
        ActiveLIST,
        withTumor,
):
        """
        Reconstruction d'image utilisant la méthode iRadon.

        :return: Image reconstruite.
        """
            
        # ======================================================
        # 1. AddSinCos (structuration)
        # ======================================================
        theta = np.radians(theta)
        F_ct_kx, theta_u, decim_u = add_sincos_cpu(R, decimation, theta)
        ScanParam = np.stack([decimation, theta], axis=1)
        _, ia, _ = np.unique(ScanParam, axis=0, return_index=True, return_inverse=True)
        ActiveLIST = np.asarray(ActiveLIST).T
        DelayLAWS = np.asarray(DelayLAWS).T
        ActiveLIST_unique = ActiveLIST[:,ia]

        # ======================================================
        # 2. FFT z
        # ======================================================
        z_gpu = cp.asarray(z)
        Fin = fourierz_gpu(z, F_ct_kx)


        dz = float(z[1] - z[0])
        fz = cp.fft.fftshift(cp.fft.fftfreq(len(z), d=dz))

        Nz, Nk = Fin.shape

        # ======================================================
        # 3. Filtrage OS exact
        # ======================================================
        decim_gpu = cp.asarray(decim_u)
        I0 = decim_gpu == 0
        F0 = Fin * I0[None, :]  # I0 broadcastée sur les lignes

        DEC, FZ = cp.meshgrid(decim_gpu, fz)

        Hinf = cp.abs(FZ) < cp.abs(DEC) * df0x
        Hsup = FZ >= 0

        Fc = 1 / Lc
        FILTER = filter_radon_gpu(fz, Fc)[:, None]

        Finf = F0 * FILTER[:, :F0.shape[1]] * Hinf[:, :F0.shape[1]]
        Fsup = Fin * FILTER * Hsup

        # ======================================================
        # 4. Retour espace z
        # ======================================================
        Finf = ifourierz_gpu(z, Finf)
        Fsup = ifourierz_gpu(z, Fsup)

        # ======================================================
        # 5. Grille image
        # ======================================================
        X_gpu = cp.asarray(X_m)
        X, Z = cp.meshgrid(X_gpu, z_gpu)
        Xc = float(np.mean(X_m))

        # ======================================================
        # 6. Calcul du centre M0 pour chaque angle
        # ======================================================
        M0 = EvalDelayLawOS_center(X_m,theta, DelayLAWS[:, ia], ActiveLIST_unique, c)
        M0_gpu = cp.asarray(M0)

        # ======================================================
        # 7. Rétroprojection
        # ======================================================
        Irec_list = []
        Irec = cp.zeros_like(X, dtype=cp.complex64)
        for i in trange(len(theta_u), desc=f"AOT-BioMaps -- Analytic Reconstruction Tomography: iFourier ({'with tumor' if withTumor else 'without tumor'})  ---- processing on single GPU ----", unit="angle"):
            th = float(theta_u[i])
            
            # calcul de T, S, h0...
            T = (X - M0_gpu[i,0]) * cp.sin(th) + (Z - M0_gpu[i,1]) * cp.cos(th) + M0_gpu[i,1]
            S = (X - Xc) * cp.cos(th) - (Z - M0_gpu[i,1]) * cp.sin(th)
            h0 = cp.exp(1j * 2*cp.pi * decim_u[i] * df0x * S)

            # interpolation GPU
            Tind = (T - z_gpu[0]) / (z_gpu[1] - z_gpu[0])
            i0 = cp.floor(Tind).astype(cp.int32)
            i1 = i0 + 1
            i0 = cp.clip(i0, 0, Nz-1)
            i1 = cp.clip(i1, 0, Nz-1)
            w = Tind - i0

            proj_sup = (1-w)*Fsup[i0, i] + w*Fsup[i1, i]
            proj_inf = (1-w)*Finf[i0, i] + w*Finf[i1, i]

            Irec += 2*h0*proj_sup + proj_inf

            # Ajouter à la liste
            Irec_list.append(cp.real(Irec).get())

        return (Irec_list, z, X_m)

