# SparkSessionの初期化とUDFに必要なインポート

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, IntegerType, FloatType, BooleanType

# SparkSessionを作成
spark = SparkSession.builder \
    .appName("UDF Example") \
    .getOrCreate()

print("SparkSession準備完了")

# 人物データを作成

# データ：(名前, 年齢, 身長cm, 体重kg)
people_data = [
    ("田中太郎", 25, 170, 65),
    ("鈴木花子", 30, 160, 50),
    ("佐藤次郎", 35, 175, 70),
    ("山田美咲", 18, 165, 55)
]

df_people = spark.createDataFrame(people_data, ["name", "age", "height", "weight"])

print("人物データ:")
df_people.show()

# 商品データを作成

# データ：(商品名, 価格, カテゴリ)
products_data = [
    ("ノートPC", 100000, "電化製品"),
    ("コーヒー", 500, "食品"),
    ("Tシャツ", 3000, "衣料品"),
    ("スマートフォン", 80000, "電化製品")
]

df_products = spark.createDataFrame(products_data, ["product", "price", "category"])
df_products.createOrReplaceTempView("products")

print("商品データ:")
df_products.show()

# ========================================
# STEP 17 サンプルデータ一括作成スクリプト
# ========================================

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, IntegerType, FloatType, BooleanType

spark = SparkSession.builder.appName("UDF Example").getOrCreate()

# 1. 人物データ
people_data = [
    ("田中太郎", 25, 170, 65), ("鈴木花子", 30, 160, 50),
    ("佐藤次郎", 35, 175, 70), ("山田美咲", 18, 165, 55)
]
df_people = spark.createDataFrame(people_data, ["name", "age", "height", "weight"])

# 2. 商品データ
products_data = [
    ("ノートPC", 100000, "電化製品"), ("コーヒー", 500, "食品"),
    ("Tシャツ", 3000, "衣料品"), ("スマートフォン", 80000, "電化製品")
]
df_products = spark.createDataFrame(products_data, ["product", "price", "category"])
df_products.createOrReplaceTempView("products")

print("✅ サンプルデータを作成しました")
print("  - df_people: 人物データ")
print("  - df_products: 商品データ（products ビュー）")

# @udfデコレータでUDFを作成

@udf(returnType=StringType())
def age_category(age):
    """年齢を年齢層に分類するUDF"""
    if age is None:
        return "不明"
    if age < 20:
        return "10代"
    elif age < 30:
        return "20代"
    elif age < 40:
        return "30代"
    else:
        return "40代以上"

# UDFを使用
result = df_people.withColumn("age_group", age_category("age"))

print("【年齢層を追加】")
result.show()

# BMIを計算するUDF

@udf(returnType=FloatType())
def calculate_bmi(height_cm, weight_kg):
    """BMIを計算（体重kg / 身長m^2）"""
    if height_cm is None or weight_kg is None:
        return None
    height_m = height_cm / 100
    return weight_kg / (height_m ** 2)

# UDFを使用
result = df_people.withColumn("bmi", F.round(calculate_bmi("height", "weight"), 1))

print("【BMIを計算】")
result.show()

# 成人かどうか判定するUDF

@udf(returnType=BooleanType())
def is_adult(age):
    """20歳以上かどうか判定"""
    if age is None:
        return None
    return age >= 20

# UDFを使用
result = df_people.withColumn("is_adult", is_adult("age"))

print("【成人判定】")
result.show()

# 消費税を計算するUDFをSQLに登録

def calc_tax(price):
    """消費税（10%）を計算"""
    if price is None:
        return None
    return int(price * 0.1)

# SQLに登録
spark.udf.register("calc_tax", calc_tax, IntegerType())

# SQLで使用
print("【消費税計算（SQL）】")
spark.sql("""
    SELECT 
        product,
        price,
        calc_tax(price) AS tax,
        price + calc_tax(price) AS total
    FROM products
""").show()

# 割引率を計算するUDF

def calculate_discount(price, category):
    """
    カテゴリ別の割引率を適用
    - 電化製品：5%引き
    - 食品：割引なし
    - 衣料品：10%引き
    """
    if price is None:
        return None
    
    discount_rates = {
        "電化製品": 0.05,
        "食品": 0,
        "衣料品": 0.10
    }
    
    rate = discount_rates.get(category, 0)
    return int(price * (1 - rate))

# SQLに登録
spark.udf.register("calc_discount", calculate_discount, IntegerType())

# SQLで使用
print("【カテゴリ別割引（SQL）】")
spark.sql("""
    SELECT 
        product,
        category,
        price AS 定価,
        calc_discount(price, category) AS 割引後価格
    FROM products
""").show()

# 標準関数 vs UDF のパフォーマンス比較

# 【悪い例】UDFで大文字変換
@udf(returnType=StringType())
def to_upper_udf(text):
    return text.upper() if text else None

# UDF版（遅い）
result_udf = df_people.withColumn("name_upper", to_upper_udf("name"))

# 【良い例】標準関数で大文字変換（速い）
result_standard = df_people.withColumn("name_upper", F.upper("name"))

print("UDF版:")
result_udf.select("name", "name_upper").show()

print("標準関数版（こちらを使うべき）:")
result_standard.select("name", "name_upper").show()