# critwire Learn: /learn/labs/s2.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 s2.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;
""")

# %% s2-kde
# q: s2-kde
# status: ran
# after: py-connect
import numpy as np
from sklearn.neighbors import KernelDensity

with connect() as conn:
    points = np.array(conn.execute(
        """SELECT ST_X(g), ST_Y(g)
           FROM (SELECT ST_Transform(o.geom, 3577) AS g
                 FROM occurrences AS o JOIN species AS s ON s.id = o.species_id
                 WHERE s.scientific_name = 'Sarcophilus harrisii') AS devils
           ORDER BY 1, 2"""
    ).fetchall())
    cells = conn.execute(
        "SELECT id, n, cx, cy, ST_Area(geom) FROM s4_cells WHERE edge_m = 5000 ORDER BY id"
    ).fetchall()

kde = KernelDensity(bandwidth=5000, kernel="gaussian").fit(points)
centroids = np.array([(c[2], c[3]) for c in cells])
per_km2 = np.exp(kde.score_samples(centroids)) * len(points) * 1e6
cell_km2 = cells[0][4] / 1e6

print(f"records: {len(points)}  cells: {len(cells)}  cell area: {cell_km2:.2f} km2")
print(f"density x area, summed over the cells: {(per_km2 * cell_km2).sum():.0f} records")
print("densest cells: id | records in cell | records per km2")
for k in np.argsort(-per_km2, kind="stable")[:5]:
    print(f"{cells[k][0]} | {cells[k][1]} | {per_km2[k]:.3f}")