import numpy as np
import pydicom
from decimal import Decimal, ROUND_HALF_UP
import traceback

# Imports relativos de la librería
from .image_utils import segment_and_draw_rois
from .physics.metrics import (
    calculate_snr, 
    calculate_sdnr, 
    detect_angle_in_horizontal_roi, 
    detect_angle_in_vertical_roi, 
    pmft_horizontal_iaea_like
)
from .physics.nnps import calculate_nnps, extract_roi
from .physics.detectability import calculate_dprime

class AtiaMamo:
    """
    Clase principal de PyATIA para el análisis de mamografías digitales (protocolo IAEA).
    """

    def process(self, dicom_path: str) -> dict:
        """
        Procesa un archivo DICOM y calcula métricas de calidad de imagen.

        Args:
            dicom_path (str): Ruta al archivo .dcm.

        Returns:
            dict: Diccionario con llaves 'success', 'metrics', 'rois', 'image_processed', 'error'.
        """
        try:
            # -----------------------------------------------------------
            # 1. CARGA
            # -----------------------------------------------------------
            ds = pydicom.dcmread(dicom_path)
            
            # Pixel Spacing (Crítico)
            px = ds.get("PixelSpacing", [0.1, 0.1])
            if isinstance(px, (list, pydicom.multival.MultiValue)) and len(px) >= 1:
                pixel_spacing_mm = float(px[0])
            else:
                pixel_spacing_mm = 0.1  # Fallback
            
            # Imagen RAW (float)
            original_image = ds.pixel_array.astype(np.float32)

            # -----------------------------------------------------------
            # 2. PREPROCESAMIENTO VISUAL (ROBUST)
            # -----------------------------------------------------------
            # Obtener valores crudos (pueden ser MultiValue/listas)
            raw_wc = ds.get("WindowCenter")
            raw_ww = ds.get("WindowWidth")

            # Valores por defecto basados en la imagen
            default_wc = np.median(original_image)
            default_ww = np.max(original_image) - np.min(original_image)

            # Lógica para Window Center (WC)
            if raw_wc is None:
                wc = float(default_wc)
            elif isinstance(raw_wc, (list, pydicom.multival.MultiValue)):
                wc = float(raw_wc[0]) # Tomar el primer valor si es lista
            else:
                wc = float(raw_wc)

            # Lógica para Window Width (WW)
            if raw_ww is None:
                ww = float(default_ww)
            elif isinstance(raw_ww, (list, pydicom.multival.MultiValue)):
                ww = float(raw_ww[0]) # Tomar el primer valor si es lista
            else:
                ww = float(raw_ww)

            # Aplicar Windowing
            min_win = wc - (ww / 2)
            max_win = wc + (ww / 2)
            image_vis = np.clip(original_image, min_win, max_win)
            image_vis = ((image_vis - min_win) / (max_win - min_win)) * 255.0
            image_vis = image_vis.astype(np.uint8)

            # -----------------------------------------------------------
            # 3. SEGMENTACIÓN DE ROIS
            # -----------------------------------------------------------
            image_with_rois, rois = segment_and_draw_rois(image_vis, pixel_spacing_mm=pixel_spacing_mm)

            if not rois or "background_roi" not in rois:
                return {"success": False, "error": "No se detectaron las ROIs del maniquí."}

            # -----------------------------------------------------------
            # 4. NORMALIZACIÓN FÍSICA (Mean Background = 1.0)
            # -----------------------------------------------------------
            x_bg, y_bg, w_bg, h_bg = rois["background_roi"]
            roi_bg_pixels = original_image[y_bg:y_bg+h_bg, x_bg:x_bg+w_bg]
            mean_background = float(np.mean(roi_bg_pixels))
            
            normalized_image = original_image / (mean_background if mean_background != 0 else 1.0)

            # -----------------------------------------------------------
            # 5. CÁLCULO DE NNPS (Ruido)
            # -----------------------------------------------------------
            # Extraer ROI centrada (fallback a 256 si 512 es muy grande)
            roi_nnps = extract_roi(normalized_image, centered=True, size=(512, 512))
            if roi_nnps.shape[0] < 512:
                roi_nnps = extract_roi(normalized_image, centered=True, size=(256, 256))

            nnps_result = calculate_nnps(roi_nnps, pixel_spacing_mm, save_csv=False)

            # -----------------------------------------------------------
            # 6. CÁLCULO DE MTF
            # -----------------------------------------------------------
            # --- MTF Horizontal (Borde Inferior) ---
            roi_h = rois["roi_vertical"]
            detH = detect_angle_in_horizontal_roi(original_image, roi_h, sigma=1.0)
            f_h, pmtf_h, metrics_h, _ = pmft_horizontal_iaea_like(
                original_image, roi_h,
                pixel_spacing_mm=pixel_spacing_mm,
                angle_tilt_deg=detH["angle_deg"],
                out_png=None 
            )

            # --- MTF Vertical (Borde Derecho, rotado) ---
            x, y, w, h = rois["roi_horizontal"]
            detV = detect_angle_in_vertical_roi(original_image, (x, y, w, h), sigma=1.0)
            roi_v_rot = np.rot90(original_image[y:y+h, x:x+w], k=1)
            f_v, p_v, metrics_v, _ = pmft_horizontal_iaea_like(
                roi_v_rot, (0, 0, roi_v_rot.shape[1], roi_v_rot.shape[0]),
                pixel_spacing_mm=pixel_spacing_mm,
                angle_tilt_deg=-detV["angle_deg"],
                out_png=None
            )

            # -----------------------------------------------------------
            # 7. DETECTABILIDAD (d') - CORREGIDA
            # -----------------------------------------------------------
            # Contraste medido del inserto real
            rx, ry, rw, rh = rois["roi_small"]
            mean_sig_norm = np.mean(normalized_image[ry:ry+rh, rx:rx+rw])
            
            # Contraste absoluto: |Signal - Bg| / Bg (donde Bg=1.0)
            contrast_measured = abs(mean_sig_norm - 1.0)
            
            # Factor de corrección IAEA (Amplitud = Contraste / 2)
            contrast_final = contrast_measured * 0.5

            dprime_res = calculate_dprime(
                nnps_result, 
                mtf_freq=f_h,  # Usamos MTF horizontal como referencia
                mtf_val=pmtf_h, 
                diameter_mm_list=[0.1, 0.25, 0.5],
                contrast=contrast_final
            )

            # -----------------------------------------------------------
            # 8. CONSTRUCCIÓN DE RESULTADOS
            # -----------------------------------------------------------
            
            # Helpers internos para limpiar floats y redondear
            def _rd(val, d=2):
                if val is None: return None
                return float(Decimal(float(val)).quantize(Decimal("1."+"0"*d), rounding=ROUND_HALF_UP))

            def _cross(freqs, mtf, level):
                idx = np.where(mtf <= level)[0]
                if idx.size == 0: return None
                j = idx[0]
                if j == 0: return float(freqs[0])
                # Interp lineal
                f0, f1 = freqs[j-1], freqs[j]
                y0, y1 = mtf[j-1], mtf[j]
                return float(f0 + (level - y0) * (f1 - f0) / (y1 - y0))

            # Valores puntuales MTF
            mtf50_h = metrics_h.get("pMTF50") or _cross(f_h, pmtf_h, 0.5)
            mtf10_h = metrics_h.get("pMTF10") or _cross(f_h, pmtf_h, 0.1)
            mtf50_v = metrics_v.get("pMTF50") or _cross(f_v, p_v, 0.5)
            mtf10_v = metrics_v.get("pMTF10") or _cross(f_v, p_v, 0.1)
            nyquist = 1.0 / (2.0 * pixel_spacing_mm)

            metrics = {
                "pixel_spacing_mm": pixel_spacing_mm,
                "snr": _rd(calculate_snr(normalized_image, rois["background_roi"])),
                "sdnr": _rd(calculate_sdnr(normalized_image, rois["roi_small"], rois["background_roi"])),
                
                # Geometría
                "angle_h_deg": _rd(detH["angle_deg"]),
                "angle_v_deg": _rd(detV["angle_deg"]),

                # MTF Resumen
                "mtf50_h": _rd(mtf50_h), "mtf10_h": _rd(mtf10_h),
                "mtf50_v": _rd(mtf50_v), "mtf10_v": _rd(mtf10_v),
                "nyquist_freq": _rd(nyquist),

                # Detectabilidad
                "dprime_0.10mm": _rd(dprime_res.get("dprime_0.1mm"), 3),
                "dprime_0.25mm": _rd(dprime_res.get("dprime_0.25mm"), 3),
                "dprime_0.50mm": _rd(dprime_res.get("dprime_0.5mm"), 3),
                
                # Datos crudos para graficar (opcional)
                "mtf_curve_h": {f"{k:.2f}": float(v) for k,v in zip(f_h, pmtf_h)},
                "mtf_curve_v": {f"{k:.2f}": float(v) for k,v in zip(f_v, p_v)},
                "nnps_radial": {f"{k:.2f}": float(v) for k,v in zip(nnps_result["radial_freq"], nnps_result["radial_nnps"])}
            }

            return {
                "success": True,
                "metrics": metrics,
                "rois": rois, # Coordenadas (x,y,w,h)
                "image_processed": image_with_rois, # Array uint8 para mostrar/guardar
                "error": None
            }

        except Exception as e:
            traceback.print_exc() # Imprime el error en consola para debug
            return {
                "success": False, 
                "error": str(e),
                "metrics": {},
                "image_processed": None
            }