#!/usr/bin/env python3
"""
Analysis of the Dr. Duke's Phytochemical and Ethnobotanical Databases
data dictionary.

Source: Dr. Duke's Phytochemical and Ethnobotanical Databases (data dictionary)
DOI: 10.15482/USDA.ADC/1239279  (licence CC0)
File: DrDukesDatabaseDataDictionary-prelim.csv
      md5 = ee486f85c5aecd88c1c27d08c47b88be

PRE-REGISTERED QUESTION
-----------------------
Within this relational database schema, is a column's declared DATA TYPE
associated with WHICH TABLE the column belongs to? Equivalently: do tables
differ in their data-type composition beyond what independent assignment
would produce?

Rationale / what is NOT tested (established during profiling, recorded here
for honesty):
  * `required` is FALSE for every one of the 199 columns -> zero variance,
    untestable.
  * `size` is a deterministic function of `type` (Short Text=255, Double=8,
    Date With Time=8, Long Integer=4; within-type variance = 0) -> degenerate,
    untestable as an independent variable.
  * Every table carries exactly one Long Integer `ID` auto-increment key ->
    a structural constant; a robustness check removes it.
The only field with genuine, non-degenerate cross-table variation is `type`.

PRE-REGISTERED TESTS (fixed before seeing results)
--------------------------------------------------
Primary (H1): Table (15 levels) x declared type (4 levels) independence.
  Test  : chi-square test of independence with a MONTE-CARLO simulated
          p-value (B = 20000, seeded). Monte Carlo is mandated a priori
          because the 15x4 table is sparse (most expected counts < 5), so the
          asymptotic chi-square reference distribution is invalid.
  Effect: Cramer's V (with bias-corrected V as a secondary read).
  Expected output: chi-square statistic, MC p-value, Cramer's V, n columns,
          the full contingency table.

Robustness / sensitivity checks (pre-registered):
  R1: Drop the structural Long Integer `ID` key from every table, then repeat
      the primary Monte-Carlo chi-square. Association must not be an artefact
      of the constant key column.
  R2: Collapse type to a binary Text vs. Non-text distinction (Short Text vs.
      {Double, Date With Time, Long Integer}) and repeat the Monte-Carlo
      chi-square. Tests whether any association survives the coarsest
      meaningful grouping.

Multiple-comparison control: the primary test plus R1 and R2 form a family of
3 p-values, corrected together with the Holm-Bonferroni method.

DISCIPLINE
----------
Nonparametric throughout (permutation / Monte-Carlo). Effect sizes and n
reported for every test. All seeds set and recorded. Results reported whether
strong, weak, or null. No causal claims.

Re-running this script reproduces every number, figure, and table.
"""

import os
import sys
import json
import hashlib
import urllib.request
import urllib.error

import numpy as np
import pandas as pd
from scipy.stats import chi2_contingency

# --------------------------------------------------------------------------
# 0. Reproducibility
# --------------------------------------------------------------------------
SEED = 20240607
RNG = np.random.default_rng(SEED)
B_MC = 20000  # Monte-Carlo replicates

HERE = os.path.dirname(os.path.abspath(__file__))
FIG_DIR = os.path.join(HERE, "figures")
TAB_DIR = os.path.join(HERE, "tables")
os.makedirs(FIG_DIR, exist_ok=True)
os.makedirs(TAB_DIR, exist_ok=True)

# --------------------------------------------------------------------------
# 1. Data acquisition: canonical download first, md5-gated local fallback
# --------------------------------------------------------------------------
EXPECTED_MD5 = "ee486f85c5aecd88c1c27d08c47b88be"
CANONICAL_URL = "https://ndownloader.figshare.com/files/43363338"
LOCAL_FALLBACK = "/Users/markhahnel/datasetpapers-handfeed/dr-duke/DrDukesDatabaseDataDictionary-prelim.csv"
WORK_COPY = os.path.join(HERE, "DrDukesDatabaseDataDictionary-prelim.csv")


def _md5(path):
    h = hashlib.md5()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 16), b""):
            h.update(chunk)
    return h.hexdigest()


def acquire_data():
    """Try canonical figshare download; fall back to md5-verified local copy.
    Hard-fail on any md5 mismatch."""
    # Attempt canonical download (works on any machine where figshare is reachable)
    try:
        print(f"[data] attempting canonical download: {CANONICAL_URL}")
        req = urllib.request.Request(CANONICAL_URL, headers={"User-Agent": "python-urllib"})
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = resp.read()
        with open(WORK_COPY, "wb") as fh:
            fh.write(data)
        got = _md5(WORK_COPY)
        if got != EXPECTED_MD5:
            raise ValueError(f"downloaded md5 {got} != expected {EXPECTED_MD5}")
        print(f"[data] canonical download OK, md5 verified {got}")
        return WORK_COPY
    except Exception as exc:  # network blocked, timeout, etc.
        print(f"[data] canonical download unavailable ({exc}); using local fallback")

    # md5-gated local fallback
    if not os.path.exists(LOCAL_FALLBACK):
        raise FileNotFoundError(
            f"canonical download failed and local fallback not found at {LOCAL_FALLBACK}"
        )
    got = _md5(LOCAL_FALLBACK)
    if got != EXPECTED_MD5:
        raise ValueError(
            f"local fallback md5 {got} != expected {EXPECTED_MD5} -- refusing to proceed"
        )
    print(f"[data] local fallback md5 verified {got}")
    return LOCAL_FALLBACK


# --------------------------------------------------------------------------
# 2. Parse the data dictionary into a clean field-level table
# --------------------------------------------------------------------------
def load_fields(csv_path):
    raw = pd.read_csv(csv_path)
    # Drop fully-empty trailing rows and the all-empty 'Unnamed:' columns
    named_cols = [c for c in raw.columns if not str(c).startswith("Unnamed")]
    raw = raw[named_cols]
    # Table-header rows: table name present, column name absent -> carry RecordCount
    rc = (raw[raw["column"].isna() & raw["table"].notna()]
          .set_index("table")["RecordCount"])
    # Forward-fill table name onto its field rows, then keep only field rows
    d = raw.copy()
    d["table"] = d["table"].ffill()
    fields = d[d["column"].notna()].copy().reset_index(drop=True)
    fields["RecordCount"] = fields["table"].map(rc)
    # Normalise whitespace in text fields
    for c in ["type", "required", "column", "table"]:
        fields[c] = fields[c].astype(str).str.strip()
    fields["size"] = pd.to_numeric(fields["size"], errors="coerce")
    return fields, rc


# --------------------------------------------------------------------------
# 3. Monte-Carlo chi-square via label permutation
# --------------------------------------------------------------------------
def mc_chi2(table_labels, type_labels, rng, B):
    """Permutation Monte-Carlo p-value for independence between two categorical
    label vectors. Statistic = Pearson chi-square. Returns (chi2_obs, p_mc,
    cramers_v, cramers_v_bc, contingency_df, dof, n)."""
    ct = pd.crosstab(pd.Series(table_labels, name="table"),
                     pd.Series(type_labels, name="type"))
    obs = ct.values.astype(float)
    chi2_obs, _, dof, _ = chi2_contingency(obs, correction=False)
    n = int(obs.sum())
    r, k = ct.shape
    # Cramer's V and bias-corrected V (Bergsma 2013)
    phi2 = chi2_obs / n
    cv = np.sqrt(phi2 / min(r - 1, k - 1)) if min(r - 1, k - 1) > 0 else np.nan
    phi2c = max(0.0, phi2 - (r - 1) * (k - 1) / (n - 1))
    rc_ = r - (r - 1) ** 2 / (n - 1)
    kc_ = k - (k - 1) ** 2 / (n - 1)
    denom = min(rc_ - 1, kc_ - 1)
    cv_bc = np.sqrt(phi2c / denom) if denom > 0 else np.nan
    # Permutation null
    type_arr = np.asarray(type_labels)
    tab_ser = pd.Series(table_labels)
    ge = 0
    for _ in range(B):
        perm = rng.permutation(type_arr)
        sim = pd.crosstab(tab_ser, pd.Series(perm)).values
        chi2_sim, _, _, _ = chi2_contingency(sim, correction=False)
        if chi2_sim >= chi2_obs - 1e-9:
            ge += 1
    p_mc = (ge + 1) / (B + 1)  # add-one (Davison & Hinkley) -> never exactly 0
    return chi2_obs, p_mc, cv, cv_bc, ct, dof, n


def holm(pvals):
    """Holm-Bonferroni corrected p-values, preserving input order."""
    p = np.asarray(pvals, float)
    m = len(p)
    order = np.argsort(p)
    corrected = np.empty(m)
    running = 0.0
    for rank, idx in enumerate(order):
        val = (m - rank) * p[idx]
        running = max(running, val)
        corrected[idx] = min(running, 1.0)
    return corrected.tolist()


# --------------------------------------------------------------------------
# 4. Figures
# --------------------------------------------------------------------------
def make_figures(fields, ct_primary, results):
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    try:
        # Optional shared style helper; script runs without it.
        from kernel import apply_figure_style  # type: ignore
        apply_figure_style()
    except Exception:
        mpl.rcParams.update({
            "figure.dpi": 150, "savefig.dpi": 150, "font.size": 9,
            "axes.spines.top": False, "axes.spines.right": False,
        })

    # Colourblind-safe qualitative palette (Okabe-Ito subset), type-threaded
    TYPE_ORDER = ["Short Text", "Double", "Long Integer", "Date With Time"]
    TYPE_COLOR = {
        "Short Text":     "#0072B2",  # blue
        "Double":         "#E69F00",  # orange
        "Long Integer":   "#009E73",  # green
        "Date With Time": "#CC79A7",  # magenta
    }
    ct = ct_primary.reindex(columns=[t for t in TYPE_ORDER if t in ct_primary.columns])
    # order tables by number of columns (descending) for readability
    order = fields.groupby("table").size().sort_values(ascending=False).index.tolist()
    ct = ct.reindex(order)

    # ---- Figure 1: stacked proportional composition per table -------------
    prop = ct.div(ct.sum(axis=1), axis=0)
    fig, ax = plt.subplots(figsize=(7.2, 5.0))
    left = np.zeros(len(prop))
    ylab = prop.index.tolist()
    ypos = np.arange(len(prop))
    for t in ct.columns:
        vals = prop[t].values
        ax.barh(ypos, vals, left=left, color=TYPE_COLOR[t], label=t,
                edgecolor="white", linewidth=0.6)
        left = left + vals
    ax.set_yticks(ypos)
    ax.set_yticklabels(ylab, fontsize=7)
    ax.invert_yaxis()
    ax.set_xlim(0, 1)
    ax.set_xlabel("Proportion of columns")
    ax.set_ylabel("Database table")
    v = results["primary"]["cramers_v"]
    p = results["primary"]["p_mc"]
    ax.set_title("Data-type composition differs across tables\n"
                 f"(Monte-Carlo \u03c7\u00b2 p = {p:.4f}, Cram\u00e9r's V = {v:.2f}, n = "
                 f"{results['primary']['n']} columns)", fontsize=9, loc="left")
    ax.legend(title="Declared type", frameon=False, fontsize=7,
              title_fontsize=7, ncol=4, loc="upper center",
              bbox_to_anchor=(0.5, -0.10))
    fig.tight_layout()
    f1 = os.path.join(FIG_DIR, "fig-1-type-composition-by-table.png")
    fig.savefig(f1, dpi=150, bbox_inches="tight")
    plt.close(fig)

    # ---- Figure 2: standardized Pearson residuals heatmap -----------------
    obs = ct.fillna(0).values.astype(float)
    row = obs.sum(1, keepdims=True)
    col = obs.sum(0, keepdims=True)
    tot = obs.sum()
    exp = row @ col / tot
    with np.errstate(divide="ignore", invalid="ignore"):
        std_resid = (obs - exp) / np.sqrt(exp * (1 - row / tot) * (1 - col / tot))
    std_resid = np.nan_to_num(std_resid, nan=0.0, posinf=0.0, neginf=0.0)
    vmax = np.max(np.abs(std_resid))
    vmax = max(vmax, 1.0)
    fig, ax = plt.subplots(figsize=(6.4, 5.2))
    im = ax.imshow(std_resid, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
    ax.set_xticks(range(len(ct.columns)))
    ax.set_xticklabels(ct.columns, rotation=30, ha="right", fontsize=7)
    ax.set_yticks(range(len(ct.index)))
    ax.set_yticklabels(ct.index, fontsize=7)
    for i in range(std_resid.shape[0]):
        for j in range(std_resid.shape[1]):
            val = std_resid[i, j]
            if abs(val) >= 0.05:
                ax.text(j, i, f"{val:.1f}", ha="center", va="center",
                        fontsize=6,
                        color="white" if abs(val) > vmax * 0.6 else "black")
    cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cb.set_label("Standardized Pearson residual\n(obs \u2212 exp, centred at 0)", fontsize=7)
    ax.set_xlabel("Declared type")
    ax.set_ylabel("Database table")
    ax.set_title("Where the association lives: residuals vs. independence",
                 fontsize=9, loc="left")
    fig.tight_layout()
    f2 = os.path.join(FIG_DIR, "fig-2-residuals-heatmap.png")
    fig.savefig(f2, dpi=150, bbox_inches="tight")
    plt.close(fig)

    # ---- Figure 3: Monte-Carlo null distribution + observed statistic -----
    fig, ax = plt.subplots(figsize=(6.4, 4.2))
    null = results["primary"]["mc_null_sample"]
    chi2_obs = results["primary"]["chi2"]
    ax.hist(null, bins=50, color="#B0B0B0", edgecolor="white", linewidth=0.3)
    ax.axvline(chi2_obs, color="#D55E00", linewidth=2,
               label=f"observed \u03c7\u00b2 = {chi2_obs:.1f}")
    ax.set_xlabel("\u03c7\u00b2 statistic under permutation null")
    ax.set_ylabel("Monte-Carlo replicates")
    ax.set_title(f"Permutation null vs. observed (B = {B_MC}, p = "
                 f"{results['primary']['p_mc']:.4f})", fontsize=9, loc="left")
    ax.legend(frameon=False, fontsize=8)
    fig.tight_layout()
    f3 = os.path.join(FIG_DIR, "fig-3-mc-null-distribution.png")
    fig.savefig(f3, dpi=150, bbox_inches="tight")
    plt.close(fig)

    return [f1, f2, f3]


# --------------------------------------------------------------------------
# 5. Main
# --------------------------------------------------------------------------
def main():
    csv_path = acquire_data()
    fields, rc = load_fields(csv_path)

    n_tables = fields["table"].nunique()
    n_fields = len(fields)
    assert n_fields == 199, f"expected 199 field rows, got {n_fields}"
    assert n_tables == 15, f"expected 15 tables, got {n_tables}"

    results = {
        "meta": {
            "seed": SEED, "mc_replicates": B_MC,
            "source_doi": "10.15482/USDA.ADC/1239279",
            "source_md5": EXPECTED_MD5,
            "n_tables": int(n_tables), "n_columns": int(n_fields),
        },
        "degenerate_fields": {},
    }

    # Record degeneracy of the fields we deliberately do NOT test
    results["degenerate_fields"]["required"] = {
        "distinct_values": fields["required"].unique().tolist(),
        "note": "all FALSE -> zero variance, untestable",
    }
    size_by_type = (fields.groupby("type")["size"]
                    .agg(["nunique", "min", "max"]).to_dict("index"))
    results["degenerate_fields"]["size"] = {
        "within_type_unique_sizes": {k: int(v["nunique"]) for k, v in size_by_type.items()},
        "note": "size is a deterministic function of type (within-type variance 0)",
    }

    # ---------- PRIMARY TEST ----------
    chi2_obs, p_mc, cv, cv_bc, ct, dof, n = mc_chi2(
        fields["table"].values, fields["type"].values, RNG, B_MC)
    # keep a small null sample for the figure (regenerate deterministically)
    rng_fig = np.random.default_rng(SEED + 1)
    null_sample = []
    type_arr = fields["type"].values
    tab_ser = pd.Series(fields["table"].values)
    for _ in range(B_MC):
        perm = rng_fig.permutation(type_arr)
        null_sample.append(chi2_contingency(pd.crosstab(tab_ser, pd.Series(perm)).values,
                                            correction=False)[0])
    results["primary"] = {
        "test": "Monte-Carlo chi-square test of independence (table x type)",
        "chi2": float(chi2_obs), "dof": int(dof), "p_mc": float(p_mc),
        "cramers_v": float(cv), "cramers_v_bias_corrected": float(cv_bc),
        "n": int(n), "n_cells": int(ct.size),
        "cells_expected_lt5": int((chi2_contingency(ct.values)[3] < 5).sum()),
        "contingency": ct.to_dict("index"),
        "mc_null_sample": null_sample,
    }

    # ---------- R1: drop structural ID key ----------
    no_id = fields[fields["column"].str.upper() != "ID"].copy()
    chi2_1, p_1, cv_1, cv_bc_1, ct1, dof1, n1 = mc_chi2(
        no_id["table"].values, no_id["type"].values,
        np.random.default_rng(SEED + 2), B_MC)
    results["R1_drop_id"] = {
        "test": "Monte-Carlo chi-square, Long Integer ID key removed",
        "chi2": float(chi2_1), "dof": int(dof1), "p_mc": float(p_1),
        "cramers_v": float(cv_1), "cramers_v_bias_corrected": float(cv_bc_1),
        "n": int(n1),
    }

    # ---------- R2: binary Text vs Non-text ----------
    bin_type = np.where(fields["type"].values == "Short Text", "Text", "Non-text")
    chi2_2, p_2, cv_2, cv_bc_2, ct2, dof2, n2 = mc_chi2(
        fields["table"].values, bin_type,
        np.random.default_rng(SEED + 3), B_MC)
    results["R2_binary_text"] = {
        "test": "Monte-Carlo chi-square, type collapsed to Text vs Non-text",
        "chi2": float(chi2_2), "dof": int(dof2), "p_mc": float(p_2),
        "cramers_v": float(cv_2), "cramers_v_bias_corrected": float(cv_bc_2),
        "n": int(n2),
    }

    # ---------- Holm correction across the family of 3 ----------
    fam = [("primary", p_mc), ("R1_drop_id", p_1), ("R2_binary_text", p_2)]
    corr = holm([p for _, p in fam])
    for (name, _), pc in zip(fam, corr):
        results[name]["p_holm"] = float(pc)
    results["multiple_comparison"] = {
        "method": "Holm-Bonferroni", "family_size": 3,
        "family": {name: {"p_mc": float(p), "p_holm": float(pc)}
                   for (name, p), pc in zip(fam, corr)},
    }

    # --------------------------------------------------------------------
    # Tables (CSV) + Frictionless datapackage
    # --------------------------------------------------------------------
    # tbl-1: table x type contingency (with margins)
    ct_out = ct.copy()
    ct_out["(all types)"] = ct_out.sum(axis=1)
    ct_out.loc["(all tables)"] = ct_out.sum(axis=0)
    t1 = os.path.join(TAB_DIR, "tbl-1-type-by-table-contingency.csv")
    ct_out.to_csv(t1)

    # tbl-2: per-table summary (n columns, record count, type counts)
    summ = fields.groupby("table").agg(
        n_columns=("column", "size"),
        record_count=("RecordCount", "first"),
    )
    tcounts = pd.crosstab(fields["table"], fields["type"])
    summ = summ.join(tcounts)
    summ = summ.sort_values("n_columns", ascending=False)
    t2 = os.path.join(TAB_DIR, "tbl-2-per-table-summary.csv")
    summ.to_csv(t2)

    # tbl-3: test results summary
    rows = []
    for name in ["primary", "R1_drop_id", "R2_binary_text"]:
        r = results[name]
        rows.append({
            "test": name, "description": r["test"],
            "chi2": r["chi2"], "dof": r["dof"], "n": r["n"],
            "p_mc": r["p_mc"], "p_holm": r["p_holm"],
            "cramers_v": r["cramers_v"],
            "cramers_v_bias_corrected": r["cramers_v_bias_corrected"],
        })
    tests_df = pd.DataFrame(rows)
    t3 = os.path.join(TAB_DIR, "tbl-3-test-results.csv")
    tests_df.to_csv(t3, index=False)

    # Frictionless datapackage.json describing the three tables
    def _fields_schema(df, index_name=None):
        fs = []
        if index_name is not None:
            fs.append({"name": index_name, "type": "string"})
        for c in df.columns:
            dt = "integer" if pd.api.types.is_integer_dtype(df[c]) else (
                "number" if pd.api.types.is_float_dtype(df[c]) else "string")
            fs.append({"name": str(c), "type": dt})
        return fs

    datapackage = {
        "name": "drduke-datadictionary-type-by-table",
        "title": "Dr. Duke data dictionary: data-type composition across tables",
        "licenses": [{"name": "CC0-1.0",
                      "path": "https://creativecommons.org/publicdomain/zero/1.0/"}],
        "profile": "tabular-data-package",
        "sources": [{"title": "Dr. Duke's Phytochemical and Ethnobotanical Databases (data dictionary)",
                     "path": "https://doi.org/10.15482/USDA.ADC/1239279"}],
        "resources": [
            {"name": "tbl-1-type-by-table-contingency",
             "path": "tbl-1-type-by-table-contingency.csv", "format": "csv",
             "mediatype": "text/csv", "profile": "tabular-data-resource",
             "schema": {"fields": _fields_schema(ct_out, "table")}},
            {"name": "tbl-2-per-table-summary",
             "path": "tbl-2-per-table-summary.csv", "format": "csv",
             "mediatype": "text/csv", "profile": "tabular-data-resource",
             "schema": {"fields": _fields_schema(summ, "table")}},
            {"name": "tbl-3-test-results",
             "path": "tbl-3-test-results.csv", "format": "csv",
             "mediatype": "text/csv", "profile": "tabular-data-resource",
             "schema": {"fields": _fields_schema(tests_df)}},
        ],
    }
    with open(os.path.join(TAB_DIR, "datapackage.json"), "w") as fh:
        json.dump(datapackage, fh, indent=2)

    # --------------------------------------------------------------------
    # Figures
    # --------------------------------------------------------------------
    fig_paths = make_figures(fields, ct, results)

    # --------------------------------------------------------------------
    # results.json (drop the large null sample before writing headline file)
    # --------------------------------------------------------------------
    results_out = json.loads(json.dumps(results))  # deep copy
    results_out["primary"].pop("mc_null_sample", None)
    with open(os.path.join(HERE, "results.json"), "w") as fh:
        json.dump(results_out, fh, indent=2)

    print("[done] tables, figures, results.json written")
    print(f"  primary: chi2={chi2_obs:.3f} p_mc={p_mc:.4f} "
          f"p_holm={results['primary']['p_holm']:.4f} V={cv:.3f} n={n}")
    print(f"  R1(no ID): p_mc={p_1:.4f} p_holm={results['R1_drop_id']['p_holm']:.4f} V={cv_1:.3f}")
    print(f"  R2(binary): p_mc={p_2:.4f} p_holm={results['R2_binary_text']['p_holm']:.4f} V={cv_2:.3f}")
    return results_out, fig_paths


if __name__ == "__main__":
    main()
