import os
import gradio as gr
from Oracle.deepfundingoracle import prepare_dataset, train_predict_weight, create_submission_csv
import pandas as pd
import matplotlib.pyplot as plt
import time
import io
from PIL import Image

def analyze_file(file, progress=gr.Progress(track_tqdm=True)):
    start_time = time.time()
    progress(0, desc="Preparing dataset...")
    df = prepare_dataset(file.name)
    progress(0.3, desc="Predicting weights...")
    df = train_predict_weight(df)
    progress(0.6, desc="Saving results to CSV...")
    csv_path = create_submission_csv(df, "submission.csv")
    progress(0.8, desc="Generating graph...")
    # Example: plot histogram of a column if exists
    fig, ax = plt.subplots()
    if 'final_weight' in df.columns:
        df['final_weight'].hist(ax=ax)
        ax.set_title('Distribution of Final Weights')
        ax.set_xlabel('Final Weight')
        ax.set_ylabel('Count')
    else:
        ax.text(0.5, 0.5, 'No final_weight column to plot', ha='center')
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close(fig)
    img = Image.open(buf)
    progress(1, desc="Done!")
    elapsed = time.time() - start_time
    preview = df.head().to_csv(index=False)
    return preview, csv_path, img, f"Analysis completed in {elapsed:.2f} seconds."

iface = gr.Interface(
    fn=analyze_file,
    inputs=gr.File(label="Upload CSV"),
    outputs=[
        gr.Textbox(label="Preview of Results"),
        gr.File(label="Download CSV"),
        gr.Image(label="Analysis Graph"),
        gr.Textbox(label="Status/Timing Info")
    ],
    title="DeepFunding Oracle",
    description="Upload a CSV of repo-parent relationships; see analysis progress, get a graph, and download results as CSV.",
    allow_flagging="never"
)

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    iface.launch(server_name="0.0.0.0", server_port=port)