from fastapi import FastAPI
from src.do_pca_on_tfidf import query_docs
from src.search_embeddings import sbert_query_docs
import polars as pl


app = FastAPI()

@app.get("/")
def default():

    return {"status": "ok", "version": 0.1}


@app.get("/sbert")
def sb(query: str):
    res_sbert = sbert_query_docs(query)

    return {"scores": str(res_sbert)}


@app.get("/tfidf")
def tf(query: str):
    res_tfidf = query_docs(query)

    return {"scores": str(res_tfidf)}


@app.get("/square")
def square(x: int):
    return {"result": x * x}

@app.get("/search")
def greet_json(query: str):
    res_tfidf = query_docs(query)
    res_sbert = sbert_query_docs(query)

    joined = res_sbert.join(res_tfidf, on='file', how = 'inner')

    res_combined = joined.with_columns((0.7 * pl.col("rank-sbert") + 0.3 * pl.col("rank-tfidf")).alias("rank-combined")).sort("rank-combined")

    


    return {"result": res_combined.to_pandas().to_html(), "shape": res_combined.shape}




@app.get("/test")
def echo(query: str):
    
    return {"echo": query}