from pathfinder_setup_fns import *

def get_df_query(query, top_k = 10, date_range = None, ec=ec):

    input_keywords = ['']
    query_keywords = get_keywords(query)
    ec.query_input_keywords = input_keywords+query_keywords
    ec.toggles = ['Keywords']
    ec.hyde = False
    ec.rerank = False
    ec.date_range = date_range
    rs, small_df = ec.retrieve(query, top_k = top_k, return_scores=True)
    formatted_df = ec.return_formatted_df(rs, small_df)
    qembd = np.array(ec.query_embedding)
    return formatted_df, qembd


sfh_query = 'What can the star formation histories of distant galaxies obtained from SED fitting tell us about their evolution?'

start_date = datetime.date(1996, 1, 1)
# end_date = datetime.date(2012, 12, 31)
end_date = datetime.datetime.now().date()
date_range = (start_date, end_date)

sfh_df, qembd = get_df_query(sfh_query, date_range=date_range, ec=ec)

print(qembd.shape)
sfh_df[['title','date']]