# An arena for 3D generations with code inspired from TTS arena

import gradio as gr
import pandas as pd
from langdetect import detect
from datasets import load_dataset
import threading, time, uuid, sqlite3, shutil, os, random, asyncio, threading
from pathlib import Path
from huggingface_hub import CommitScheduler, delete_file, hf_hub_download
from gradio_client import Client
from detoxify import Detoxify
import os
import tempfile
import PIL.Image

toxicity = Detoxify('original')

####################################
# Constants
####################################
AVAILABLE_MODELS = {
    'TripoSR': 'TripoSR',
    'LGM-Mini': 'LGM-Mini',
    'CRM': 'CRM'
}

SPACE_ID = os.getenv('SPACE_ID')
MAX_SAMPLE_TXT_LENGTH = 300
MIN_SAMPLE_TXT_LENGTH = 10
DB_DATASET_ID = os.getenv('DATASET_ID')
DB_NAME = "database.db"

# If /data available => means local storage is enabled => let's use it!
DB_PATH = f"/data/{DB_NAME}" if os.path.isdir("/data") else DB_NAME
print(f"Using {DB_PATH}")

####################################
# Functions
####################################

def create_db_if_missing():
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS model (
            name TEXT UNIQUE,
            upvote INTEGER,
            downvote INTEGER
        );
    ''')
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS vote (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT,
            model TEXT,
            vote INTEGER,
            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    ''')
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS votelog (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT,
            chosen TEXT,
            rejected TEXT,
            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    ''')
    
def get_db():
    return sqlite3.connect(DB_PATH)

####################################
# Space initialization
####################################

# Download existing DB
if not os.path.isfile(DB_PATH):
    print("Downloading DB...")
    try:
        cache_path = hf_hub_download(repo_id=DB_DATASET_ID, repo_type='dataset', filename=DB_NAME)
        shutil.copyfile(cache_path, DB_PATH)
        print("Downloaded DB")
    except Exception as e:
        print("Error while downloading DB:", e)

# Create DB table (if doesn't exist)
create_db_if_missing()
    
# Sync local DB with remote repo every 5 minute (only if a change is detected)
scheduler = CommitScheduler(
    repo_id=DB_DATASET_ID,
    repo_type="dataset",
    folder_path=Path(DB_PATH).parent,
    every=5,
    allow_patterns=DB_NAME,
)

####################################
# Router API
####################################
router = Client("RamAnanth1/3D-Arena-Router", hf_token=os.getenv('HF_TOKEN'))
####################################
# Gradio app
####################################
MUST_BE_LOGGEDIN = "Please login with Hugging Face to participate in the 3D Arena."
DESCR = """
# ⚔️3D Arena: Benchmarking Image-to-3D models

Vote to help the community find the best Image-to-3D model!
""".strip()

INSTR = """
## 🗳️ Vote

* Input image to generate a 3D reconstruction.
* View the responses of the models, one after the other.
* Vote on which model made a better reconstruction.
* _Note: Model names are revealed after the vote is cast._

Note: It may take up to 60 seconds to get a response.
""".strip()
request = ''
if SPACE_ID:
    request = f"""
### Request a model

Please [create a Discussion](https://huggingface.co/spaces/{SPACE_ID}/discussions/new) to request a model.
"""
ABOUT = f"""
## 📄 About

The 3D Arena evaluates leading 3D generation model. It is inspired primarily by [TTS-Arena](https://huggingface.co/spaces/TTS-AGI/TTS-Arena) and LMsys's [Chatbot Arena](https://chat.lmsys.org/).

### The Arena

The leaderboard allows a user to input an image, for which a 3D reconstruction be synthesized by two models. After viewing each sample, the user can vote on which model works better. Due to the risks of human bias and abuse, model names are revealed only after a vote is submitted.

{request}


""".strip()
LDESC = """
## 🏆 Leaderboard

Vote to help the community find the best Image-to-3D model!

The leaderboard displays models in descending order of how suitable the models are (based on votes cast by the community).

Important: In order to help keep results fair, the leaderboard hides results by default until the number of votes passes a threshold. Tick the `Reveal preliminary results` to show models without sufficient votes. Please note that preliminary results may be inaccurate.
""".strip()

def del_db(txt):
    if not txt.lower() == 'delete db':
        raise gr.Error('You did not enter "delete db"')

    # Delete local + remote
    os.remove(DB_PATH)
    delete_file(path_in_repo=DB_NAME, repo_id=DB_DATASET_ID, repo_type='dataset')

    # Recreate
    create_db_if_missing()
    return 'Delete DB'

theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    radius_size=gr.themes.sizes.radius_sm,
    font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],
)
css = ".generating {visibility: hidden}"

model_names = {
    'TripoSR': 'TripoSR',
    'LGM-Mini': 'LGM-Mini',
    'CRM':'CRM'
}
model_licenses = {
    'TripoSR': 'MIT License',
    'LGM-Mini': 'MIT License',
    'CRM': 'MIT License'
    
}
model_links = {
    'TripoSR': 'https://github.com/VAST-AI-Research/TripoSR',
    'LGM-Mini': 'https://github.com/3DTopia/LGM',
    'CRM': 'https://github.com/thu-ml/CRM'
}

def model_license(name):
    print(name)
    for k, v in AVAILABLE_MODELS.items():
        if k == name:
            if v in model_licenses:
                return model_licenses[v]
    print('---')
    return 'Unknown'
def get_leaderboard(reveal_prelim = False):
    conn = get_db()
    cursor = conn.cursor()
    sql = 'SELECT name, upvote, downvote FROM model'
    # if not reveal_prelim: sql += ' WHERE EXISTS (SELECT 1 FROM model WHERE (upvote + downvote) > 750)'
    if not reveal_prelim: sql += ' WHERE (upvote + downvote) > 500'
    cursor.execute(sql)
    data = cursor.fetchall()
    df = pd.DataFrame(data, columns=['name', 'upvote', 'downvote'])
    # df['license'] = df['name'].map(model_license)
    df['name'] = df['name'].replace(model_names)
    df['votes'] = df['upvote'] + df['downvote']
    # df['score'] = round((df['upvote'] / df['votes']) * 100, 2) # Percentage score

    ## ELO SCORE
    df['score'] = 1200
    for i in range(len(df)):
        for j in range(len(df)):
            if i != j:
                expected_a = 1 / (1 + 10 ** ((df['score'][j] - df['score'][i]) / 400))
                expected_b = 1 / (1 + 10 ** ((df['score'][i] - df['score'][j]) / 400))
                actual_a = df['upvote'][i] / df['votes'][i]
                actual_b = df['upvote'][j] / df['votes'][j]
                df.at[i, 'score'] += 32 * (actual_a - expected_a)
                df.at[j, 'score'] += 32 * (actual_b - expected_b)
    df['score'] = round(df['score'])
    ## ELO SCORE
    df = df.sort_values(by='score', ascending=False)
    df['order'] = ['#' + str(i + 1) for i in range(len(df))]
    # df = df[['name', 'score', 'upvote', 'votes']]
    # df = df[['order', 'name', 'score', 'license', 'votes']]
    df = df[['order', 'name', 'score', 'votes']]
    return df

def mkuuid(uid):
    if not uid:
        uid = uuid.uuid4()
    return uid

def upvote_model(model, uname):
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('UPDATE model SET upvote = upvote + 1 WHERE name = ?', (model,))
    if cursor.rowcount == 0:
        cursor.execute('INSERT OR REPLACE INTO model (name, upvote, downvote) VALUES (?, 1, 0)', (model,))
    cursor.execute('INSERT INTO vote (username, model, vote) VALUES (?, ?, ?)', (uname, model, 1,))
    with scheduler.lock:
        conn.commit()
    cursor.close()

def downvote_model(model, uname):
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('UPDATE model SET downvote = downvote + 1 WHERE name = ?', (model,))
    if cursor.rowcount == 0:
        cursor.execute('INSERT OR REPLACE INTO model (name, upvote, downvote) VALUES (?, 0, 1)', (model,))
    cursor.execute('INSERT INTO vote (username, model, vote) VALUES (?, ?, ?)', (uname, model, -1,))
    with scheduler.lock:
        conn.commit()
    cursor.close()

def a_is_better(model1, model2, userid):
    userid = mkuuid(userid)
    if model1 and model2:
        conn = get_db()
        cursor = conn.cursor()
        cursor.execute('INSERT INTO votelog (username, chosen, rejected) VALUES (?, ?, ?)', (str(userid), model1, model2,))
        with scheduler.lock:
            conn.commit()
            cursor.close()
        upvote_model(model1, str(userid))
        downvote_model(model2, str(userid))
    return reload(model1, model2, userid, chose_a=True)
def b_is_better(model1, model2, userid):
    userid = mkuuid(userid)
    if model1 and model2:
        conn = get_db()
        cursor = conn.cursor()
        cursor.execute('INSERT INTO votelog (username, chosen, rejected) VALUES (?, ?, ?)', (str(userid), model2, model1,))
        with scheduler.lock:
            conn.commit()
            cursor.close()
        upvote_model(model2, str(userid))
        downvote_model(model1, str(userid))
    return reload(model1, model2, userid, chose_b=True)
def both_bad(model1, model2, userid):
    userid = mkuuid(userid)
    if model1 and model2:
        downvote_model(model1, str(userid))
        downvote_model(model2, str(userid))
    return reload(model1, model2, userid)
def both_good(model1, model2, userid):
    userid = mkuuid(userid)
    if model1 and model2:
        upvote_model(model1, str(userid))
        upvote_model(model2, str(userid))
    return reload(model1, model2, userid)
def reload(chosenmodel1=None, chosenmodel2=None, userid=None, chose_a=False, chose_b=False):
    out = [
        gr.update(interactive=False, visible=False),
        gr.update(interactive=False, visible=False)
    ]
    if chose_a == True:
        out.append(gr.update(value=f'Your vote: {chosenmodel1}', interactive=False, visible=True))
        out.append(gr.update(value=f'{chosenmodel2}', interactive=False, visible=True))
    else:
        out.append(gr.update(value=f'{chosenmodel1}', interactive=False, visible=True))
        out.append(gr.update(value=f'Your vote: {chosenmodel2}', interactive=False, visible=True))
    out.append(gr.update(visible=True))
    return out

with gr.Blocks() as leaderboard:
    gr.Markdown(LDESC)
    # df = gr.Dataframe(interactive=False, value=get_leaderboard())
    df = gr.Dataframe(interactive=False, min_width=0, wrap=True, column_widths=[30, 200, 50, 50])
    with gr.Row():
        reveal_prelim = gr.Checkbox(label="Reveal preliminary results", info="Show all models, including models with very few human ratings.", scale=1)
        reloadbtn = gr.Button("Refresh", scale=3)
    reveal_prelim.input(get_leaderboard, inputs=[reveal_prelim], outputs=[df])
    leaderboard.load(get_leaderboard, inputs=[reveal_prelim], outputs=[df])
    reloadbtn.click(get_leaderboard, inputs=[reveal_prelim], outputs=[df])
    
def synthandreturn(image):
    file_path = "temp.png"
    image.save(file_path)
    # Get two random models
    mdl1, mdl2 = random.sample(list(AVAILABLE_MODELS.keys()), 2)
    print("[debug] Using", mdl1, mdl2)
    def predict_and_update_result(image, model, result_storage):
        try:
            if model in AVAILABLE_MODELS:
                result = router.predict(file_path, AVAILABLE_MODELS[model].lower(), api_name="/synthesize")
            else:
                result = router.predict(file_path, model.lower(), api_name="/synthesize")
        except:
            raise gr.Error('Unable to call API, please try again :)')
        print('Done with', model)
        result_storage[model] = result
    
    results = {}
    thread1 = threading.Thread(target=predict_and_update_result, args=(image, mdl1, results))
    thread2 = threading.Thread(target=predict_and_update_result, args=(image, mdl2, results))
    thread1.start()
    thread2.start()
    thread1.join()
    thread2.join()
    
    return (
        image,
        "Generate",
        gr.update(visible=True), # r2
        mdl1, # model1
        mdl2, # model2
        gr.update(visible=True, value=results[mdl1]),
        gr.update(visible=True, value=results[mdl2]),
        gr.update(visible=True, interactive=True),
        gr.update(visible=True, interactive=True),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False), #nxt round btn
    )
    
def clear_stuff():
    return "", "Generate", gr.update(visible=False), '', '', gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

with gr.Blocks() as vote:
    useridstate = gr.State()
    gr.Markdown(INSTR)
    with gr.Group():
        with gr.Row():
            image = gr.Image(container=False, type='pil', show_label=False,  min_width=0)
        btn = gr.Button("Generate", variant='primary')
    model1 = gr.Textbox(interactive=False, lines=1, max_lines=1, visible=False)
    model2 = gr.Textbox(interactive=False, lines=1, max_lines=1, visible=False)
    with gr.Row(visible=False) as r2:
        with gr.Column():
            with gr.Group():
                threed1 = gr.Model3D(show_label=False)
                abetter = gr.Button("A is better", variant='primary')
                prevmodel1 = gr.Textbox(interactive=False, show_label=False, container=False, value="Vote to reveal model A", text_align="center", lines=1, max_lines=1, visible=False)
        with gr.Column():
            with gr.Group():
                threed2 = gr.Model3D(show_label=False)
                bbetter = gr.Button("B is better", variant='primary')
                prevmodel2 = gr.Textbox(interactive=False, show_label=False, container=False, value="Vote to reveal model B", text_align="center", lines=1, max_lines=1, visible=False)
    nxtroundbtn = gr.Button('Next round', visible=False)
    outputs = [image, btn, r2, model1, model2, threed1, threed2, abetter, bbetter, prevmodel1, prevmodel2, nxtroundbtn]
    btn.click(synthandreturn, inputs=[image], outputs=outputs)
    nxtroundbtn.click(clear_stuff, outputs=outputs)

    
    nxt_outputs = [abetter, bbetter, prevmodel1, prevmodel2, nxtroundbtn]
    abetter.click(a_is_better, outputs=nxt_outputs, inputs=[model1, model2, useridstate])
    bbetter.click(b_is_better, outputs=nxt_outputs, inputs=[model1, model2, useridstate])

with gr.Blocks() as about:
    gr.Markdown(ABOUT)
with gr.Blocks(theme=theme, css=css, title="3D Arena") as demo:
    gr.Markdown(DESCR)
    gr.TabbedInterface([vote, leaderboard, about], ['🗳️ Vote', '🏆 Leaderboard', '📄 About'])
    
demo.queue(api_open=False, default_concurrency_limit=40).launch(show_api=False)