"""Shared helpers for reproducing the numerical experiments and figures of
Bishop (2006), Pattern Recognition and Machine Learning, Chapters 9, 10, 12, 14.

Every chapter module (ch09_mixture_em.py, ch10_approximate_inference.py,
ch12_continuous_latent.py, ch14_combining_models.py) imports this file for

* data loaders: Old Faithful (bundled CSV), the oil flow data and MNIST
  (downloaded once and cached), and two public-domain photographs;
* a common plotting style and colour palette close to the book's figures;
* savefig(), which writes ../img/chNN/figX-Y.png next to the lecture notes;
* run(), a small command-line driver: python ch09_mixture_em.py --fig 9.1

Data are cached in ~/.cache/prml-repro (set PRML_DATA to change this) and
figures go to ../img (set PRML_OUT, or pass --out, to change this).
"""
import argparse
import gzip
import logging
import os
import struct
import sys
import tarfile
import time
import urllib.request

import numpy as np
import matplotlib

matplotlib.use("Agg")
# fontTools (used to embed fonts in the PDFs) logs a harmless note about font timestamps
logging.getLogger("fontTools").setLevel(logging.ERROR)
import matplotlib.pyplot as plt  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(HERE, "data")
CACHE = os.environ.get("PRML_DATA", os.path.join(os.path.expanduser("~"), ".cache", "prml-repro"))
OUT_DIR = os.environ.get("PRML_OUT", os.path.normpath(os.path.join(HERE, "..", "img")))
USER_AGENT = "prml-repro/1.0 (teaching code)"

# Colours used throughout, chosen to match the book's red / blue / green figures.
RED = "#d62728"
BLUE = "#1f5fbf"
GREEN = "#2ca02c"
MAGENTA = "#c2185b"
CYAN = "#17becf"
ORANGE = "#ff7f0e"
PURPLE = "#7b4fbf"
GREY = "#7f7f7f"
CLASS_COLOURS = (RED, BLUE, GREEN)


# --------------------------------------------------------------------------- style
def setup_style():
    """A clean, book-like matplotlib style (white background, boxed axes)."""
    plt.rcParams.update({
        "figure.dpi": 100,
        "savefig.dpi": 130,
        "font.size": 10,
        "axes.titlesize": 10.5,
        "axes.labelsize": 10,
        "legend.fontsize": 9,
        "xtick.labelsize": 9,
        "ytick.labelsize": 9,
        "axes.linewidth": 0.8,
        "lines.linewidth": 1.6,
        "figure.facecolor": "white",
        "axes.facecolor": "white",
        "savefig.facecolor": "white",
        "savefig.bbox": "tight",
        "savefig.pad_inches": 0.06,
        "mathtext.fontset": "cm",
        "svg.hashsalt": "prml-repro",   # stable ids, so re-running gives identical SVG files
        "pdf.fonttype": 42,             # embed TrueType fonts: the PDFs stay editable in Illustrator/Inkscape
    })


setup_style()


def panel_grid(nrows, ncols, size=(2.6, 2.4), **kw):
    """fig, axes for a grid of equally sized panels (axes is always 2-D)."""
    fig, axes = plt.subplots(nrows, ncols, figsize=(size[0] * ncols, size[1] * nrows), squeeze=False, **kw)
    return fig, axes


def panel_label(ax, label, loc="upper left"):
    """Write (a), (b), ... inside a panel, as in the book."""
    x, ha = (0.03, "left") if "left" in loc else (0.97, "right")
    y, va = (0.97, "top") if "upper" in loc else (0.03, "bottom")
    ax.text(x, y, "(%s)" % label, transform=ax.transAxes, ha=ha, va=va, fontsize=10,
            bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.8))


def savefig(fig, chapter, name, out_dir=None, formats=("png", "svg", "pdf")):
    """Save fig as <out_dir>/chNN/<name>.png, .svg and .pdf and return the PNG path.

    name is the figure id in the book with a dash, e.g. 'fig9-1' for Figure 9.1.
    The SVG is what the web pages show (sharp at any zoom); the PDF is for slides
    and papers; the PNG is a light fallback and the thumbnail in the code gallery.
    """
    out = os.path.join(out_dir or OUT_DIR, "ch%02d" % int(chapter))
    os.makedirs(out, exist_ok=True)
    paths = {}
    for ext in formats:
        path = os.path.join(out, name + "." + ext)
        meta = {"Date": None} if ext in ("svg", "pdf") else None   # no timestamps: reproducible files
        if ext == "pdf":
            meta["CreationDate"] = None
            del meta["Date"]
        fig.savefig(path, metadata=meta)
        paths[ext] = path
    plt.close(fig)
    print("wrote", os.path.relpath(paths.get("png", next(iter(paths.values()))), os.getcwd()), "(+ " + ", ".join(e for e in formats if e != "png") + ")")
    return paths.get("png")


# --------------------------------------------------------------------------- downloads
def _download(urls, dest):
    """Download the first URL that works to dest (atomic rename)."""
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    last = None
    for url in urls:
        try:
            req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
            with urllib.request.urlopen(req, timeout=120) as r:
                data = r.read()
            tmp = dest + ".part"
            with open(tmp, "wb") as f:
                f.write(data)
            os.replace(tmp, dest)
            return dest
        except Exception as e:  # try the next mirror
            last = e
            time.sleep(1)
    raise RuntimeError("could not download %s (last error: %s)" % (os.path.basename(dest), last))


# --------------------------------------------------------------------------- data sets
def old_faithful(standardize=False):
    """Old Faithful geyser data (Bishop 2006, Appendix A): 272 x 2 array.

    Column 0 is the eruption duration (minutes) and column 1 the waiting time to
    the next eruption (minutes). The same 272 observations ship with R as
    datasets::faithful (Azzalini and Bowman 1990; Haerdle 1991).
    With standardize=True each column has mean 0 and standard deviation 1, the
    're-scaled' version used in Figures 9.1, 9.8 and 10.6.
    """
    X = np.loadtxt(os.path.join(DATA_DIR, "faithful.csv"), delimiter=",", skiprows=1)
    if standardize:
        X = (X - X.mean(0)) / X.std(0)
    return X


_OIL_URLS = [
    "https://web.archive.org/web/2019id_/http://staffwww.dcs.shef.ac.uk/people/N.Lawrence/resources/3PhData.tar.gz",
    "http://staffwww.dcs.shef.ac.uk/people/N.Lawrence/resources/3PhData.tar.gz",
]


def oil_flow(split="trn"):
    """Three-phase oil flow data (Bishop and James 1993; Appendix A).

    Returns X (1000 x 12 gamma-densitometer measurements) and labels in
    {0, 1, 2} = {'homogeneous', 'annular', 'laminar'} in the column order of the
    1-of-3 label file, for split in {'trn', 'vdn', 'tst'}. The data
    (copyright Markus Svensen 1996) are downloaded once from a public mirror.
    """
    folder = os.path.join(CACHE, "oil")
    tag = {"trn": "Trn", "vdn": "Vdn", "tst": "Tst"}[split]
    fx, fl = os.path.join(folder, "Data%s.txt" % tag), os.path.join(folder, "Data%sLbls.txt" % tag)
    if not (os.path.exists(fx) and os.path.exists(fl)):
        tgz = _download(_OIL_URLS, os.path.join(folder, "3PhData.tar.gz"))
        with tarfile.open(tgz) as t:
            for m in t.getmembers():
                if m.isfile() and "/" not in m.name and m.name.endswith(".txt"):
                    with open(os.path.join(folder, m.name), "wb") as f:
                        f.write(t.extractfile(m).read())
    X = np.loadtxt(fx)
    L = np.loadtxt(fl)
    return X, L.argmax(1)


OIL_CLASS_NAMES = ("homogeneous", "annular", "laminar")

_MNIST_MIRRORS = [
    "https://ossci-datasets.s3.amazonaws.com/mnist/",
    "https://storage.googleapis.com/cvdf-datasets/mnist/",
]


def _mnist_file(name):
    path = os.path.join(CACHE, "mnist", name)
    if not os.path.exists(path):
        _download([m + name for m in _MNIST_MIRRORS], path)
    return path


def mnist(train=True, digits=None):
    """MNIST digits (LeCun et al. 1998), the 'off-line digits' of Appendix A.

    Returns images as a float array in [0, 1] of shape (N, 784) and integer labels.
    digits, e.g. (2, 3, 4), keeps only those classes.
    """
    pre = "train" if train else "t10k"
    with gzip.open(_mnist_file(pre + "-images-idx3-ubyte.gz"), "rb") as f:
        _, n, r, c = struct.unpack(">IIII", f.read(16))
        X = np.frombuffer(f.read(), dtype=np.uint8).reshape(n, r * c)
    with gzip.open(_mnist_file(pre + "-labels-idx1-ubyte.gz"), "rb") as f:
        f.read(8)
        y = np.frombuffer(f.read(), dtype=np.uint8)
    if digits is not None:
        keep = np.isin(y, digits)
        X, y = X[keep], y[keep]
    return X.astype(float) / 255.0, y.astype(int)


def photographs():
    """Two public-domain colour photographs for the image-segmentation figure.

    Grace Hopper (U.S. Navy photograph, shipped with matplotlib) and, when
    scikit-image is installed, the astronaut Eileen Collins (NASA). Returns a
    list of (name, H x W x 3 float array in [0, 1]).
    """
    out = []
    path = os.path.join(matplotlib.get_data_path(), "sample_data", "grace_hopper.jpg")
    img = plt.imread(path)
    out.append(("grace_hopper", img[..., :3].astype(float) / (255.0 if img.dtype == np.uint8 else 1.0)))
    try:
        from skimage import data as skdata
        a = skdata.astronaut()
        out.append(("astronaut", a.astype(float) / 255.0))
    except Exception:
        pass
    return out


# --------------------------------------------------------------------------- maths helpers
def logsumexp(a, axis=None, keepdims=False):
    a = np.asarray(a, dtype=float)
    m = np.max(a, axis=axis, keepdims=True)
    m = np.where(np.isfinite(m), m, 0.0)
    s = np.log(np.sum(np.exp(a - m), axis=axis, keepdims=True)) + m
    return s if keepdims else np.squeeze(s, axis=axis) if axis is not None else s.item()


def gauss_logpdf(X, mu, Sigma):
    """log N(x | mu, Sigma) for each row of X."""
    X = np.atleast_2d(X)
    D = X.shape[1]
    L = np.linalg.cholesky(Sigma)
    Z = np.linalg.solve(L, (X - mu).T)
    return -0.5 * (D * np.log(2 * np.pi) + 2 * np.log(np.diag(L)).sum() + (Z ** 2).sum(0))


def ellipse_xy(mu, Sigma, nsd=1.0, n=200):
    """Points on the nsd-standard-deviation contour of N(mu, Sigma) in 2-D."""
    vals, vecs = np.linalg.eigh(Sigma)
    t = np.linspace(0, 2 * np.pi, n)
    circ = np.stack([np.cos(t), np.sin(t)])
    pts = vecs @ (np.sqrt(np.maximum(vals, 0))[:, None] * circ) * nsd
    return pts[0] + mu[0], pts[1] + mu[1]


def sigmoid(a):
    return 0.5 * (1.0 + np.tanh(0.5 * np.asarray(a, dtype=float)))


# --------------------------------------------------------------------------- driver
def run(figures, chapter, argv=None):
    """Command-line driver shared by the chapter modules.

    figures maps the book's figure number ('9.1') to a function f(out_dir) that
    builds and saves that figure. With no --fig option every figure is made.
    """
    p = argparse.ArgumentParser(description="Reproduce Bishop (2006) Chapter %d figures." % chapter)
    p.add_argument("--fig", nargs="*", help="figure numbers, e.g. 9.1 9.8 (default: all)")
    p.add_argument("--out", default=None, help="output folder (default: ../img)")
    p.add_argument("--list", action="store_true", help="list the reproducible figures")
    a = p.parse_args(argv)
    if a.list:
        for k, f in figures.items():
            print(k, "-", (f.__doc__ or "").strip().split("\n")[0])
        return
    keys = a.fig or list(figures)
    for k in keys:
        if k not in figures:
            sys.exit("unknown figure %s; use --list" % k)
        t0 = time.time()
        figures[k](a.out)
        print("  Figure %s done in %.1f s" % (k, time.time() - t0))