# src/feature_extractor.py

import cv2
import numpy as np
from typing import Dict, List
from langsmith import traceable


# ─── Yardımcı kategorileştirici fonksiyonlar ─────────────────────────────────

def _size_category(max_dim_mm: float) -> str:
    """EAU klinik eşiklerine göre boyut kategorisi."""
    if max_dim_mm < 5:
        return "small"
    elif max_dim_mm < 10:
        return "medium"
    elif max_dim_mm < 20:
        return "large"
    else:
        return "very_large"


def _density_category(mean_intensity: float) -> str:
    """0-255 piksel skalasına göre relative density kategorisi."""
    if mean_intensity > 230:
        return "very_high"
    elif mean_intensity > 200:
        return "high"
    elif mean_intensity > 160:
        return "moderate"
    else:
        return "low"


def _shape_from_mask(circularity: float, eccentricity: float) -> str:
    """Circularity + eccentricity'den şekil kategorisi."""
    if circularity >= 0.85:
        return "round"
    elif circularity >= 0.65 and eccentricity < 0.7:
        return "oval"
    elif circularity >= 0.40:
        return "irregular"
    else:
        return "highly_irregular"


def _empty_morphology() -> Dict:
    return {
        "area_px": 0, "area_mm2": 0, "perimeter_px": 0,
        "circularity": 0, "solidity": 0, "eccentricity": 0,
        "orientation_deg": 0, "equivalent_diameter_mm": 0,
        "size_category": "unknown", "shape_category": "unknown"
    }


# ─── Bbox-tabanlı fonksiyonlar (fallback — mask yoksa kullanılır) ─────────────

def extract_basic_features(detection_result: Dict) -> Dict:
    return {
        "stone_detected": detection_result["num_detections"] > 0,
        "count": detection_result["num_detections"],
        "image_dimensions": {
            "height": detection_result["image_shape"][0],
            "width": detection_result["image_shape"][1]
        }
    }


def estimate_size_mm(bbox, pixel_spacing_mm: float = 0.7) -> Dict:
    """Bbox'tan yaklaşık boyut (mask yoksa fallback)."""
    x1, y1, x2, y2 = bbox
    width_mm = (x2 - x1) * pixel_spacing_mm
    height_mm = (y2 - y1) * pixel_spacing_mm
    max_dim_mm = max(width_mm, height_mm)
    return {
        "width_mm": round(width_mm, 1),
        "height_mm": round(height_mm, 1),
        "max_dimension_mm": round(max_dim_mm, 1),
        "size_category": _size_category(max_dim_mm)
    }


def get_location(bbox, image_shape) -> Dict:
    """Bbox center'dan konum (mask yoksa fallback)."""
    x1, y1, x2, y2 = bbox
    cx = (x1 + x2) / 2
    cy = (y1 + y2) / 2
    H, W = image_shape
    vertical = "upper" if cy < H / 2 else "lower"
    horizontal = "left" if cx < W / 2 else "right"
    return {
        "quadrant": f"{vertical}-{horizontal}",
        "normalized_center": {
            "x": round(cx / W, 3),
            "y": round(cy / H, 3)
        }
    }


def get_shape(bbox) -> str:
    """Bbox aspect ratio'dan şekil (mask yoksa fallback)."""
    x1, y1, x2, y2 = bbox
    w = x2 - x1
    h = y2 - y1
    if w == 0 or h == 0:
        return "unknown"
    aspect_ratio = max(w, h) / min(w, h)
    if aspect_ratio < 1.3:
        return "round"
    elif aspect_ratio < 2.0:
        return "oval"
    else:
        return "elongated"


def get_density(image, bbox) -> Dict:
    """Bbox ROI'sinden density (mask yoksa fallback)."""
    x1, y1, x2, y2 = [int(c) for c in bbox]
    roi = image[y1:y2, x1:x2]
    if roi.size == 0:
        return {"category": "unknown", "mean_intensity": 0}
    roi_gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if len(roi.shape) == 3 else roi
    bright = roi_gray[roi_gray > 150]
    mean_int = float(bright.mean()) if len(bright) > 0 else float(roi_gray.mean())
    return {
        "category": _density_category(mean_int),
        "mean_intensity": round(mean_int, 1)
    }


# ─── Mask-tabanlı fonksiyonlar ────────────────────────────────────────────────

def get_mask_morphology(mask_polygon: List,
                        pixel_spacing_mm: float = 0.7) -> Dict:
    """
    Mask polygon'dan morfometrik analiz.
    Döndürür: alan, çevre, circularity, solidity, eccentricity,
              orientation, equivalent_diameter, şekil kategorisi.
    """
    pts = np.array(mask_polygon, dtype=np.float32)
    if pts.ndim == 1:
        pts = pts.reshape(-1, 2)
    if len(pts) < 3:
        return _empty_morphology()

    contour_f = pts.reshape((-1, 1, 2))
    contour_i = contour_f.astype(np.int32)

    area_px = float(cv2.contourArea(contour_f))
    if area_px < 1:
        return _empty_morphology()

    perimeter_px = float(cv2.arcLength(contour_f, closed=True))

    # Circularity (1.0 = mükemmel daire)
    circularity = (
        min((4 * np.pi * area_px) / (perimeter_px ** 2), 1.0)
        if perimeter_px > 0 else 0.0
    )

    # Solidity: alan / convex hull alanı
    hull = cv2.convexHull(contour_i)
    hull_area = float(cv2.contourArea(hull))
    solidity = area_px / hull_area if hull_area > 0 else 0.0

    # Eccentricity + orientation (fitEllipse ≥ 5 nokta gerektirir)
    eccentricity = 0.0
    orientation_deg = 0.0
    if len(contour_f) >= 5:
        try:
            _, (minor_ax, major_ax), angle = cv2.fitEllipse(contour_f)
            if major_ax > 0:
                eccentricity = float(np.sqrt(max(0.0, 1 - (minor_ax / major_ax) ** 2)))
            orientation_deg = float(angle)
        except cv2.error:
            pass

    equiv_diameter_mm = round(
        float(np.sqrt(4 * area_px / np.pi)) * pixel_spacing_mm, 1
    )
    area_mm2 = round(area_px * (pixel_spacing_mm ** 2), 2)

    return {
        "area_px": round(area_px, 1),
        "area_mm2": area_mm2,
        "perimeter_px": round(perimeter_px, 1),
        "circularity": round(circularity, 3),
        "solidity": round(solidity, 3),
        "eccentricity": round(eccentricity, 3),
        "orientation_deg": round(orientation_deg, 1),
        "equivalent_diameter_mm": equiv_diameter_mm,
        "size_category": _size_category(equiv_diameter_mm),
        "shape_category": _shape_from_mask(circularity, eccentricity)
    }


def get_density_profile(image, mask_polygon: List, image_shape: tuple) -> Dict:
    """
    Mask içindeki piksellerden yoğunluk profili.
    Sadece taş piksellerini kullanır — bbox ROI'sinden daha doğru.
    Not: Değerler 0-255 PNG piksel skalasında, gerçek HU değil.
    """
    H, W = image_shape[:2]
    mask_img = np.zeros((H, W), dtype=np.uint8)
    pts = np.array(mask_polygon, dtype=np.int32)
    if pts.ndim == 1:
        pts = pts.reshape(-1, 2)
    cv2.fillPoly(mask_img, [pts], 255)

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
    masked_pixels = gray[mask_img > 0]

    if masked_pixels.size == 0:
        return {
            "category": "unknown", "mean_intensity": 0,
            "std_intensity": 0, "homogeneity": "unknown",
            "contrast_to_surrounding": 0
        }

    mean_int = float(masked_pixels.mean())
    std_int = float(masked_pixels.std())

    if std_int < 20:
        homogeneity = "homogeneous"
    elif std_int < 40:
        homogeneity = "moderately_heterogeneous"
    else:
        homogeneity = "heterogeneous"

    # Çevre kontrast: mask bbox'u etrafındaki dış pikseller
    ys, xs = np.where(mask_img > 0)
    contrast = 0.0
    if len(ys) > 0:
        pad = 10
        y1s = max(0, int(ys.min()) - pad)
        y2s = min(H, int(ys.max()) + pad)
        x1s = max(0, int(xs.min()) - pad)
        x2s = min(W, int(xs.max()) + pad)
        outer_pixels = gray[y1s:y2s, x1s:x2s][mask_img[y1s:y2s, x1s:x2s] == 0]
        if outer_pixels.size > 0:
            contrast = round(mean_int - float(outer_pixels.mean()), 1)

    return {
        "category": _density_category(mean_int),
        "mean_intensity": round(mean_int, 1),
        "std_intensity": round(std_int, 1),
        "homogeneity": homogeneity,
        "contrast_to_surrounding": contrast
    }


def get_location_from_mask(mask_polygon: List, image_shape: tuple) -> Dict:
    """Mask centroid'den konum — bbox center'dan daha doğru."""
    pts = np.array(mask_polygon, dtype=np.float32)
    if pts.ndim == 1:
        pts = pts.reshape(-1, 2)
    contour = pts.reshape((-1, 1, 2)).astype(np.int32)

    M = cv2.moments(contour)
    H, W = image_shape
    if M["m00"] > 0:
        cx = M["m10"] / M["m00"]
        cy = M["m01"] / M["m00"]
    else:
        cx = float(pts[:, 0].mean())
        cy = float(pts[:, 1].mean())

    vertical = "upper" if cy < H / 2 else "lower"
    horizontal = "left" if cx < W / 2 else "right"
    return {
        "quadrant": f"{vertical}-{horizontal}",
        "normalized_center": {
            "x": round(cx / W, 3),
            "y": round(cy / H, 3)
        }
    }


# ─── Ana çıkarım fonksiyonu ───────────────────────────────────────────────────

@traceable(name="Feature Extraction", run_type="tool")
def extract_full_features(detection_result: Dict,
                          pixel_spacing_mm: float = 0.7) -> Dict:
    """
    YOLO26-seg detection/segmentation result'tan RAG için komple feature object üret.
    mask_polygon varsa mask-based özellikler kullanılır, yoksa bbox fallback.
    """
    image = cv2.imread(detection_result["image_path"])
    features = extract_basic_features(detection_result)

    if not features["stone_detected"]:
        features["message"] = "Bu görüntüde taş tespit edilmedi."
        features["stones"] = []
        return features

    features["stones"] = []
    for i, det in enumerate(detection_result["detections"]):
        bbox = det["bbox"]
        mask_polygon = det.get("mask_polygon")
        has_mask = mask_polygon is not None and len(mask_polygon) >= 3

        if has_mask:
            morph = get_mask_morphology(mask_polygon, pixel_spacing_mm)
            density_info = get_density_profile(
                image, mask_polygon, detection_result["image_shape"]
            )
            location_info = get_location_from_mask(
                mask_polygon, detection_result["image_shape"]
            )
            size_info = {
                "equivalent_diameter_mm": morph["equivalent_diameter_mm"],
                "max_dimension_mm": morph["equivalent_diameter_mm"],
                "area_mm2": morph["area_mm2"],
                "size_category": morph["size_category"]
            }
            shape = morph["shape_category"]
        else:
            morph = None
            size_info = estimate_size_mm(bbox, pixel_spacing_mm)
            location_info = get_location(bbox, detection_result["image_shape"])
            shape = get_shape(bbox)
            density_info = get_density(image, bbox)

        stone_features = {
            "stone_id": i + 1,
            "detection_confidence": round(det["confidence"], 2),
            "size": size_info,
            "location": location_info,
            "shape": shape,
            "density": density_info,
            "bbox_pixels": [round(c, 1) for c in bbox]
        }

        if has_mask and morph:
            stone_features["morphology"] = {
                "area_px": morph["area_px"],
                "area_mm2": morph["area_mm2"],
                "perimeter_px": morph["perimeter_px"],
                "circularity": morph["circularity"],
                "solidity": morph["solidity"],
                "eccentricity": morph["eccentricity"],
                "orientation_deg": morph["orientation_deg"]
            }

        features["stones"].append(stone_features)

    if features["count"] > 1:
        features["largest_stone_mm"] = max(
            s["size"]["max_dimension_mm"] for s in features["stones"]
        )
        features["multiple_stones"] = True
        areas = [
            s["size"].get("area_mm2") for s in features["stones"]
            if s["size"].get("area_mm2") is not None
        ]
        if areas:
            features["total_stone_area_mm2"] = round(sum(areas), 2)
    else:
        features["multiple_stones"] = False

    return features