import functools

import numpy as np

from worldhistory.world import World, neighbours_from_ids

DEFAULTS = dict(ocean=False, elevation_m=100.0, T_mean=15.0, P_ann=1000.0, holdridge=20, landform=1, river=False,
                strahler=0, lake=False, ice=0, coal_potential=False, iron_potential=False, vent_potential=0.0,
                seabed_type=0, lithology=1, ground=0, gravity_g=1.0, o2_fraction=0.21, pressure_bar=1.0,
                bottom_temp_c=2.0, plate=0)


def fields(n, **over):
    out = {}
    for k, v in {**DEFAULTS, **over}.items():
        a = np.asarray(v)
        out[k] = np.broadcast_to(a, (n,)).copy() if a.ndim == 0 else a.copy()
    return out


def line_world(n, spacing_km=100.0, radius_km=6371.0, **over):
    """Cells on the equator in a row, each linked to the next: 1-D worlds with hand-checkable answers."""
    lon = np.arange(n) * np.degrees(spacing_km / radius_km)
    lat = np.zeros(n)
    xyz = np.stack([np.cos(np.radians(lon)), np.sin(np.radians(lon)), np.zeros(n)], 1)
    nb = np.full((n, 6), -1, np.int64)
    nb[1:, 0] = np.arange(n - 1)
    nb[:-1, 1] = np.arange(1, n)
    return World(nb=nb, lat=lat, lon=lon, xyz=xyz, area=np.full(n, spacing_km ** 2), radius_km=radius_km,
                 base=fields(n, **over))


@functools.lru_cache(maxsize=4)
def _h3_grid(res):
    import h3.api.basic_int as h3
    ids = sorted(c for r0 in h3.get_res0_cells() for c in h3.cell_to_children(r0, res))
    ll = np.array([h3.cell_to_latlng(c) for c in ids])
    area = np.array([h3.cell_area(c, unit="rads^2") for c in ids])
    return np.array(ids, np.uint64), ll, area


def globe_world(res=1, radius_km=6371.0):
    """A whole H3 globe (res 1 = 842 cells) with default fields; tests edit world.base afterwards."""
    ids, ll, area = _h3_grid(res)
    lat, lon = ll[:, 0].copy(), ll[:, 1].copy()
    la, lo = np.radians(lat), np.radians(lon)
    xyz = np.stack([np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)], 1)
    return World(nb=neighbours_from_ids(ids), lat=lat, lon=lon, xyz=xyz, area=area * radius_km ** 2,
                 radius_km=radius_km, base=fields(len(ids)))


from worldhistory.config import parse_history, parse_race


def make_race(id="a", **over):
    """A land race with density 1/km² at q=1, growth 0.1/step, no Allee, founder 1, no mobility; sections override."""
    d = {"id": id,
         "demography": {"density": 1.0, "growth": 0.1, "allee": 0.0, "founder": 1.0, "mobility": 0.0},
         "habitat": {"realm": "land", "terms": [{"p": "land", "w": 1.0}]}}
    for k, v in over.items():
        d[k] = {**d.get(k, {}), **v} if isinstance(v, dict) and isinstance(d.get(k), dict) else v
    return parse_race(d, f"test:{id}")


def make_history(**over):
    d = {"run": {"years": 100, "step": 10, "snapshot_every": 50, "seed": 1}}
    d.update(over)
    return parse_history(d, "test:history")