import gradio as gr
import pandas as pd
import numpy as np
import pickle
import os
from sentence_transformers import SentenceTransformer
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, confusion_matrix

MODEL_OUTPUT = "tree_house_model.pkl"
BERT_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"

# ============================================================
# HELPER: CREATE TRAINING TEXT
# ============================================================
def create_training_text(row):
    parts = []
    for field in ["title", "characters", "setting", "plot", "problem", "ending"]:
        if field in row and pd.notna(row[field]):
            parts.append(f"{field.capitalize()}: {row[field]}")
    return " | ".join(parts)

# ============================================================
# TRAINING FUNCTION
# ============================================================
def train_model(tree_file, house_file):
    logs = []

    if tree_file is None or house_file is None:
        return "❌ Please upload both CSV files.", None

    logs.append("📂 Loading datasets...")
    tree_df = pd.read_csv(tree_file.name)
    house_df = pd.read_csv(house_file.name)

    tree_df["category"] = "tree"
    house_df["category"] = "house"

    full_df = pd.concat([tree_df, house_df], ignore_index=True)

    logs.append(f"✓ Tree stories: {len(tree_df)}")
    logs.append(f"✓ House stories: {len(house_df)}")

    full_df["training_text"] = full_df.apply(create_training_text, axis=1)

    X_texts = full_df["training_text"].tolist()
    y_labels = full_df["category"].tolist()

    logs.append("🤖 Loading SentenceTransformer...")
    bert_model = SentenceTransformer(BERT_MODEL_NAME)

    logs.append("🧠 Creating embeddings...")
    X_embeddings = bert_model.encode(
        X_texts,
        show_progress_bar=True,
        normalize_embeddings=True  # 🔥 improves accuracy
    )

    X_train, X_test, y_train, y_test = train_test_split(
        X_embeddings,
        y_labels,
        test_size=0.2,
        random_state=42,
        stratify=y_labels
    )

    logs.append("🎯 Training Random Forest...")

    classifier = RandomForestClassifier(
        n_estimators=500,          # 🔥 more trees
        max_depth=30,              # 🔥 deeper trees
        min_samples_split=2,
        min_samples_leaf=1,
        class_weight="balanced",   # 🔥 handles imbalance
        random_state=42,
        n_jobs=-1
    )

    classifier.fit(X_train, y_train)

    y_pred = classifier.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    cm = confusion_matrix(y_test, y_pred, labels=["tree", "house"])

    logs.append(f"✅ Accuracy: {accuracy * 100:.2f}%")
    logs.append("📊 Confusion Matrix:")
    logs.append(str(cm))

    model_package = {
        "classifier": classifier,
        "bert_model_name": BERT_MODEL_NAME,
        "categories": ["tree", "house"],
        "accuracy": accuracy,
        "feature_dim": X_embeddings.shape[1],
        "confusion_matrix": cm.tolist(),
        "trained_on": pd.Timestamp.now().strftime("%Y-%m-%d %H:%M:%S"),
        "tree_stories": tree_df.to_dict("records"),
        "house_stories": house_df.to_dict("records"),
        "total_stories": len(full_df),
    }

    with open(MODEL_OUTPUT, "wb") as f:
        pickle.dump(model_package, f)

    logs.append(f"💾 Model saved as {MODEL_OUTPUT}")
    logs.append("⬇️ Download using the button below")

    return "\n".join(logs), MODEL_OUTPUT

# ============================================================
# GRADIO UI
# ============================================================
with gr.Blocks(title="Tree vs House Model Trainer") as demo:
    gr.Markdown(
        """
        # 🌳🏠 Tree vs House Story Model Trainer
        Upload your datasets and train a high-accuracy model.
        """
    )

    with gr.Row():
        tree_file = gr.File(label="🌳 Tree Dataset CSV", file_types=[".csv"])
        house_file = gr.File(label="🏠 House Dataset CSV", file_types=[".csv"])

    train_button = gr.Button("🚀 Train Model")

    output_logs = gr.Textbox(label="Training Logs", lines=20)
    download_model = gr.File(label="⬇️ Download Trained Model")

    train_button.click(
        fn=train_model,
        inputs=[tree_file, house_file],
        outputs=[output_logs, download_model]
    )

if __name__ == "__main__":
    demo.launch()