from data import (
    FB15KDataset,
    WN18Dataset,
    WN18RRDataset,
    YAGO310Dataset,
)
from metrics.wlcrec import WLCREC


def main() -> None:
    datasets = {
        "YAGO3-10": YAGO310Dataset(split="train"),
        # "WN18": WN18Dataset(split="train"),
        # "WN18RR": WN18RRDataset(split="train"),
        # "FB15K": FB15KDataset(split="train"),
        # "Hetionet": HetionetDataset(split="train"),
        # "OpenBioLink": OpenBioLinkDataset(split="train"),
        # "OpenKEWiki": OpenKEWikiDataset(split="train"),
    }

    for name, dataset in datasets.items():
        wl_crec = WLCREC(dataset)
        entropy, c_ratio, c_nwlec, h_cond, d_ratio, d_nwlec = wl_crec.compute(H=20)
        print(
            f"\nDataset: {name}",
            f"\nResults (H={20})",
            f"\n  • avg WLEC        : {entropy:.6f} nats",
            f"\n  • C_ratio         : {c_ratio:.6f}",
            f"\n  • C_NWLEC         : {c_nwlec:.6f}",
            f"\n  • H_cond(R|S_H)   : {h_cond:.6f} nats",
            f"\n  • D_ratio         : {d_ratio:.6f}",
            f"\n  • D_NWLEC         : {d_nwlec:.6f}",
            sep="",
        )


if __name__ == "__main__":
    main()