# -*- coding: utf-8 -*-
"""PS13.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/141X72rA2VxgIbA93KoMkwHzHSc1fsbty

# Agglomerative hierarchical clustering

### Setup
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotly.figure_factory as ff
import plotly.graph_objects as go
import seaborn as sns
from scipy.cluster.hierarchy import dendrogram, linkage
from sklearn.cluster import AgglomerativeClustering
from sklearn.preprocessing import LabelEncoder

csv_path = "/content/drive/MyDrive/Academic_Resources/Semester 08/20XW87/PS13/mall_customers_dataset.csv"

df = pd.read_csv(csv_path)
display(df)

"""### Remove CustomerID column"""

df = df.drop(columns=["CustomerID"])
display(df)

"""### Check for missing values"""

missing_counts = df.isna().sum()
print(missing_counts)

missing_counts = df.isna().sum()
missing_percent = (missing_counts / len(df)) * 100
print(missing_percent)

print("Any missing values?", df.isna().values.any())

for col in df.columns:
    if df[col].isna().sum() == 0:
        print(f"Column '{col}' has no missing values.")
    else:
        print(f"Column '{col}' has {df[col].isna().sum()} missing values.")

"""### Convert categorical variable value of gender to numerical (Male-1, Female-0)"""

le = LabelEncoder()
df["gender_encoded"] = le.fit_transform(df["Gender"])

print("Classes:", le.classes_)

display(df)

"""### Display male, female ratio as pie chart"""

gender_counts = df["gender_encoded"].value_counts()
labels = ["Female", "Male"]

plt.figure(figsize=(5, 5))
plt.pie(
    gender_counts,
    labels=labels,
    autopct="%1.1f%%",
    startangle=90,
    colors=["#66b3ff", "#ff9999"],
)
plt.title("Gender Distribution")
plt.show()

"""### Display age, annual income as bar graph"""

values = [df["Age"].mean(), df["Annual Income (k$)"].mean()]
labels = ["Age", "Annual Income (k$)"]

plt.bar(labels, values, color=["skyblue", "orange"])
plt.title("Age and Annual Income")
plt.ylabel("Value")
plt.show()

df["Age_group"] = pd.cut(df["Age"], bins=[18, 25, 35, 45, 55, 65], right=False)

plt.figure(figsize=(8, 5))
sns.barplot(x="Age_group", y="Annual Income (k$)", data=df)

plt.title("Average Income by Age Group")
plt.xticks(rotation=45)
plt.show()

sns.barplot(x="Age", y="Annual Income (k$)", data=df, estimator="mean")

import plotly.express as px

fig = px.bar(
    df.groupby("Age")["Annual Income (k$)"].mean().reset_index(),
    x="Age",
    y="Annual Income (k$)",
    title="Mean Income by Age",
)
fig.show()

"""### Perform agglomerative clustering with ward algorithm as linkage."""

agg_clust = AgglomerativeClustering(n_clusters=2, linkage="ward")
df["Cluster"] = agg_clust.fit_predict(df[["Age", "Annual Income (k$)"]])

print(df)

linked = linkage(df[["Age", "Annual Income (k$)"]], method="ward")

plt.figure(figsize=(20, 5))
dendrogram(
    linked, labels=df.index.tolist(), distance_sort="descending", show_leaf_counts=True
)
plt.title("Dendrogram")
plt.xlabel("Sample index")
plt.ylabel("Distance")
plt.show()

fig = ff.create_dendrogram(
    df[["Age", "Annual Income (k$)"]].values,
    labels=df.index.tolist(),
    linkagefun=lambda x: linkage(x, "ward"),
)
fig.update_layout(width=1200, height=600)
fig.show()

"""## Custom

### Impl
"""


def euclidean(a, b):
    return np.sqrt(np.sum((a - b) ** 2))


def ward_distance(cluster_i, cluster_j, data):
    """
    Ward's formula:
    d(i,j) = sqrt( (2 * ni * nj) / (ni + nj) ) * ||mean_i - mean_j||
    """
    ni, nj = len(cluster_i), len(cluster_j)
    mean_i = np.mean(data[cluster_i], axis=0)
    mean_j = np.mean(data[cluster_j], axis=0)
    return np.sqrt((2 * ni * nj) / (ni + nj)) * euclidean(mean_i, mean_j)


def custom_agglomerative(data, n_clusters=2):
    n = len(data)
    clusters = {i: [i] for i in range(n)}
    linkage_matrix = []
    next_id = n

    while len(clusters) > 1:
        keys = list(clusters.keys())
        min_dist = float("inf")
        merge_a, merge_b = None, None

        for i in range(len(keys)):
            for j in range(i + 1, len(keys)):
                d = ward_distance(clusters[keys[i]], clusters[keys[j]], data)
                if d < min_dist:
                    min_dist = d
                    merge_a, merge_b = keys[i], keys[j]

        new_members = clusters[merge_a] + clusters[merge_b]

        id_a = float(min(merge_a, merge_b))
        id_b = float(max(merge_a, merge_b))
        prev_dist = linkage_matrix[-1][2] if linkage_matrix else 0.0
        safe_dist = float(max(min_dist, prev_dist))

        linkage_matrix.append([id_a, id_b, safe_dist, float(len(new_members))])

        clusters[next_id] = new_members
        del clusters[merge_a]
        del clusters[merge_b]
        next_id += 1

    Z = np.array(linkage_matrix, dtype=float)

    from scipy.cluster.hierarchy import fcluster

    labels = fcluster(Z, t=n_clusters, criterion="maxclust") - 1

    return labels, Z


features = df[["Age", "Annual Income (k$)"]].values

labels, Z = custom_agglomerative(features, n_clusters=2)
df["Cluster"] = labels
print(df)

"""### Visualizations"""

plt.figure(figsize=(8, 5))
for c in np.unique(labels):
    mask = labels == c
    plt.scatter(features[mask, 0], features[mask, 1], label=f"Cluster {c}", s=60)
plt.xlabel("Age")
plt.ylabel("Annual Income (k$)")
plt.title("Custom Agglomerative Clustering (Ward)")
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(12, 5))
dendrogram(linked, truncate_mode="lastp", p=20, leaf_rotation=45, leaf_font_size=10)
plt.title("Dendrogram (Ward Linkage)")
plt.xlabel("Cluster index / size")
plt.ylabel("Ward distance")
plt.tight_layout()
plt.show()

ddata = dendrogram(Z, truncate_mode="lastp", p=20, no_plot=True)

fig = go.Figure()
for x_seg, y_seg in zip(ddata["icoord"], ddata["dcoord"]):
    fig.add_trace(
        go.Scatter(
            x=x_seg,
            y=y_seg,
            mode="lines",
            line=dict(color="royalblue", width=1.5),
            showlegend=False,
            hoverinfo="skip",
        )
    )

fig.update_layout(
    title="Dendrogram (Custom Ward Linkage)",
    xaxis=dict(
        title="Cluster index / size",
        tickangle=45,
        tickfont=dict(size=10),
        showgrid=False,
    ),
    yaxis=dict(title="Ward Distance", showgrid=True, gridcolor="lightgrey"),
    width=1200,
    height=500,
    plot_bgcolor="white",
    paper_bgcolor="white",
    showlegend=False,
)
fig.show()

"""## Tests"""

test_data = np.array([[1, 2], [1, 4], [5, 8], [5, 9], [8, 1], [8, 2]], dtype=float)

# Custom
_, Z_custom = custom_agglomerative(test_data, n_clusters=2)

# Scipy
from scipy.cluster.hierarchy import linkage

Z_scipy = linkage(test_data, method="ward")

print("Custom:\n", Z_custom.round(3))
print("\nScipy:\n", Z_scipy.round(3))
