import gzip
import json


def read_selected_torsions(input_json):
    """ Read data generated by select_torsions.py

    Returns
    -------
    selected_torsions: dict
        Dictionary for selected torsions, has this structure:
        {
            canonical_torsion_index1: {
                'initial_molecules': [ Molecule1a, Molecule1b, .. ],
                'atom_indices': [ (0,1,2,3) ],
                'attributes': {'canonical_explicit_hydrogen_smiles': .., 'canonical_isomeric_smiles': .., ..}
            },
            ..
        }
    """
    if input_json.endswith(".gz"):
        with gzip.open(input_json, 'r') as infile:
            selected_torsions = json.load(infile)
    else:
        with open(input_json) as infile:
            selected_torsions = json.load(infile)
    return selected_torsions