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

# %% s5-gistar-pysal
# q: s5-gistar-pysal
# status: ran
# after: py-connect
import esda
import numpy as np
from libpysal.weights import DistanceBand
from scipy.stats import norm

with connect() as conn:
    rows = conn.execute("SELECT id, y, cx, cy FROM s4_cells WHERE edge_m = 5000 ORDER BY id").fetchall()
ids = [r[0] for r in rows]
y = np.array([r[1] for r in rows])
w = DistanceBand(np.array([(r[2], r[3]) for r in rows]), threshold=10000, binary=True, ids=ids)

g = esda.G_Local(y, w, transform="B", star=True, permutations=0)
p = 2 * norm.sf(np.abs(g.Zs))
cut = esda.fdr(p, 0.05)
print("cells | hot_05 | cold_05 | hot_fdr | cold_fdr | max_z | min_z")
print(f"{len(y)} | {np.sum((g.Zs > 0) & (p <= 0.05))} | {np.sum((g.Zs < 0) & (p <= 0.05))} | "
      f"{np.sum((g.Zs > 0) & (p <= cut))} | {np.sum((g.Zs < 0) & (p <= cut))} | {g.Zs.max():.6f} | {g.Zs.min():.6f}")

# %% s5-lisa
# q: s5-lisa
# status: ran
# after: py-connect
import esda
import numpy as np
from libpysal.weights import DistanceBand

with connect() as conn:
    rows = conn.execute("SELECT id, y, cx, cy FROM s4_cells WHERE edge_m = 5000 ORDER BY id").fetchall()
ids = [r[0] for r in rows]
y = np.array([r[1] for r in rows])
w = DistanceBand(np.array([(r[2], r[3]) for r in rows]), threshold=10000, binary=True, ids=ids)

lisa = esda.Moran_Local(y, w, transformation="r", permutations=999, seed=2026)
cut = esda.fdr(lisa.p_sim, 0.05)
print(f"false discovery rate cut: p <= {cut:.4f}")
print("quadrant | cells | p <= 0.05 | after FDR")
for q, name in [(1, "high-high"), (2, "low-high"), (3, "low-low"), (4, "high-low")]:
    in_q = lisa.q == q
    print(f"{name} | {in_q.sum()} | {np.sum(in_q & (lisa.p_sim <= 0.05))} | {np.sum(in_q & (lisa.p_sim <= cut))}")