from smv import *
from pyspark.sql.functions import *
from pyspark.sql.types import *

from airlineapp.feature import cmprofilefeat
from airlineapp.feature import cmtravelfeat
from airlineapp.core import conf

__all__ = ['CustomerAllFeature']

class CustomerAllFeatureJoin(SmvPyModule, SmvPyOutput):
    """Combine all features of each customer"""

    def requiresDS(self):
        return [cmprofilefeat.CustomerProfileFeature, cmtravelfeat.CustomerTravelFeature]

    def run(self, i):
        feat_df1 = i[cmprofilefeat.CustomerProfileFeature]
        feat_df2 = i[cmtravelfeat.CustomerTravelFeature]

        return feat_df1.smvJoinByKey(feat_df2, ["CUST_ID"], "leftouter")


class CustomerAllFeatureTreat(SmvPyModule, SmvPyOutput):
    """Treatments of features after join"""

    def requiresDS(self):
        return [CustomerAllFeatureJoin]

    def run(self, i):
        df = i[CustomerAllFeatureJoin]
        feat_df = df.smvSelectPlus(
            avg(col("cmint_age_now")).over(conf.WindowConf.win_allcm).alias("cmint_avg_age_now")
        )

        return feat_df.smvSelectPlus(
            when(col("cmstr_gender_cd") == "F", 1).otherwise(0).alias("cmind_gender_f"),
            when(col("cmstr_gender_cd") == "M", 1).otherwise(0).alias("cmind_gender_m"),
            when(col("cmstr_tier_now") == "0", 1).otherwise(0).alias("cmind_tier0_now"),
            when(col("cmstr_tier_now") == "1", 1).otherwise(0).alias("cmind_tier1_now"),
            when(col("cmstr_tier_now") == "2", 1).otherwise(0).alias("cmind_tier2_now"),
            when(col("cmstr_tier_now") == "3", 1).otherwise(0).alias("cmind_tier3_now")
        ).smvSelectPlus(
            coalesce(col("cmint_age_now"), col("cmint_avg_age_now")).alias("cmint_age_now2")
        ).smvSelectMinus(
            "cmint_age_now",
            "cmint_avg_age_now"
        ).smvRenameField(
            ("cmint_age_now2", "cmint_age_now")
        ).fillna({"cmcnt_seg_tvl_p1y": 0.0, "cmcnt_lh_seg_tvl_p1y": 0.0, "cmsum_rev_tvl_p1y": 0.0})
