import sys
import os
sys.path.append(os.path.abspath("src"))

import torch
import numpy as np
from sklearn.metrics import precision_recall_curve, f1_score
from dataset import create_dataloaders
from model import SEMViTAutoencoder
from train_eval import train_model, evaluate_anomaly
import json
from safetensors.torch import save_file

# Configuration
DATA_DIR = "/content"  #"./data"
BATCH_SIZE = 8 # Small batch for 512p on CPU
NUM_WORKERS = 2
EPOCHS = 15
DEVICE = torch.device("cuda") #if torch.cuda.is_available() else "cpu")

def main():
    print(f"--- Running on {DEVICE} ---")
    print("--- Loading Data ---")
    # Ensure create_dataloaders handles the /content/ path
    train_loader, test_loader = create_dataloaders(DATA_DIR, BATCH_SIZE, NUM_WORKERS)

    print("--- Initializing Model ---")
    model = SEMViTAutoencoder().to(DEVICE)

    print("--- Starting Training (Normal Images Only) ---")
    # Pass DEVICE to your training function if it's not handled internally
    train_model(model, train_loader, epochs=EPOCHS)

    print("--- Evaluating & Finding Threshold ---")
    model.eval()
    errors, labels = evaluate_anomaly(model, test_loader)

    target_labels = 1 - np.array(labels)

    # Calculate best threshold
    precisions, recalls, thresholds = precision_recall_curve(target_labels, errors)
    f1_scores = 2 * (precisions * recalls) / (precisions + recalls + 1e-8)

    # Safety check: if no defects were found, argmax might fail
    best_idx = np.argmax(f1_scores)
    best_threshold = thresholds[best_idx]

    print(f"Optimal Anomaly Threshold: {best_threshold:.6f}")
    print(f"Max F1-Score: {f1_scores[best_idx]:.4f}")

    # Save Safetensors
    save_file(model.state_dict(), "model.safetensors")

    # Save Config with extra metadata for the App
    config = {
        "threshold": float(best_threshold),
        "model_type": "ViT-Autoencoder",
        "resolution": 512
    }

    with open("config.json", "w") as f:
        json.dump(config, f, indent=4)

    print("✅ Pipeline Complete. Download 'model.safetensors' and 'config.json' from the file pane.")


if __name__ == "__main__":
    main()