import pandas as pd


class Selector:
    def select(self, codes: [str], df: pd.DataFrame) -> [str]:
        return codes


class Strategy:
    def __init__(self, selectors: [Selector]):
        self.selectors = selectors

    def select(self, codes: [str], df: pd.DataFrame) -> [str]:
        return list(map(lambda code: self.selectors.select(code, df), codes))


class PeriodSelector(Selector):
    def __init__(self, period: int = 5):
        self.__period = period

    def select(self, codes: [str], df: pd.DataFrame) -> [str]:
        size_df = df.groupby("code").size()
        return list(filter(lambda code: size_df[code] > self.__period, codes))


class PyramidSelector(Selector):
    def select(self, codes: [str], df: pd.DataFrame) -> [str]:
        target_df = df[df["code"].isin(codes)]
        target_df["score"] = 0
        group_df = target_df.groupby("code")

        target_df["prev_total_stockholder_interest"] = group_df["total_stockholder_interest"].shift(1)
        target_df["roe"] = target_df["net_income"] / ((target_df["prev_total_stockholder_interest"] + target_df["total_stockholder_interest"]) / 2)
        target_df["average_roe"] = target_df["roe"].mean()
        target_df[target_df["average_roe"] >= 35] = target_df["score"] + 550
        target_df[(target_df["average_roe"] < 35) & (target_df["average_roe"] >= 30)] = target_df["score"] + 500
        target_df[(target_df["average_roe"] < 30) & (target_df["average_roe"] >= 25)] = target_df["score"] + 450
        target_df[(target_df["average_roe"] < 25) & (target_df["average_roe"] >= 15)] = target_df["score"] + 300
        target_df[(target_df["average_roe"] < 15) & (target_df["average_roe"] >= 10)] = target_df["score"] + 250

        target_df["prev_total_assets"] = group_df["total_assets"].shift(1)
        target_df["roa"] = target_df["net_income"] / ((target_df["prev_total_assets"] + target_df["total_assets"]) / 2)

        return super().select(codes, df)