from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import precision_recall_fscore_support

from petri_cv import add_annotation_centers, candidate_points, features_for_points, label_candidates, prepare_image

DATASET = Path("AGAR_representative")
MODEL_PATH = Path("models/petri_candidate_classifier.joblib")


def main() -> None:
    features, targets = [], []
    annotations = sorted(DATASET.rglob("*.json"))
    for annotation_path in annotations:
        data = json.loads(annotation_path.read_text())
        image_path = annotation_path.with_suffix(".jpg")
        if not image_path.exists():
            continue
        image = prepare_image(image_path)
        points = add_annotation_centers(candidate_points(image), data["labels"], image.scale)
        labels = label_candidates(points, data["labels"], image.scale)
        # keep a useful but not overwhelming background-to-colony ratio
        positive = np.where(labels == 1)[0]
        negative = np.where(labels == 0)[0]
        rng = np.random.default_rng(abs(hash(annotation_path.stem)) % (2**32))
        negative = rng.choice(negative, size=min(len(negative), max(50, len(positive) * 5)), replace=False)
        selected = np.concatenate([positive, negative])
        matrix = features_for_points(image, [points[index] for index in selected])
        features.append(matrix)
        targets.append(labels[selected])
        print(f"{image_path.name}: {len(positive)} positives, {len(negative)} negatives")
    x = np.vstack(features)
    y = np.concatenate(targets)
    model = RandomForestClassifier(
        n_estimators=280, max_depth=14, min_samples_leaf=2, class_weight="balanced_subsample", n_jobs=-1, random_state=42
    )
    model.fit(x, y)
    predicted = model.predict(x)
    precision, recall, f1, _ = precision_recall_fscore_support(y, predicted, average="binary", zero_division=0)
    MODEL_PATH.parent.mkdir(exist_ok=True)
    joblib.dump({"model": model, "feature_version": 1}, MODEL_PATH)
    print(f"Saved {MODEL_PATH}; training candidate precision={precision:.3f}, recall={recall:.3f}, f1={f1:.3f}; rows={len(y)}")


if __name__ == "__main__":
    main()