"""Piper diagram in modern Python.

A line-for-line counterpart to piper.R in this directory: same geometry, same
tick and label placement, same default colours, so the two produce visually
equivalent figures from the same CSV. Requires only pandas and matplotlib.

Geometry (all units in plot space):
    cation ternary   (0,0) (100,0) (50, 86.603)
    anion ternary    (120,0) (220,0) (170, 86.603)
    diamond          bottom (110, 17.3206)  left  (60, 103.9236)
                     right  (160, 103.9236) top   (110, 190.5266)
The diamond vertices are derived as the intersections of the two ternaries'
projection lines at +/- tan(60 deg).

Usage:
    import pandas as pd, piper
    raw = pd.read_csv("piper_example.csv")
    df  = piper.to_percent(raw)
    fig, ax = piper.plot(piper.transform(df), groups=df["Formation"])
    fig.savefig("piper.png", dpi=120)
"""

from __future__ import annotations

import json
import math

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.lines import Line2D
from matplotlib.ticker import MaxNLocator

GRAD = math.tan(math.radians(60))   # 1.7320508
APEX = 50 * GRAD                    # 86.60254
OFFSET = 120                        # x shift of the anion ternary

CATIONS = ["Ca", "Mg", "Na", "K"]
ANIONS = ["Cl", "SO4", "CO3", "HCO3"]

# The ggplot2 hue palette for three groups, so the R and Python figures use
# the same colours.
GGPLOT_HUE3 = ["#F8766D", "#00BA38", "#619CFF"]


# ---- data preparation ------------------------------------------------------

def to_percent(df: pd.DataFrame, cations=None, anions=None) -> pd.DataFrame:
    """Convert milliequivalent concentrations to percent of cation/anion sums."""
    cations = cations or CATIONS
    anions = anions or ANIONS
    missing = [c for c in (*cations, *anions) if c not in df.columns]
    if missing:
        raise KeyError(f"missing columns: {missing}")

    out = df.copy()
    cat_sum = out[cations].sum(axis=1)
    an_sum = out[anions].sum(axis=1)
    if (cat_sum == 0).any() or (an_sum == 0).any():
        raise ValueError("a sample has a zero cation or anion sum")

    out[cations] = 100 * out[cations].div(cat_sum, axis=0)
    out[anions] = 100 * out[anions].div(an_sum, axis=0)
    return out


def transform(df: pd.DataFrame) -> pd.DataFrame:
    """Project percentage data into Piper plot coordinates."""
    ca, mg, cl, so4 = df["Ca"], df["Mg"], df["Cl"], df["SO4"]

    cx = 100 * (1 - ca / 100 - mg / 200)
    cy = mg * GRAD / 2

    ax_ = OFFSET + cl + 0.5 * so4
    ay = so4 * GRAD / 2

    # diamond: up-slope from the cation point meets the down-slope from the anion point
    dx = (GRAD * cx + GRAD * ax_ + ay - cy) / (2 * GRAD)
    dy = GRAD * (dx - cx) + cy

    return pd.DataFrame(
        {"cation_x": cx, "cation_y": cy, "anion_x": ax_, "anion_y": ay,
         "diamond_x": dx, "diamond_y": dy}
    )


# ---- plot furniture --------------------------------------------------------

def _ternary_grid(dx: float = 0.0):
    """Internal 20/40/60/80 grid for one ternary, as (x, y, xend, yend) tuples."""
    seg = []
    for t in (20, 40, 60, 80):
        h = t * GRAD / 2
        seg.append((dx + t / 2, h, dx + 100 - t / 2, h))                       # left -> right edge
        seg.append((dx + t, 0, dx + t + (100 - t) / 2, (100 - t) * GRAD / 2))  # base -> right edge
        seg.append((dx + t, 0, dx + t / 2, t * GRAD / 2))                      # base -> left edge
    return seg


def _diamond_grid():
    seg = []
    for t in (20, 40, 60, 80):
        h = t * GRAD / 2
        seg.append((110 - t / 2, 17.3206 + h, 160 - t / 2, 103.9236 + h))
        seg.append((110 + t / 2, 17.3206 + h, 60 + t / 2, 103.9236 + h))
    return seg


def _draw_frame(ax):
    for x, y, xe, ye in (*_ternary_grid(0), *_ternary_grid(OFFSET), *_diamond_grid()):
        ax.plot([x, xe], [y, ye], ls=(0, (5, 4)), lw=0.6, color="0.6", zorder=1)

    tri_l = ([0, 100, 50, 0], [0, 0, APEX, 0])
    tri_r = ([x + OFFSET for x in (0, 100, 50, 0)], [0, 0, APEX, 0])
    diam = ([110, 60, 110, 160, 110], [17.3206, 103.9236, 190.5266, 103.9236, 17.3206])
    for xs, ys in (tri_l, tri_r, diam):
        ax.plot(xs, ys, lw=1.2, color="black", zorder=2)


def _draw_ticks(ax):
    """Tick labels. The two outer near-vertical edges (Mg, SO4) read
    horizontally; every other edge is rotated onto its own grid family."""
    t = [20, 40, 60, 80]
    rt = t[::-1]
    pad = 4.6
    dx, dy = pad * 0.87, pad * 0.5
    common = dict(fontsize=7, color="0.2", ha="center", va="center", zorder=3)

    rows = []
    # cation ternary
    rows += [(v / 2 - dx - 1.6, v * GRAD / 2 + dy, lab, 0) for v, lab in zip(t, t)]
    rows += [(v, -pad, lab, 60) for v, lab in zip(t, rt)]
    rows += [(100 - v / 2 + dx, v * GRAD / 2 + dy, lab, -60) for v, lab in zip(t, rt)]
    # anion ternary
    rows += [(OFFSET + v / 2 - dx, v * GRAD / 2 + dy, lab, 60) for v, lab in zip(t, rt)]
    rows += [(OFFSET + v, -pad, lab, 60) for v, lab in zip(t, t)]
    rows += [(OFFSET + 100 - v / 2 + dx + 1.6, v * GRAD / 2 + dy, lab, 0) for v, lab in zip(t, t)]
    # diamond
    rows += [(60 + v / 2 - dx, 103.9236 + v * GRAD / 2 + dy, lab, 60) for v, lab in zip(t, t)]
    rows += [(160 - v / 2 + dx, 103.9236 + v * GRAD / 2 + dy, lab, -60) for v, lab in zip(t, t)]

    for x, y, lab, rot in rows:
        ax.text(x, y, str(lab), rotation=rot, **common)


def _draw_titles(ax):
    # Positions are computed on each edge's outward normal. Identical to piper.R.
    ann = [
        (76.2, 158.1, r"$\mathrm{SO_4^{2-} + Cl^-}$", 60),
        (143.8, 158.1, r"$\mathrm{Ca^{2+} + Mg^{2+}}$", -60),
        (21.5, 59.2, r"$\mathrm{Mg^{2+}}$", 60),
        (50, -13, r"$\mathrm{Ca^{2+}}$", 0),
        (78.5, 59.2, r"$\mathrm{Na^+ + K^+}$", -60),
        (OFFSET + 21.5, 59.2, r"$\mathrm{Alkalinity\ as\ HCO_3^-}$", 60),
        (OFFSET + 50, -13, r"$\mathrm{Cl^-}$", 0),
        (OFFSET + 78.5, 59.2, r"$\mathrm{SO_4^{2-}}$", -60),
    ]
    for x, y, lab, rot in ann:
        ax.text(x, y, lab, rotation=rot, rotation_mode="anchor",
                fontsize=9, ha="center", va="center", zorder=3)


# ---- location map ----------------------------------------------------------

M_PER_DEG_LAT = 111132


def basin_polygon(path):
    """Read a single-polygon GeoJSON into (lons, lats).

    Avoids a GDAL/geopandas dependency for the one polygon this needs.
    """
    with open(path) as f:
        g = json.load(f)
    geom = g["features"][0]["geometry"]
    ring = geom["coordinates"][0] if geom["type"] == "Polygon" else geom["coordinates"][0][0]
    return [p[0] for p in ring], [p[1] for p in ring]


def basin_map(basin, wells, groups=None, km=5, colors=None, ax=None,
              legend_title="Formation"):
    """Location map with graticule labels, a north arrow and a scale bar."""
    lons, lats = basin
    lat_mid = (min(lats) + max(lats)) / 2
    # One degree of longitude shortens with latitude; fixing the aspect this
    # way keeps distances true without reprojecting.
    aspect = 1 / math.cos(math.radians(lat_mid))

    if ax is None:
        _, ax = plt.subplots(figsize=(5.2, 6.2))

    ax.fill(lons, lats, facecolor="#b8b8b8", edgecolor="0.35", linewidth=0.8, zorder=1)

    palette = colors or GGPLOT_HUE3
    if groups is None:
        ax.scatter(wells["Longitude"], wells["Latitude"], s=16, color="black", zorder=3)
    else:
        g = pd.Series(list(groups))
        for i, lv in enumerate(sorted(pd.unique(g))):
            m = (g == lv).to_numpy()
            ax.scatter(wells["Longitude"][m], wells["Latitude"][m], s=16,
                       color=palette[i % len(palette)], linewidths=0, zorder=3)

    xr = (min(lons), max(lons))
    yr = (min(lats), max(lats))
    padx = (xr[1] - xr[0]) * 0.10
    pady = (yr[1] - yr[0]) * 0.10
    ax.set_xlim(xr[0] - padx, xr[1] + padx)
    ax.set_ylim(yr[0] - pady, yr[1] + pady)

    # scale bar, bottom-left
    deg_per_km = 1000 / (M_PER_DEG_LAT * math.cos(math.radians(lat_mid)))
    bx0 = xr[0] - padx * 0.3
    bx1 = bx0 + km * deg_per_km
    by = yr[0] - pady * 0.45
    tick = pady * 0.10
    ax.plot([bx0, bx1], [by, by], color="black", lw=1.2, zorder=4)
    for bx in (bx0, bx1):
        ax.plot([bx, bx], [by - tick, by + tick], color="black", lw=1.2, zorder=4)
    ax.text((bx0 + bx1) / 2, by + tick * 2.4, f"{km} km",
            ha="center", va="bottom", fontsize=8, zorder=4)

    # north arrow, top-right
    nx = xr[1] + padx * 0.45
    ny0 = yr[1] - (yr[1] - yr[0]) * 0.13
    ny1 = yr[1] - (yr[1] - yr[0]) * 0.02
    ax.annotate("", xy=(nx, ny1), xytext=(nx, ny0),
                arrowprops=dict(arrowstyle="-|>", color="black", lw=1.1))
    ax.text(nx, ny1 + (yr[1] - yr[0]) * 0.035, "N",
            ha="center", va="bottom", fontsize=9, fontweight="bold")

    ax.set_aspect(aspect)
    # Degree labels are wide; cap the tick count so they cannot collide when
    # the map is rendered as a narrow panel.
    ax.xaxis.set_major_locator(MaxNLocator(nbins=4, prune="both"))
    ax.yaxis.set_major_locator(MaxNLocator(nbins=5))
    ax.xaxis.set_major_formatter(lambda v, _: f"{abs(v):.2f}°W")
    ax.yaxis.set_major_formatter(lambda v, _: f"{abs(v):.2f}°N")
    ax.tick_params(labelsize=7, colors="0.25")
    ax.grid(True, color="0.88", lw=0.4)
    ax.set_axisbelow(True)
    for s in ax.spines.values():
        s.set_color("0.3")
    return ax


def plot(pts: pd.DataFrame, groups=None, legend_title="Formation",
         colors=None, figsize=(9, 7.2)):
    """Draw a full Piper diagram on a new figure. `pts` is from transform()."""
    fig, ax = plt.subplots(figsize=figsize)
    draw(pts, ax, groups=groups, legend_title=legend_title, colors=colors)
    fig.tight_layout()
    return fig, ax


def draw(pts: pd.DataFrame, ax, groups=None, legend_title="Formation",
         colors=None):
    """Draw a Piper diagram onto an existing axis."""
    _draw_frame(ax)
    _draw_ticks(ax)
    _draw_titles(ax)

    xs = pd.concat([pts["cation_x"], pts["anion_x"], pts["diamond_x"]])
    ys = pd.concat([pts["cation_y"], pts["anion_y"], pts["diamond_y"]])

    if groups is None:
        ax.scatter(xs, ys, s=22, color="black", alpha=0.9, zorder=4)
    else:
        g = pd.concat([pd.Series(list(groups))] * 3, ignore_index=True)
        levels = sorted(pd.unique(pd.Series(list(groups))))
        palette = colors or GGPLOT_HUE3
        handles = []
        for i, lv in enumerate(levels):
            m = (g == lv).to_numpy()
            c = palette[i % len(palette)]
            ax.scatter(xs[m], ys[m], s=22, color=c, alpha=0.9,
                       linewidths=0, zorder=4)
            handles.append(Line2D([], [], marker="o", ls="", color=c,
                                  markersize=5, label=str(lv)))
        leg = ax.legend(handles=handles, title=legend_title, loc="upper left",
                        frameon=False, fontsize=7, title_fontsize=8,
                        handletextpad=0.4, borderpad=0.2)
        leg.get_title().set_fontweight("bold")
        leg._legend_box.align = "left"

    ax.set_aspect("equal")
    ax.set_xlim(-16, 236)
    ax.set_ylim(-22, 206)
    ax.axis("off")
    return ax