# critwire Learn: /learn/labs/s3.py
# Run it after the lab loader (/learn/data/postgis/learn_curriculum.sql), on the same database, in the environment
# /learn/labs/requirements.txt pins (the setup page, /learn/setup/, gives the uv commands):
#   python s3.py
# It connects with libpq's PG* variables (PGHOST, PGPORT, PGDATABASE, PGUSER, PGPASSWORD) and reads the lab data from
# LEARN_BASE (default https://critwire.com). Each block runs once, in page order, after py-connect and the boot setup;
# an untested block is commented out under its reason.

# %% py-connect
# q: py-connect
# status: ran
import psycopg


def connect():
    """A connection to the lab database: plain table names find learn_curriculum's tables, then PostGIS in public."""
    return psycopg.connect(options="-c search_path=learn_curriculum,public")

# %% boot: s4-cells
with connect() as conn:
    conn.execute("""
-- q: s4-cells
-- status: ran
-- engine: postgis
-- boot
DROP VIEW IF EXISTS s4_neighbours;
DROP TABLE IF EXISTS s4_cells;
CREATE TABLE s4_cells AS
WITH island AS (
  SELECT (part).geom
  FROM (SELECT ST_Dump(ST_Transform(geom, 3577)) AS part
        FROM admin1_shapes WHERE country_code = 'AU' AND name = 'Tasmania') AS parts
  ORDER BY ST_Area((part).geom) DESC
  LIMIT 1
), devils AS (
  SELECT ST_Transform(o.geom, 3577) AS geom
  FROM occurrences AS o JOIN species AS s ON s.id = o.species_id
  WHERE s.scientific_name = 'Sarcophilus harrisii'
), grid AS (
  SELECT e.edge_m, h.i, h.j, h.geom
  FROM island, (VALUES (5000), (10000)) AS e (edge_m), ST_HexagonGrid(e.edge_m, island.geom) AS h
  WHERE ST_Intersects(h.geom, island.geom)
)
SELECT row_number() OVER (ORDER BY g.edge_m, g.i, g.j)::integer AS id,
       g.edge_m, g.geom,
       ST_X(ST_Centroid(g.geom)) AS cx,
       ST_Y(ST_Centroid(g.geom)) AS cy,
       count(d.geom)::integer AS n,
       ln(1 + count(d.geom)::float8) AS y
FROM grid AS g LEFT JOIN devils AS d ON ST_Contains(g.geom, d.geom)
GROUP BY g.edge_m, g.i, g.j, g.geom;
ALTER TABLE s4_cells ADD PRIMARY KEY (id);
CREATE VIEW s4_neighbours AS
SELECT a.edge_m, a.id AS i, b.id AS j
FROM s4_cells AS a
JOIN s4_cells AS b
  ON b.edge_m = a.edge_m AND b.id <> a.id
 AND ST_DWithin(ST_Point(a.cx, a.cy, 3577), ST_Point(b.cx, b.cy, 3577), 2 * a.edge_m);
SELECT c.edge_m, count(*) AS cells, sum(c.n) AS records_in_cells,
       (SELECT count(*) FROM occurrences AS o JOIN species AS s ON s.id = o.species_id
        WHERE s.scientific_name = 'Sarcophilus harrisii') - sum(c.n) AS records_in_no_cell,
       max(c.n) AS max_n
FROM s4_cells AS c
GROUP BY c.edge_m
ORDER BY c.edge_m;
""")

# %% s3-pysal
# q: s3-pysal
# status: ran
# after: py-connect
import geopandas as gpd
import numpy as np
import shapely
from libpysal.weights import KNN, DistanceBand, Queen
from scipy.spatial import cKDTree

with connect() as conn:
    rows = conn.execute(
        "SELECT id, cx, cy, ST_AsBinary(geom) FROM s4_cells WHERE edge_m = 5000 ORDER BY id"
    ).fetchall()
ids = [r[0] for r in rows]
xy = np.array([(r[1], r[2]) for r in rows])
cells = gpd.GeoDataFrame({"id": ids}, geometry=shapely.from_wkb([bytes(r[3]) for r in rows]), crs=3577)

queen = Queen.from_dataframe(cells, ids="id")
band = DistanceBand(xy, threshold=10000, binary=True, ids=ids)
same = all(set(queen.neighbors[i]) == set(band.neighbors[i]) for i in ids)
print(f"queen = distance band: {same}  links: {queen.s0:.0f}  islands: {queen.islands}")
print("neighbours: cells | " + "  ".join(f"{k}: {c}" for k, c in queen.histogram))

queen.transform = "r"
rows_sum_to_one = np.allclose(np.asarray(queen.sparse.sum(axis=1)).ravel(), 1)
lone = next(i for i in ids if queen.cardinalities[i] == 1)
inner = next(i for i in ids if queen.cardinalities[i] == 6)
print(f"rows sum to 1: {rows_sum_to_one}  cell {lone} (1 neighbour): {queen.weights[lone]}  "
      f"cell {inner} (6): {[round(v, 4) for v in queen.weights[inner]]}")

knn = KNN(xy, k=4, ids=ids, silence_warnings=True)
one_way = sum(1 for i in ids for j in knn.neighbors[i] if i not in knn.neighbors[j])
d, _ = cKDTree(xy).query(xy, k=6)
tied = int(np.sum(np.round(d[:, 4], 2) == np.round(d[:, 5], 2)))
print(f"k = 4: links: {knn.s0:.0f}  one-way: {one_way}  cells with the 4th and 5th tied: {tied}")