import pandas as pd
import os
from crabnet.utils.composition import _fractional_composition


def norm_form(formula):
    comp = _fractional_composition(formula)
    form = ""
    for key, value in comp.items():
        form += f"{key}{str(value)[0:9]}"
    return form


def count_elems(string):
    count = 0
    switch = 1
    for c in string:
        if c.isalpha():
            count += switch
            switch = 0
        if c.isnumeric():
            switch = 1
    return count


# %%
if __name__ == "__main__":
    print("processing all model predictions and calculating metrics")
    print("this will take a few minutes...")

    # %%
    results_path = "publication_predictions"
    benchmark_path = "data/benchmark_data"
    test_directories = os.listdir(results_path)
    benchmark_props = os.listdir(benchmark_path)

    benchmark_test_directories = [
        test for test in test_directories if "benchmark" in test
    ]

    dataset_results = {}
    dataset_preds = {}
    dataset_acts = {}
    test_maes = pd.DataFrame()
    df_stats = pd.DataFrame()
    for benchmark in benchmark_props:
        df_compositions = pd.DataFrame()
        df_preds = pd.DataFrame()
        df_acts = pd.DataFrame()
        models = []
        for directory in benchmark_test_directories:
            df_train_orig = pd.read_csv(
                f"{benchmark_path}/{benchmark}/train.csv",
                keep_default_na=False,
                na_values=[""],
            )
            df_val = pd.read_csv(
                f"{benchmark_path}/{benchmark}/val.csv",
                keep_default_na=False,
                na_values=[""],
            )
            df_train = pd.concat([df_train_orig, df_val], ignore_index=True)
            df_train["formula"] = [
                norm_form(formula) for formula in df_train["formula"]
            ]
            df_train.index = df_train["formula"]
            files = os.listdir(f"{results_path}\{directory}")
            file = [file for file in files if benchmark in file and "test" in file]
            if len(file) > 0:
                models.append(directory.split("_")[0])
                file = file[0]
                df = pd.read_csv(
                    f"{results_path}\{directory}\{file}",
                    keep_default_na=False,
                    na_values=[""],
                )
                composition = df["formula"]
                pred = df["predicted"]
                act = df["actual"]
                print(f"processing {benchmark} {models[-1]}")
                df_compositions = pd.concat([df_compositions, composition], axis=1)
                df_preds = pd.concat([df_preds, pred], axis=1)
                df_acts = pd.concat([df_acts, act], axis=1)

                n_total = act.count() + df_val.shape[0] + df_train_orig.shape[0]
                df_stats.at[benchmark, "mean_test"] = act.mean()
                df_stats.at[benchmark, "std_test"] = act.std()
                df_stats.at[benchmark, "n_test"] = act.count()

                df_stats.at[benchmark, "mean_train"] = df_train["target"].mean()
                df_stats.at[benchmark, "std_train"] = df_train["target"].std()
                df_stats.at[benchmark, "n_train"] = df_train_orig.shape[0]

                df_stats.at[benchmark, "n_val"] = df_val.shape[0]
                df_stats.at[benchmark, "n_total"] = n_total

                df_stats.at[benchmark, "prop_train"] = df_train_orig.shape[0] / n_total
                df_stats.at[benchmark, "prop_val"] = df_val.shape[0] / n_total
                df_stats.at[benchmark, "prop_test"] = act.count() / n_total

        df_compositions.columns = models
        df_preds.columns = models
        df_acts.columns = models
        df_diff = df_preds - df_acts
        df_mae = df_diff.abs().mean()
        test_maes[benchmark] = df_mae

        dataset_results[benchmark] = df_compositions
        dataset_preds[benchmark] = df_preds
        dataset_acts[benchmark] = df_acts

    maes = test_maes.T
    model_names = ["roost", "mat2vec", "onehot", "elemnet", "rf"]
    out_1 = maes[model_names]
    out = pd.concat([out_1, df_stats], axis=1)
    df_benchmark = out.copy()

    # %%
    results_path = "publication_predictions"
    matbench_path = "data/matbench_cv"
    test_directories = os.listdir(results_path)
    matbench_props = os.listdir(matbench_path)

    matbench_test_directories = [
        test for test in test_directories if "matbench" in test
    ]

    dataset_results = {}
    dataset_preds = {}
    dataset_acts = {}
    test_maes = pd.DataFrame()
    df_stats = pd.DataFrame()
    for matbench in matbench_props:
        df_compositions = pd.DataFrame()
        df_preds = pd.DataFrame()
        df_acts = pd.DataFrame()
        models = []
        for directory in matbench_test_directories:
            train_files = os.listdir(f"{matbench_path}/{matbench}")
            train_files = [file for file in train_files if "train" in file]
            test_files = os.listdir(f"{results_path}/{directory}")
            test_files = [
                file for file in test_files if matbench in file and "test" in file
            ]
            for i, (train_file, test_file) in enumerate(zip(train_files, test_files)):
                df_train_orig = pd.read_csv(
                    f"{matbench_path}/{matbench}/{train_file}",
                    keep_default_na=False,
                    na_values=[""],
                )
                df_val = pd.read_csv(
                    f'{matbench_path}/{matbench}/{train_file.replace("train", "val")}',
                    keep_default_na=False,
                    na_values=[""],
                )
                df_train = pd.concat([df_train_orig, df_val], ignore_index=True)
                df_train["formula"] = [
                    norm_form(formula) for formula in df_train["formula"]
                ]
                df_train.index = df_train["formula"]
                if len(file) > 0:
                    models.append(directory.split("_")[0] + f"_{i}")
                    file = file[0]
                    df = pd.read_csv(
                        f"{results_path}\{directory}\{test_file}",
                        keep_default_na=False,
                        na_values=[""],
                    )
                    df.index = df["formula"].values
                    composition = df["formula"]
                    pred = df["predicted"]
                    act = df["actual"]
                    print(f"processing {matbench} {models[-1]}")
                    df_compositions = pd.concat([df_compositions, composition], axis=1)
                    df_preds = pd.concat([df_preds, pred], axis=1)
                    df_acts = pd.concat([df_acts, act], axis=1)

                    n_total = act.count() + df_val.shape[0] + df_train_orig.shape[0]
                    df_stats.at[matbench, "mean_test"] = act.mean()
                    df_stats.at[matbench, "std_test"] = act.std()
                    df_stats.at[matbench, "n_test"] = act.count()

                    df_stats.at[matbench, "mean_train"] = df_train["target"].mean()
                    df_stats.at[matbench, "std_train"] = df_train["target"].std()
                    df_stats.at[matbench, "n_train"] = df_train_orig.shape[0]

                    df_stats.at[matbench, "n_val"] = df_val.shape[0]
                    df_stats.at[matbench, "n_total"] = n_total

                    df_stats.at[matbench, "prop_train"] = (
                        df_train_orig.shape[0] / n_total
                    )
                    df_stats.at[matbench, "prop_val"] = df_val.shape[0] / n_total
                    df_stats.at[matbench, "prop_test"] = act.count() / n_total

        df_compositions.columns = models
        df_preds.columns = models
        df_acts.columns = models
        df_diff = df_preds - df_acts
        df_mae_cv = df_diff.abs()
        df_mae = pd.DataFrame()
        model_names = []
        _ = [model_names.append(x[:-2]) for x in models if x[:-2] not in model_names]
        for j, model in enumerate(model_names):
            df_mae.loc[model, 0] = (
                df_mae_cv.iloc[:, (j) * 5 : (j + 1) * 5].max(axis=1).mean()
            )

        test_maes[matbench] = df_mae[0]
        dataset_results[matbench] = df_compositions
        dataset_preds[matbench] = df_preds
        dataset_acts[matbench] = df_acts

    maes = test_maes.T
    model_names = ["roost", "mat2vec", "onehot", "elemnet", "rf"]
    out_1 = maes[model_names]
    out = pd.concat([out_1, df_stats], axis=1)
    df_matbench = out.copy()

    # %%
    df_all = pd.concat([df_benchmark, df_matbench], axis=0, ignore_index=False)

    rename_dict = {"mat2vec": "crabnet", "onehot": "hotcrab"}
    df_all = df_all.rename(columns=rename_dict)
    df_all.to_csv("metrics/all_metrics.csv", index_label="property")
