import logging


def format_mesh_terms(df):
    """
    Removes unrequired columns and pivots the mesh terms data into a dictionary.

    Args:
        df (dataframe): mesh terms as returned from retrieve_mesh_terms

    Returns:
        (dict): document_id: list of mesh terms
    """
    logging.info("Formatting mesh terms")
    # remove PRC rows
    df = df.drop(df[df.term == 'PRC'].index, axis=0)

    # remove invalid error rows
    df = df.drop(df[df.doc_id.astype(str).str.contains('ERROR.*ERROR', na=False)].index, axis=0)

    # pivot and remove unrequired columns
    doc_terms = {doc_id: list(grouped.term) for doc_id, grouped in df.groupby("doc_id")}
    return doc_terms