import gradio as gr
import pandas as pd
import numpy as np
from datasets import Dataset
from sklearn.ensemble import RandomForestClassifier
import torch
from diffusers import StableDiffusionPipeline

print("System Initialization...")
np.random.seed(42)
ozellikler = np.random.rand(200, 14) * 100 
etiketler = np.random.choice([0, 1], 200) 
df = pd.DataFrame(ozellikler)
df['durum'] = etiketler

hf_dataset = Dataset.from_pandas(df)
X = hf_dataset.to_pandas().iloc[:, :-1]
y = hf_dataset.to_pandas().iloc[:, -1]

clf = RandomForestClassifier(n_estimators=50, random_state=42)
clf.fit(X, y)

pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-tiny", torch_dtype=torch.float32)
if torch.cuda.is_available():
    pipe = pipe.to("cuda")

def generate_neuro_image():
    yeni_sinyal = np.random.rand(1, 14) * 100
    tahmin = clf.predict(yeni_sinyal)[0]
    
    if tahmin == 0:
        durum_adi = "Awake & Focused (Eyes Open)"
        prompt = "A highly focused futuristic brain, glowing neon blue nodes, cyberpunk city background, ultra detailed, 8k"
    else:
        durum_adi = "Calm & Resting (Eyes Closed)"
        prompt = "A peaceful zen garden, glowing soft green tree, calm relaxing atmosphere, digital art, highly detailed"
        
    image = pipe(prompt, num_inference_steps=20).images[0]
    return durum_adi, prompt, image

# Minimalist ve Profesyonel UI (Monochrome Theme)
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# CognitiveDiffusion: Edge-Optimized Neuro-AI Pipeline")
    gr.Markdown("This Proof of Concept simulates an active 14-channel EEG data stream, categorizes the real-time cognitive state using a lightweight classifier, and visually renders the state utilizing a pruned diffusion model (bk-sdm-tiny). Designed for resource-constrained edge environments.")
    
    with gr.Row():
        with gr.Column():
            btn = gr.Button("Simulate EEG Signal & Generate Output", variant="primary")
            state_out = gr.Textbox(label="Detected Cognitive State", lines=1)
            prompt_out = gr.Textbox(label="Generated Gen-AI Prompt", lines=2)
        with gr.Column():
            image_out = gr.Image(label="Neuro-AI Output")
            
    btn.click(fn=generate_neuro_image, inputs=[], outputs=[state_out, prompt_out, image_out])

demo.launch()