app.py

import os
import gradio as gr
import random
import time
import logging
import google.generativeai as genai

import torch
import numpy as np
from diffusers import DiffusionPipeline
from transformers import pipeline as hf_pipeline
import re

##############################################################################
# 1) ZeroGPU Environment Setup + Device and Dtype Configuration
##############################################################################
try:
    import zerogpu
    zerogpu.init()
    print("ZeroGPU initialized successfully")
    device = "cuda" if torch.cuda.is_available() else "cpu"
except ImportError:
    print("ZeroGPU package not installed, continuing without it")
    if os.getenv("ZERO_GPU"):
        print("ZeroGPU environment variable is set but zerogpu package is not installed.")
    device = "cuda" if torch.cuda.is_available() else "cpu"
except Exception as e:
    print(f"Error initializing ZeroGPU: {e}")
    print("Continuing without ZeroGPU")
    device = "cuda" if torch.cuda.is_available() else "cpu"

dtype = torch.bfloat16 if device == "cuda" else torch.float32

print(f"Using device: {device}, dtype: {dtype}")

##############################################################################
# 2) Load Models: Translation Model, Diffusion Pipeline
##############################################################################
try:
    translator = hf_pipeline(
        "translation",
        model="Helsinki-NLP/opus-mt-ko-en",
        device=0 if device == "cuda" else -1
    )

    pipe = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell",
        torch_dtype=dtype
    ).to(device)
    
    print("Models loaded successfully")
except Exception as e:
    print(f"Error loading models: {e}")
    def dummy_translator(text):
        return [{'translation_text': text}]
    
    class DummyPipe:
        def __call__(self, **kwargs):
            from PIL import Image
            import numpy as np
            dummy_img = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
            class DummyResult:
                def __init__(self, img):
                    self.images = [img]
            return DummyResult(dummy_img)
    
    translator = dummy_translator
    pipe = DummyPipe()

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

##############################################################################
# Korean detection and input text cleaning functions
##############################################################################
def contains_korean(text):
    for char in text:
        if ord('가') <= ord(char) <= ord('힣'):
            return True
    return False

def clean_input_text(text):
    """
    Allows only Korean, English, numbers, whitespace and common punctuation marks.
    Adjust allowed characters as needed.
    """
    allowed = re.compile(r'[^ㄱ-ㅎ가-힣a-zA-Z0-9\s\.\,\!\?\-\:\;\'\"]')
    cleaned_text = allowed.sub('', text)
    return cleaned_text

def log_unexpected_characters(text):
    allowed = re.compile(r'[ㄱ-ㅎ가-힣a-zA-Z0-9\s\.\,\!\?\-\:\;\'\"]')
    unexpected_chars = [char for char in text if not allowed.match(char)]
    if unexpected_chars:
        print("Unexpected characters found:", set(unexpected_chars))
    else:
        print("No unexpected characters in the input text.")

##############################################################################
# Image Generation Function
##############################################################################
def generate_design_image(prompt, seed=42, randomize_seed=True, width=1024, height=1024, num_inference_steps=4):
    original_prompt = prompt
    translated = False
    
    # Clean the input text
    prompt = clean_input_text(prompt)
    
    # Pre-process: if input is too long, trim to 1000 characters
    if len(prompt) > 1000:
        prompt = prompt[:1000]
    
    if contains_korean(prompt):
        # When calling translation, add max_length and truncation options to avoid length issues
        translation = translator(prompt, max_length=400, truncation=True)
        prompt = translation[0]['translation_text']
        translated = True

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    generator = torch.Generator(device=device).manual_seed(seed)
    
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=0.0
    ).images[0]
    
    return image

##############################################################################
# Logging Setup
##############################################################################
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("api_debug.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("idea_generator")

##############################################################################
# Gemini API Key
##############################################################################
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
genai.configure(api_key=GEMINI_API_KEY)

##############################################################################
# Optional Transformation Choice Function
##############################################################################
def choose_alternative(transformation):
    if "/" not in transformation:
        return transformation
    parts = transformation.split("/")
    if len(parts) != 2:
        return random.choice([part.strip() for part in parts])
    left = parts[0].strip()
    right = parts[1].strip()
    if " " in left:
        tokens = left.split(" ", 1)
        prefix = tokens[0]
        if not right.startswith(prefix):
            option1 = left
            option2 = prefix + " " + right
        else:
            option1 = left
            option2 = right
        return random.choice([option1, option2])
    else:
        return random.choice([left, right])

##############################################################################
# Transformation Categories Dictionaries
##############################################################################
# Korean version
physical_transformation_categories = {
    "센서 기능": [
        "시각 센서/감지", "청각 센서/감지", "촉각 센서/감지", "미각 센서/감지", "후각 센서/감지",
        "온도 센서/감지", "습도 센서/감지", "압력 센서/감지", "가속도 센서/감지", "회전 센서/감지",
        "근접 센서/감지", "위치 센서/감지", "운동 센서/감지", "가스 센서/감지", "적외선 센서/감지",
        "자외선 센서/감지", "방사선 센서/감지", "자기장 센서/감지", "전기장 센서/감지", "화학물질 센서/감지",
        "생체신호 센서/감지", "진동 센서/감지", "소음 센서/감지", "빛 세기 센서/감지", "빛 파장 센서/감지",
        "기울기 센서/감지", "pH 센서/감지", "전류 센서/감지", "전압 센서/감지", "이미지 센서/감지",
        "거리 센서/감지", "깊이 센서/감지", "중력 센서/감지", "속도 센서/감지", "흐름 센서/감지",
        "수위 센서/감지", "탁도 센서/감지", "염도 센서/감지", "금속 감지", "압전 센서/감지",
        "광전 센서/감지", "열전대 센서/감지", "홀 효과 센서/감지", "초음파 센서/감지", "레이더 센서/감지",
        "라이다 센서/감지", "터치 센서/감지", "제스처 센서/감지", "심박 센서/감지", "혈압 센서/감지"
    ],
    "크기와 형태 변화": [
        "부피 늘어남/줄어듦", "길이 늘어남/줄어듦", "너비 늘어남/줄어듦", "높이 늘어남/줄어듦",
        "밀도 변화", "무게 증가/감소", "모양 변형", "상태 변화", "불균등 변형",
        "복잡한 형태 변형", "비틀림/꼬임", "불균일한 확장/축소", "모서리 둥글게/날카롭게",
        "깨짐/갈라짐", "여러 조각 나눠짐", "물 저항", "먼지 저항", "찌그러짐/복원",
        "접힘/펼쳐짐", "압착/팽창", "늘어남/수축", "구겨짐/평평해짐", "뭉개짐/단단해짐",
        "말림/펴짐", "꺾임/구부러짐"
    ],
    "표면 및 외관 변화": [
        "색상 변화", "질감 변화", "투명/불투명 변화", "반짝임/무광 변화",
        "빛 반사 정도 변화", "무늬 변화", "각도에 따른 색상 변화", "빛에 따른 색상 변화",
        "온도에 따른 색상 변화", "홀로그램 효과", "표면 각도별 빛 반사", "표면 모양 변형",
        "초미세 표면 구조 변화", "자가 세정 효과", "얼룩/패턴 생성", "흐림/선명함 변화",
        "광택/윤기 변화", "색조/채도 변화", "발광/형광", "빛 산란 효과",
        "빛 흡수 변화", "반투명 효과", "그림자 효과 변화", "자외선 반응 변화",
        "야광 효과"
    ],
    "물질의 상태 변화": [
        "고체/액체/기체 전환", "결정화/용해", "산화/부식", "딱딱해짐/부드러워짐",
        "특수 상태 전환", "무정형/결정형 전환", "성분 분리", "미세 입자 형성/분해",
        "젤 형성/풀어짐", "준안정 상태 변화", "분자 자가 정렬/분해", "상태변화 지연 현상",
        "녹음", "굳음", "증발/응축", "승화/증착", "침전/부유", "분산/응집",
        "건조/습윤", "팽윤/수축", "동결/해동", "풍화/침식", "충전/방전",
        "결합/분리", "발효/부패"
    ],
    "움직임 특성 변화": [
        "가속/감속", "일정 속도 유지", "진동/진동 감소", "부딪힘/튕김",
        "회전 속도 증가/감소", "회전 방향 변화", "불규칙 움직임", "멈췄다 미끄러지는 현상",
        "공진/반공진", "유체 속 저항/양력 변화", "움직임 저항 변화", "복합 진동 움직임",
        "특수 유체 속 움직임", "회전-이동 연계 움직임", "관성 정지", "충격 흡수",
        "충격 전달", "운동량 보존", "마찰력 변화", "관성 탈출", "불안정 균형",
        "동적 안정성", "흔들림 감쇠", "경로 예측성", "회피 움직임"
    ],
    "구조적 변화": [
        "부품 추가/제거", "조립/분해", "접기/펴기", "변형/원상복구", "최적 구조 변화",
        "자가 재배열", "자연 패턴 형성/소멸", "규칙적 패턴 변화", "모듈식 변형",
        "복잡성 증가 구조", "원래 모양 기억 효과", "시간에 따른 형태 변화", "부분 제거",
        "부분 교체", "결합", "분리", "분할/통합", "중첩/겹침", "내부 구조 변화",
        "외부 구조 변화", "중심축 이동", "균형점 변화", "계층 구조 변화", "지지 구조 변화",
        "응력 분산 구조", "충격 흡수 구조", "그리드/매트릭스 구조 변화", "상호 연결성 변화"
    ],
    "공간 이동": [
        "앞/뒤 이동", "좌/우 이동", "위/아래 이동", "세로축 회전(고개 끄덕임)",
        "가로축 회전(고개 젓기)", "길이축 회전(옆으로 기울임)", "원 운동", "나선형 이동",
        "관성에 의한 미끄러짐", "회전축 변화", "불규칙 회전", "흔들림 운동", "포물선 이동",
        "무중력 부유", "수면 위 부유", "점프/도약", "슬라이딩", "롤링", "자유 낙하",
        "왕복 운동", "탄성 튕김", "관통", "회피 움직임", "지그재그 이동", "스윙 운동"
    ],
    "시간 관련 변화": [
        "노화/풍화", "마모/부식", "색 바램/변색", "손상/회복", "수명 주기 변화",
        "사용자 상호작용에 따른 적응", "학습 기반 형태 최적화", "시간에 따른 물성 변화",
        "집단 기억 효과", "문화적 의미 변화", "지연 반응", "이전 상태 의존 변화",
        "점진적 시간 변화", "진화적 변화", "주기적 재생", "계절 변화 적응",
        "생체리듬 변화", "생애 주기 단계", "성장/퇴화", "자기 복구/재생",
        "자연 순환 적응", "지속성/일시성", "기억 효과", "지연된 작용", "누적 효과"
    ],
    "빛과 시각 효과": [
        "발광/소등", "빛 투과/차단", "빛 산란/집중", "색상 스펙트럼 변화", "빛 회절",
        "빛 간섭", "홀로그램 생성", "레이저 효과", "빛 편광", "형광/인광",
        "자외선/적외선 발광", "광학적 착시", "빛 굴절", "그림자 생성/제거",
        "색수차 효과", "무지개 효과", "글로우 효과", "플래시 효과", "조명 패턴",
        "빔 효과", "광 필터 효과", "빛의 방향성 변화", "투영 효과", "빛 감지/반응",
        "광도 변화"
    ],
    "소리와 진동 효과": [
        "소리 발생/소멸", "소리 높낮이 변화", "소리 크기 변화", "음색 변화",
        "공명/반공명", "음향 진동", "초음파/저음파 발생", "음향 집중/분산",
        "음향 반사/흡수", "음향 도플러 효과", "음파 간섭", "음향 공진",
        "진동 패턴 변화", "타악 효과", "음향 피드백", "음향 차폐/증폭",
        "소리 지향성", "음향 왜곡", "비트 생성", "하모닉스 생성", "주파수 변조",
        "음향 충격파", "음향 필터링"
    ]
}

# English version
physical_transformation_categories_en = {
    "Sensor Functions": [
        "Visual sensor/detection", "Auditory sensor/detection", "Tactile sensor/detection", "Taste sensor/detection", "Olfactory sensor/detection",
        "Temperature sensor/detection", "Humidity sensor/detection", "Pressure sensor/detection", "Acceleration sensor/detection", "Rotational sensor/detection",
        "Proximity sensor/detection", "Position sensor/detection", "Motion sensor/detection", "Gas sensor/detection", "Infrared sensor/detection",
        "Ultraviolet sensor/detection", "Radiation sensor/detection", "Magnetic sensor/detection", "Electric field sensor/detection", "Chemical sensor/detection",
        "Biosignal sensor/detection", "Vibration sensor/detection", "Noise sensor/detection", "Light intensity sensor/detection", "Light wavelength sensor/detection",
        "Tilt sensor/detection", "pH sensor/detection", "Current sensor/detection", "Voltage sensor/detection", "Image sensor/detection",
        "Distance sensor/detection", "Depth sensor/detection", "Gravity sensor/detection", "Speed sensor/detection", "Flow sensor/detection",
        "Water level sensor/detection", "Turbidity sensor/detection", "Salinity sensor/detection", "Metal detection", "Piezoelectric sensor/detection",
        "Photovoltaic sensor/detection", "Thermocouple sensor/detection", "Hall effect sensor/detection", "Ultrasonic sensor/detection", "Radar sensor/detection",
        "Lidar sensor/detection", "Touch sensor/detection", "Gesture sensor/detection", "Heart rate sensor/detection", "Blood pressure sensor/detection"
    ],
    "Size and Shape Change": [
        "Volume increase/decrease", "Length increase/decrease", "Width increase/decrease", "Height increase/decrease",
        "Density change", "Weight increase/decrease", "Shape deformation", "State change", "Uneven deformation",
        "Complex shape deformation", "Twisting/entwining", "Non-uniform expansion/contraction", "Rounded/sharpened edges",
        "Cracking/splitting", "Fragmentation", "Water resistance", "Dust resistance", "Denting/recovery",
        "Folding/unfolding", "Compression/expansion", "Stretching/contraction", "Wrinkling/flattening", "Crushing/hardening",
        "Rolling/unrolling", "Bending/curving"
    ],
    "Surface and Appearance Change": [
        "Color change", "Texture change", "Transparency change", "Glossy/matte change",
        "Light reflection variation", "Pattern change", "Angle-dependent color change", "Light-induced color change",
        "Temperature-dependent color change", "Holographic effect", "Angle-specific light reflection", "Surface shape alteration",
        "Nano-scale surface structure change", "Self-cleaning effect", "Stain/pattern formation", "Blurriness/clarity change",
        "Luster/shine change", "Hue/saturation change", "Luminescence/fluorescence", "Light scattering effect",
        "Light absorption change", "Translucency effect", "Shadow effect change", "UV response change",
        "Glow effect"
    ],
    "Material State Change": [
        "Solid/liquid/gas transition", "Crystallization/dissolution", "Oxidation/corrosion", "Hardening/softening",
        "Special state transition", "Amorphous/crystalline transition", "Component separation", "Particle formation/disintegration",
        "Gel formation/dissolution", "Metastable state change", "Molecular self-assembly/disintegration", "Delayed state change",
        "Melting", "Solidification", "Evaporation/condensation", "Sublimation/deposition", "Precipitation/suspension", "Dispersion/aggregation",
        "Drying/moistening", "Swelling/shrinkage", "Freezing/thawing", "Weathering/erosion", "Charging/discharging",
        "Bonding/separation", "Fermentation/decay"
    ],
    "Movement Characteristics Change": [
        "Acceleration/deceleration", "Maintaining constant speed", "Vibration/vibration reduction", "Collision/bouncing",
        "Increase/decrease in rotational speed", "Change in rotational direction", "Irregular movement", "Stop-and-slide phenomenon",
        "Resonance/anti-resonance", "Resistance/lift change in fluid", "Change in movement resistance", "Complex vibrational movement",
        "Movement in special fluid", "Rotational-translational movement", "Inertial stoppage", "Shock absorption",
        "Shock transfer", "Conservation of momentum", "Friction change", "Overcoming inertia", "Unstable equilibrium",
        "Dynamic stability", "Damping of oscillation", "Path predictability", "Evasive movement"
    ],
    "Structural Change": [
        "Addition/removal of components", "Assembly/disassembly", "Folding/unfolding", "Deformation/recovery", "Optimal structural change",
        "Self-rearrangement", "Natural pattern formation/disappearance", "Regular pattern change", "Modular transformation",
        "Increased structural complexity", "Memory of original shape effect", "Shape change over time", "Partial removal",
        "Partial replacement", "Bonding", "Separation", "Division/integration", "Overlaying", "Internal structure change",
        "External structure change", "Shift of center axis", "Balance point change", "Hierarchical structure change", "Support structure change",
        "Stress distribution structure", "Shock absorption structure", "Grid/matrix structure change", "Interconnectivity change"
    ],
    "Spatial Movement": [
        "Forward/backward movement", "Left/right movement", "Up/down movement", "Vertical axis rotation (nodding)",
        "Horizontal axis rotation (shaking head)", "Longitudinal axis rotation (tilting sideways)", "Circular motion", "Spiral movement",
        "Slipping due to inertia", "Change of rotation axis", "Irregular rotation", "Shaking movement", "Parabolic motion",
        "Zero-gravity floating", "Floating on water surface", "Jump/leap", "Sliding", "Rolling", "Free fall",
        "Reciprocating motion", "Elastic bouncing", "Penetration", "Evasive movement", "Zigzag movement", "Swinging movement"
    ],
    "Time-Related Change": [
        "Aging/weathering", "Wear/corrosion", "Fading/discoloration", "Damage/recovery", "Lifecycle change",
        "Adaptation through user interaction", "Learning-based shape optimization", "Property change over time",
        "Collective memory effect", "Cultural significance change", "Delayed response", "History-dependent change",
        "Gradual time change", "Evolutionary change", "Periodic regeneration", "Seasonal adaptation",
        "Circadian rhythm change", "Lifecycle stage", "Growth/decline", "Self-repair/regeneration",
        "Natural cycle adaptation", "Persistence/transience", "Memory effect", "Delayed effect", "Cumulative effect"
    ],
    "Light and Visual Effects": [
        "Illumination/shutdown", "Light transmission/blocking", "Light scattering/concentration", "Color spectrum change", "Light diffraction",
        "Light interference", "Hologram creation", "Laser effect", "Light polarization", "Fluorescence/phosphorescence",
        "UV/IR emission", "Optical illusion", "Light refraction", "Shadow creation/removal",
        "Chromatic aberration", "Rainbow effect", "Glow effect", "Flash effect", "Lighting pattern",
        "Beam effect", "Light filter effect", "Change in light direction", "Projection effect", "Light detection/response",
        "Luminance change"
    ],
    "Sound and Vibration Effects": [
        "Sound generation/cessation", "Pitch change", "Volume change", "Timbre change",
        "Resonance/antiresonance", "Acoustic vibration", "Ultrasonic/infrasonic emission", "Sound concentration/distribution",
        "Sound reflection/absorption", "Acoustic Doppler effect", "Sound wave interference", "Acoustic resonance",
        "Vibration pattern change", "Percussive effect", "Audio feedback", "Sound shielding/amplification",
        "Directional sound", "Sound distortion", "Beat generation", "Harmonics generation", "Frequency modulation",
        "Acoustic shockwave", "Sound filtering"
    ],
    "Thermal Changes": [
        "Temperature rise/fall", "Thermal expansion/contraction", "Heat transfer/blocking", "Pressure increase/decrease",
        "Magnetization due to heat change", "Entropy change", "Thermoelectric effect", "Magnetic-induced thermal change",
        "Heat storage/release during phase change", "Thermal stress buildup/release", "Impact of rapid temperature change",
        "Radiative cooling/heating", "Exothermic/endothermic", "Heat distribution change", "Heat reflection/absorption",
        "Cooling condensation", "Thermal activation", "Thermal discoloration", "Coefficient of thermal expansion change", "Thermal stability change",
        "Heat resistance/cold resistance", "Self-heating", "Thermal equilibrium/imbalance", "Thermal deformation", "Heat dispersion/concentration"
    ],
    "Electrical and Magnetic Changes": [
        "Magnetism creation/cessation", "Charge increase/decrease", "Electric field creation/cessation", "Magnetic field creation/cessation",
        "Superconducting transition", "Ferroelectric property change", "Quantum state change", "Plasma formation/cessation",
        "Spin wave transmission", "Electricity generation by light", "Electricity generation by pressure", "Current change in magnetic field",
        "Electrical resistance change", "Electrical conductivity change", "Static electricity generation/discharge", "Electromagnetic induction",
        "Electromagnetic wave emission/absorption", "Capacitance change", "Magnetic hysteresis", "Electrical polarization",
        "Electron flow direction change", "Electrical resonance", "Electrical shielding/exposure", "Magnetic shielding/exposure",
        "Magnetic field alignment"
    ],
    "Chemical Change": [
        "Surface coating change", "Material composition change", "Chemical reaction change", "Catalytic action start/stop",
        "Light-induced chemical reaction", "Electricity-induced chemical reaction", "Monolayer formation", "Molecular-level structural change",
        "Biomimetic surface change", "Environmentally responsive material change", "Periodic chemical reaction", "Oxidation", "Reduction",
        "Polymerization", "Water splitting", "Compound formation", "Radiation effects", "Acid-base reaction", "Neutralization reaction",
        "Ionization", "Chemical adsorption/desorption", "Catalytic efficiency change", "Enzyme activity change", "Colorimetric reaction",
        "pH change", "Chemical equilibrium shift", "Bond formation/breakage", "Solubility change"
    ],
    "Biological Change": [
        "Growth/shrinkage", "Cell division/death", "Bioluminescence", "Metabolic change", "Immune response",
        "Hormone secretion", "Neural response", "Genetic expression", "Adaptation/evolution", "Circadian rhythm change",
        "Regeneration/healing", "Aging/maturation", "Biomimetic change", "Biofilm formation", "Biological degradation",
        "Enzyme activation/inactivation", "Biological signaling", "Stress response", "Thermoregulation", "Biological clock change",
        "Extracellular matrix change", "Biomechanical response", "Cell motility", "Cell polarity change", "Nutritional status change"
    ],
    "Environmental Interaction": [
        "Temperature response", "Humidity response", "Pressure response", "Gravity response", "Magnetic field response",
        "Light response", "Sound response", "Chemical detection", "Mechanical stimulus detection", "Electrical stimulus response",
        "Radiation response", "Vibration detection", "pH response", "Solvent response", "Gas exchange",
        "Pollution response", "Weather response", "Seasonal response", "Circadian response", "Ecosystem interaction",
        "Symbiotic/competitive interaction", "Predator/prey relationship", "Swarm formation", "Territorial behavior", "Migration/settlement pattern"
    ]
}

##############################################################################
# Gemini API Call Function (Language Independent)
##############################################################################
def query_gemini_api(prompt):
    try:
        model = genai.GenerativeModel('gemini-2.0-flash-thinking-exp-01-21')
        response = model.generate_content(prompt)
        try:
            if hasattr(response, 'text'):
                return response.text
            if hasattr(response, 'candidates') and response.candidates:
                candidate = response.candidates[0]
                if hasattr(candidate, 'content'):
                    content = candidate.content
                    if hasattr(content, 'parts') and content.parts:
                        if len(content.parts) > 0:
                            return content.parts[0].text
            if hasattr(response, 'parts') and response.parts:
                if len(response.parts) > 0:
                    return response.parts[0].text
            return "Unable to generate a response. API response structure is different than expected."
        except Exception as inner_e:
            logger.error(f"Error processing response: {inner_e}")
            return f"An error occurred while processing the response: {str(inner_e)}"
    except Exception as e:
        logger.error(f"Error calling Gemini API: {e}")
        if "API key not valid" in str(e):
            return "API key is not valid. Please check your GEMINI_API_KEY environment variable."
        return f"An error occurred while calling the API: {str(e)}"

##############################################################################
# Description Expansion Functions (LLM) - Korean and English Versions
##############################################################################
def enhance_with_llm(base_description, obj_name, category):
    prompt = f"""
다음은 '{obj_name}'의 '{category}' 관련 간단한 설명입니다:
"{base_description}"
위 내용을 보다 구체화하여, 
1) 창의적인 모델/컨셉/형상의 변화에 대한 이해, 
2) 혁신 포인트와 기능성 등을 중심으로 
3~4문장의 아이디어로 확장해 주세요.
"""
    return query_gemini_api(prompt)

def enhance_with_llm_en(base_description, obj_name, category):
    prompt = f"""
Below is a brief description related to '{category}' for '{obj_name}':
"{base_description}"
Please expand the above content into a more detailed explanation, focusing on:
1) Creative transformation of the model/concept/shape,
2) Innovative aspects and functionality,
in 3-4 sentences.
"""
    return query_gemini_api(prompt)

##############################################################################
# Transformation Idea Generation Functions for Both Languages
##############################################################################
def generate_single_object_transformation_for_category_lang(obj, selected_category, categories_dict, lang="ko"):
    transformations = categories_dict.get(selected_category)
    if not transformations:
        return {}
    transformation = choose_alternative(random.choice(transformations))
    if lang == "ko":
        base_description = f"{obj}이(가) {transformation} 현상을 보인다"
    else:
        base_description = f"{obj} exhibits {transformation}"
    return {selected_category: {"base": base_description, "enhanced": None}}

def generate_two_objects_interaction_for_category_lang(obj1, obj2, selected_category, categories_dict, lang="ko"):
    transformations = categories_dict.get(selected_category)
    if not transformations:
        return {}
    transformation = choose_alternative(random.choice(transformations))
    if lang == "ko":
        template = random.choice([
            "{obj1}이(가) {obj2}에 결합하여 {change}가 발생했다",
            "{obj1}과(와) {obj2}이(가) 충돌하면서 {change}가 일어났다"
        ])
    else:
        template = random.choice([
            "{obj1} combined with {obj2} resulted in {change}",
            "A collision between {obj1} and {obj2} led to {change}"
        ])
    base_description = template.format(obj1=obj1, obj2=obj2, change=transformation)
    return {selected_category: {"base": base_description, "enhanced": None}}

def generate_three_objects_interaction_for_category_lang(obj1, obj2, obj3, selected_category, categories_dict, lang="ko"):
    transformations = categories_dict.get(selected_category)
    if not transformations:
        return {}
    transformation = choose_alternative(random.choice(transformations))
    if lang == "ko":
        template = random.choice([
            "{obj1}, {obj2}, {obj3}이(가) 삼각형 구조로 결합하여 {change}가 발생했다",
            "{obj1}이(가) {obj2}와(과) {obj3} 사이에서 매개체 역할을 하며 {change}를 촉진했다"
        ])
    else:
        template = random.choice([
            "{obj1}, {obj2}, and {obj3} formed a triangular structure resulting in {change}",
            "{obj1} acted as an intermediary between {obj2} and {obj3}, facilitating {change}"
        ])
    base_description = template.format(obj1=obj1, obj2=obj2, obj3=obj3, change=transformation)
    return {selected_category: {"base": base_description, "enhanced": None}}

def enhance_descriptions_lang(results, objects, lang="ko"):
    obj_name = " 및 ".join([obj for obj in objects if obj]) if lang=="ko" else " and ".join([obj for obj in objects if obj])
    for category, result in results.items():
        if lang == "ko":
            result["enhanced"] = enhance_with_llm(result["base"], obj_name, category)
        else:
            result["enhanced"] = enhance_with_llm_en(result["base"], obj_name, category)
    return results

def generate_transformations_lang(text1, text2, text3, selected_category, categories_dict, lang="ko"):
    if text2 and text3:
        results = generate_three_objects_interaction_for_category_lang(text1, text2, text3, selected_category, categories_dict, lang)
        objects = [text1, text2, text3]
    elif text2:
        results = generate_two_objects_interaction_for_category_lang(text1, text2, selected_category, categories_dict, lang)
        objects = [text1, text2]
    else:
        results = generate_single_object_transformation_for_category_lang(text1, selected_category, categories_dict, lang)
        objects = [text1]
    return enhance_descriptions_lang(results, objects, lang)

def format_results_lang(results, lang="ko"):
    formatted = ""
    if lang == "ko":
        for category, result in results.items():
            formatted += f"## {category}\n**기본 아이디어**: {result['base']}\n\n**확장된 아이디어**: {result['enhanced']}\n\n---\n\n"
    else:
        for category, result in results.items():
            formatted += f"## {category}\n**Base Idea**: {result['base']}\n\n**Expanded Idea**: {result['enhanced']}\n\n---\n\n"
    return formatted

def process_inputs_lang(text1, text2, text3, selected_category, categories_dict, lang="ko", progress=gr.Progress()):
    text1 = text1.strip() if text1 else None
    text2 = text2.strip() if text2 else None
    text3 = text3.strip() if text3 else None
    
    if not text1:
        return "오류: 최소 하나의 키워드를 입력해주세요." if lang=="ko" else "Error: Please enter at least one keyword."
    
    if lang == "ko":
        progress(0.05, desc="아이디어 생성 준비 중...")
        time.sleep(0.3)
        progress(0.1, desc="창의적인 아이디어 생성 시작...")
    else:
        progress(0.05, desc="Preparing idea generation...")
        time.sleep(0.3)
        progress(0.1, desc="Generating creative idea...")
    
    results = generate_transformations_lang(text1, text2, text3, selected_category, categories_dict, lang)
    
    if lang == "ko":
        progress(0.8, desc="결과 포맷팅 중...")
        formatted = format_results_lang(results, lang)
        progress(1.0, desc="완료!")
    else:
        progress(0.8, desc="Formatting results...")
        formatted = format_results_lang(results, lang)
        progress(1.0, desc="Done!")
    return formatted

def process_all_lang(text1, text2, text3, selected_category, categories_dict, lang="ko", progress=gr.Progress()):
    idea_result = process_inputs_lang(text1, text2, text3, selected_category, categories_dict, lang, progress)
    image_result = generate_design_image(
        idea_result, 
        seed=42, 
        randomize_seed=True, 
        width=1024, 
        height=1024, 
        num_inference_steps=4
    )
    return idea_result, image_result

##############################################################################
# Warning Message Function for API Key (Language Specific)
##############################################################################
def get_warning_message_lang(lang="ko"):
    if not GEMINI_API_KEY:
        return "⚠️ 환경 변수 GEMINI_API_KEY가 설정되지 않았습니다. Gemini API 키를 설정하세요." if lang=="ko" else "⚠️ The GEMINI_API_KEY environment variable is not set. Please set your Gemini API key."
    return ""


##############################################################################
# Gradio UI with Two Tabs: English (Main Home) and Korean
##############################################################################
with gr.Blocks(
    title="Idea Transformer",
    theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")
) as demo:
    
    gr.HTML("""
    <style>
    body { 
        background: linear-gradient(135deg, #e0eafc, #cfdef3); 
        font-family: 'Arial', sans-serif; 
    }
    .gradio-container { 
        padding: 20px; 
    }
    h1, h2 { 
        text-align: center; 
    }
    h1 { 
        color: #333; 
    }
    h2 { 
        color: #555; 
    }
    .output { 
        background-color: #ffffff; 
        padding: 15px; 
        border-radius: 8px; 
    }
    .gr-button { 
        background-color: #4CAF50; 
        color: white; 
        border: none; 
        border-radius: 4px; 
        padding: 8px 16px; 
    }
    .progress-message { 
        color: #2196F3; 
        font-weight: bold; 
        margin-top: 10px; 
    }
    </style>
    """)
    
    with gr.Tabs():
        with gr.Tab(label="English"):
            gr.Markdown("# 🚀 Idea Transformer")
            gr.Markdown("Based on up to **three keywords** and a **selected category**, this tool generates a creative transformation idea and a design image using the expanded idea as a prompt.")
            
            warning_en = gr.Markdown(get_warning_message_lang("en"))
            
            with gr.Row():
                with gr.Column(scale=1):
                    text_input1_en = gr.Textbox(label="Keyword 1 (required)", placeholder="e.g., Smartphone")
                    text_input2_en = gr.Textbox(label="Keyword 2 (optional)", placeholder="e.g., Artificial Intelligence")
                    text_input3_en = gr.Textbox(label="Keyword 3 (optional)", placeholder="e.g., Healthcare")
                    category_radio_en = gr.Radio(
                        label="Select Category",
                        choices=list(physical_transformation_categories_en.keys()),
                        value=list(physical_transformation_categories_en.keys())[0],
                        info="Select a category."
                    )
                    status_msg_en = gr.Markdown("💡 Click the 'Generate Idea' button to create an idea and design image based on the selected category.")
                    
                    processing_indicator_en = gr.HTML("""
                    <div style="display: flex; justify-content: center; align-items: center; margin: 10px 0;">
                        <div style="border: 5px solid #f3f3f3; border-top: 5px solid #3498db; border-radius: 50%; width: 30px; height: 30px; animation: spin 2s linear infinite;"></div>
                        <p style="margin-left: 10px; font-weight: bold; color: #3498db;">Processing...</p>
                    </div>
                    <style>
                    @keyframes spin {
                        0% { transform: rotate(0deg); }
                        100% { transform: rotate(360deg); }
                    }
                    </style>
                    """, visible=False)
                    
                    submit_button_en = gr.Button("Generate Idea", variant="primary")
                
                with gr.Column(scale=2):
                    idea_output_en = gr.Markdown(label="Idea Output")
                    generated_image_en = gr.Image(label="Generated Design Image", type="pil")
            
            gr.Examples(
                examples=[
                    ["Smartphone", "", "", "Sensor Functions"],
                    ["Car", "", "", "Size and Shape Change"],
                    ["Car", "Artificial Intelligence", "", "Surface and Appearance Change"],
                    ["Drone", "Artificial Intelligence", "", "Material State Change"],
                    ["Sneakers", "Wearable", "Health", "Structural Change"],
                ],
                inputs=[text_input1_en, text_input2_en, text_input3_en, category_radio_en],
            )
            
            def show_processing_indicator_en():
                return gr.update(visible=True)
            
            def hide_processing_indicator_en():
                return gr.update(visible=False)
            
            submit_button_en.click(
                fn=show_processing_indicator_en,
                inputs=None,
                outputs=processing_indicator_en
            ).then(
                fn=process_all_lang,
                inputs=[text_input1_en, text_input2_en, text_input3_en, category_radio_en, gr.State(physical_transformation_categories_en), gr.State("en")],
                outputs=[idea_output_en, generated_image_en]
            ).then(
                fn=hide_processing_indicator_en,
                inputs=None,
                outputs=processing_indicator_en
            )
        
        with gr.Tab(label="한국어"):
            gr.Markdown("# 🚀 아이디어 트랜스포머")
            gr.Markdown("입력한 **키워드**(최대 3개)와 **카테고리**를 바탕으로, 창의적인 모델/컨셉/형상 변화 아이디어를 생성하고, 해당 확장 아이디어를 프롬프트로 하여 디자인 이미지를 생성합니다.")
            
            warning_ko = gr.Markdown(get_warning_message_lang("ko"))
            
            with gr.Row():
                with gr.Column(scale=1):
                    text_input1_ko = gr.Textbox(label="키워드 1 (필수)", placeholder="예: 스마트폰")
                    text_input2_ko = gr.Textbox(label="키워드 2 (선택)", placeholder="예: 인공지능")
                    text_input3_ko = gr.Textbox(label="키워드 3 (선택)", placeholder="예: 헬스케어")
                    category_radio_ko = gr.Radio(
                        label="카테고리 선택",
                        choices=list(physical_transformation_categories.keys()),
                        value=list(physical_transformation_categories.keys())[0],
                        info="출력할 카테고리를 선택하세요."
                    )
                    status_msg_ko = gr.Markdown("💡 '아이디어 생성하기' 버튼을 클릭하면 선택한 카테고리에 해당하는 아이디어와 디자인 이미지가 생성됩니다.")
                    
                    processing_indicator_ko = gr.HTML("""
                    <div style="display: flex; justify-content: center; align-items: center; margin: 10px 0;">
                        <div style="border: 5px solid #f3f3f3; border-top: 5px solid #3498db; border-radius: 50%; width: 30px; height: 30px; animation: spin 2s linear infinite;"></div>
                        <p style="margin-left: 10px; font-weight: bold; color: #3498db;">처리 중입니다...</p>
                    </div>
                    <style>
                    @keyframes spin {
                        0% { transform: rotate(0deg); }
                        100% { transform: rotate(360deg); }
                    }
                    </style>
                    """, visible=False)
                    
                    submit_button_ko = gr.Button("아이디어 생성하기", variant="primary")
                
                with gr.Column(scale=2):
                    idea_output_ko = gr.Markdown(label="아이디어 결과")
                    generated_image_ko = gr.Image(label="생성된 디자인 이미지", type="pil")
            
            gr.Examples(
                examples=[
                    ["스마트폰", "", "", "센서 기능"],
                    ["자동차", "", "", "크기와 형태 변화"],
                    ["자동차", "인공지능", "", "표면 및 외관 변화"],
                    ["드론", "인공지능", "", "물질의 상태 변화"],
                    ["운동화", "웨어러블", "건강", "구조적 변화"],
                ],
                inputs=[text_input1_ko, text_input2_ko, text_input3_ko, category_radio_ko],
            )
            
            def show_processing_indicator_ko():
                return gr.update(visible=True)
            
            def hide_processing_indicator_ko():
                return gr.update(visible=False)
            
            submit_button_ko.click(
                fn=show_processing_indicator_ko,
                inputs=None,
                outputs=processing_indicator_ko
            ).then(
                fn=process_all_lang,
                inputs=[text_input1_ko, text_input2_ko, text_input3_ko, category_radio_ko, gr.State(physical_transformation_categories), gr.State("ko")],
                outputs=[idea_output_ko, generated_image_ko]
            ).then(
                fn=hide_processing_indicator_ko,
                inputs=None,
                outputs=processing_indicator_ko
            )

if __name__ == "__main__":
    demo.launch(debug=True)