import requests
import zipfile
import io
import pandas as pd

from pyspark.sql.types import DoubleType, StringType, LongType
from synapse.ml.recommendation import SAR

url = "http://files.grouplens.org/datasets/movielens/ml-25m.zip"
response = requests.get(url)

with zipfile.ZipFile(io.BytesIO(response.content)) as z:
    with z.open('ml-25m/ratings.csv') as csvfile:
        pdf_ratings = pd.read_csv(csvfile)

# 明示的評価を暗黙的評価にするために全部に1.0を代入する
pdf_ratings["rating"] = 1.0

# pandas DataFrameをSpark DataFrameに変換
spark_df_ratings = spark.createDataFrame(pdf_ratings)

# 各列のデータ型を表示して確認
print("Before casting:")
spark_df_ratings.printSchema()

# データ型を明示的に変換
spark_df_ratings = spark_df_ratings.withColumn("userId", spark_df_ratings["userId"].cast(StringType()))
spark_df_ratings = spark_df_ratings.withColumn("movieId", spark_df_ratings["movieId"].cast(StringType()))
spark_df_ratings = spark_df_ratings.withColumn("rating", spark_df_ratings["rating"].cast(DoubleType()))
spark_df_ratings = spark_df_ratings.withColumn("timestamp", spark_df_ratings["timestamp"].cast(LongType()))

# 各列のデータ型を再度表示して確認
print("After casting:")
spark_df_ratings.printSchema()

# SARモデルの設定
sar = SAR(
    userCol="userId",
    itemCol="movieId",
    ratingCol="rating",
    timeCol="timestamp",
    implicitPrefs=True,
    activityTimeFormat="epoch"
)

# モデルのトレーニング
model = sar.fit(spark_df_ratings)