from typing import Counter


def aligned_print(ref, hyp, steps):
    """This funcition is to print the result of comparing reference and
    hypothesis sentences in an aligned way.

    Args:
        ref: list of words obtained by splitting reference sentence string
        hyp: list of words obtained by splitting hypothesis sentence string
        steps: list of edit steps with elements "corr", "sub", "ins" or "del".

    Return:
        out_str: aligned reference and hypothesis string with edit steps.
    """

    assert isinstance(ref, list) and isinstance(hyp, list)
    assert isinstance(steps, list)

    if len(steps) == 0:  # in case both ref and hyp are empty
        assert len(ref) == 0 and len(hyp) == 0
        out_str = "REF: \nHYP: \nSTP: \nWER: {:.2f}%\n\n".format(0.0)
        return out_str

    out_str = "REF: "
    for i in range(len(steps)):
        delim = " " if i < len(steps) - 1 else "\n"
        if steps[i] == "sub":
            ref_idx = i - steps[: i].count("ins")
            hyp_idx = i - steps[: i].count("del")
            if len(ref[ref_idx]) < len(hyp[hyp_idx]):
                out_str += (
                    ref[ref_idx] + " " * (len(hyp[hyp_idx]) - len(ref[ref_idx])) + delim
                )
            else:
                out_str += ref[ref_idx] + delim
        elif steps[i] == "ins":
            idx = i - steps[: i].count("del")
            out_str += " " * len(hyp[idx]) + delim
        else:
            assert steps[i] == "del" or steps[i] == "corr"
            idx = i - steps[: i].count("ins")
            out_str += ref[idx] + delim

    out_str += "HYP: "
    for i in range(len(steps)):
        delim = " " if i < len(steps) - 1 else "\n"
        if steps[i] == "sub":
            ref_idx = i - steps[: i].count("ins")
            hyp_idx = i - steps[: i].count("del")
            if len(ref[ref_idx]) > len(hyp[hyp_idx]):
                out_str += (
                    hyp[hyp_idx] + " " * (len(ref[ref_idx]) - len(hyp[hyp_idx])) +
                    delim
                )
            else:
                out_str += hyp[hyp_idx] + delim
        elif steps[i] == "del":
            idx = i - steps[: i].count("ins")
            out_str += " " * len(ref[idx]) + delim
        else:
            assert steps[i] == "ins" or steps[i] == "corr"
            idx = i - steps[: i].count("del")
            out_str += hyp[idx] + delim

    out_str += "STP: "
    for i in range(len(steps)):
        delim = " " if i < len(steps) - 1 else "\n"
        if steps[i] == "sub":
            ref_idx = i - steps[: i].count("ins")
            hyp_idx = i - steps[: i].count("del")
            if len(ref[ref_idx]) > len(hyp[hyp_idx]):
                out_str += "S" + " " * (len(ref[ref_idx]) - 1) + delim
            else:
                out_str += "S" + " " * (len(hyp[hyp_idx]) - 1) + delim
        elif steps[i] == "ins":
            idx = i - steps[: i].count("del")
            out_str += "I" + " " * (len(hyp[idx]) - 1) + delim
        else:
            assert steps[i] == "del" or steps[i] == "corr"
            idx = i - steps[: i].count("ins")
            sym = "D" if steps[i] == "del" else " "
            out_str += sym + " " * (len(ref[idx]) - 1) + delim

    counter = Counter(steps)
    wer = (
        float(counter["sub"] + counter["ins"] + counter["del"]) / len(ref) * 100
        if len(ref) > 0
        else 0.0
    )
    out_str += "WER: " + "{:.2f}%".format(wer) + "\n"
    out_str += "\n"

    return out_str