from typing import Dict, Tuple, List
import pandas as pd
from investigation.database.house import scan_cases
from investigation.database.dataframe import format_and_Q_eq_in


def filter_entry(
    commit_hash: str = "",
    batch_hash: List[str] = [],
    **kwargs,
):
    job_dict = scan_cases(commit_hash, batch_hash)

    # if (len(job_dict) == 0 or len(kwargs) == 0):
    #     return job_dict

    data_df = pd.DataFrame.from_dict(job_dict, orient='index')
    data_df.index.name = "job_name"
    data_df.reset_index(inplace=True)
    query_df = data_df.fillna("NaN")

    query_str, local_dict = format_and_Q_eq_in(kwargs)
    if (query_str == ""):
        data_df.set_index("job_name", inplace=True)
        data_df.index.name = "job_name"
        return data_df, job_dict

    # data_df = data_df.replace({np.nan: None})  # Replace NaN with None for query compatibility
    # input(query_df)
    filtered_df = query_df.query(query_str, local_dict=local_dict)

    filtered_df = data_df.loc[filtered_df.index]
    filtered_df.set_index("job_name", inplace=True)
    filtered_df.index.name = "job_name"

    job_dict = {k: v for k, v in job_dict.items() if k in filtered_df.index}
    return filtered_df, job_dict


def extrac_unique_config(
    commit_hash: str = "",
    batch_hash: str = "",
    **kwargs,
):
    cases_df, cases = filter_entry(
        commit_hash,
        batch_hash,
        **kwargs,
    )
    non_unique_cols = [col for col in cases_df.columns if len(cases_df[col].dropna().unique()) > 1]
    non_unique_cols = list(
        set(non_unique_cols) - set(['num_cpus', 'num_gpus', 'time', 'requeue', "compile_model", "batch_hash"]))
    return cases_df[non_unique_cols], {k: v for k, v in cases.items() if k in cases_df.index}
