import dash
from dash import dcc, html, Input, Output, State
import dash_bootstrap_components as dbc
from transformers import pipeline
import plotly.graph_objects as go
import os

# Initialize Hugging Face pipelines
try:
    sentiment_pipeline = pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment-latest")
    text_generator = pipeline("text-generation", model="gpt2", max_length=100)
except Exception as e:
    print(f"Error loading models: {e}")
    sentiment_pipeline = None
    text_generator = None

# Initialize Dash app
app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
app.title = "Hugging Face Dash Demo"

# Define the layout
app.layout = dbc.Container([
    dbc.Row([
        dbc.Col([
            html.H1("🤗 Hugging Face + Dash Demo", className="text-center mb-4"),
            html.Hr(),
        ])
    ]),
    
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H4("Sentiment Analysis", className="card-title"),
                    dcc.Textarea(
                        id='sentiment-input',
                        placeholder='Enter text to analyze sentiment...',
                        style={'width': '100%', 'height': 100},
                        className="mb-3"
                    ),
                    dbc.Button("Analyze Sentiment", id="sentiment-btn", color="primary", className="mb-3"),
                    html.Div(id='sentiment-output')
                ])
            ])
        ], width=6),
        
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H4("Text Generation", className="card-title"),
                    dcc.Textarea(
                        id='generation-input',
                        placeholder='Enter prompt for text generation...',
                        style={'width': '100%', 'height': 100},
                        className="mb-3"
                    ),
                    dbc.Button("Generate Text", id="generation-btn", color="success", className="mb-3"),
                    html.Div(id='generation-output')
                ])
            ])
        ], width=6)
    ], className="mb-4"),
    
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H4("Sentiment Score Visualization", className="card-title"),
                    dcc.Graph(id='sentiment-graph')
                ])
            ])
        ])
    ])
], fluid=True)

# Callback for sentiment analysis
@app.callback(
    [Output('sentiment-output', 'children'),
     Output('sentiment-graph', 'figure')],
    [Input('sentiment-btn', 'n_clicks')],
    [State('sentiment-input', 'value')]
)
def analyze_sentiment(n_clicks, text):
    if not n_clicks or not text or not sentiment_pipeline:
        return "Enter text and click 'Analyze Sentiment'", {}
    
    try:
        result = sentiment_pipeline(text)
        label = result[0]['label']
        score = result[0]['score']
        
        # Create output
        output = dbc.Alert([
            html.H5(f"Sentiment: {label}"),
            html.P(f"Confidence: {score:.2%}")
        ], color="info")
        
        # Create visualization
        colors = {'POSITIVE': 'green', 'NEGATIVE': 'red', 'NEUTRAL': 'orange'}
        fig = go.Figure(data=[
            go.Bar(x=[label], y=[score], marker_color=colors.get(label, 'blue'))
        ])
        fig.update_layout(
            title="Sentiment Analysis Result",
            xaxis_title="Sentiment",
            yaxis_title="Confidence Score",
            yaxis=dict(range=[0, 1])
        )
        
        return output, fig
        
    except Exception as e:
        return dbc.Alert(f"Error: {str(e)}", color="danger"), {}

# Callback for text generation
@app.callback(
    Output('generation-output', 'children'),
    [Input('generation-btn', 'n_clicks')],
    [State('generation-input', 'value')]
)
def generate_text(n_clicks, prompt):
    if not n_clicks or not prompt or not text_generator:
        return "Enter a prompt and click 'Generate Text'"
    
    try:
        result = text_generator(prompt, max_length=len(prompt.split()) + 50, num_return_sequences=1)
        generated_text = result[0]['generated_text']
        
        return dbc.Alert([
            html.H5("Generated Text:"),
            html.P(generated_text)
        ], color="success")
        
    except Exception as e:
        return dbc.Alert(f"Error: {str(e)}", color="danger")

# Run the app
if __name__ == '__main__':
    # Hugging Face Spaces requires the app to run on port 7860
    app.run_server(host='0.0.0.0', port=7860, debug=False)