"""
Glasgow Public 5G Performance Dataset (2025)
Figure Generation Script
Reproduces all 7 figure

Usage:
    python FigureGeneration.py

Requires:
    pip install pandas openpyxl matplotlib seaborn geopandas shapely contextily pillow numpy

Output:
    GRAPHS/ directory containing all figure files.
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib.colors as mcolors
import seaborn as sns
import geopandas as gpd
from shapely.geometry import Point

# ── Configuration ─────────────────────────────────────────────────────────────
DATASET_PATH = "Glasgow5GDataSet.xlsx"
OUTPUT_DIR   = "GRAPHS"
DPI          = 150

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ── Load data ─────────────────────────────────────────────────────────────────
sheets = ["glasgow_5g_2025-04-06", "glasgow_5g_2025-04-07", "glasgow_5g_2025-04-08"]
df = pd.concat(
    [pd.read_excel(DATASET_PATH, sheet_name=s) for s in sheets],
    ignore_index=True
)

df.columns = [c.strip() for c in df.columns]
df = df.rename(columns={
    "Signal Strength (dBm)": "Signal",
    "Download Speed (Mbps)": "Download",
    "Upload Speed (Mbps)":   "Upload",
    "Ping (ms)":             "Ping",
    "Network Provider":      "Provider",
})

avg_df = df.groupby("Location", as_index=False).agg(
    Download=("Download", "mean"),
    Upload=("Upload",   "mean"),
    Ping=("Ping",       "mean"),
    Signal=("Signal",   "mean"),
)

# ── Figure 1: Download speed distribution per provider (box plot) ─────────────
fig, ax = plt.subplots(figsize=(10, 6))
provider_order = ["EE", "Vodafone", "O2", "Sky Mobile"]
palette = ["#3B4D8C", "#4E8FA2", "#3D8F6F", "#8FB04E"]

sns.boxplot(
    data=df,
    x="Provider", y="Download",
    hue="Provider",
    order=provider_order,
    palette=palette,
    width=0.5,
    linewidth=1.2,
    flierprops=dict(marker="", linestyle="none"),
    legend=False,
    ax=ax,
)
ax.set_title("Download Speed Distribution per Network Provider", fontsize=13)
ax.set_xlabel("Network Provider", fontsize=11)
ax.set_ylabel("Download Speed (Mbps)", fontsize=11)
ax.set_ylim(0, 1300)
sns.despine()
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/download_speed_distribution.png", dpi=DPI)
plt.close()
print("Figure 1 saved.")

# ── Figure 2: Average download speed by neighbourhood (bar chart) ─────────────
sorted_dl = avg_df.sort_values("Download")
n = len(sorted_dl)
colours = sns.color_palette("mako", n)

fig, ax = plt.subplots(figsize=(10, 7))
bars = ax.barh(sorted_dl["Location"], sorted_dl["Download"], color=colours)
ax.set_title("Average Download Speed by Area (Mbps)", fontsize=13)
ax.set_xlabel("Average Download Speed (Mbps)", fontsize=11)
ax.set_ylabel("Location", fontsize=11)
ax.set_xlim(0, 750)
sns.despine()
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/avg_download_by_area.png", dpi=DPI)
plt.close()
print("Figure 2 saved.")

# ── Figure 4: Average upload speed by neighbourhood ───────────────────────────
sorted_ul = avg_df.sort_values("Upload")
n = len(sorted_ul)
colours_ul = sns.color_palette("mako", n)

fig, ax = plt.subplots(figsize=(10, 7))
ax.barh(sorted_ul["Location"], sorted_ul["Upload"], color=colours_ul)
ax.set_title("Average Upload Speed by Area", fontsize=13)
ax.set_xlabel("Average Upload Speed (Mbps)", fontsize=11)
ax.set_ylabel("Location", fontsize=11)
ax.set_xlim(0, 200)
sns.despine()
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/avg_upload_by_area.png", dpi=DPI)
plt.close()
print("Figure 4 saved.")

# ── Figure 5: Average ping latency by neighbourhood ───────────────────────────
sorted_ping = avg_df.sort_values("Ping", ascending=False)
n = len(sorted_ping)
colours_ping = sns.color_palette("RdPu_r", n)

fig, ax = plt.subplots(figsize=(10, 7))
ax.barh(sorted_ping["Location"], sorted_ping["Ping"], color=colours_ping)
ax.set_title("Average Ping by Area (ms)", fontsize=13)
ax.set_xlabel("Average Ping (ms)", fontsize=11)
ax.set_ylabel("Location", fontsize=11)
ax.set_xlim(0, 26)
sns.despine()
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/avg_ping_by_area.png", dpi=DPI)
plt.close()
print("Figure 5 saved.")

# ── Figure 6: SS-RSRP signal strength by provider (box plot) ─────────────────
sig_palette = ["#7B8CC4", "#A8C4D4", "#D4B8A8", "#C48888"]

fig, ax = plt.subplots(figsize=(10, 6))
sns.boxplot(
    data=df,
    x="Provider", y="Signal",
    hue="Provider",
    order=provider_order,
    palette=sig_palette,
    width=0.5,
    linewidth=1.2,
    flierprops=dict(marker="", linestyle="none"),
    legend=False,
    ax=ax,
)
ax.set_title("Signal Strength by Provider (lower dBm = stronger)", fontsize=13)
ax.set_xlabel("Network Provider", fontsize=11)
ax.set_ylabel("Signal Strength (dBm)", fontsize=11)
sns.despine()
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/signal_strength_by_provider.png", dpi=DPI)
plt.close()
print("Figure 6 saved.")

# ── Figure 7: Area × Provider download heat map ───────────────────────────────
heatmap_df = df.groupby(["Location", "Provider"])["Download"].mean().unstack()
heatmap_df = heatmap_df[["EE", "O2", "Sky Mobile", "Vodafone"]]
heatmap_df = heatmap_df.sort_index()

fig, ax = plt.subplots(figsize=(12, 9))
sns.heatmap(
    heatmap_df,
    annot=True, fmt=".1f",
    cmap="YlGnBu_r",
    linewidths=0.5,
    linecolor="white",
    cbar_kws={"label": "Avg Download Speed (Mbps)"},
    ax=ax,
)
ax.set_title("Avg Download Speed: Area vs Network Provider", fontsize=13)
ax.set_xlabel("Network Provider", fontsize=11)
ax.set_ylabel("Location", fontsize=11)
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/avg_download_heatmap.png", dpi=DPI)
plt.close()
print("Figure 7 saved.")

print(f"\nAll figures saved to ./{OUTPUT_DIR}/")


# ── Figure 3: Geospatial heat map of average download speed ───────────────────
import contextily as ctx
from PIL import Image

area_coords = {
    "Glasgow City Centre": (55.8609, -4.2514),
    "Merchant City":       (55.8590, -4.2440),
    "Dennistoun":          (55.8650, -4.2160),
    "Partick":             (55.8699, -4.3125),
    "Shawlands":           (55.8292, -4.2923),
    "Govan":               (55.8615, -4.3083),
    "Hillhead":            (55.8748, -4.2928),
    "Maryhill":            (55.8910, -4.2930),
    "Pollok":              (55.8340, -4.3460),
    "Easterhouse":         (55.8670, -4.1260),
    "Bearsden":            (55.9190, -4.3320),
    "Springburn":          (55.8830, -4.2280),
    "Govanhill":           (55.8360, -4.2580),
    "Drumchapel":          (55.9040, -4.3620),
    "Cathcart":            (55.8160, -4.2610),
}

# Build GeoDataFrame from averaged data
coord_df = pd.DataFrame.from_dict(
    area_coords, orient="index", columns=["Latitude", "Longitude"]
).reset_index().rename(columns={"index": "Location"})

map_df = avg_df.merge(coord_df, on="Location", how="inner")
gdf = gpd.GeoDataFrame(
    map_df,
    geometry=gpd.points_from_xy(map_df.Longitude, map_df.Latitude),
    crs="EPSG:4326"
).to_crs(epsg=3857)

# Map bounds with padding
bounds  = gdf.total_bounds
padding = 5000  # metres
x_min, y_min = bounds[0] - padding, bounds[1] - padding
x_max, y_max = bounds[2] + padding, bounds[3] + padding

fig, ax = plt.subplots(figsize=(12, 12))
ax.set_xlim(x_min, x_max)
ax.set_ylim(y_min, y_max)

# Greyscale CartoDB basemap
ctx.add_basemap(ax, source=ctx.providers.CartoDB.Positron)
for im in ax.get_images():
    if im.get_array().shape[-1] in [3, 4]:
        basemap_pil  = Image.fromarray((im.get_array() * 255).astype(np.uint8))
        basemap_grey = np.array(basemap_pil.convert("L")) / 255.0
        im.set_array(np.stack(
            [basemap_grey, basemap_grey, basemap_grey,
             np.ones_like(basemap_grey) * 0.7], axis=-1
        ))

# Scatter points coloured by download speed
scatter = ax.scatter(
    gdf.geometry.x, gdf.geometry.y,
    s=500,
    c=gdf["Download"],
    cmap="viridis",
    vmin=600, vmax=710,
    alpha=1.0,
    edgecolors="none",
    zorder=3,
)

# Labels
for _, row in gdf.iterrows():
    ax.text(
        row.geometry.x, row.geometry.y + 500,
        row["Location"],
        fontsize=8, ha="center", color="black",
        bbox=dict(facecolor="white", alpha=0.7, edgecolor="none", boxstyle="round,pad=0.2"),
        zorder=4,
    )

# Colorbar
cbar = fig.colorbar(
    plt.cm.ScalarMappable(cmap=plt.cm.viridis, norm=plt.Normalize(vmin=600, vmax=710)),
    ax=ax, label="Download Speed (Mbps)"
)
cbar.set_ticks(np.linspace(600, 710, 5))

ax.set_axis_off()
plt.title("5G Average Download Speeds Across Glasgow", fontsize=18, weight="bold")
plt.tight_layout()
plt.savefig(f"{OUTPUT_DIR}/HeatMap2.png", dpi=DPI)
plt.close()
print("Figure 3 saved.")