# app.py
import os
import io
import zipfile
import json
import shutil
from pathlib import Path
from PIL import Image
import torch
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader, random_split
import torch.nn as nn
import torch.optim as optim
import gradio as gr
import time

ROOT = Path(".")
DATA_ZIP_NAME = "dataset.zip"   # upload your Roboflow export here
WORK_DIR = ROOT / "roboflow_dataset"
CLASSIFY_DIR = ROOT / "classification_data"
MODEL_PATH = ROOT / "model.pth"
CLASSES_JSON = ROOT / "classes.json"

# Training config (tweak if needed)
BATCH_SIZE = 16
IMG_SIZE = 224
NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", 3))   # small default for Spaces CPU
LR = 1e-3
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def safe_mkdir(p: Path):
    p.mkdir(parents=True, exist_ok=True)


def extract_zip_to_workdir(zip_path: Path, out_dir: Path):
    if out_dir.exists():
        shutil.rmtree(out_dir)
    safe_mkdir(out_dir)
    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(out_dir)


def find_classes_mapping(workdir: Path):
    # Roboflow usually includes a data.yaml or classes.txt or a names list.
    # Try common locations.
    data_yaml = workdir / "data.yaml"
    classes_txt = workdir / "classes.txt"
    # Sometimes Roboflow includes a folder "labels" and a file "labels.names" or "classes.txt"
    if classes_txt.exists():
        names = [x.strip() for x in classes_txt.read_text().splitlines() if x.strip()]
        return names
    if data_yaml.exists():
        import yaml  # note: pyyaml must be in requirements if needed
        try:
            parsed = yaml.safe_load(data_yaml.read_text())
            if "names" in parsed:
                # could be list or dict
                n = parsed["names"]
                if isinstance(n, dict):
                    return [n[k] for k in sorted(n.keys(), key=lambda x: int(x))]
                elif isinstance(n, list):
                    return n
        except Exception:
            pass
    # fallback: try to find a file named "classes.txt" or "labels.names"
    for candidate in workdir.rglob("classes.txt"):
        names = [x.strip() for x in candidate.read_text().splitlines() if x.strip()]
        if names:
            return names
    for candidate in workdir.rglob("labels.names"):
        names = [x.strip() for x in candidate.read_text().splitlines() if x.strip()]
        if names:
            return names
    # last resort: scan label files to get max class index, produce numeric names
    max_idx = -1
    for lbl in workdir.rglob("labels/*.txt"):
        for line in lbl.read_text().splitlines():
            parts = line.strip().split()
            if len(parts) >= 1:
                try:
                    idx = int(float(parts[0]))
                    max_idx = max(max_idx, idx)
                except:
                    pass
    if max_idx >= 0:
        return [f"class_{i}" for i in range(max_idx + 1)]
    return []


def convert_roboflow_detection_to_classification(workdir: Path, outdir: Path):
    """
    Creates a folder-structured classification dataset:
    outdir/train/<class_name>/*.jpg
    outdir/valid/<class_name>/*.jpg
    It uses label files (YOLO txt) to assign the main class for each image.
    If bounding box info is available, it crops the bbox; otherwise it copies the image.
    """
    if outdir.exists():
        shutil.rmtree(outdir)
    safe_mkdir(outdir)

    # Try common image and label folders
    images_dirs = []
    labels_dirs = []
    for p in workdir.iterdir():
        if p.is_dir():
            if p.name.lower() in ("images", "image", "images/train", "train", "valid", "test"):
                images_dirs.append(p)
            if p.name.lower() in ("labels", "annotations"):
                labels_dirs.append(p)

    # simpler approach: look for 'images' and 'labels' in any depth
    images_all = list(workdir.rglob("images/*")) + list(workdir.rglob("images/*/*"))
    if not images_all:
        # fallback to all popular image file types in workdir
        images_all = [p for p in workdir.rglob("*") if p.suffix.lower() in (".jpg", ".jpeg", ".png")]

    # mapping of image filename (no path) to its full path
    img_map = {p.name: p for p in images_all}

    # find label files
    label_files = list(workdir.rglob("labels/*.txt")) + list(workdir.rglob("labels/*/*.txt"))
    if not label_files:
        # some exports put labels alongside images with same base name but different extension
        label_files = [p for p in workdir.rglob("*.txt") if p.stem in img_map]

    # find class names
    classes = find_classes_mapping(workdir)
    if not classes:
        # if not available, default to single class "unknown"
        classes = ["class_0"]

    # prepare train/valid split target folders (Roboflow often has train/valid folders; try to preserve)
    # We'll just create train and valid
    train_out = outdir / "train"
    valid_out = outdir / "valid"
    safe_mkdir(train_out)
    safe_mkdir(valid_out)

    # Load label->image mapping from label_files
    # We'll assume label files mirror the image names: e.g., images/train/img1.jpg and labels/train/img1.txt
    img_to_labels = {}
    for lbl in label_files:
        name = lbl.stem
        if name in img_map:
            img_to_labels[name] = lbl

    # If Roboflow has images split into train/valid dirs, detect them
    # Otherwise we'll create a split based on filenames (80/20)
    # Build a dataset list
    dataset_rows = []
    for img_name, img_path in img_map.items():
        lbl = img_to_labels.get(Path(img_name).stem)
        # Determine main class for this image (first label line)
        main_class = None
        bbox = None
        if lbl and lbl.exists():
            lines = [l for l in lbl.read_text().splitlines() if l.strip()]
            if lines:
                parts = lines[0].split()
                try:
                    cls_idx = int(float(parts[0]))
                    main_class = classes[cls_idx] if cls_idx < len(classes) else f"class_{cls_idx}"
                    if len(parts) >= 5:
                        # YOLO format: cls x_center y_center width height (normalized)
                        bbox = tuple(float(x) for x in parts[1:5])
                except Exception:
                    pass
        if not main_class:
            # fallback: mark as unknown
            main_class = "unknown"
            if "unknown" not in classes:
                classes.append("unknown")
        dataset_rows.append((img_path, main_class, bbox))

    # do deterministic split
    dataset_rows.sort(key=lambda x: x[0].name)
    split_idx = int(0.8 * len(dataset_rows))
    train_rows = dataset_rows[:split_idx]
    valid_rows = dataset_rows[split_idx:]

    def save_rows(rows, dest_folder):
        for img_path, cls_name, bbox in rows:
            dest_cls = dest_folder / cls_name
            safe_mkdir(dest_cls)
            try:
                img = Image.open(img_path).convert("RGB")
                if bbox:
                    # bbox are normalized; convert to pixel coords
                    w, h = img.size
                    xc, yc, bw, bh = bbox
                    left = int((xc - bw / 2) * w)
                    right = int((xc + bw / 2) * w)
                    top = int((yc - bh / 2) * h)
                    bottom = int((yc + bh / 2) * h)
                    # clamp
                    left = max(0, left); right = min(w, right)
                    top = max(0, top); bottom = min(h, bottom)
                    if right - left > 10 and bottom - top > 10:
                        img = img.crop((left, top, right, bottom))
                # save with a unique name
                dest_path = dest_cls / img_path.name
                img.save(dest_path)
            except Exception as e:
                print("Skipping", img_path, "due to", e)

    save_rows(train_rows, train_out)
    save_rows(valid_rows, valid_out)

    # Save classes json
    with open(CLASSES_JSON, "w") as f:
        json.dump(classes, f)
    return classes


def build_model(num_classes):
    model = models.resnet18(pretrained=True)
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, num_classes)
    return model


def train_model(data_dir: Path, classes):
    print("Starting training. This may take some time on CPU.")
    num_classes = len(classes)
    model = build_model(num_classes).to(DEVICE)

    transform_train = transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
    ])
    transform_valid = transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
    ])

    dataset_train = datasets.ImageFolder(str(data_dir / "train"), transform=transform_train)
    dataset_valid = datasets.ImageFolder(str(data_dir / "valid"), transform=transform_valid)

    # If ImageFolder class mapping differs from classes list, use folder names.
    # Dataloaders
    if len(dataset_train) == 0:
        raise RuntimeError("No training images found. Please check dataset structure.")

    loader_train = DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
    loader_valid = DataLoader(dataset_valid, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LR)

    best_val = 0.0
    for epoch in range(NUM_EPOCHS):
        model.train()
        running = 0.0
        for imgs, labels in loader_train:
            imgs = imgs.to(DEVICE)
            labels = labels.to(DEVICE)
            optimizer.zero_grad()
            outputs = model(imgs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running += loss.item()
        # validation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for imgs, labels in loader_valid:
                imgs = imgs.to(DEVICE)
                labels = labels.to(DEVICE)
                outputs = model(imgs)
                _, preds = torch.max(outputs, 1)
                correct += (preds == labels).sum().item()
                total += labels.size(0)
        acc = correct / total if total > 0 else 0.0
        print(f"Epoch {epoch+1}/{NUM_EPOCHS}, loss={running:.4f}, val_acc={acc:.4f}")
        if acc > best_val:
            best_val = acc
            # save best
            torch.save({
                "model_state": model.state_dict(),
                "classes": classes
            }, MODEL_PATH)
    print("Training complete. Best val acc:", best_val)
    # final save if not saved
    if not MODEL_PATH.exists():
        torch.save({
            "model_state": model.state_dict(),
            "classes": classes
        }, MODEL_PATH)
    return MODEL_PATH.exists()


def load_saved_model(path: Path):
    data = torch.load(path, map_location=DEVICE)
    classes = data.get("classes", None)
    if not classes and Path(CLASSES_JSON).exists():
        classes = json.loads(Path(CLASSES_JSON).read_text())
    if not classes:
        classes = [f"class_{i}" for i in range(2)]
    model = build_model(len(classes))
    model.load_state_dict(data["model_state"])
    model.to(DEVICE).eval()
    return model, classes


# Prepare model at startup
MODEL = None
BREEDS = None

def startup():
    global MODEL, BREEDS
    # If model exists, load directly
    if Path(MODEL_PATH).exists():
        try:
            MODEL, BREEDS = load_saved_model(Path(MODEL_PATH))
            print("Loaded existing model with classes:", BREEDS)
            return
        except Exception as e:
            print("Failed to load existing model:", e)

    # If dataset.zip exists, extract and convert, then train
    if Path(DATA_ZIP_NAME).exists():
        print("dataset.zip found. Extracting and preparing...")
        extract_zip_to_workdir(Path(DATA_ZIP_NAME), WORK_DIR)
        classes = convert_roboflow_detection_to_classification(WORK_DIR, CLASSIFY_DIR)
        print("Prepared classification dataset with classes:", classes)
        # train (may be slow on CPU)
        try:
            trained = train_model(CLASSIFY_DIR, classes)
            if trained:
                MODEL, BREEDS = load_saved_model(Path(MODEL_PATH))
        except Exception as e:
            print("Training failed:", e)
    else:
        print("No dataset.zip found. Please upload dataset.zip to the Space root or upload a model.pth")


# Prediction function
transform_predict = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])

def predict_image(pil_img):
    global MODEL, BREEDS
    if MODEL is None:
        return {"error": "Model not ready. Upload dataset.zip to train, or model.pth to load."}
    img = pil_img.convert("RGB")
    x = transform_predict(img).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        out = MODEL(x)
        probs = torch.nn.functional.softmax(out[0], dim=0).cpu().numpy()
    # top 3
    indices = probs.argsort()[::-1][:3]
    return {BREEDS[int(i)]: float(probs[int(i)]) for i in indices}

# Run startup (this will attempt to load or train)
start_time = time.time()
startup()
print("Startup complete in", time.time() - start_time, "seconds")

# Build Gradio app
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    examples=[],
    title="Cow Breed Classifier",
    description="Upload a cow image. If you uploaded Roboflow dataset.zip to the Space root, the Space will auto-train on start (small number of epochs). If you already have a trained model.pth, upload that instead to skip training."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)