import matplotlib.pyplot as plt
import firebase_admin
from firebase_admin import credentials, firestore
import gradio as gr
import io
from PIL import Image
import os
import json

# Initialize Firebase if not already initialized
if not firebase_admin._apps:
    cred = credentials.Certificate("firebase_key.json")
    firebase_admin.initialize_app(cred)

db = firestore.client()

# ✅ This stays as-is: Firebase Feedback Summary
def update_dashboard_plot():
    logs_ref = db.collection("evo_feedback")
    docs = logs_ref.stream()

    count_1 = 0
    count_2 = 0
    for doc in docs:
        data = doc.to_dict()
        winner = data.get("winner", "")
        if winner == "1":
            count_1 += 1
        elif winner == "2":
            count_2 += 1

    # Generate a bar chart
    fig, ax = plt.subplots()
    ax.bar(["Solution 1", "Solution 2"], [count_1, count_2], color=["blue", "green"])
    ax.set_ylabel("Votes")
    ax.set_title("EvoTransformer Feedback Summary")

    # Convert matplotlib figure to PIL Image
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)
    img = Image.open(buf)
    return img

# ✅ NEW: Accuracy Plot from Local Log
def evolution_accuracy_plot():
    try:
        log_path = "trained_model/evolution_log.json"
        if not os.path.exists(log_path):
            fig, ax = plt.subplots()
            ax.text(0.5, 0.5, "No evolution log found", ha="center", va="center")
            return fig

        with open(log_path, "r") as f:
            log_data = json.load(f)

        generations = list(range(1, len(log_data) + 1))
        accuracies = [entry.get("accuracy", 0) for entry in log_data]

        fig, ax = plt.subplots()
        ax.plot(generations, accuracies, marker="o", linestyle="-")
        ax.set_xlabel("Generation")
        ax.set_ylabel("Accuracy")
        ax.set_title("EvoTransformer Evolution Accuracy")
        ax.grid(True)

        return fig
    except Exception as e:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, f"Error loading plot: {e}", ha="center")
        return fig