dimensions = ["A", "B-S", "B-O", "C-S", "C-O", "D"]
dimension_to_layer_name_raw = { # in raw annotation files (Ayal's outputs)
    "A": "Segment Patient Affect (A)",
    "B-S": "Segment Patient B-S",
    "B-O": "Segment Patient B-O",
    "C-S": "Segment Patient C-S",
    "C-O": "Segment Patient C-O",
    "D": "Segment Patient Desire (D)",
}
dimension_to_layer_name_cleaned = { # in cleaned annotation files (Yael's outputs)
    "A": "Segment_Patient_A",
    "B-S": "Segment_Patient_B-S",
    "B-O": "Segment_Patient_B-O",
    "C-S": "Segment_Patient_C-S",
    "C-O": "Segment_Patient_C-O",
    "D": "Segment_Patient_D",
}

        
def df_to_self_states_json(df, doc_name, annotator = None):
    """Convert a dataframe into a json object that can be more easily used for visualization."""
    # df is the dataframe of annotations
    # doc_name is the name of the document
    # annotator is the name of the annotator (optional)
    # select dimension_to_layer_name
    if 'Segment_Patient_A' in df.layer.unique():
        dimension_to_layer_name = dimension_to_layer_name_cleaned
    else:
        dimension_to_layer_name = dimension_to_layer_name_raw
    
    def get_evidence_obj(evidence_df):
        "Assume that the evidence_df is a partial dataframe including only annotation of a single evidence span."
        evidence_obj = {k: v.value.iloc[0] for k, v in evidence_df.groupby("feature")}
        # evidence_obj["text"] = evidence_df.span_text.iloc[0]
        # evidence_obj["span_index_begin"] = evidence_df.begin.iloc[0]
        # evidence_obj["span_index_end"] = evidence_df.end.iloc[0]
        return evidence_obj

    doc_object = {"document": doc_name, "annotator": annotator}
    doc_object["segments"] = []
    doc_df = df[df.document == doc_name]
    if annotator:
        doc_df = doc_df[doc_df.annotator == annotator]
        
    # now add the segments
    for segment_index, segment_group in df[df.document == doc_name].groupby("segment"):
        # add Segment Summary features into segment object
        segment_object = {"segment": segment_index}
        segment_summary_df = segment_group[segment_group.layer == "Segment Summary"]
        # # if not post-summary, skip this post
        # if segment_summary_df.empty:
        #     continue
        segment_object["Segment Summary"] = {k: v.value.iloc[0] for k, v in segment_summary_df.groupby("feature")}
        
        state1_df = segment_group[segment_group.self_state_index == 1]
        state2_df = segment_group[segment_group.self_state_index == 2]
        states_list = list()
        state1_obj = dict()
        state2_obj = dict()
        # set is_adaptive for each state
        if not state1_df.empty:
            state1_obj["is_adaptive"] = state1_df.is_adaptive.dropna().iloc[0]
            states_list.append(state1_obj)
        if not state2_df.empty:
            state2_obj["is_adaptive"] = state2_df.is_adaptive.dropna().iloc[0]
            states_list.append(state2_obj)     
        # collect elements per dimension
        for dimension in dimensions:
            segment_dim_layer_name = dimension_to_layer_name[dimension]
            state1_dimension_df = state1_df[state1_df.layer == segment_dim_layer_name]
            state2_dimension_df = state2_df[state2_df.layer == segment_dim_layer_name]
            # search for evidence layers that match the same segment and dimension
            dim_evidence_rows = segment_group[segment_group.layer == f"Patient_{dimension}_evidence"]
            if not state1_dimension_df.empty:
                state1_obj[dimension] = {k: v.value.iloc[0] for k, v in state1_dimension_df.groupby("feature")} # "Category", "Adaptivity", "Presence"
                evidences_obj = []
                # for _, evidence_df in dim_evidence_rows.groupby("span_text"):
                for _, evidence_df in dim_evidence_rows.groupby(["begin", "end"]):
                    # take only the evidence that matches the category of the segment-level element of the same dimension
                    if not evidence_df.empty and evidence_df[evidence_df.feature == "Category"].value.iloc[0] == state1_obj[dimension]["Category"]:
                        evidences_obj.append(get_evidence_obj(evidence_df))
                if evidences_obj:
                    state1_obj[dimension]["evidences"] = evidences_obj
            if not state2_dimension_df.empty:
                state2_obj[dimension] = {k: v.value.iloc[0] for k, v in state2_dimension_df.groupby("feature")} # "Category", "Adaptivity", "Presence"
                evidences_obj = []
                # for _, evidence_df in dim_evidence_rows.groupby("span_text"):
                for _, evidence_df in dim_evidence_rows.groupby(["begin", "end"]):
                    # take only the evidence that matches the category of the segment-level element of the same dimension
                    if not evidence_df.empty and evidence_df[evidence_df.feature == "Category"].value.iloc[0] == state2_obj[dimension]["Category"]:
                        evidences_obj.append(get_evidence_obj(evidence_df))
                if evidences_obj:
                    state2_obj[dimension]["evidences"] = evidences_obj
        
        segment_object["self-states"] = states_list
        
        # add the post object to the document object
        doc_object["segments"].append(segment_object)
    return doc_object

element_short_desc_map = {
    'A:(11) Proud': 'A:(11) Proud',
    'B-O:(1) Relating behavior': 'B-O:(1) Relating',
    'C-S:(1) Self-acceptance and compassion': 'C-S:(1) Self-compassion',
    'D:(1) Relatedness': 'D:(1) Relatedness',
    'A:(4) Depressed, despair, hopeless': 'A:(4) Depressed',
    'C-O:(4) Perception of the other as blocking autonomy needs': 'C-O:(4) Other blocks autonomy',
    'C-S:(2) Self criticism': 'C-S:(2) Self-criticism',
    'C-O:(2) Perception of the other as detached or over attached': 'C-O:(2) Other detached/overattached',
    'C-O:(1) Perception of the other as related': 'C-O:(1) Other related',
    'A:(3) Sad, emotional pain, grieving': 'A:(3) Sadness',
    'B-O:(2) Fight or flight behavior': 'B-O:(2) Fight/flight',
    'A:(14) Feel lonely': 'A:(14) Lonely',
    'D:(2) Expectation that relatedness needs will not be met': 'D:(2) Relatedness (-)',
    'B-S:(2) Self harm, neglect and avoidance': 'B-S:(2) Self-harm',
    'A:(10) Angry (aggression), disgust, contempt': 'A:(10) Angry/Aggressive',
    'A:(8) Apathic, don’t care, blunted': 'A:(8) Apathetic',
    'B-S:(1) Self care and improvement': 'B-S:(1) Self-care',
    'D:(5) Competence, self esteem, self-care': 'D:(5) Competence',
    'D:(6) Expectation that competence needs will not be met': 'D:(6) Competence (-)',
    'C-O:(3) Perception of the other as facilitating autonomy needs': 'C-O:(3) Other supports autonomy',
    'A:(2) Anxious/ fearful/ tense': 'A:(2) Anxious',
    'A:(12) Ashamed, guilty': 'A:(12) Ashamed/Guilty',
    'B-O:(4) Over controlled or controlling behavior': 'B-O:(4) Controlling',
    'A:(1) Calm/ laid back': 'A:(1) Calm',
    'D:(4) Expectation that autonomy needs will not be met': 'D:(4) Autonomy (-)',
    'D:(3) Autonomy and adaptive control': 'D:(3) Autonomy',
    'A:(5) Content, happy, joy, hopeful': 'A:(5) Happy',
    'B-O:(3) Autonomous or adaptive control behavior': 'B-O:(3) Adaptive control',
    'A:(9) Justifiable anger/ assertive anger, justifiable outrage': 'A:(9) Justified anger',
    'A:(13) Feel loved, belong': 'A:(13) Loved/Belonging',
    'A:(7) Vigor / energetic': 'A:(7) Vigor'
}