# docker build -t reward-simulator .docker run -p 7860:7860 -v $(pwd)/data:/app/data reward-simulator

from flask import Flask, request, jsonify, render_template, send_from_directory
from PIL import Image
import numpy as np
import io
import torch

from request import get_ft  # get_ft(model, image) doit retourner un np.ndarray

app = Flask(__name__)

# Global model
model = None

def load_model():
    """Load DINOv2 model"""
    torch.hub.set_dir("/tmp/torch_cache")  # Dossier temporaire autorisé
    model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
    model.eval()
    model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
    return model

def init_model():
    global model
    model = load_model()

@app.route('/')
def home():
    return render_template('index.html')  # Si tu as un front-end intégré

@app.route('/static/<path:filename>')
def serve_static(filename):
    return send_from_directory('static', filename)

@app.route('/process', methods=['POST'])
def process_images():
    if 'image1' not in request.files or 'image2' not in request.files:
        return jsonify({'error': 'Two images must be provided (image1 and image2)'}), 400

    try:
        image1 = Image.open(io.BytesIO(request.files['image1'].read())).convert('RGB')
        image2 = Image.open(io.BytesIO(request.files['image2'].read())).convert('RGB')

        features1 = get_ft(model, image1)
        features2 = get_ft(model, image2)

        distance = float(np.linalg.norm(features1 - features2))
        return jsonify({'distance': distance})

    except Exception as e:
        print(f"Erreur back-end: {e}")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    init_model()
    app.run(host='0.0.0.0', port=7860)