import duckdb
import numpy as np
import json
from pyserini.analysis import Analyzer, get_lucene_analyzer

conn = duckdb.connect("../duck.db")
# analyzer = Analyzer(get_lucene_analyzer())
# corpus_path = 'indexes/nfcorpus.bge-base-en-v1.5/corpus_embeddings.jsonl'
# query_path = 'indexes/nfcorpus.bge-base-en-v1.5/query_embeddings.jsonl'

# def load_jsonl_to_table(file_path, table_name):
#     conn.execute(f"drop table if exists {table_name}")
#     conn.execute(f"""create table {table_name} (id varchar primary key, contents varchar, embedding double[{768}])""")
#     with open(file_path, 'r') as file:
#         for line in file:
#             row = json.loads(line.strip())
#             conn.execute(f"""insert into {table_name} (id, contents, embedding) values (?, ?, ?)""", (row['id'], ' '.join(analyzer.analyze(row['contents'])), row['vector']))
# load_jsonl_to_table(corpus_path, "corpus")
# load_jsonl_to_table(query_path, "query")

# conn.execute("PRAGMA create_fts_index(corpus, id, contents, stemmer = 'none', stopwords = 'none', ignore = 'a^', strip_accents = 0, lower = 0, overwrite = 1)")

def rrf(query_id, query_string, k=60, top_n=5):
    #         ORDER BY array_cosine_similarity(corpus.embedding, query_embedding.embedding) desc
    # ORDER BY COALESCE(fts_main_corpus.match_bm25(id, ?), 0) DESC
    query = f"""
    WITH 
    embd AS (
        SELECT corpus.id,
            ROW_NUMBER() OVER (ORDER BY array_cosine_similarity(corpus.embedding, q.embedding) DESC) AS sim_rank
        FROM corpus, (SELECT embedding FROM query WHERE id = ?) AS q
        limit 1000
    ),
    fts AS (
        SELECT id, 
            ROW_NUMBER() OVER (ORDER BY COALESCE(fts_main_corpus.match_bm25(id, ?), 0) DESC) AS fts_rank
        FROM corpus
        limit 1000
    ),
    combined_results AS (
        SELECT 
            COALESCE(s.id, f.id) AS id,
            COALESCE(1.0 / ({k} + s.sim_rank), 0) + COALESCE(1.0 / ({k} + f.fts_rank), 0) AS rrf_score
        FROM embd s
        FULL OUTER JOIN fts f ON s.id = f.id
    )
    SELECT id, rrf_score
    FROM combined_results
    ORDER BY rrf_score DESC
    LIMIT ?
    """
    return conn.execute(query, [query_id, query_string, top_n]).fetchall()

from tqdm import tqdm
queries = conn.execute("SELECT id, contents, embedding FROM query").fetchall()
run_tag = "rrf_duckdb"

all_results = []

for query_id, query_string, query_embedding in tqdm(queries, desc=f"Processing {run_tag}", unit="query"):
    results = rrf(query_id, query_string, top_n=1000)
    for rank, (doc_id, score) in enumerate(results, 1):
        all_results.append((query_id, doc_id, score, rank))

all_results.sort(key=lambda x: (x[0], x[3])) # sort by queryid, then rank

with open("runs/duckdb_rrf_nfcorpus.txt", "w") as f:
    for query_id, doc_id, score, rank in all_results:
        a = f.write(f"{query_id} Q0 {doc_id} {rank} {score} {run_tag}\n")

# python -m pyserini.eval.trec_eval -c -m ndcg_cut.10 collections/nfcorpus/qrels/test.qrels runs/duckdb_rrf_nfcorpus.txt
