import os
import pandas as pd
from pandas.testing import assert_frame_equal


REFERENCE_FOLDER = "/g/scb2/bork/ustick/nfixplanet_pipeline/01_annotation_pipeline/data/nfixplanet_test_full/QC_sequences"
GENERATED_FOLDER = "tmp_full"

def get_names() -> list[str]:
    names = []
    for f in os.listdir(GENERATED_FOLDER):
        if f.endswith(".tsv"):
            name = f.split(".")[0]
            names.append(name)
    return names


def compare_dfs(name: str):
    ref_file = f"{REFERENCE_FOLDER}/04_{name}_final.out"
    gen_file = f"{GENERATED_FOLDER}/{name}.tsv"

    ref_df = pd.read_csv(
        ref_file,
        sep="\t",
        header=None,
        names=[
            "query_name",
            "accession",
            "target_name",
            "full_accession",
            "full_evalue",
            "full_score",
            "full_bias",
            "domain_evalue",
            "domain_score",
            "domain_bias",
            "exp",
            "reg",
            "clu",
            "ov",
            "env",
            "dom",
            "rep",
            "inc",
        ],
    )
    ref_df = ref_df.sort_values(by=ref_df.columns[0]).reset_index(drop=True)
    gen_df = pd.read_csv(gen_file, sep="\t")
    gen_df = gen_df.sort_values(by=gen_df.columns[0]).reset_index(drop=True)

    # print(ref_df)
    # print(gen_df)
    try:
        assert_frame_equal(ref_df, gen_df, check_dtype=False)
    except Exception as e:
        print(f"failed: {name}")

        # # NOTE Rounding to ignore differences for now
        # # TODO: we have more results now, check if we are correct instead of the old results
        # float_cols = gen_df.select_dtypes('float').columns
        # gen_df[float_cols] = gen_df[float_cols].round(0).astype(int)
        # ref_df[float_cols] = ref_df[float_cols].round(0).astype(int)


        # gen_df.to_csv("gen_df.tsv", sep="\t", index=False)
        # ref_df.to_csv("ref_df.tsv", sep="\t", index=False)
        # # TODO: check 153;SAMN17048824;k141_265651_14 in hmm
        
        # raise e

print(f"Running on folder: {GENERATED_FOLDER}")
names = get_names()
for name in names:
    print(f"Checking {name}")
    compare_dfs(name)
