"""
# multigaussian kriging

multigaussian kriging estimates a full distribution at every target from a single variogram. the samples take a
normal-score transform, and the scores are simple-kriged about 0. if the scores are multivariate gaussian, the unknown
score at a target is gaussian with the kriged mean and the kriging variance. back-transformed into data units, that
distribution gives the E-type mean, its variance, quantiles and the probability above any cutoff.
"""

# %% [hidden]
import sys
from pathlib import Path

HERE = Path(__file__).parent if "__file__" in globals() else Path.cwd()
sys.path.insert(0, str(HERE.parents[1]))

# %%
import boitata as bt
import matplotlib.pyplot as plt
import numpy as np
from common import ACCENT, INK, map_axes, save
from matplotlib.colors import PowerNorm

samples = bt.datasets.walker_lake()
truth = bt.datasets.walker_lake_exhaustive()["V"].reshape(300, 260)
xy, v = samples.coords, samples["V"]
azimuth = 170.0

# %% [markdown]
# ## scores and their variogram
#
# the declustering weights set the transform, so the E-type follows the declustered histogram instead of the clustered
# high-grade samples. some samples tie at 0 ppm; `despike=True` in `fit` breaks the ties as
# [`despike`](../../03-exploratory-analysis/06-despiking/README.md) does, so they get distinct scores. the scores'
# variogram, along N170° (the major axis found in
# [variogram fitting](../../05-spatial-continuity/02-variogram-fitting/README.md)) and N260°, is scaled to a unit sill,
# which makes the kriging variance the variance of the conditional gaussian.

# %%
weights = bt.cell_declustering(xy, v, sizes=np.arange(2.5, 102.5, 2.5)).weights
print(f"{np.sum(v == 0)} samples at 0 ppm")
y = bt.NormalScore().fit_transform(bt.despike(xy, v), weights=weights)

lag, max_lag = 10.0, 120.0
major = bt.experimental_variogram(xy, y, lag, max_lag, azimuth=azimuth).fit("spherical")
minor = bt.experimental_variogram(xy, y, lag, max_lag, azimuth=azimuth + 90).fit("spherical")
a_major = major.structures[0].range
gaussian = major.standardized().with_anisotropy(
    (azimuth, 0, 0), (min(minor.structures[0].range / a_major, 1.0), 1.0)
)
print(gaussian)

# %% [markdown]
# ## distributions at points
#
# `predict` returns an `IndicatorSummary`, the same summary multiple indicator kriging gives: here the E-type mean, the
# 10th and 90th percentiles and the probability above 500 ppm on a 5 m grid. the tails bound the back-transform to 0 and
# the largest sample.

# %%
grid = bt.BlockModel(origin=(0.5, 0.5), size=(5, 5), count=(52, 60))
mg = bt.MultigaussianKriging(gaussian, bt.Search(radius=100, max_samples=24), tails=(0.0, v.max()))
mg.fit(samples, "V", weights=weights, despike=True)
summary = mg.predict(grid, cutoffs=[500.0], quantiles=[0.1, 0.9])
etype, (p10, p90), p500 = summary.mean, summary.quantile_values.T, summary.probability_above[:, 0]

nodes = grid.coords.astype(int)
true_at_nodes = truth[nodes[:, 1] - 1, nodes[:, 0] - 1]
print(f"mean E-type {etype.mean():.0f} ppm, true {true_at_nodes.mean():.0f}")
print(f"correlation with the truth {bt.compare(etype, true_at_nodes)['correlation']:.2f}")
print(
    f"truth inside the 80% interval at {np.mean((true_at_nodes >= p10) & (true_at_nodes <= p90)):.0%} of nodes"
)
print(f"expected area above 500 ppm {p500.mean():.1%}, true {np.mean(true_at_nodes > 500):.1%}")

# %% [markdown]
# the E-type is smooth like kriging. the probability map says where grade above 500 ppm is likely, and the interval
# width says where the samples leave the grade uncertain.

# %%
extent = (0.5, 260.5, 0.5, 300.5)
norm = PowerNorm(0.5, vmin=0, vmax=1500)
fig, axes = plt.subplots(1, 4, figsize=(15, 4.4), layout="constrained")
for ax, image, title in ((axes[0], true_at_nodes, "True V"), (axes[1], etype, "E-type")):
    im = ax.imshow(grid.grid(image)[0], origin="lower", extent=extent, norm=norm)
    map_axes(ax, title)
fig.colorbar(im, ax=axes[:2], shrink=0.8, label="V (ppm)")
width = axes[2].imshow(grid.grid(p90 - p10)[0], origin="lower", extent=extent, cmap="Greys")
map_axes(axes[2], "P90 − P10")
fig.colorbar(width, ax=axes[2], shrink=0.8, label="ppm")
p = axes[3].imshow(grid.grid(p500)[0], origin="lower", extent=extent, vmin=0, vmax=1)
map_axes(axes[3], "P(V > 500 ppm)")
axes[3].scatter(xy[:, 0], xy[:, 1], s=1, color=INK, linewidths=0)
fig.colorbar(p, ax=axes[3], shrink=0.8, label="probability")
save(fig, "maps")

# %% [markdown]
# ## panels
#
# with `discretization`, each 20 m panel takes the average of the distributions at 4 × 4 points within it: the
# distribution of the point grades inside the panel. its probability above 500 ppm is then the expected share of the
# panel above the cutoff, compared here with the share of the 400 exhaustive values in each panel.

# %%
panels = bt.BlockModel(origin=(0.0, 0.0), size=(20, 20), count=(13, 15))
blocks = mg.predict(panels, cutoffs=[500.0], discretization=(4, 4, 1))
cells = truth.reshape(15, 20, 13, 20)
true_share = (cells > 500).mean(axis=(1, 3)).ravel()
true_mean = cells.mean(axis=(1, 3)).ravel()
print(f"panel means: r {bt.compare(blocks.mean, true_mean)['correlation']:.2f}")
print(
    f"panel shares above 500 ppm: r {bt.compare(blocks.probability_above[:, 0], true_share)['correlation']:.2f}"
)

fig, (a, b) = plt.subplots(1, 2, figsize=(9, 4), layout="constrained")
bt.plot.scatter(true_at_nodes, etype, ax=a, s=4, color=ACCENT, alpha=0.4)
a.set(xlim=(0, 1600), ylim=(0, 1600), xlabel="True V (ppm)", ylabel="E-type (ppm)", title="Points")
bt.plot.scatter(true_share, blocks.probability_above[:, 0], ax=b, s=10, color=ACCENT)
b.set(xlabel="True share above 500 ppm", ylabel="P(V > 500 ppm)", title="20 m panels")
for ax in (a, b):
    ax.set_aspect("equal")
    ax.legend(loc="upper left")
save(fig, "checks")

# %% [markdown]
# ## checks
#
# over the whole area the E-type averages to about the declustered mean. cross-validation re-estimates every sample's
# distribution from the others: `accuracy(p)` is the share of samples inside their central `p` interval, close to `p`
# when the distributions have the right spread.

# %%
print(f"mean E-type {etype.mean():.0f} ppm, declustered sample mean {np.average(v, weights=weights):.0f} ppm")
cv = mg.cross_validate()
print(f"cross-validation: ME {cv.mean_error:.1f}  RMSE {cv.rmse:.1f}  goodness {cv.goodness:.2f}")
for p, a in zip([0.5, 0.8, 0.9], cv.accuracy([0.5, 0.8, 0.9])):
    print(f"inside the central {p:.0%} interval: {a:.0%}")

# %% [markdown]
# ## compared with multiple indicator kriging
#
# [multiple indicator kriging](../../06-kriging/11-multiple-indicator-kriging/README.md) builds the distribution from
# indicators kriged at a set of thresholds, one variogram each, then corrects order relations and fills in between the
# thresholds and in the tails. it lets each grade range have its own continuity, such as high grades less continuous
# than low ones. multigaussian kriging needs one variogram and no order-relation correction and gives smooth
# distributions at any cutoff. it assumes the gaussian model, though: high and low scores are equally continuous, and
# extremes are disconnected. where the data show connected high grades, indicators follow them better.