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

from PIL import Image
import numpy as np
import io
import faiss

import requests
import torch

from request import get_ft, get_topk
from flickrapi import FlickrAPI

from flask import Flask, request, render_template, jsonify, send_from_directory
app = Flask(__name__)

PRESET_IMAGES = {
    1: "static/1.webp",
    2: "static/2.webp",
    3: "static/3.webp"
}

# Add Flickr configuration
FLICKR_API_KEY = '80ef21a6f7eb0984ea613c316a89ca69'
FLICKR_API_SECRET = '4d0e8ce6734f4b3f'
flickr = FlickrAPI(FLICKR_API_KEY, FLICKR_API_SECRET, format='parsed-json', store_token=False)

def get_photo_id(url):
    """Extract photo ID from Flickr URL"""
    try:
        return url.split('/')[-1].split('_')[0]
    except:
        return None

def get_other_info(url):
    """Get author information from Flickr"""
    try:
        photo_id = get_photo_id(url)
        if photo_id:
            photo_info = flickr.photos.getInfo(photo_id=photo_id)
            license = photo_info['photo']['license']
            owner = photo_info['photo']['owner']
            flickr_url = f"https://www.flickr.com/photos/{owner.get('nsid', '')}/{photo_id}"
            return {
                'username': owner.get('username', ''),
                'realname': owner.get('realname', ''),
                'nsid': owner.get('nsid', ''),
                'flickr_url': flickr_url,
                'license': license
            }
    except:
        pass
    return {
        'username': 'Unknown',
        'realname': 'Unknown',
        'nsid': '',
        'flickr_url': '',
        'license': 'Unknown'
    }


def load_model():
    """Load DINOv2 model once and cache it"""
    torch.hub.set_dir('static')
    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 load_index(index_path):
    """Load FAISS index once and cache it"""
    return faiss.read_index(index_path)

def distance_to_similarity(distances, temp=1e-4):
    """Convert distance to similarity"""
    for ii in range(len(distances)):
        contribs = distances[ii].max() - distances[ii]
        contribs = contribs / temp
        sum_contribs = np.exp(contribs).sum()
        distances[ii] = np.exp(contribs) / sum_contribs
    return distances

def calculate_rewards(subscription, num_generations, author_share, ro_share, num_users_k, similarities, num_authors=1800):
    """Calculate rewards based on user inputs and similarities"""
    num_users = num_users_k * 1000
    
    # Monthly revenue allocated to authors
    authors_monthly_revenue = subscription * num_users * (author_share / 100)
    
    rewards = []
    for sim in similarities[0]:
        # Attribution bonus based on similarity score and number of neighbors
        attribution_bonus = sim * len(similarities[0])
        
        # Calculate monthly rewards
        author_month_reward = (authors_monthly_revenue / num_authors) * attribution_bonus
        ro_month_reward = author_month_reward / (author_share / 100) * (ro_share / 100)
        
        rewards.append({
            'paid_per_month': f"{subscription:.0f}€",
            'attribution': f"{sim*100:.0f}%",
            'author_month_reward': f"{author_month_reward:.0f}€",
            'ro_month_reward': f"{ro_month_reward:.0f}€"
            # 'paid_per_month': f"{subscription:.0f}€",
            # 'paid_per_gen': f"{paid_per_gen:.2f}€",
            # 'aro_share': f"{aro_share:.2f}c€",
            # 'attribution': f"{sim*100:.0f}%",
            # 'training_data_reward': f"{training_data_reward:.2f}c€",
            # 'author_month_reward': f"{author_month_reward:.0f}€",
            # 'ro_month_reward': f"{ro_month_reward:.0f}€"
        })
    return rewards

# Global variables for model and index
model = None
index = None
urls = None

def init_model():
    global model, index, urls
    model = load_model()
    index = load_index("data/openimages_index.bin")
    with open("data/openimages_urls.txt", "r") as f:
        urls = f.readlines()

@app.route('/')
def home():
    return render_template('index.html')

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

DEFAULT_PARAMS = {
    'subscription': 12,
    'num_generations': 60,
    'author_share': 5,
    'ro_share': 10,
    'num_users_k': 500,
    'num_neighbors': 10,
    'num_authors': 2000
}

@app.route('/select_preset/<int:preset_id>')
def select_preset(preset_id):
    if preset_id not in PRESET_IMAGES:
        return jsonify({'error': 'Invalid preset ID'}), 400
    
    try:
        image_path = PRESET_IMAGES[preset_id]
        image = Image.open(image_path).convert('RGB')
        
        # Use default parameters for presets
        params = DEFAULT_PARAMS.copy()
        
        # Get features and search
        features = get_ft(model, image)
        distances, indices = get_topk(index, features, topk=params['num_neighbors'])
        
        # Collect valid results first
        valid_results = []
        valid_similarities = []
        for i in range(params['num_neighbors']):
            image_url = urls[indices[0][i]].strip()
            try:
                response = requests.head(image_url)
                if response.status_code == 200:
                    valid_results.append({
                        'index': i,
                        'url': image_url
                    })
                    valid_similarities.append(distances[0][i])
            except requests.RequestException:
                continue
        
        # Renormalize similarities for valid results
        if valid_similarities:
            similarities = distance_to_similarity(np.array([valid_similarities]), temp=1e-5)
            
            # Calculate rewards with renormalized similarities
            rewards = calculate_rewards(
                params['subscription'],
                params['num_generations'],
                params['author_share'],
                params['ro_share'],
                params['num_users_k'],
                similarities,
                params['num_authors']
            )
            
            # Build final results
            results = []
            for i, result in enumerate(valid_results):
                other_info = get_other_info(result['url'])
                results.append({
                    'image_url': result['url'],
                    'rewards': rewards[i],
                    'other': other_info
                })
            
            return jsonify({'results': results})
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/process', methods=['POST'])
def process_image():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400
    
    try:
        image_file = request.files['image']
        image = Image.open(io.BytesIO(image_file.read())).convert('RGB')
        
        # Use default parameters if none provided
        params = DEFAULT_PARAMS.copy()
        if request.form:
            params.update({
                'subscription': float(request.form.get('subscription', params['subscription'])),
                'num_generations': int(request.form.get('num_generations', params['num_generations'])),
                'author_share': float(request.form.get('author_share', params['author_share'])),
                'ro_share': float(request.form.get('ro_share', params['ro_share'])),
                'num_users_k': int(request.form.get('num_users_k', params['num_users_k'])),
                'num_neighbors': int(request.form.get('num_neighbors', params['num_neighbors'])),
                'num_authors': int(request.form.get('num_authors', DEFAULT_PARAMS['num_authors'])),
            })
        
        # Process image
        features = get_ft(model, image)
        distances, indices = get_topk(index, features, topk=params['num_neighbors'])
        
        # Collect valid results first
        valid_results = []
        valid_similarities = []
        for i in range(params['num_neighbors']):
            image_url = urls[indices[0][i]].strip()
            try:
                response = requests.head(image_url)
                if response.status_code == 200:
                    valid_results.append({
                        'index': i,
                        'url': image_url
                    })
                    valid_similarities.append(distances[0][i])
            except requests.RequestException:
                continue
        
        # Renormalize similarities for valid results
        if valid_similarities:
            similarities = distance_to_similarity(np.array([valid_similarities]), temp=1e-5)
            
            # Calculate rewards with renormalized similarities
            rewards = calculate_rewards(
                params['subscription'],
                params['num_generations'],
                params['author_share'],
                params['ro_share'],
                params['num_users_k'],
                similarities,
                params['num_authors']
            )
            
            # Build final results
            results = []
            for i, result in enumerate(valid_results):
                other_info = get_other_info(result['url'])
                results.append({
                    'image_url': result['url'],
                    'rewards': rewards[i],
                    'other': other_info
                })
            
            return jsonify({'results': results})
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

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