from typing import Dict, Any

import numpy as np
from monai.transforms.utils import ndimage

from src.quality.quality_scorer import QualityScorer

class MedicalQualityScorer(QualityScorer):
    def __init__(self, anatomy_type="spine"):
        super().__init__()
        self.anatomy_type = anatomy_type  # "spine", "femur", "whole_body"

    def evaluate(self, image, segmentation) -> Dict[str, Any]:
        # Базовые проверки
        base = super().evaluate(image, segmentation)

        # Специфические проверки
        if self.anatomy_type == "spine":
            medical_check = self.check_spine_segmentation(segmentation)
        elif self.anatomy_type == "femur":
            medical_check = self.check_femur_segmentation(segmentation)

        return {
            **base,
            "medical_validation": medical_check
        }

    def check_spine_segmentation(self, segmentation: np.ndarray) -> Dict[str, Any]:
        """
        Комплексная проверка сегментации позвоночника

        Args:
            segmentation: бинарная маска [H, W]

        Returns:
            Словарь с результатами проверки
        """
        results = {
            "valid": True,
            "issues": [],
            "metrics": {},
            "recommendations": []
        }

        # 1. Проверка количества позвонков
        labeled, num_features = ndimage.label(segmentation)
        results["metrics"]["num_vertebrae"] = num_features

        if num_features < 3:
            results["valid"] = False
            results["issues"].append({
                "type": "TOO_FEW_VERTEBRAE",
                "details": f"Найдено только {num_features} позвонков, ожидается минимум 3",
                "severity": "HIGH"
            })
            results["recommendations"].append("Повторите исследование, захватите больше позвонков")
            return results  # Дальнейшие проверки не имеют смысла

        # 2. Собираем центры позвонков
        centers = []
        sizes = []
        for i in range(1, num_features + 1):
            y, x = np.where(labeled == i)
            if len(y) > 0:
                centers.append((np.mean(y), np.mean(x)))
                sizes.append(len(y))

        # 3. Проверка выравнивания
        x_coords = [c[1] for c in centers]
        x_std = np.std(x_coords)
        results["metrics"]["alignment_deviation"] = float(x_std)

        if x_std > 20:
            results["valid"] = False
            results["issues"].append({
                "type": "POOR_ALIGNMENT",
                "details": f"Позвонки не выровнены (отклонение {x_std:.1f} пикселей)",
                "severity": "MEDIUM"
            })
            results["recommendations"].append("Проверьте укладку пациента")

        # 4. Проверка расстояний между позвонками
        if len(centers) >= 2:
            centers_sorted = sorted(centers, key=lambda c: c[0])
            distances = []
            for i in range(len(centers_sorted) - 1):
                y1, _ = centers_sorted[i]
                y2, _ = centers_sorted[i + 1]
                distances.append(abs(y2 - y1))

            if len(distances) > 0:
                mean_dist = np.mean(distances)
                std_dist = np.std(distances)
                results["metrics"]["spacing_variation"] = float(std_dist / mean_dist)

                if std_dist / mean_dist > 0.3:
                    results["valid"] = False
                    results["issues"].append({
                        "type": "IRREGULAR_SPACING",
                        "details": f"Неравномерное расстояние между позвонками",
                        "severity": "MEDIUM"
                    })
                    results["recommendations"].append("Возможно, пропущен позвонок или артефакт")

        # 5. Проверка размера позвонков
        if len(sizes) > 1:
            mean_size = np.mean(sizes)
            results["metrics"]["mean_vertebra_size"] = float(mean_size)

            for i, size in enumerate(sizes):
                if abs(size - mean_size) / mean_size > 0.5:
                    results["valid"] = False
                    results["issues"].append({
                        "type": "ABNORMAL_SIZE",
                        "details": f"Позвонок {i + 1} аномального размера",
                        "severity": "HIGH"
                    })
                    results["recommendations"].append("Проверьте сегментацию, возможно артефакт")
                    break

        # Итоговая оценка
        if not results["issues"]:
            results["status"] = "GOOD"
            results["recommendations"].append("Исследование выполнено качественно")
        else:
            results["status"] = "WARNING" if len(results["issues"]) <= 2 else "POOR"

        return results

    def check_spine_alignment(self, segmentation):
        """Проверка, что позвонки выровнены по вертикали"""
        labeled, num_features = ndimage.label(segmentation)

        # Собираем центры всех позвонков
        centers = []
        for i in range(1, num_features + 1):
            y, x = np.where(labeled == i)
            if len(y) > 0:
                center_y = np.mean(y)
                center_x = np.mean(x)
                centers.append((center_y, center_x))

        # Проверяем, что центры находятся примерно на одной вертикальной линии
        x_coords = [c[1] for c in centers]
        x_std = np.std(x_coords)  # стандартное отклонение по X

        # Если отклонение большое → позвоночник искривлен
        # Если позвонки смещены влево-вправо → пациент лежал неправильно.
        if x_std > 20:  # порог в пикселях
            return {
                "valid": False,
                "issue": "позвонки не выровнены",
                "deviation": x_std,
                "recommendation": "проверьте укладку пациента"
            }

        return {"valid": True, "alignment": "good"}

    def check_vertebrae_spacing(self, centers):
        """Проверка равномерности расстояния между позвонками"""
        # Сортируем по Y (сверху вниз)
        centers_sorted = sorted(centers, key=lambda c: c[0])

        # Вычисляем расстояния между соседними позвонками
        distances = []
        for i in range(len(centers_sorted) - 1):
            y1, _ = centers_sorted[i]
            y2, _ = centers_sorted[i + 1]
            distances.append(abs(y2 - y1))

        # Если расстояния сильно различаются → проблема
        if len(distances) > 0:
            mean_dist = np.mean(distances)
            std_dist = np.std(distances)

            if std_dist / mean_dist > 0.3:  # больше 30% вариации
                return {
                    "valid": False,
                    "issue": "неравномерное расстояние между позвонками",
                    "distances": distances,
                    "variation": std_dist / mean_dist
                }

        return {"valid": True}

    def check_vertebrae_size(labeled, num_features):
        """Проверка, что все позвонки примерно одного размера"""
        sizes = []
        for i in range(1, num_features + 1):
            size = np.sum(labeled == i)
            sizes.append(size)

        # Если один позвонок сильно отличается по размеру
        if len(sizes) > 1:
            mean_size = np.mean(sizes)
            std_size = np.std(sizes)

            # Проверяем каждый позвонок
            issues = []
            for i, size in enumerate(sizes):
                if abs(size - mean_size) / mean_size > 0.5:  # отклонение > 50%
                    issues.append({
                        "vertebra": i + 1,
                        "size": size,
                        "expected": mean_size,
                        "issue": "позвонок аномального размера"
                    })

            if issues:
                return {
                    "valid": False,
                    "issue": "обнаружены позвонки аномального размера",
                    "details": issues,
                    "recommendation": "возможно, артефакт или неправильная сегментация"
                }

        return {"valid": True}