import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline  # Для генерации изображений
import torch
from PIL import Image
import requests
from io import BytesIO

# Загружаем модели для анализа тональности, суммаризации текста, генерации подписей к изображениям, ответов на вопросы, перевода текста, определения эмоций, автодополнения кода, определения фейковых новостей, NER, классификации изображений, генерации кода и исправления кода
sentiment_pipeline = pipeline("sentiment-analysis")
summarization_pipeline = pipeline("summarization")
image_captioning_pipeline = pipeline("image-to-text")
qa_pipeline = pipeline("question-answering")
translation_pipeline = pipeline("translation_en_to_ru", model="Helsinki-NLP/opus-mt-en-ru")
emotion_pipeline = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion")
code_completion_pipeline = pipeline("text-generation", model="Salesforce/codegen-350M-mono")
fake_news_pipeline = pipeline("text-classification", model="roberta-base-openai-detector")
ner_pipeline = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", grouped_entities=True)
image_classification_pipeline = pipeline("image-classification", model="google/vit-base-patch16-224")
code_generation_pipeline = pipeline("text-generation", model="deepseek-ai/deepseek-coder-1.3b-instruct")
code_fix_pipeline = pipeline("text-generation", model="deepseek-ai/deepseek-coder-1.3b-instruct")

# Загружаем модель для генерации изображений (Stable Diffusion)
if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

image_generation_pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1").to(device)

# Загружаем модель для расширения изображений (Real-ESRGAN)
def load_esrgan_model():
    from realesrgan import RealESRGANer
    model = RealESRGANer(
        scale=4,
        model_path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
        model=None,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=False,
    )
    return model

esrgan_model = load_esrgan_model()

# Функция для анализа тональности текста
def analyze_sentiment(text):
    result = sentiment_pipeline(text)[0]
    return f"Label: {result['label']}, Confidence: {result['score']:.4f}"

# Функция для суммаризации текста
def summarize_text(text):
    result = summarization_pipeline(text, max_length=50, min_length=25, do_sample=False)
    return result[0]['summary_text']

# Функция для генерации подписи к изображению
def generate_caption(image):
    result = image_captioning_pipeline(image)
    return result[0]['generated_text']

# Функция для ответов на вопросы
def answer_question(context, question):
    result = qa_pipeline(question=question, context=context)
    return f"Answer: {result['answer']}, Confidence: {result['score']:.4f}"

# Функция для перевода текста
def translate_text(text):
    result = translation_pipeline(text)
    return result[0]['translation_text']

# Функция для определения эмоций
def detect_emotion(text):
    result = emotion_pipeline(text)[0]
    return f"Emotion: {result['label']}, Confidence: {result['score']:.4f}"

# Функция для автодополнения кода
def complete_code(code):
    result = code_completion_pipeline(code, max_length=50, num_return_sequences=1)
    return result[0]['generated_text']

# Функция для определения фейковых новостей
def detect_fake_news(text):
    result = fake_news_pipeline(text)[0]
    return f"Label: {result['label']}, Confidence: {result['score']:.4f}"

# Функция для распознавания именованных сущностей (NER)
def recognize_entities(text):
    result = ner_pipeline(text)
    entities = []
    for entity in result:
        entities.append(f"Entity: {entity['word']}, Label: {entity['entity_group']}, Confidence: {entity['score']:.4f}")
    return "\n".join(entities)

# Функция для классификации изображений
def classify_image(image):
    result = image_classification_pipeline(image)
    classifications = []
    for item in result:
        classifications.append(f"Label: {item['label']}, Confidence: {item['score']:.4f}")
    return "\n".join(classifications)

# Функция для генерации кода по запросу
def generate_code(prompt):
    result = code_generation_pipeline(prompt, max_length=100, num_return_sequences=1)
    return result[0]['generated_text']

# Функция для исправления кода
def fix_code(error, problem, solution, example):
    prompt = f"""
**Ошибка:** {error}
**Проблема:** {problem}
**Решение:** {solution}
**Пример:** {example}
"""
    result = code_fix_pipeline(prompt, max_length=200, num_return_sequences=1)
    return result[0]['generated_text']

# Функция для расширения изображений (Flux Extend Image)
def extend_image(image):
    # Преобразуем изображение в формат, подходящий для модели
    img = Image.open(image).convert("RGB")
    output, _ = esrgan_model.enhance(img, outscale=4)
    return output

# Функция для генерации изображений (Image Generation)
def generate_image(prompt):
    with torch.autocast("cuda"):
        image = image_generation_pipeline(prompt).images[0]
    return image

# Примеры текстов для анализа тональности
sentiment_examples = [
    "I love programming, it's so much fun!",
    "This movie was terrible, I hated it.",
    "The weather is nice today.",
    "I feel so frustrated with this project.",
    "Gradio is an amazing tool for building ML demos!"
]

# Примеры текстов для суммаризации
summarization_examples = [
    "Gradio is a powerful tool for building machine learning demos. It allows developers to quickly create interactive interfaces for their models.",
    "The weather today is sunny with a slight breeze. It's a perfect day to go outside and enjoy nature.",
    "Artificial intelligence is transforming industries by automating tasks and providing insights from large datasets."
]

# Примеры изображений для генерации подписей
image_examples = [
    "https://a.d-cd.net/b977306s-1920.jpg",  # Пример 1
    "https://i.pinimg.com/originals/ba/bd/6d/babd6d37eb2dd965c7f1dfb516d54094.jpg",  # Пример 2
    "https://get.wallhere.com/photo/sea-bay-water-beach-coast-swimming-pool-resort-island-lagoon-Caribbean-vacation-estate-leisure-ocean-tropics-2560x1440-px-geographical-feature-atoll-554636.jpg"  # Пример 3
]

# Примеры для ответов на вопросы
qa_examples = [
    ["Gradio is a Python library for building machine learning demos. It allows developers to quickly create interactive interfaces for their models.", "What is Gradio?"],
    ["The weather today is sunny with a slight breeze. It's a perfect day to go outside and enjoy nature.", "What is the weather like today?"],
    ["Artificial intelligence is transforming industries by automating tasks and providing insights from large datasets.", "How is AI transforming industries?"]
]

# Примеры текстов для перевода
translation_examples = [
    "Hello, how are you?",
    "I love machine learning and artificial intelligence.",
    "The weather is beautiful today."
]

# Примеры текстов для определения эмоций
emotion_examples = [
    "I am so happy today!",
    "I feel really sad about what happened.",
    "This situation makes me angry.",
    "I am scared of the dark.",
    "I am surprised by the results."
]

# Примеры кода для автодополнения
code_examples = [
    "def factorial(n):",
    "import numpy as np",
    "for i in range(10):"
]

# Примеры текстов для определения фейковых новостей
fake_news_examples = [
    "A new study shows that eating chocolate every day can make you live longer.",
    "The government has secretly been working on time travel technology for decades.",
    "Scientists have discovered a new planet in our solar system that is inhabited by aliens."
]

# Примеры текстов для распознавания именованных сущностей (NER)
ner_examples = [
    "My name is John Doe and I live in New York.",
    "Apple is looking at buying a startup in the UK for $1 billion.",
    "Elon Musk is the CEO of Tesla and SpaceX."
]

# Примеры изображений для классификации
classification_examples = [
    "https://a.d-cd.net/b977306s-1920.jpg",  # Пример 1
    "https://i.pinimg.com/originals/ba/bd/6d/babd6d37eb2dd965c7f1dfb516d54094.jpg",  # Пример 2
    "https://get.wallhere.com/photo/sea-bay-water-beach-coast-swimming-pool-resort-island-lagoon-Caribbean-vacation-estate-leisure-ocean-tropics-2560x1440-px-geographical-feature-atoll-554636.jpg"  # Пример 3
]

# Примеры запросов для генерации кода
code_generation_examples = [
    "Write a Python function to calculate the factorial of a number.",
    "Create a JavaScript function to reverse a string.",
    "Generate a SQL query to find all users older than 30."
]

# Примеры для исправления кода
code_fix_examples = [
    ["SyntaxError: invalid syntax", "Missing colon at the end of the if statement", "Add a colon at the end of the if statement", "if x == 5\n    print('Hello')"],
    ["NameError: name 'x' is not defined", "Variable 'x' is not defined before use", "Define the variable 'x' before using it", "print(x)\nx = 10"],
    ["IndentationError: unexpected indent", "Incorrect indentation in the code", "Fix the indentation to match the correct level", "def foo():\n  print('Hello')\n    print('World')"]
]

# Создаем интерфейс Gradio с вкладками
with gr.Blocks() as demo:
    with gr.Tab("Sentiment Analysis"):
        gr.Interface(
            fn=analyze_sentiment,
            inputs=gr.Textbox(lines=2, placeholder="Введите текст для анализа тональности..."),
            outputs="text",
            title="Анализ тональности текста",
            description="Введите текст, чтобы определить его тональность.",
            examples=sentiment_examples,
            examples_per_page=3
        )
    with gr.Tab("Text Summarization"):
        gr.Interface(
            fn=summarize_text,
            inputs=gr.Textbox(lines=5, placeholder="Введите текст для суммаризации..."),
            outputs="text",
            title="Суммаризация текста",
            description="Введите текст, чтобы получить его краткое содержание.",
            examples=summarization_examples,
            examples_per_page=2
        )
    with gr.Tab("Image Captioning"):
        gr.Interface(
            fn=generate_caption,
            inputs=gr.Image(type="pil", label="Загрузите изображение"),
            outputs="text",
            title="Генерация подписи к изображению",
            description="Загрузите изображение, чтобы сгенерировать его описание.",
            examples=image_examples,
            examples_per_page=2
        )
    with gr.Tab("Question Answering"):
        gr.Interface(
            fn=answer_question,
            inputs=[
                gr.Textbox(lines=5, placeholder="Введите контекст..."),
                gr.Textbox(lines=2, placeholder="Введите вопрос...")
            ],
            outputs="text",
            title="Ответы на вопросы",
            description="Введите контекст и вопрос, чтобы получить ответ.",
            examples=qa_examples,
            examples_per_page=2
        )
    with gr.Tab("Language Translation"):
        gr.Interface(
            fn=translate_text,
            inputs=gr.Textbox(lines=2, placeholder="Введите текст на английском..."),
            outputs="text",
            title="Перевод текста (английский → русский)",
            description="Введите текст на английском, чтобы перевести его на русский.",
            examples=translation_examples,
            examples_per_page=2
        )
    with gr.Tab("Emotion Detection"):
        gr.Interface(
            fn=detect_emotion,
            inputs=gr.Textbox(lines=2, placeholder="Введите текст для определения эмоции..."),
            outputs="text",
            title="Определение эмоций",
            description="Введите текст, чтобы определить эмоцию.",
            examples=emotion_examples,
            examples_per_page=2
        )
    with gr.Tab("Code Completion"):
        gr.Interface(
            fn=complete_code,
            inputs=gr.Textbox(lines=5, placeholder="Введите начало кода..."),
            outputs="text",
            title="Автодополнение кода",
            description="Введите начало кода, чтобы получить его продолжение.",
            examples=code_examples,
            examples_per_page=2
        )
    with gr.Tab("Fake News Detection"):
        gr.Interface(
            fn=detect_fake_news,
            inputs=gr.Textbox(lines=5, placeholder="Введите текст новости..."),
            outputs="text",
            title="Определение фейковых новостей",
            description="Введите текст новости, чтобы определить, является ли она фейковой.",
            examples=fake_news_examples,
            examples_per_page=2
        )
    with gr.Tab("Named Entity Recognition (NER)"):
        gr.Interface(
            fn=recognize_entities,
            inputs=gr.Textbox(lines=5, placeholder="Введите текст для распознавания сущностей..."),
            outputs="text",
            title="Распознавание именованных сущностей (NER)",
            description="Введите текст, чтобы извлечь из него именованные сущности.",
            examples=ner_examples,
            examples_per_page=2
        )
    with gr.Tab("Image Classification"):
        gr.Interface(
            fn=classify_image,
            inputs=gr.Image(type="pil", label="Загрузите изображение"),
            outputs="text",
            title="Классификация изображений",
            description="Загрузите изображение, чтобы классифицировать его.",
            examples=classification_examples,
            examples_per_page=2
        )
    with gr.Tab("Code Generation"):
        gr.Interface(
            fn=generate_code,
            inputs=gr.Textbox(lines=5, placeholder="Введите запрос для генерации кода..."),
            outputs="text",
            title="Генерация кода по запросу",
            description="Введите текстовый запрос, чтобы сгенерировать код.",
            examples=code_generation_examples,
            examples_per_page=2
        )
    with gr.Tab("Code Fix"):
        gr.Interface(
            fn=fix_code,
            inputs=[
                gr.Textbox(lines=2, placeholder="Ошибка..."),
                gr.Textbox(lines=2, placeholder="Проблема..."),
                gr.Textbox(lines=2, placeholder="Решение..."),
                gr.Textbox(lines=2, placeholder="Пример...")
            ],
            outputs="text",
            title="Исправление кода",
            description="Введите ошибку, проблему, решение и пример, чтобы получить исправленный код.",
            examples=code_fix_examples,
            examples_per_page=2
        )
    with gr.Tab("Flux Extend Image"):
        gr.Interface(
            fn=extend_image,
            inputs=gr.Image(type="pil", label="Загрузите изображение"),
            outputs=gr.Image(type="pil", label="Расширенное изображение"),
            title="Расширение изображений (Flux Extend Image)",
            description="Загрузите изображение, чтобы увеличить его разрешение.",
            examples=image_examples,
            examples_per_page=2
        )
    with gr.Tab("Image Generation"):
        gr.Interface(
            fn=generate_image,
            inputs=gr.Textbox(lines=2, placeholder="Введите текстовый запрос..."),
            outputs=gr.Image(type="pil", label="Сгенерированное изображение"),
            title="Генерация изображений (Image Generation)",
            description="Введите текстовый запрос, чтобы сгенерировать изображение.",
            examples=[
                "A futuristic cityscape at night",
                "A beautiful landscape with mountains and a lake",
                "An astronaut riding a horse in space"
            ],
            examples_per_page=2
        )

# Запускаем интерфейс
demo.launch()