import re

"""
These helper functions can convert an arbitrary graph of otherwise json-compatible nested dicts and lists
into a serializable tree, by replacing cyclic references with {"#": <index>} dicts.
"""


def _deflate(obj, _sofar: dict[int, dict], _cur_ref: list[int]):
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, list):
        return [_deflate(v, _sofar, _cur_ref) for v in obj]
    if not isinstance(obj, dict):
        raise TypeError(f"Unexpected type: {type(obj)}")

    if id(obj) not in _sofar:
        _sofar[id(obj)] = {}
    else:
        flat = _sofar[id(obj)]
        if "#" not in flat:
            _cur_ref[0] += 1
            flat["#"] = _cur_ref[0]
        return {"#": flat["#"]}
    
    cleaned = {('#' + k if re.fullmatch('#+', k) else k): _deflate(v, _sofar, _cur_ref) for k, v in obj.items()}
    _sofar[id(obj)].update(cleaned)
    return _sofar[id(obj)]


def deflate(obj):
    """
    Given a json-compatible nested set of dicts and lists, replace self-references with {"#ref": <index>}
    """
    return _deflate(obj, {}, [0])


def _inflate(obj, _refs: dict):
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, list):
        return [_inflate(v, _refs) for v in obj]
    if not isinstance(obj, dict):
        raise TypeError(f"Unexpected type: {type(obj)}")
    
    ref_id = None
    if "#" in obj:
        ref_id = obj["#"]
        if ref_id not in _refs:
            _refs[ref_id] = {}

    cleaned = {(k[1:] if re.fullmatch("#+", k) else k): _inflate(v, _refs) for k, v in obj.items() if k != "#"}

    if ref_id is None:
        return cleaned

    _refs[ref_id].update(cleaned)
    return _refs[ref_id]


def inflate(obj):
    """
    Inverse operation as deflate: convert {"#ref": <index>} to self-references.
    """
    return _inflate(obj, {})


def test():
    import json
    import collections

    a = {"#": "collide", "##": "collide"}
    a["b"] = a
    a["c"] = [a, a]
    a["d"] = {"e": a, "f": [a, a]}
    a["g"] = {"h": {"i": a, "j": [a, a]}}
    a2 = inflate(json.loads(json.dumps(deflate(a), indent=2)))
    assert repr(a) == repr(a2)


if __name__ == "__main__":
    try:
        test()
    except:
        import traceback
        traceback.print_exc()
        import pdb
        pdb.post_mortem()