import re
import gradio as gr

from chessfenbot.chessboard_finder import findGrayscaleTilesInImage
from chessfenbot.tensorflow_chessbot import ChessboardPredictor
from chessfenbot.helper_functions import shortenFEN


def predict(img, active="w"):
    """
    main predict function for gradio.
    Predict a chessboard FEN.
    Wraps model from https://github.com/Elucidation/tensorflow_chessbot/tree/chessfenbot

    Args:
        img (PIL image): input image of a chess board
        active (str): defaults to "w"
    """

    # Look for chessboard in image, get corners and split chessboard into tiles
    tiles, corners = findGrayscaleTilesInImage(img)

    # Initialize predictor, takes a while, but only needed once
    predictor = ChessboardPredictor(frozen_graph_path='chessfenbot/saved_models/frozen_graph.pb')
    fen, tile_certainties = predictor.getPrediction(tiles)
    predictor.close()
    short_fen = shortenFEN(fen)
    # Use the worst case certainty as our final uncertainty score
    certainty = tile_certainties.min()

    print('Per-tile certainty:')
    print(tile_certainties)
    print("Certainty range [%g - %g], Avg: %g" % (
    tile_certainties.min(), tile_certainties.max(), tile_certainties.mean()))

    # predicted FEN
    fen_out = f"{short_fen} {active} - - 0 1"
    # certainty
    certainty = "%.1f%%" % (certainty*100)
    # link to analysis board on Lichess
    lichess_link = f'https://lichess.org/analysis/standard/{re.sub(" ", "_", fen_out)}'
     
    return fen_out, certainty, lichess_link


gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Upload chess board", type="pil"),
    outputs=[
        gr.Textbox(label="FEN"), 
        gr.Textbox(label="certainty"), 
        gr.Textbox(label="Link to Lichess analysis board (copy and paste into URL)"), 
        ],
    title="Chess FEN bot",
    examples=["chessfenbot/example_input.png"],
    description="Simple wrapper around TensorFlow Chessbot (https://github.com/Elucidation/tensorflow_chessbot)"
).launch()