from pathlib import Path
from datetime import datetime
from urllib.parse import quote
from urllib.request import urlretrieve
import os
import sys

try:
    from IPython.display import display as _ipy_display
except Exception:
    _ipy_display = None

def show(value):
    if _ipy_display is not None:
        _ipy_display(value)
    else:
        print(value)

# --- Project metadata ---
CAPSTONE = 4
SESSION_TITLE = 'Session 4: Data Visualization Techniques'
IS_COLAB = 'google.colab' in sys.modules
RAW_BASE = os.environ.get('FRANCISBURNET_RAW_BASE', 'https://raw.githubusercontent.com/FrancisBurnet/francisburnet/main')

print(f"Capstone: {CAPSTONE} | Session: {SESSION_TITLE}")
print("Run timestamp:", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
print("Runtime:", "Google Colab" if IS_COLAB else "Local / notebook runtime")

CWD = Path.cwd()
if IS_COLAB:
    BASE_DIR = Path('/content/francisburnet_capstone_4')
elif (CWD / f"Capstone {CAPSTONE}").exists():
    BASE_DIR = CWD / f"Capstone {CAPSTONE}"
elif CWD.name == f"Capstone {CAPSTONE}":
    BASE_DIR = CWD
elif (CWD / "Incremental_Capstone" / f"Capstone {CAPSTONE}").exists():
    BASE_DIR = CWD / "Incremental_Capstone" / f"Capstone {CAPSTONE}"
else:
    BASE_DIR = CWD

INPUT_DIR = BASE_DIR / "inputs"
OUTPUT_DIR = BASE_DIR / "outputs"
PLOTS_DIR = OUTPUT_DIR / "plots"
INPUT_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
PLOTS_DIR.mkdir(parents=True, exist_ok=True)

# --- Helpers ---
def first_existing_path(candidates):
    for candidate in candidates:
        candidate = Path(candidate).expanduser()
        if candidate.exists():
            return candidate
    return None

def github_raw_url(relative_path: str) -> str:
    normalized = relative_path.replace('\\', '/')
    return RAW_BASE.rstrip('/') + '/' + quote(normalized, safe='/')

def stage_colab_asset(relative_path: str, destination_name=None) -> Path:
    destination = INPUT_DIR / (destination_name or Path(relative_path).name)
    if not destination.exists():
        asset_url = github_raw_url(relative_path)
        urlretrieve(asset_url, destination)
        print('Downloaded:', destination.name, 'from', asset_url)
    return destination

def resolve_dataset_path(default_filename: str) -> Path:
    """Resolve dataset path from local folders or the Colab staging folder."""
    path = first_existing_path([
        INPUT_DIR / default_filename,
        BASE_DIR / default_filename,
        CWD / default_filename,
        CWD / "Incremental_Capstone" / f"Capstone {CAPSTONE}" / default_filename,
    ])
    if path is None:
        raise FileNotFoundError(f"Dataset not found: {default_filename}")
    return path

print("Base directory:", BASE_DIR)
print("Output directory:", OUTPUT_DIR)
print("Plots directory:", PLOTS_DIR)

import pandas as pd  # DataFrames + aggregation/pivot/correlation preparation
import matplotlib.pyplot as plt  # Plot creation and export to outputs/plots

DEFAULT_DATASET = "NSMES1988updated.csv"

try:
    dataset_path = resolve_dataset_path(DEFAULT_DATASET)
except FileNotFoundError:
    fallback = first_existing_path([
        BASE_DIR.parent / "Capstone 2" / "outputs" / "NSMES1988updated.csv",
        CWD / "Capstone 2" / "outputs" / "NSMES1988updated.csv",
        CWD / "Incremental_Capstone" / "Capstone 2" / "outputs" / "NSMES1988updated.csv",
    ])
    if fallback is None and IS_COLAB:
        fallback = stage_colab_asset(
            'Incremental Capstones/Applied Data Science with Python/Capstone 2/outputs/NSMES1988updated.csv'
        )
    if fallback is None:
        raise
    dataset_path = fallback
    print("Local default not found; using fallback:", dataset_path)

df = pd.read_csv(dataset_path)
if "age_years" not in df.columns and "age" in df.columns:
    df["age_years"] = (df["age"] * 10).round(0)
if "income_dollars" not in df.columns and "income" in df.columns:
    df["income_dollars"] = (df["income"] * 10000).round(0)

print("Loaded:", dataset_path)
print("Shape:", df.shape)
show(df.head())

expected_cols = [
    "visits", "nvisits", "ovisits", "novisits", "emergency", "hospital",
    "health", "chronic", "adl", "region", "age", "gender",
    "married", "school", "income", "employed", "insurance", "medicaid",
    "age_years", "income_dollars"
]

missing_cols = [c for c in expected_cols if c not in df.columns]
print("Missing expected columns:", missing_cols)

print("\nDtypes:")
show(df.dtypes)

print("\nMissing values (count):")
na_counts = df.isna().sum().sort_values(ascending=False)
show(na_counts[na_counts > 0] if (na_counts > 0).any() else na_counts.head())

print("Plotting library: matplotlib")
print("Reason: direct pandas compatibility, flexible formatting, and reliable PNG export for grading artifacts.")

# Plot 1: counts by region
if "region" in df.columns:
    region_counts = df["region"].value_counts()
    fig, ax = plt.subplots(figsize=(7, 4.5))
    region_counts.sort_index().plot(kind="bar", title="Count by Region", ax=ax)
    ax.set_xlabel("Region")
    ax.set_ylabel("Count")
    out = PLOTS_DIR / "region_counts.png"
    fig.tight_layout()
    fig.savefig(out, dpi=150)
    plt.show()
    print("Saved:", out)

# Plot 2: counts by health
if "health" in df.columns:
    fig, ax = plt.subplots(figsize=(7, 4.5))
    df["health"].value_counts().sort_index().plot(kind="bar", title="Count by Health Category", ax=ax)
    ax.set_xlabel("Health")
    ax.set_ylabel("Count")
    out = PLOTS_DIR / "health_counts.png"
    fig.tight_layout()
    fig.savefig(out, dpi=150)
    plt.show()
    print("Saved:", out)

# Plot 3: mean visits by region and health
if set(["region", "health", "visits"]).issubset(df.columns):
    pivot_visits = pd.pivot_table(df, index="region", columns="health", values="visits", aggfunc="mean")
    fig, ax = plt.subplots(figsize=(8, 5))
    pivot_visits.plot(kind="bar", ax=ax)
    ax.set_title("Mean Visits by Region and Health")
    ax.set_xlabel("Region")
    ax.set_ylabel("Mean Visits")
    out = PLOTS_DIR / "mean_visits_region_health.png"
    fig.tight_layout()
    fig.savefig(out, dpi=150)
    plt.show()
    print("Saved:", out)

num_df = df.select_dtypes(include=["number"])
corr = num_df.corr(numeric_only=True)
show(corr)

fig, ax = plt.subplots(figsize=(10, 8))
cax = ax.imshow(corr, cmap="coolwarm", vmin=-1, vmax=1)
fig.colorbar(cax, ax=ax, fraction=0.046, pad=0.04)
ax.set_xticks(range(len(corr.columns)))
ax.set_yticks(range(len(corr.columns)))
ax.set_xticklabels(corr.columns, rotation=90)
ax.set_yticklabels(corr.columns)
ax.set_title("Correlation Matrix (Numeric Features)")

out = PLOTS_DIR / "correlation_matrix.png"
fig.tight_layout()
fig.savefig(out, bbox_inches="tight", dpi=150)
plt.show()
print("Saved:", out)

scatter_specs = [
    ("income_dollars", "visits", "scatter_income_vs_visits.png", "Income vs Visits"),
    ("age_years", "visits", "scatter_age_vs_visits.png", "Age vs Visits"),
    ("income_dollars", "emergency", "scatter_income_vs_emergency.png", "Income vs Emergency Visits")
]

for x, y, fname, title in scatter_specs:
    if x in df.columns and y in df.columns:
        fig, ax = plt.subplots(figsize=(6.5, 4.5))
        ax.scatter(df[x], df[y], alpha=0.3, s=12)
        ax.set_title(title)
        ax.set_xlabel(x)
        ax.set_ylabel(y)
        out = PLOTS_DIR / fname
        fig.tight_layout()
        fig.savefig(out, dpi=150)
        plt.show()
        print("Saved:", out)

print("Capstone 4 completed: C4-T3 to C4-T5")
print("Artifacts saved under outputs/plots")