import gradio as gr
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
import torch

# Load a smaller, optimized model
model_name = "google/flan-t5-base"  # Switch to a smaller model for faster inference
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Load model onto CPU with optimization
strategy_generator = pipeline(
    "text2text-generation",
    model=model,
    tokenizer=tokenizer,
    device=0 if torch.cuda.is_available() else -1,  # Use GPU if available
)

# Function to generate actionable steps
def generate_steps(industry, challenge, goals):
    prompt = f"""
    You are a business consultant with expertise in the {industry} industry.
    The company faces the following challenge: {challenge}.
    The company's goal is to achieve: {goals}.
    Provide three to five actionable steps to help the company achieve this goal.
    Focus on specific, realistic, and innovative strategies relevant to the industry.
    """
    try:
        response = strategy_generator(prompt, max_length=200, num_return_sequences=1, temperature=0.7, top_p=0.9)
        return response[0]['generated_text']
    except Exception as e:
        return f"Error generating steps: {e}"

# Function to combine rationale ("why") and implementation ("how")
def expand_step(step):
    prompt = f"""
    You are a business consultant. For the following strategy:
    "{step}"
    Provide:
    - Why this step is recommended.
    - How to implement this step effectively.
    """
    try:
        response = strategy_generator(prompt, max_length=150, num_return_sequences=1, temperature=0.7, top_p=0.9)
        return response[0]['generated_text']
    except Exception as e:
        return f"Error expanding step: {e}"

# Combined function to generate detailed strategy
def generate_strategy(industry, challenge, goals):
    # Generate initial steps
    steps = generate_steps(industry, challenge, goals)
    if "Error" in steps:
        return steps
    
    # Split steps and expand each
    steps_list = steps.split("\n")
    detailed_steps = []
    for step in steps_list:
        if step.strip():
            expanded = expand_step(step)
            detailed_steps.append(f"{step}\n{expanded}")
    
    return "\n\n".join(detailed_steps)

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# AI Business Strategy Generator")
    gr.Markdown("Generate actionable business strategies and SWOT analyses using AI.")

    # Tab 1: Generate Business Strategy
    with gr.Tab("Generate Strategy"):
        gr.Markdown("### Input Information to Generate a Business Strategy")
        industry_input = gr.Textbox(label="Industry", placeholder="E.g., E-commerce, Healthcare")
        challenge_input = gr.Textbox(label="Key Challenge", placeholder="E.g., Low customer retention")
        goals_input = gr.Textbox(label="Goals", placeholder="E.g., Increase sales by 20% in 6 months")
        strategy_button = gr.Button("Generate Strategy")
        strategy_output = gr.Textbox(label="Generated Strategy", lines=10)

        strategy_button.click(
            generate_strategy,
            inputs=[industry_input, challenge_input, goals_input],
            outputs=[strategy_output]
        )

# Launch the Gradio app
demo.launch()