#!/usr/bin/env python3
"""
04_grd_explore.py — Load and explore GRD (hospital discharges) data.

IELE756 · Week 0 · March 6, 2026
"""

import zipfile
import pandas as pd

# ── 1. Unzip and load (key columns only) ──
cols_grd = ["COMUNA", "NACIONALIDAD", "SEXO", "DIAGNOSTICO1",
            "FECHA_INGRESO", "FECHAALTA",
            "IR_29301_SEVERIDAD", "IR_29301_COD_GRD"]

with zipfile.ZipFile("materials/grd/GRD_PUBLICO_2024.zip") as z:
    with z.open("GRD_PUBLICO_2024.txt") as f:
        grd = pd.read_csv(f, sep="|", usecols=cols_grd,
                          encoding="latin-1", low_memory=False)

print(f"Total discharges 2024: {len(grd):,}")

# ── 2. Quick look ──
print("\n── First rows ──")
print(grd.head())

# ── 3. Filter to Tarapacá (comuna names) ──
comunas_tarapaca = ["IQUIQUE", "ALTO HOSPICIO", "POZO ALMONTE",
                    "HUARA", "CAMIÑA", "COLCHANE", "PICA"]
grd_tar = grd[grd["COMUNA"].isin(comunas_tarapaca)]
print(f"\nTarapacá discharges: {len(grd_tar):,}")

# ── 4. Join to ICD-10 lookup ──
cie10 = pd.read_excel("materials/grd/CIE-10.xlsx",
                       sheet_name="CIE 10", engine="openpyxl")
print(f"CIE-10 codes: {len(cie10):,}")

grd_tar = grd_tar.merge(
    cie10[["Código", "Descripción", "Capítulo"]],
    left_on="DIAGNOSTICO1", right_on="Código", how="left"
)

# ── 5. Top 10 diagnoses ──
print("\n── Top 10 diagnoses ──")
print(grd_tar["Descripción"].value_counts().head(10))

# ── 6. Severity distribution ──
# 0 = Sin gravedad, 1 = Menor, 2 = Moderada, 3 = Mayor
print("\n── Severity distribution ──")
print(grd_tar["IR_29301_SEVERIDAD"].value_counts().sort_index())

# ── 7. Nationality ──
print("\n── Nationality ──")
print(grd_tar["NACIONALIDAD"].value_counts())