import gradio as gr
import json
import os
from datetime import datetime

# File to store model data
MODEL_FILE = "models.json"

# Load existing data or create a new file if it doesn't exist
def load_models():
    if os.path.exists(MODEL_FILE):
        with open(MODEL_FILE, "r") as f:
            return json.load(f)
    return []

# Save data to the JSON file
def save_models(models):
    with open(MODEL_FILE, "w") as f:
        json.dump(models, f, indent=4)

# Helper function to calculate time differences between releases
def calculate_time_difference(models):
    for i in range(1, len(models)):
        prev_date = datetime.strptime(models[i - 1]['dateOfRelease'], "%Y-%m-%d")
        curr_date = datetime.strptime(models[i]['dateOfRelease'], "%Y-%m-%d")
        models[i]['time_difference'] = (curr_date - prev_date).days
    return models

# Function to display all models with optional filters
def display_models(developer=None, use_case=None, year_range=None):
    models = load_models()
    if not models:
        return "No models have been added yet."

    # Apply filters if provided
    if developer:
        models = [m for m in models if m.get('developer') == developer]
    if use_case:
        models = [m for m in models if m.get('use_case') == use_case]
    if year_range:
        start_year, end_year = year_range
        models = [m for m in models if start_year <= int(m['dateOfRelease'][:4]) <= end_year]

    models = sorted(models, key=lambda x: x['dateOfRelease'])  # Sort by release date
    models = calculate_time_difference(models)  # Calculate time differences

    output = "<div style='display: flex; flex-direction: column; align-items: flex-start; gap: 20px;'>"
    for model in models:
        time_gap = model.get('time_difference', 0) * 2  # Scale time gap for visualization
        output += f"<div style='margin-top: {time_gap}px; border: 1px solid #ccc; border-radius: 10px; padding: 10px; width: 90%; background: linear-gradient(135deg, #f9f9f9, #e9e9e9); box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.1);'>"
        output += f"<h3 style='color: #2d89ef;'>{model['name']} ({model['dateOfRelease']})</h3>"
        output += f"<p><strong>Description:</strong> {model['description']}</p>"
        output += f"<p><strong>Developer:</strong> {model.get('developer', 'Unknown')}</p>"
        output += f"<p><strong>Use Case:</strong> {model.get('use_case', 'General')}</p>"
        output += f"<p><strong>Impact:</strong> {model.get('impact', 'Not specified')}</p>"
        output += "</div>"
    output += "</div>"
    return output

# Function to add a new model
def add_model(name, description, dateOfRelease, developer, use_case, impact):
    models = load_models()
    models.append({
        "name": name,
        "description": description,
        "dateOfRelease": dateOfRelease,
        "developer": developer,
        "use_case": use_case,
        "impact": impact
    })
    save_models(models)
    return "Model added successfully!", display_models()

# Function to edit an existing model
def edit_model(index, name, description, dateOfRelease, developer, use_case, impact):
    models = load_models()
    if index < 1 or index > len(models):
        return "Invalid index. Please provide a valid model number.", display_models()
    models[index - 1] = {
        "name": name,
        "description": description,
        "dateOfRelease": dateOfRelease,
        "developer": developer,
        "use_case": use_case,
        "impact": impact
    }
    save_models(models)
    return "Model updated successfully!", display_models()

# Gradio interface
with gr.Blocks(css="""
    body {
        font-family: 'Arial', sans-serif;
        background-color: #121212;  /* Dark background color */
        color: white;  /* White text color for better contrast */
        margin: 0;
        padding: 0;
    }
    .gradio-container {
        width: 80%;
        margin: auto;
        padding: 20px;
    }
    .gradio-markdown {
        color: white;  /* Ensure markdown text is white */
    }
    .gr-button {
        background-color: #333;  /* Dark button background */
        color: white;  /* White button text */
    }
    .gr-button:hover {
        background-color: #555;  /* Lighter shade on hover */
    }
    .gr-dropdown, .gr-textbox {
        background-color: #333;  /* Dark input fields */
        color: white;  /* White text */
        border: 1px solid #555;  /* Lighter border */
    }
""") as app:
    gr.Markdown("# AI Timeline\n\nVisualize the development of AI models through an interactive timeline.")

    with gr.Tab("View Timeline"):
        with gr.Row():
            developer_filter = gr.Dropdown(label="Filter by Developer", choices=["All"] + list(set([m['developer'] for m in load_models()])), value="All")
            use_case_filter = gr.Dropdown(label="Filter by Use Case", choices=["All"] + list(set([m['use_case'] for m in load_models()])), value="All")
            year_range_filter = gr.Slider(label="Filter by Year Range", minimum=2000, maximum=2025, value=(2000, 2025), step=1)
        filter_button = gr.Button("Apply Filters")
        view_output = gr.HTML()

        def apply_filters(developer, use_case, year_range):
            dev = None if developer == "All" else developer
            use = None if use_case == "All" else use_case
            return display_models(developer=dev, use_case=use, year_range=year_range)

        filter_button.click(apply_filters, inputs=[developer_filter, use_case_filter, year_range_filter], outputs=view_output)

    with gr.Tab("Add Model"):
        with gr.Row():
            name_input = gr.Textbox(label="Model Name", placeholder="Enter model name")
            year_input = gr.Textbox(label="Publication Date", placeholder="Enter date of publication (YYYY-MM-DD)")
        description_input = gr.Textbox(label="Description", placeholder="Enter a short description")
        developer_input = gr.Textbox(label="Developer", placeholder="Enter the developer or organization")
        use_case_input = gr.Textbox(label="Use Case", placeholder="Enter the primary use case")
        impact_input = gr.Textbox(label="Impact", placeholder="Enter the model's impact")
        add_button = gr.Button("Add Model")
        add_output = gr.Markdown()
        add_button.click(
            add_model,
            inputs=[name_input, description_input, year_input, developer_input, use_case_input, impact_input],
            outputs=[add_output, view_output]
        )

    with gr.Tab("Edit Model"):
        edit_index = gr.Number(label="Model Number", precision=0)
        with gr.Row():
            edit_name = gr.Textbox(label="New Model Name", placeholder="Enter new model name")
            edit_year = gr.Textbox(label="New Publication Date", placeholder="Enter new date of publication (YYYY-MM-DD)")
        edit_description = gr.Textbox(label="New Description", placeholder="Enter new description")
        edit_developer = gr.Textbox(label="New Developer", placeholder="Enter new developer or organization")
        edit_use_case = gr.Textbox(label="New Use Case", placeholder="Enter new primary use case")
        edit_impact = gr.Textbox(label="New Impact", placeholder="Enter new impact")
        edit_button = gr.Button("Edit Model")
        edit_output = gr.Markdown()
        edit_button.click(
            edit_model,
            inputs=[edit_index, edit_name, edit_description, edit_year, edit_developer, edit_use_case, edit_impact],
            outputs=[edit_output, view_output]
        )

# Run the app
app.launch()