import torch
from diffusers import (
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
    DDIMScheduler,
    PNDMScheduler,
    EulerDiscreteScheduler
)
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np
from typing import Optional, List, Union

class InteriorDesignerPro:
    """AI модель для профессионального дизайна интерьеров"""
    
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.dtype = torch.float16 if self.device.type == 'cuda' else torch.float32
        
        # Определяем мощность GPU
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0).lower()
            self.is_powerful_gpu = any(x in gpu_name for x in ['a100', 'h100', 'rtx 4090', 'rtx 3090', 'v100', 'h200'])
        else:
            self.is_powerful_gpu = False
            
        # Основная модель для дизайна
        model_id = "runwayml/stable-diffusion-v1-5"
        self.model_name = "SD 1.5 Professional"
        
        print(f"🚀 Loading {self.model_name} on {self.device}")
        print(f"💪 Powerful GPU detected: {self.is_powerful_gpu}")
        
        # Загружаем пайплайны
        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            model_id,
            torch_dtype=self.dtype,
            safety_checker=None,
            requires_safety_checker=False
        ).to(self.device)
        
        # Пайплайн для inpainting
        self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "runwayml/stable-diffusion-inpainting",
            torch_dtype=self.dtype,
            safety_checker=None,
            requires_safety_checker=False
        ).to(self.device)
        
        # Оптимизация памяти
        if self.device.type == 'cuda':
            self.pipe.enable_attention_slicing()
            self.inpaint_pipe.enable_attention_slicing()
            
            # xFormers если доступно
            try:
                self.pipe.enable_xformers_memory_efficient_attention()
                self.inpaint_pipe.enable_xformers_memory_efficient_attention()
                print("✅ xFormers optimization enabled")
            except:
                pass
        
        # Настройка scheduler для старой версии
        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
        self.inpaint_pipe.scheduler = DDIMScheduler.from_config(self.inpaint_pipe.scheduler.config)
        
    def apply_style_pro(self, image: Image.Image, style: str, room_type: str, 
                       strength: float = 0.75, quality: str = "balanced") -> Image.Image:
        """Применяет стиль к изображению с учетом типа комнаты"""
        from design_styles import DESIGN_STYLES, get_detailed_prompt
        
        # Получаем детальный промпт
        style_info = DESIGN_STYLES.get(style, DESIGN_STYLES["Современный минимализм"])
        
        # Базовый промпт + специфика комнаты
        base_prompt = style_info["prompt"]
        room_specific = style_info.get("room_specific", {}).get(room_type, "")
        
        # Комбинируем промпты
        if room_specific:
            full_prompt = f"{base_prompt}, {room_specific}"
        else:
            full_prompt = f"{base_prompt}, {room_type} interior"
            
        # Добавляем усилители качества
        quality_boost = {
            "fast": "",
            "balanced": "high quality, professional photo, ",
            "ultra": "masterpiece, best quality, ultra detailed, 8k uhd, professional photography, "
        }
        
        full_prompt = quality_boost.get(quality, "") + full_prompt
        
        # Параметры генерации для старой версии diffusers
        params = {
            "fast": {"num_inference_steps": 20, "guidance_scale": 7.5},
            "balanced": {"num_inference_steps": 35, "guidance_scale": 8.5},
            "ultra": {"num_inference_steps": 50, "guidance_scale": 10}
        }
        
        # Генерация
        result = self.pipe(
            prompt=full_prompt,
            image=image,
            strength=strength,
            negative_prompt=style_info.get("negative", ""),
            **params[quality]
        ).images[0]
        
        return result
        
    def create_variations(self, image: Image.Image, num_variations: int = 4) -> List[Image.Image]:
        """Создает вариации дизайна"""
        variations = []
        base_seed = torch.randint(0, 1000000, (1,)).item()
        
        for i in range(num_variations):
            # Устанавливаем seed для вариации
            torch.manual_seed(base_seed + i)
            
            # Случайный стиль и сила
            from design_styles import DESIGN_STYLES
            styles = list(DESIGN_STYLES.keys())
            random_style = styles[i % len(styles)]
            
            strength = 0.6 + (i * 0.05)  # От 0.6 до 0.75
            
            variation = self.apply_style_pro(
                image, 
                random_style, 
                "living room",
                strength=strength,
                quality="fast"
            )
            variations.append(variation)
            
        return variations
        
    def create_hdr_lighting(self, image: Image.Image, intensity: float = 0.3) -> Image.Image:
        """Добавляет HDR эффект освещения"""
        # Конвертируем в numpy
        img_array = np.array(image).astype(np.float32) / 255.0
        
        # Применяем tone mapping
        gamma = 1.0 + intensity
        img_hdr = np.power(img_array, 1.0 / gamma)
        
        # Увеличиваем контраст в светлых областях
        img_hdr = img_hdr * (1.0 + intensity * 0.5)
        img_hdr = np.clip(img_hdr, 0, 1)
        
        # Обратно в PIL
        img_hdr = (img_hdr * 255).astype(np.uint8)
        result = Image.fromarray(img_hdr)
        
        # Дополнительная обработка
        enhancer = ImageEnhance.Color(result)
        result = enhancer.enhance(1.0 + intensity * 0.3)
        
        return result
        
    def enhance_details(self, image: Image.Image) -> Image.Image:
        """Улучшает детализацию изображения"""
        # Увеличиваем резкость
        enhanced = image.filter(ImageFilter.UnsharpMask(radius=2, percent=150, threshold=3))
        
        # Улучшаем контраст
        enhancer = ImageEnhance.Contrast(enhanced)
        enhanced = enhancer.enhance(1.15)
        
        # Насыщенность
        enhancer = ImageEnhance.Color(enhanced)
        enhanced = enhancer.enhance(1.1)
        
        return enhanced
        
    def change_element(self, image: Image.Image, element: str, value: str, 
                      strength: float = 0.5) -> Image.Image:
        """Изменяет отдельный элемент интерьера"""
        from design_styles import ROOM_ELEMENTS
        
        element_info = ROOM_ELEMENTS.get(element, {})
        prompt_add = element_info.get("prompt_add", element)
        
        # Формируем промпт
        prompt = f"interior with {value} {prompt_add}, high quality, detailed"
        
        # Применяем изменение
        result = self.pipe(
            prompt=prompt,
            image=image,
            strength=strength,
            negative_prompt="low quality, blurry",
            num_inference_steps=30,
            guidance_scale=8.0
        ).images[0]
        
        return result
        
    def create_style_comparison(self, image: Image.Image, styles: List[str], 
                               room_type: str = "living room") -> Image.Image:
        """Создает сравнение разных стилей"""
        from utils import ImageProcessor
        processor = ImageProcessor()
        
        styled_images = []
        for style in styles:
            styled = self.apply_style_pro(
                image, 
                style, 
                room_type, 
                strength=0.75,
                quality="fast"
            )
            styled_images.append(styled)
            
        # Создаем сетку
        comparison = processor.create_grid(styled_images, titles=styles)
        return comparison

# Добавляем метод для создания сетки если его нет
def _create_comparison_grid(images: List[Image.Image], titles: List[str]) -> Image.Image:
    """Создает сетку из изображений"""
    if not images:
        return None
        
    # Определяем размер сетки
    n = len(images)
    cols = min(3, n)
    rows = (n + cols - 1) // cols
    
    # Размер одного изображения
    img_width, img_height = images[0].size
    grid_width = img_width * cols
    grid_height = img_height * rows
    
    # Создаем сетку
    grid = Image.new('RGB', (grid_width, grid_height), 'white')
    
    for idx, (img, title) in enumerate(zip(images, titles)):
        row = idx // cols
        col = idx % cols
        x = col * img_width
        y = row * img_height
        grid.paste(img, (x, y))
        
    return grid

# Патчим метод если его нет
if not hasattr(InteriorDesignerPro, '_create_comparison_grid'):
    InteriorDesignerPro._create_comparison_grid = _create_comparison_grid