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

# %% e2-file
# q: e2-file
# status: ran
HEADER = ["gbifID", "datasetKey", "species", "countryCode", "locality", "decimalLatitude", "decimalLongitude",
          "coordinateUncertaintyInMeters", "year", "month", "taxonKey", "speciesKey", "basisOfRecord", "license",
          "recordedBy"]
ROWS = [
    ["1000000001", "00000000-0000-0000-0000-000000000001", "Sarcophilus harrisii", "AU",
     r"Cradle Mountain\Dove Lake track", "-41.6512", "145.9601", "30", "2019", "3", "2435451", "2435451",
     "HUMAN_OBSERVATION", "CC_BY_4_0", 'A. "Sam" Lee'],
    ["1000000002", "00000000-0000-0000-0000-000000000001", "Smutsia temminckii", "ZA",
     r"Pafuri\north bank", "-22.4310", "31.2120", "", "2021", "", "5786560", "5219630",
     "HUMAN_OBSERVATION", "CC0_1_0", ""],
    ["1000000003", "00000000-0000-0000-0000-000000000002", "Tapirus bairdii", "CR",
     "Sirena station", "8.4800", "-83.5900", "250", "", "", "2440897", "8395500",
     "MACHINE_OBSERVATION", "CC_BY_4_0", r"\N"],
]


def gbif_file(rows):
    """Records in GBIF's SIMPLE_CSV form: a header line, tab-separated, no quoting, an empty field for "no value"."""
    return ("\n".join("\t".join(row) for row in [HEADER, *rows]) + "\n").encode()


GBIF_FILE = gbif_file(ROWS)  # three made-up records

# Every column text: the file is read exactly as delivered, and typed in SQL afterwards.
CREATE_STAGING = "CREATE TEMP TABLE gbif_staging (" + ", ".join(f'"{h}" text' for h in HEADER) + ")"
# CSV mode with a quote character that never occurs, so backslashes and double quotes stay data; '' is NULL.
COPY_STAGING = ("COPY gbif_staging FROM STDIN "
                "(FORMAT csv, DELIMITER E'\\t', QUOTE E'\\b', NULL '', HEADER true)")


def stage(cur, data, copy_sql=COPY_STAGING):
    """Load the file's bytes into gbif_staging, streamed: COPY reads what it is sent, in pieces of any size."""
    cur.execute(CREATE_STAGING)
    with cur.copy(copy_sql) as copy:
        for start in range(0, len(data), 1 << 20):
            copy.write(data[start:start + (1 << 20)])

# %% e2-copy
# q: e2-copy
# status: ran
# after: py-connect, e2-file
TEXT_FORMAT = "COPY gbif_staging FROM STDIN (FORMAT text, HEADER true)"  # the default format: wrong for this file


def show(value):
    """NULL, '' and a line break made visible."""
    return "NULL" if value is None else "''" if value == "" else value.replace("\n", "↵")


with connect() as conn, conn.cursor() as cur:
    read = {}
    for name, sql in [("csv form", COPY_STAGING), ("text format", TEXT_FORMAT)]:
        stage(cur, GBIF_FILE, sql)
        read[name] = cur.execute('SELECT "gbifID", locality, "recordedBy" FROM gbif_staging ORDER BY "gbifID"').fetchall()
        cur.execute("DROP TABLE gbif_staging")
    conn.rollback()

sent = {row[0]: (row[4], row[14] or None) for row in ROWS}
print("COPY | gbifID | locality | recordedBy | read")
for name, rows in read.items():
    for gbif_id, locality, recorded_by in rows:
        same = (locality, recorded_by) == sent[gbif_id]
        print(f"{name} | {gbif_id} | {show(locality)} | {show(recorded_by)} | {'as sent' if same else 'CHANGED'}")

# %% e2-copy-fails
# q: e2-copy-fails
# status: ran
# after: py-connect, e2-file
import re

import psycopg

# Three ways a delivered file breaks COPY itself, before any type is involved: each damages one record of GBIF_FILE.
BROKEN = [
    ("a tab inside a field", GBIF_FILE.replace(b"Sirena station", b"Sirena\tstation")),
    ("a NUL byte", GBIF_FILE.replace(b"Pafuri", b"Paf\x00uri")),
    ("Latin-1, not UTF-8", GBIF_FILE.replace('A. "Sam" Lee'.encode(), 'A. "Sam" Lée'.encode("latin-1"))),
]

with connect() as conn, conn.cursor() as cur:
    cur.execute(CREATE_STAGING)
    print("file | SQLSTATE | line | error")
    for name, data in BROKEN:
        try:
            with conn.transaction():  # a savepoint: the failed COPY is undone, the table stays
                with cur.copy(COPY_STAGING) as copy:
                    copy.write(data)
        except psycopg.Error as error:
            line = re.search(r"\bline (\d+)", error.diag.context or "")  # COPY's context names the file's line
            print(f"{name} | {error.sqlstate} | {line[1] if line else '?'} | {error.diag.message_primary}")
        else:
            raise SystemExit(f"{name}: COPY accepted it")
    failed = cur.execute("SELECT count(*) FROM gbif_staging").fetchone()[0]
    with cur.copy(COPY_STAGING) as copy:
        copy.write(GBIF_FILE)
    delivered = cur.execute("SELECT count(*) FROM gbif_staging").fetchone()[0]
    conn.rollback()

print(f"rows staged: {failed} after the three failed COPYs, {delivered} from the file as delivered")

# %% e2-typed
# q: e2-typed
# status: ran
# after: py-connect, e2-file
TYPED = """
CREATE TEMP TABLE occurrences_new (
  id                       bigint PRIMARY KEY,
  species_id               integer NOT NULL,
  dataset_key              uuid NOT NULL,
  license                  text NOT NULL,
  basis_of_record          text NOT NULL,
  longitude                double precision NOT NULL CHECK (longitude BETWEEN -180 AND 180),
  latitude                 double precision NOT NULL CHECK (latitude BETWEEN -90 AND 90),
  coordinate_uncertainty_m double precision,
  country_code_gbif        text,
  year                     integer,
  month                    integer CHECK (month BETWEEN 1 AND 12),
  geom                     geometry(Point, 4326) GENERATED ALWAYS AS (ST_Point(longitude, latitude, 4326)) STORED
)"""

# The species whose GBIF key is the record's speciesKey, else its taxonKey (a synonym keeps our key in taxonKey).
# The license: this file's two codes renamed, for illustration only. The lab sample keeps a record only if its
# effective license, the more restrictive of its own and its dataset's, is open; that needs the dataset's too.
INSERT_TYPED = """
INSERT INTO occurrences_new
SELECT g."gbifID"::bigint, coalesce(by_species.id, by_taxon.id), g."datasetKey"::uuid,
       replace(replace(g.license, 'CC_BY_4_0', 'CC BY 4.0'), 'CC0_1_0', 'CC0 1.0'), g."basisOfRecord",
       g."decimalLongitude"::float8, g."decimalLatitude"::float8, g."coordinateUncertaintyInMeters"::float8,
       g."countryCode", g.year::integer, g.month::integer
FROM gbif_staging AS g
LEFT JOIN species AS by_species ON by_species.gbif_id = g."speciesKey"::bigint
LEFT JOIN species AS by_taxon ON by_taxon.gbif_id = g."taxonKey"::bigint
"""

# One row per check, with the rows that break it. Any non-zero count rolls the whole load back.
CHECKS = """
SELECT 'rows typed = rows staged', (SELECT count(*) FROM gbif_staging) - (SELECT count(*) FROM occurrences_new)
UNION ALL SELECT 'speciesKey and taxonKey name two species', count(*) FROM gbif_staging AS g
  JOIN species AS a ON a.gbif_id = g."speciesKey"::bigint
  JOIN species AS b ON b.gbif_id = g."taxonKey"::bigint
  WHERE a.id <> b.id
UNION ALL SELECT 'license not CC0 1.0 or CC BY 4.0', count(*) FROM occurrences_new
  WHERE license NOT IN ('CC0 1.0', 'CC BY 4.0')
"""

with connect() as conn, conn.cursor() as cur:
    with conn.transaction():  # all or nothing: an exception inside rolls every statement back
        stage(cur, GBIF_FILE)
        cur.execute(TYPED)
        cur.execute(INSERT_TYPED)  # a record that maps to no species fails NOT NULL here, and so the load
        failed = [(name, n) for name, n in cur.execute(CHECKS).fetchall() if n]
        if failed:
            raise SystemExit(f"load refused, nothing kept: {failed}")
        by_species_only = cur.execute(
            'SELECT count(*) FROM gbif_staging AS g JOIN species AS s ON s.gbif_id = g."speciesKey"::bigint').fetchone()[0]
        rows = cur.execute("""
            SELECT o.id, s.scientific_name, o.license, o.coordinate_uncertainty_m, o.year, ST_AsText(o.geom)
            FROM occurrences_new AS o JOIN species AS s ON s.id = o.species_id ORDER BY o.id""").fetchall()

print(f"joined on speciesKey alone: {by_species_only} of {len(rows)} records")
print("id | species | license | uncertainty_m | year | geom")
for row in rows:
    print(" | ".join("NULL" if v is None else str(v) for v in row))

# %% e2-casts
# q: e2-casts
# status: ran
# after: py-connect, e2-file
import psycopg

# The type each staged column gets in e2-typed's INSERT_TYPED.
CASTS = {"gbifID": "bigint", "datasetKey": "uuid", "decimalLatitude": "float8", "decimalLongitude": "float8",
         "coordinateUncertaintyInMeters": "float8", "year": "integer", "month": "integer", "taxonKey": "bigint",
         "speciesKey": "bigint"}
# Every staged value its cast would refuse, with its record's key and its column: what a failed cast doesn't say.
BAD_VALUES = " UNION ALL ".join(
    f"""SELECT "gbifID", '{column}', "{column}", '{type_}' FROM gbif_staging
        WHERE NOT pg_input_is_valid("{column}", '{type_}')""" for column, type_ in CASTS.items()) + " ORDER BY 1, 2"

rows = [row.copy() for row in ROWS]
rows[0][HEADER.index("coordinateUncertaintyInMeters")] = "30 m"
rows[1][HEADER.index("decimalLatitude")] = "-22,4310"

with connect() as conn, conn.cursor() as cur:
    stage(cur, gbif_file(rows))  # staging takes both: every column is text
    try:
        with conn.transaction():
            cur.execute('SELECT "decimalLatitude"::float8, "coordinateUncertaintyInMeters"::float8 FROM gbif_staging')
    except psycopg.Error as error:
        print(f"the cast alone: {error.sqlstate} {error.diag.message_primary}")
    bad = cur.execute(BAD_VALUES).fetchall()  # run this before INSERT_TYPED; refuse the load if it returns a row
    conn.rollback()

print("gbifID | column | value | type")
for row in bad:
    print(" | ".join(row))
print(f"load refused before INSERT_TYPED: {len(bad)} values in {len({row[0] for row in bad})} of {len(rows)} records")

# %% e2-geopandas
# q: e2-geopandas
# status: ran
# after: py-connect
import geopandas as gpd
from sqlalchemy import create_engine, text

# SQLAlchemy on py-connect's connection: the same database, role and search path.
engine = create_engine("postgresql+psycopg://", creator=connect)

with engine.connect() as conn:
    devils = gpd.read_postgis(
        text("SELECT o.id, o.year, o.geom FROM learn_curriculum.occurrences AS o "
             "JOIN learn_curriculum.species AS s ON s.id = o.species_id "
             "WHERE s.scientific_name = 'Sarcophilus harrisii' ORDER BY o.id"),
        conn, geom_col="geom")
    print(f"read: {len(devils)} rows, CRS {devils.crs.to_string()}, geometry {devils.geom_type.iloc[0]}")

    metres = devils.to_crs(3577)  # GDA94 / Australian Albers, an equal-area projection in metres
    metres["geom"] = metres.buffer(1000)  # a 1 km circle around each record
    # The schema named, not left to the search path; the write stays in this connection's transaction.
    metres.to_postgis("devil_buffers", conn, schema="learn_curriculum", if_exists="fail", index=False)
    n, srid, kind, km2 = conn.execute(text(
        "SELECT count(*), min(ST_SRID(geom)), min(GeometryType(geom)), round(sum(ST_Area(geom))::numeric / 1e6, 1) "
        "FROM learn_curriculum.devil_buffers")).one()
    print(f"written: {n} rows, SRID {srid}, {kind}, {km2} km² in all")
    schema = conn.execute(text(
        "SELECT relnamespace::regnamespace::text FROM pg_class "
        "WHERE oid = to_regclass('learn_curriculum.devil_buffers')")).scalar_one()
    assert schema == "learn_curriculum", f"devil_buffers was written to {schema}"
    print(f"devil_buffers is in schema {schema}, until the rollback")
    conn.rollback()  # the lab leaves learn_curriculum as it found it
    left = conn.execute(text("SELECT to_regclass('learn_curriculum.devil_buffers')")).scalar_one()
    print(f"after the rollback: {left or 'no table'}")

# %% e2-ex
# q: e2-ex
# status: ran
# after: py-connect, e2-file
# e2-typed's casts and CHECKs, per staged column: its type, and the range a typed value must fall in (None: any).
RULES = [
    ("datasetKey", "uuid", None),
    ("decimalLatitude", "float8", (-90, 90)),
    ("decimalLongitude", "float8", (-180, 180)),
    ("year", "integer", None),
    ("month", "integer", (1, 12)),
]


def problem(column, type_, bounds):
    """A value's problem, or NULL: the cast's own message, else the range it breaks. CASE tests in order,
    so a value is cast only once pg_input_is_valid has said it can be."""
    out_of_range = (f"""WHEN NOT ("{column}"::{type_} BETWEEN {bounds[0]} AND {bounds[1]})
                    THEN 'outside {bounds[0]} to {bounds[1]}'""" if bounds else "")
    return f"""SELECT "gbifID", '{column}' AS column_name, "{column}" AS value,
                      CASE WHEN NOT pg_input_is_valid("{column}", '{type_}')
                           THEN (pg_input_error_info("{column}", '{type_}')).message
                           {out_of_range} END AS problem
               FROM gbif_staging"""


PROBLEMS = ("SELECT * FROM (" + " UNION ALL ".join(problem(*rule) for rule in RULES) + ") AS p "
            "WHERE problem IS NOT NULL ORDER BY 1, 2")

rows = [row.copy() for row in ROWS]
rows[0][HEADER.index("month")] = "13"
rows[1][HEADER.index("decimalLatitude")] = "-122.4310"
rows[2][HEADER.index("datasetKey")] = "00000000-0000-0000-0000-00000000002"
rows[2][HEADER.index("year")] = "c. 2015"

with connect() as conn, conn.cursor() as cur:
    stage(cur, gbif_file(rows))
    problems = cur.execute(PROBLEMS).fetchall()
    conn.rollback()

print("gbifID | column | value | problem")
for row in problems:
    print(" | ".join(row))
print(f"load refused: {len(problems)} problems in {len({row[0] for row in problems})} of {len(rows)} records")