import spaces
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "Fralet/DDeduPModelv7"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")


prompts = """You are a Design Developer of Educational Programs (DDeduP). Below is an instruction that describes a task.

### Instruction:
{}

### Input:
{}

### Response:
{}"""

def generate(instruction, program, number):
    prompt = prompts.format(instruction, program, "")

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=number
        )

    text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    if "### Response:" in text:
        text = text.split("### Response:")[1].strip()

    return text


@spaces.GPU
def generate_goal(program):
    instruction="Write an educational program goal for this educational program in English. The goal should be as concise and specific as possible within the context of the professional field."
    return generate(instruction, program, 500)


@spaces.GPU
def generate_learning_outcomes(program):
    instruction="Write 10-12 learning outcomes for this educational program in English in JSON. The learning outcomes should follow Bloom's taxonomy (knowledge, comprehension, application, analysis, synthesis, evaluation)."
    return generate(instruction, program, 4000)

@spaces.GPU
def generate_course_description(program):
    instruction="Provide a concise and informative description of the course in English, highlighting its key objectives, topics covered, and intended outcomes."
    return generate(instruction, program, 1000)

@spaces.GPU
def generate_recommendation(program):
    instruction="Recommend appropriate courses that align with the educational program. Provide course titles in English, enclosed in brackets and separated by semicolons."
    return generate(instruction, program, 4000)

@spaces.GPU
def generate_course_lo_mapping(program):
    instruction = (
        "Match each course with all relevant learning outcome codes from the educational program. "
        "Format the output as a JSON object, ensuring accurate alignment between course content "
        "and program-level learning outcomes."
    )
    return generate(instruction, program, 8000)
    
@spaces.GPU
def generate_professional_standards(program):
    instruction = (
        "Extract all professional standards relevant to this educational program "
        "and present them in JSON format in the English language."
    )
    return generate(instruction, program, 500)
    
@spaces.GPU
def generate_emerging_professions(program):
    instruction = (
        "Extract relevant emerging professions from the Atlas of New Professions "
        "based on the educational program and its stated goal. "
        "Present the results in JSON format, with all content written in Russian."
    )
    return generate(instruction, program, 500)

@spaces.GPU
def generate_graduate_professions(program):
    instruction = (
        "List potential professions that graduates of the educational program "
        "may pursue upon completion in English. "
        "Present the output as a JSON array."
    )
    return generate(instruction, program, 500)
    
with gr.Blocks(title="DDeduP API") as demo:

    gr.Markdown("# DDeduP API")

    with gr.Tab("Generate Goal"):
        goal_input = gr.Textbox(lines=5, label="Educational Program")
        goal_output = gr.Textbox(lines=8, label="Goal")
        goal_btn = gr.Button("Generate Goal")

        goal_btn.click(
            fn=generate_goal,
            inputs=goal_input,
            outputs=goal_output,
            api_name="generate_goals"
        )


    with gr.Tab("Generate Learning Outcomes"):
        lo_input = gr.Textbox(lines=5, label="Educational Program")
        lo_output = gr.Textbox(lines=15, label="Learning Outcomes JSON")
        lo_btn = gr.Button("Generate Learning Outcomes")

        lo_btn.click(
            fn=generate_learning_outcomes,
            inputs=lo_input,
            outputs=lo_output,
            api_name="generate_learning_outcomes"


        )

    with gr.Tab("Generate Course Descriptions"):
        lo_input = gr.Textbox(lines=5, label="Educational Program")
        lo_output = gr.Textbox(lines=15, label="Course Descriptions")
        lo_btn = gr.Button("Generate Course Descriptions")

        lo_btn.click(
            fn=generate_course_description,
            inputs=lo_input,
            outputs=lo_output,
            api_name="generate_course_descriptions"


        )
    with gr.Tab("Generate Recommendations"):
        lo_input = gr.Textbox(lines=5, label="Educational Program")
        lo_output = gr.Textbox(lines=15, label="Recommendations")
        lo_btn = gr.Button("Generate Recommendations")

        lo_btn.click(
            fn=generate_recommendation,
            inputs=lo_input,
            outputs=lo_output,
            api_name="generate_recommendation"


        )
    with gr.Tab("Course ↔ Learning Outcomes"):
        input_box = gr.Textbox(lines=5, label="Educational Program")
        output_box = gr.Textbox(lines=20, label="Course Mapping JSON")
        btn = gr.Button("Generate Mapping")
    
        btn.click(
            fn=generate_course_lo_mapping,
            inputs=input_box,
            outputs=output_box,
            api_name="generate_course_lo_mapping"
        )

    with gr.Tab("Professional Standards"):
        input_box = gr.Textbox(lines=5, label="Educational Program")
        output_box = gr.Textbox(lines=20, label="Professional Standards JSON")
        btn = gr.Button("Generate Standards")
    
        btn.click(
            fn=generate_professional_standards,
            inputs=input_box,
            outputs=output_box,
            api_name="generate_professional_standards"
        )

    with gr.Tab("Emerging Professions"):
        input_box = gr.Textbox(lines=5, label="Educational Program")
        output_box = gr.Textbox(lines=20, label="Emerging Professions JSON")
        btn = gr.Button("Generate Emerging Professions")
    
        btn.click(
            fn=generate_emerging_professions,
            inputs=input_box,
            outputs=output_box,
            api_name="generate_emerging_professions"
        )

    with gr.Tab("Graduate Professions"):
        input_box = gr.Textbox(lines=5, label="Educational Program")
        output_box = gr.Textbox(lines=20, label="Professions JSON")
        btn = gr.Button("Generate Professions")
    
        btn.click(
            fn=generate_graduate_professions,
            inputs=input_box,
            outputs=output_box,
            api_name="generate_graduate_professions"
        )
demo.launch()