"""Mineral deposits: where each ore, fuel and industrial mineral can be mined, by
geological rules, placed at an Earth-like share of the land. Everything a society needs to reach steam, rifled guns
and aluminium airships is somewhere; the good stuff (high-grade iron ore and the alloy metals) lies mostly deep in
mountains. Each deposit is a bit of `deposits` (uint32); `deposit_main` names the rarest one in a cell (a map layer).

A deposit's cells: the land cells (not sea, not under lakes) where its rule allows it, scored by the rule × a district
noise of its own (deposits cluster into mining districts). Half the `share` is taken province by province (≈ 1500 km
cells, each its share of its own land, best-scored first): every region has its local bog iron, coal pits and small
mines where its rocks allow. The rest goes to the best-scored land world-wide (the great districts)."""
from __future__ import annotations

import numpy as np

from scipy.spatial import cKDTree

from .noise import fbm, name_seed

# name, label, share of land (fraction). Shares are rough Earth-like extents of workable districts, not grades.
DEPOSITS = [
    ("bog_iron", "bog / laterite iron (low grade)", 0.030),
    ("ironstone", "ironstone (sedimentary iron)", 0.025),
    ("iron_high", "high-grade iron (magnetite, banded)", 0.006),
    ("coal", "coal (incl. coking)", 0.050),
    ("copper", "copper", 0.012),
    ("molybdenum", "molybdenum", 0.002),
    ("tin", "tin", 0.003),
    ("tungsten", "tungsten", 0.002),
    ("lead_zinc", "lead & zinc", 0.010),
    ("silver", "silver", 0.005),
    ("gold", "gold", 0.008),
    ("manganese", "manganese", 0.004),
    ("chromium", "chromium", 0.0015),
    ("nickel", "nickel", 0.003),
    ("cobalt", "cobalt", 0.001),
    ("vanadium", "vanadium & titanium", 0.002),
    ("platinum", "platinum metals", 0.0005),
    ("mercury", "mercury (cinnabar)", 0.0012),
    ("sulfur", "sulfur", 0.004),
    ("saltpetre", "saltpetre (nitrates)", 0.004),
    ("bauxite", "bauxite (aluminium)", 0.010),
    ("fluorite", "fluorite / cryolite (flux)", 0.0025),
    ("rock_salt", "rock salt & evaporites", 0.018),
    ("potash", "potash", 0.004),
    ("phosphate", "phosphate", 0.005),
    ("oil_gas", "oil & gas", 0.025),
    ("helium", "helium (with gas)", 0.0025),
    ("diamond", "diamonds (kimberlite)", 0.0004),
    ("magnesite", "magnesite / dolomite (magnesium)", 0.003),
    ("kaolin", "kaolin / fire clay", 0.010),
    ("glass_sand", "glass sand (quartz)", 0.015),
    ("graphite", "graphite", 0.002),
]
NAMES = [d[0] for d in DEPOSITS]
BIT = {n: 1 << i for i, n in enumerate(NAMES)}
LEGEND = ["—"] + [d[1] for d in DEPOSITS]
# the good stuff: mostly in deep mountains
DEEP = ("iron_high", "tungsten", "molybdenum", "chromium", "cobalt", "vanadium", "platinum", "manganese", "tin")


PROVINCE_KM = 1500.0   # spacing of the provinces that each take their local share
LOCAL = 0.5            # part of a deposit's share taken province by province
LOCAL_MIN = 0.25       # ... only on ground at least this good relative to the world's marginal district


def provinces(g, seed):
    """Province index of each cell: the nearest of random centres ≈ PROVINCE_KM apart."""
    n = max(1, int(round(4 * np.pi * g.radius_km ** 2 / PROVINCE_KM ** 2)))
    c = np.random.default_rng(name_seed(seed, "deposit-provinces")).normal(size=(n, 3))
    c /= np.linalg.norm(c, axis=1)[:, None]
    return cKDTree(c).query(g.xyz / np.linalg.norm(g.xyz, axis=1)[:, None])[1]


def _rank_select(score, share, land, prov):
    """`share` × land cells where score > 0: LOCAL of it the best of each province (its share of its own land, on
    ground ≥ LOCAL_MIN × the world's k-th best score), the rest the best world-wide."""
    ok = np.flatnonzero((score > 0) & land)
    k = min(len(ok), int(round(share * land.sum())))
    out = np.zeros(len(score), bool)
    if k == 0:
        return out
    cut = LOCAL_MIN * np.partition(score[ok], len(ok) - k)[len(ok) - k]     # the k-th best world-wide
    ok_l = ok[score[ok] >= cut]
    quota = np.floor(LOCAL * share * np.bincount(prov[land], minlength=prov.max() + 1) + 0.5).astype(np.int64)
    o = ok_l[np.lexsort((-score[ok_l], prov[ok_l]))]           # by province, best first
    p = prov[o]
    first = np.r_[0, np.flatnonzero(p[1:] != p[:-1]) + 1]
    rank = np.arange(len(o)) - np.repeat(first, np.diff(np.r_[first, len(o)]))
    local = o[rank < quota[p]]
    out[local[np.argsort(-score[local], kind="stable")[:k]]] = True
    rest = ok[~out[ok]]
    m = k - int(out.sum())
    if m > 0:
        out[rest[np.argpartition(-score[rest], m - 1)[:m]]] = True
    return out


def rules(f):
    """Each deposit's suitability (≥ 0, 0 = impossible) from the geology/climate fields in f."""
    from .crust import CRATON, LIP, POST_OROGEN, PRE_OROGEN, RIFT, VOLCANO
    from .environment import (LF_ARC, LF_MASSIF, LF_MOUNTAINS, LI_ANDESITE, LI_BASALT, LI_GRANITE, LI_LIMESTONE,
                              LI_METAMORPHIC, LI_SANDSTONE, LF_DUNES, LF_PLAIN)
    z, lit, age, lf, rel = f["z"], f["lit"], f["age"], f["lf"], f["relief"]
    t, p = f["t_mean"], f["p_ann"]
    b = lambda m: np.asarray(m, dtype=np.float64)
    deep = np.clip((z - 800.0) / 2000.0, 0, 1) * np.clip(rel / 1500.0, 0.3, 1.0)   # deep-mountain mining ground
    deep = np.maximum(deep, 0.6 * b(np.isin(lf, [LF_MOUNTAINS, LF_MASSIF])) * np.clip(z / 2500.0, 0, 1))
    oro = b(np.isin(age, [PRE_OROGEN, POST_OROGEN]))
    craton = b(age == CRATON)
    rift = b(age == RIFT)
    arc = b((lit == LI_ANDESITE) | (lf == LF_ARC))
    volc = b(age == VOLCANO) + arc
    gran = b(lit == LI_GRANITE)
    meta = b(lit == LI_METAMORPHIC)
    lime = b(lit == LI_LIMESTONE)
    sand = b(lit == LI_SANDSTONE)
    basalt = b((lit == LI_BASALT) | (age == LIP))
    sed = lime + sand
    low = b(z < 800)
    hot, wet = b(t > 20), b(p > 1200)
    arid = b(p < 250) * b(t > 10)
    suture = b(f["d_coll"] < 400) * (oro + meta)
    return {
        "bog_iron": b(p > 700) * low * (b(lf == LF_PLAIN) + 0.5) + hot * wet * low,
        "ironstone": sed * b(z < 1200),
        "iron_high": deep * (1.0 + craton + meta + oro),
        "coal": b(f["coal"]),
        "copper": arc * (1 + deep) + 0.4 * sand * rift + 0.3 * basalt,
        "molybdenum": arc * deep,
        "tin": (gran + meta) * oro * (0.3 + deep),
        "tungsten": (gran + meta) * (oro + 0.3) * deep,
        "lead_zinc": lime + 0.3 * sand * (1 + rift),
        "silver": arc + 0.3 * lime + 0.3 * deep * oro,
        "gold": (oro + craton * 0.7 + meta) * (0.3 + deep) + 0.3 * arc,
        "manganese": (sed + craton) * (0.2 + deep),
        "chromium": suture * deep + 0.3 * craton * basalt * deep,
        "nickel": basalt * hot * wet + (craton + suture) * deep,
        "cobalt": (basalt * hot * wet + craton * deep + arc * 0.3) * deep,
        "vanadium": (basalt + craton) * deep,
        "platinum": craton * (basalt + meta + 0.3) * deep,
        "mercury": volc * (0.5 + 0.5 * b(f["t_mean"] > -30)),
        "sulfur": volc + 0.5 * b(f["salt"]) + 0.3 * sed * arid,
        "saltpetre": arid * b(lf != LF_MOUNTAINS) + 0.3 * lime * b(p > 800),
        "bauxite": hot * wet * (gran + basalt + 0.3) * b(rel < 800) + 0.4 * lime * b(t > 14) * b(p > 600),
        "fluorite": (gran + lime * 0.4) * (rift + oro + 0.2),
        "rock_salt": b(f["salt"]) + 0.6 * b(f["endo"]) + 0.3 * sed * b(p < 500),
        "potash": b(f["salt"]) + 0.3 * b(f["endo"]) + 0.2 * sed * b(p < 400),
        "phosphate": sed * low * b(f["dist_ocean"] < 400) + 0.3 * lime,
        "oil_gas": sed * low * (1 + b(f["d_coll"] < 1200) + rift),
        "helium": sed * low * (craton + gran),
        "diamond": craton * (1 + deep),
        "magnesite": (lime + meta * 0.5) * (0.3 + deep) + 0.3 * b(f["salt"]),
        "kaolin": (gran + sand * 0.4) * b(p > 900) * b(t > 8),
        "glass_sand": sand * (1 + b(lf == LF_DUNES) + b(f["dist_ocean"] < 150)),
        "graphite": meta * (0.3 + deep),
    }


def place(g, f, seed):
    """(deposits uint32 bitmask, deposit_main uint8 legend index) over the grid."""
    land = np.asarray(f["land"], bool)
    R = rules(f)
    prov = provinces(g, seed)
    bits = np.zeros(g.n, np.uint32)
    main = np.zeros(g.n, np.uint8)
    best = np.full(g.n, np.inf)
    for i, (name, _, share) in enumerate(DEPOSITS):
        nz = 0.5 + 0.5 * fbm(g.xyz, name_seed(seed, "deposit:" + name), 3, 6.0)       # mining districts
        sel = _rank_select(R[name] * (0.25 + nz) ** 2, share, land, prov)
        bits[sel] |= np.uint32(1 << i)
        rarer = sel & (share < best)
        main[rarer], best[rarer] = i + 1, share
    return bits, main