import sys
sys.path.insert(0, "..")

import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer

from common.datasets import churn_pelanggan, segmentasi_pelanggan

ADA_SPARK = False
try:
    from pyspark.sql import SparkSession, functions as F
    spark = (SparkSession.builder.master("local[*]").appName("lab14")
             .config("spark.sql.shuffle.partitions", "4").getOrCreate())
    ADA_SPARK = True
    print("Spark", spark.version, "siap")
except Exception as e:
    print("PySpark belum tersedia ->", type(e).__name__)
    print("Install: pip install pyspark (+ Java 17), atau jalankan di Google Colab.")

df = churn_pelanggan()
FITUR = ["lama_langganan", "biaya_bulanan", "jumlah_keluhan", "jam_pemakaian"]
print(df.shape, "| proporsi churn:", df.churn.mean())
df.head()

X = df[FITUR].to_numpy(float)
y = df["churn"].to_numpy(float)
partisi = np.array_split(X, 8)                     # bayangkan 8 executor

stat = [(len(p), p.sum(0), (p ** 2).sum(0)) for p in partisi]      # map
n = sum(s[0] for s in stat); S1 = sum(s[1] for s in stat); S2 = sum(s[2] for s in stat)  # reduce
mean = S1 / n
std = np.sqrt(S2 / n - mean ** 2)

ref = StandardScaler().fit(X)
pd.DataFrame({"mean_terdistribusi": mean, "mean_sklearn": ref.mean_,
              "std_terdistribusi": std, "std_sklearn": ref.scale_}, index=FITUR)

def sigmoid(z):
    return 1 / (1 + np.exp(-np.clip(z, -500, 500)))

Xs = (X - mean) / std
Xb = np.column_stack([np.ones(len(Xs)), Xs])

def latih(Xb, y, n_partisi=1, lr=0.5, epoch=300, catat=False):
    parts = list(zip(np.array_split(Xb, n_partisi), np.array_split(y, n_partisi)))
    w = np.zeros(Xb.shape[1]); riwayat = []
    for _ in range(epoch):
        hasil = [(Xp.T @ (sigmoid(Xp @ w) - yp), len(yp)) for Xp, yp in parts]  # map (di executor)
        g = sum(h[0] for h in hasil); nn = sum(h[1] for h in hasil)            # treeAggregate
        w = w - lr * g / nn                                                     # update di driver
        if catat:
            p = sigmoid(Xb @ w)
            riwayat.append(-np.mean(y * np.log(p + 1e-12) + (1 - y) * np.log(1 - p + 1e-12)))
    return w, riwayat

w1, loss = latih(Xb, y, n_partisi=1, catat=True)
w8, _ = latih(Xb, y, n_partisi=8)
print("w (1 mesin)   :", w1.round(4))
print("w (8 partisi) :", w8.round(4))
print("Selisih maksimum:", np.abs(w1 - w8).max())

plt.plot(loss); plt.xlabel("epoch"); plt.ylabel("log-loss"); plt.title("Konvergensi GD"); plt.show()

proba = sigmoid(Xb @ w8)
print(f"Akurasi: {accuracy_score(y, proba >= 0.5):.3f} | AUC: {roc_auc_score(y, proba):.3f}")
pd.Series(w8[1:], index=FITUR).plot.barh(title="Koefisien (fitur terstandarisasi)"); plt.show()

Xk, _ = make_blobs(n_samples=900, centers=3, cluster_std=1.0, random_state=7)
parts = np.array_split(Xk, 6)
c = Xk[[0, 1, 2]].astype(float)
jejak = [c.copy()]
for it in range(15):
    jumlah = np.zeros_like(c); banyak = np.zeros(len(c))
    for p in parts:                                               # map per partisi
        lab = np.linalg.norm(p[:, None] - c[None], axis=2).argmin(1)
        np.add.at(jumlah, lab, p); banyak += np.bincount(lab, minlength=len(c))
    baru = np.where(banyak[:, None] > 0, jumlah / np.maximum(banyak, 1)[:, None], c)  # driver
    if np.abs(baru - c).max() < 1e-6:
        break
    c = baru; jejak.append(c.copy())
print("Konvergen setelah", it, "iterasi")

jejak = np.array(jejak)
plt.scatter(Xk[:, 0], Xk[:, 1], s=5, alpha=0.3)
for j in range(3):
    plt.plot(jejak[:, j, 0], jejak[:, j, 1], "o-", c="r")
plt.title("Pergerakan centroid (K-Means terdistribusi)"); plt.show()

train, test = train_test_split(df, test_size=0.2, random_state=42, stratify=df.churn)
pre = ColumnTransformer([("num", StandardScaler(), FITUR),
                         ("kat", OneHotEncoder(), ["kontrak"])])
pipe = Pipeline([("pre", pre), ("lr", LogisticRegression(max_iter=200))])
t = time.time(); pipe.fit(train, train.churn)
print(f"Latih {time.time() - t:.3f}s | AUC uji: {roc_auc_score(test.churn, pipe.predict_proba(test)[:, 1]):.3f}")

if ADA_SPARK:
    from pyspark.ml import Pipeline as SparkPipeline
    from pyspark.ml.feature import StringIndexer, OneHotEncoder as SOHE, VectorAssembler, StandardScaler as SScaler
    from pyspark.ml.classification import LogisticRegression as SLR
    from pyspark.ml.evaluation import BinaryClassificationEvaluator

    sdf = spark.createDataFrame(df)
    s_train, s_test = sdf.randomSplit([0.8, 0.2], seed=42)

    stages = [
        StringIndexer(inputCol="kontrak", outputCol="kontrak_idx"),
        SOHE(inputCols=["kontrak_idx"], outputCols=["kontrak_ohe"]),
        VectorAssembler(inputCols=FITUR + ["kontrak_ohe"], outputCol="fitur_mentah"),
        SScaler(inputCol="fitur_mentah", outputCol="features", withMean=True),
        SLR(featuresCol="features", labelCol="churn", maxIter=50, regParam=0.01),
    ]
    model = SparkPipeline(stages=stages).fit(s_train)
    pred = model.transform(s_test)
    pred.select("churn", "probability", "prediction").show(5, truncate=False)
    auc = BinaryClassificationEvaluator(labelCol="churn").evaluate(pred)
    print("AUC Spark:", round(auc, 3))
    print("Koefisien:", model.stages[-1].coefficients)
else:
    print("(dilewati — PySpark belum tersedia)")

if ADA_SPARK:
    from pyspark.ml.tuning import ParamGridBuilder, CrossValidator
    lr_stage = stages[-1]
    grid = ParamGridBuilder().addGrid(lr_stage.regParam, [0.001, 0.01, 0.1]).build()
    cv = CrossValidator(estimator=SparkPipeline(stages=stages), estimatorParamMaps=grid,
                        evaluator=BinaryClassificationEvaluator(labelCol="churn"),
                        numFolds=3, parallelism=2)
    cvm = cv.fit(s_train)
    print("AUC rata-rata per regParam:", [round(a, 4) for a in cvm.avgMetrics])
else:
    print("(dilewati — PySpark belum tersedia)")

seg = segmentasi_pelanggan()
FS = ["recency_hari", "frekuensi", "monetary_ribu"]
if ADA_SPARK:
    from pyspark.ml.clustering import KMeans as SKMeans
    from pyspark.ml.evaluation import ClusteringEvaluator
    ssdf = spark.createDataFrame(seg)
    km_pipe = SparkPipeline(stages=[
        VectorAssembler(inputCols=FS, outputCol="raw"),
        SScaler(inputCol="raw", outputCol="features", withMean=True),
        SKMeans(k=4, seed=1, featuresCol="features"),
    ]).fit(ssdf)
    hasil = km_pipe.transform(ssdf)
    print("Silhouette:", round(ClusteringEvaluator().evaluate(hasil), 3))
    hasil.groupBy("prediction").agg(*[F.round(F.avg(c), 1).alias(c) for c in FS]).orderBy("prediction").show()
else:
    Xs_seg = StandardScaler().fit_transform(seg[FS])
    seg["klaster"] = KMeans(4, n_init=10, random_state=1).fit_predict(Xs_seg)
    print("(versi scikit-learn karena Spark belum tersedia)")
    print(seg.groupby("klaster")[FS].mean().round(1))

# 1.

# 2.

# 3.

if ADA_SPARK:
    spark.stop()