import pandas as pd

def transform_problem(raw_df:pd.DataFrame, problem_unique:set=None):
    """
    与えられたDataFrameに対して、problemsを分割して、各問題が含まれているかどうかを特徴量として追加する

    Parameters
    ----------
    raw_df : pd.DataFrame
        problemsを含むDataFrame
    problem_unique : set
        problemsに含まれる問題の集合(testデータを処理する場合には、trainデータの問題の集合を渡す)
    Returns
    -------
    pd.DataFrame
        problemsを分割して、各問題が含まれているかどうかを特徴量として追加したDataFrame
    problem_unique : set
        problemsに含まれる問題の集合(testデータを処理する場合にはNoneを返す)


    Examples
    --------
    >>> import pandas as pd
    >>> train = pd.read_csv("data/official/train.csv", index_col=0)
    >>> test = pd.read_csv("data/official/test.csv", index_col=0)
    >>> processed_train, problem_unique = transform_problem(train)
    >>> processed_test, _ = transform_problem(test, problem_unique)
    """
    def _split_string_by_uppercase(input_string):
        result = []
        current_word = ""
        if not type(input_string)==str:
            return result

        # Escape String "Other"
        for char in input_string:
            if char.isupper() and current_word:
                result.append(current_word)
                current_word = char
            else:
                current_word += char

        if current_word:
            result.append(current_word)

        return result

    df = raw_df.copy()
    df.loc[:,"problems_list"] = df.problems.apply(_split_string_by_uppercase)
    is_train = False
    if problem_unique is None:
        is_train = True
        problem_unique = set()
        for problems in df.loc[:,"problems"].unique():
            for problem in _split_string_by_uppercase(problems):
                problem_unique.add(problem)
    for unique_problem in problem_unique:
        df.loc[:,f"is_problem_{unique_problem}"] = df.loc[:,"problems_list"].apply(lambda x: unique_problem in x)
        df.loc[:,f"is_problem_{unique_problem}"] = df.loc[:,f"is_problem_{unique_problem}"].fillna(False)
    df.loc[:,"problems_count"] = df.loc[:,"problems_list"].apply(len)
    df = df.drop(columns=["problems_list"])
    df = df.drop(columns=["problems"])
    if is_train:
        return df, problem_unique
    else:
        return df, None

parts_list = [
    "Branch",
    "Trunk",
    "Root",
]
for parts in parts_list:
    train.loc[:,f"is_problem_{parts}"] = train.loc[:,"problems"].str.contains(parts)

cause_list = [
        'Grates',
        'Lights',
        'Metal',
        'Other',
        'Rope',
        'Sneakers',
        'Stones',
        'Wires'
        ]
for cause in cause_list:
    train.loc[:,f"is_cause_{cause}"] = train.loc[:,"problems"].str.contains(cause)
train.loc[:,[f"is_cause_{cause}" for cause in cause_list]] = train.loc[:,[f"is_cause_{cause}" for cause in cause_list]].fillna(False)

precision = TP/(TP+FP)
recall = TP/(TP+FN)