from __future__ import annotations

import io
from pathlib import Path

import joblib
import numpy as np
from flask import Flask, jsonify, request, send_from_directory

from petri_cv import candidate_points, features_for_points, prepare_image

ROOT = Path(__file__).resolve().parent
MODEL_PATH = ROOT / "models/petri_candidate_classifier.joblib"
app = Flask(__name__)
model_bundle = joblib.load(MODEL_PATH) if MODEL_PATH.exists() else None


@app.get("/")
def index():
    return send_from_directory(ROOT, "index.html")


@app.get("/<path:path>")
def assets(path: str):
    return send_from_directory(ROOT, path)


@app.post("/api/analyze")
def analyze():
    if model_bundle is None:
        return jsonify(error="Модель не обучена. Запустите train_petri_model.py."), 503
    uploaded = request.files.get("image")
    if uploaded is None or not uploaded.mimetype.startswith("image/"):
        return jsonify(error="Передайте изображение PNG, JPG или WEBP."), 400
    try:
        prepared = prepare_image(io.BytesIO(uploaded.read()))
    except Exception:
        return jsonify(error="Не удалось прочитать изображение."), 400
    points = candidate_points(prepared)
    if not points:
        return jsonify(colonies=[], candidates=0, model="RandomForest candidate ranker")
    probabilities = model_bundle["model"].predict_proba(features_for_points(prepared, points))[:, 1]
    ranked = sorted(zip(probabilities, points), reverse=True)
    selected = []
    # Non-maximum suppression avoids several markers on a single colony.
    for score, (x, y) in ranked:
        if score < .54:
            continue
        if all((x - candidate["x_px"]) ** 2 + (y - candidate["y_px"]) ** 2 > 36 ** 2 for candidate in selected):
            selected.append({"x": round(x / prepared.gray.shape[1], 5), "y": round(y / prepared.gray.shape[0], 5),
                             "x_px": x, "y_px": y, "score": round(float(score), 3)})
        if len(selected) >= 300:
            break
    for item in selected:
        item.pop("x_px")
        item.pop("y_px")
    return jsonify(colonies=selected, candidates=len(points), model="RandomForest candidate ranker")


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=8000, debug=False)