import psycopg2
import json
from pyserini.analysis import Analyzer, get_lucene_analyzer

analyzer = Analyzer(get_lucene_analyzer())
conn = psycopg2.connect(dbname='beir_datasets', user='postgres')
cur = conn.cursor()

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):
    cur.execute(f"drop table if exists {table_name};")
    cur.execute(f"""create table {table_name} (id text, contents text, embedding vector({768}));""")
    with open(file_path, 'r') as file:
        for line in file:
            row = json.loads(line.strip())
            cur.execute(f"""insert into {table_name} (id, contents, embedding) values (%s, %s, %s)""", (row['id'], ' <-> '.join(analyzer.analyze(row['contents'])), row['vector']))
    conn.commit()
load_jsonl_to_table(corpus_path, "corpus")
load_jsonl_to_table(query_path, "query")
cur.execute('''CREATE INDEX "corpus_contents_gin" ON "corpus" USING gin(to_tsvector('simple', contents))''')

import re
def _clean_query(query_string):
    cleaned = re.sub(r'[^\w\s]', ' ', query_string)
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned

def rrf(query_string, query_embedding, k=60):
    cleaned_query = _clean_query(query_string)
    ts_query = " | ".join(cleaned_query.split())
    sql = """
    WITH semantic_search AS (
        SELECT id, RANK () OVER (ORDER BY embedding <=> %(embedding)s) AS rank
        FROM corpus
        LIMIT 1000
    ),
    keyword_search AS (
        SELECT id, RANK () OVER (ORDER BY ts_rank(to_tsvector('simple', contents), query) DESC) as rank
        FROM corpus, to_tsquery('simple', %(query)s) query
        WHERE to_tsvector('simple', contents) @@ query
        LIMIT 1000
    )
    SELECT
        COALESCE(semantic_search.id, keyword_search.id) AS id,
        COALESCE(1.0 / (%(k)s + semantic_search.rank), 0.0) +
        COALESCE(1.0 / (%(k)s + keyword_search.rank), 0.0) AS score
    FROM semantic_search
    FULL OUTER JOIN keyword_search ON semantic_search.id = keyword_search.id
    ORDER BY score DESC
    LIMIT 1000
    """
    cur.execute(sql, {'query': ts_query, 'embedding': query_embedding, 'k': k})
    results = cur.fetchall()
    return results

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

all_results = []

for query_id, query_contents, query_embedding in tqdm(queries, desc=f"Processing {run_tag}", unit="query"):
  results = rrf(query_contents, query_embedding)
  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]))

with open("runs/psql_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/psql_rrf_nfcorpus.txt
