from esm_preprocess import *
from esm_JCQ import reverse_jcq_demand_control_scoring

requested_scales = snakemake.params["scales"]

df_esm = pd.read_csv(snakemake.input[0])
df_esm_preprocessed = preprocess_esm(df_esm)

if not all([scale in QUESTIONNAIRE_IDS for scale in requested_scales]):
    unknown_scales = set(requested_scales) - set(QUESTIONNAIRE_IDS.keys())
    print("The requested questionnaire name should be one of the following:")
    print(QUESTIONNAIRE_IDS.keys())
    raise ValueError("You requested scales not collected: ", unknown_scales)

df_esm_clean = clean_up_esm(df_esm_preprocessed)
df_esm_clean["esm_user_score"] = df_esm_clean["esm_user_answer_numeric"]

for scale in requested_scales:
    questionnaire_id = QUESTIONNAIRE_IDS[scale]
    mask = df_esm_clean["questionnaire_id"] == questionnaire_id
    if scale.startswith("JCQ"):
        df_esm_clean.loc[mask] = reverse_jcq_demand_control_scoring(df_esm_clean.loc[mask])
    #TODO Reverse other questionnaires if needed and/or adapt esm_user_score to original scoring.

df_esm_clean.to_csv(snakemake.output[0], index=False)