import re

import pandas as pd


def postprocess(output: str) -> dict:
    """
    Convert the model's output string to a dictionary.
    """
    output = output.replace("<unk>", "#").removeprefix("<s>")
    pattern = (
        r"(<LYSDAT>(?P<lysdat>.*?))?"
        r"(<VIGDAT>(?P<vigdat>.*?))?"
        r"((?P<kon1><MAN>|<KVINNA>)(?P<person1>.*?))"
        r"((?P<kon2><MAN>|<KVINNA>)(?P<person2>.*?))"
        r"(?=(<LYSDAT>|<VIGDAT>|</s>))"       
    )
    
    people = []
    rad = 0
    for m in re.finditer(pattern, output):
        res = m.groupdict()
        lysdat = res["lysdat"].strip() if res["lysdat"] else ""
        vigdat = res["vigdat"].strip() if res["vigdat"] else ""

        for i in range(2):
            person = f"person{i + 1}"
            kon = f"kon{i + 1}"
            p = {}
            for n in re.finditer(r"<(?P<key>.*?)>(?P<value>.*?)(?=<|$)", res[person]):
                key, value = n.groups()
                p["KON"] = res[kon].strip("<>").title()
                if key in ("OGIFT", "TROLOVAD", "ÄNKA_ÄNKLING", "SKILD"):
                    p["CIV"] = key
                else:
                    p[key] = value.strip()
            people.append({"LYSDAT": lysdat, "VIGDAT": vigdat, "RAD": rad} | p)
        rad += 1
    return pd.DataFrame(people)