#!/usr/bin/env python3
# digitize.py: read the fish osmolality points off two published figures, and write digitized.json.
#   pip install opencv-python-headless numpy && python3 digitize.py
#
# Why digitize: the values behind Yancey et al. 2014 Fig. 3 and Yancey 2025 Fig. 1 are printed only as
# plots (the 2014 paper's body and any table could not be fetched; see sources/). So each marker's
# centre is located in the image, and pixel positions are mapped to depth and osmolality through the
# axis tick marks, which are located the same way. The check on this is external: a straight line
# refitted to the digitized 2014 points must reproduce the two equations printed in the figure
# (326 + 0.0916 x and 320 + 0.0953 x), and the 2025 figure's open square at the 2014 snailfish's place
# (taken to be that point replotted; its legend style is "Various fish 1996-2007") must land on
# (7000 m, 991 mOsm/kg). Both are asserted below and again by the page's verifier.
#
# Two kinds of pick: 2014 circles are found by a Hough transform (listed pixel centres are its output,
# rounded to its 0.5 px grid, legend hits dropped); 2025 markers overlapping near 7,000-7,900 m could
# not be separated automatically and were placed by eye on a 4x enlargement (marked "by eye").
import json, cv2, numpy as np

def calib(ticks):
    vs, ps = zip(*ticks)
    return lambda p: float(np.interp(p, ps, vs)) if ps[0] < ps[-1] else float(np.interp(-p, [-q for q in ps], vs))

# Yancey et al. 2014, Fig. 3 (680 x 514 px). Axis ticks located from dark-pixel columns/rows:
# x ticks at px 73, 182.5, 292.5, 402, 512, 621.5 = 0..10,000 m by 2,000; y ticks every 41.14 px = 100 mOsm/kg from 457.5.
X14 = calib([(0, 73), (2000, 182.5), (4000, 292.5), (6000, 402), (8000, 512), (10000, 621.5)])
Y14 = calib([(0, 457.5), (1000, 457.5 - 411.4)])
fig3 = cv2.imread('sources/yancey2014_pnas.1322003111fig03.jpg', 0)
circles = cv2.HoughCircles(cv2.GaussianBlur(fig3, (3, 3), 0), cv2.HOUGH_GRADIENT, dp=1, minDist=4,
                           param1=100, param2=10, minRadius=4, maxRadius=9)[0]
PICKS14 = [(100.5, 297.5), (100.5, 308.5), (116.5, 296.5), (144.5, 267.5), (144.5, 278.5), (174.5, 249.5),
           (187.5, 242.5), (188.5, 248.5), (188.5, 253.5), (193.5, 235.5), (230.5, 226.5), (237.5, 195.5)]
found = {(float(x), float(y)) for x, y, r in circles}
for p in PICKS14:
    assert p in found, f'Hough no longer finds {p}'
pre2014 = [dict(depth=round(X14(x)), osm=round(Y14(y)), px=[x, y]) for x, y in PICKS14]

# Yancey 2025 (figshare 10.6084/m9.figshare.28692455.v7), Fig. 1 (1702 x 1326 px), "published fish data
# modified from Linley et al. 2016". Blue tick marks located from blue pixels.
X25 = calib([(0, 285), (1000, 402.5), (2000, 521), (3000, 638.5), (4000, 756.5), (5000, 873.5), (6000, 991.5),
             (7000, 1109.5), (8000, 1232.5), (9000, 1350.5), (10000, 1468)])
Y25 = calib([(0, 1141), (200, 984), (400, 827.5), (600, 670.5), (800, 513), (1000, 356), (1200, 199)])
EYE = lambda zx, zy: (1020 + zx / 4, 280 + zy / 4)   # 4x enlargement of the crop (1020,280)-(1250,390)
later = [
  ('Kermadec snailfish', 'legend: 2014', EYE(100, 355)), ('Mariana snailfish', 'legend: 2016', EYE(330, 270)),
  ('Mariana snailfish', 'legend: 2016', EYE(400, 250)), ('Kermadec snailfish', 'legend: 2014', EYE(455, 225)),
  ('Mariana snailfish', 'legend: 2016', EYE(595, 180)), ('Kermadec snailfish', 'legend: 2014', EYE(600, 105)),
  ('Mariana snailfish', 'legend: 2016', EYE(800, 120)),
]
check2014 = EYE(360, 340)   # the open square: the 2014 Kermadec snailfish, replotted
# other families, isolated markers: bounding-box centres of dark blobs (opened with a 9x9 kernel)
others = [
  ('Kermadec rattail', 'legend: 2014', (705.0, 680.3)), ('Kermadec rattail', 'legend: 2014', (780.2, 571.8)),
  ('Mariana rattail', 'legend: 2016', (808.0, 623.1)), ('Mariana rattail', 'legend: 2016', (903.9, 549.3)),
  ('Kermadec eelpout (zoarcid)', 'legend: 2014', (852.7, 598.1)),
  ('Kermadec eelpout (zoarcid)', 'legend: 2014', (880.2, 465.0)), ('Kermadec eelpout (zoarcid)', 'legend: 2014', (880.2, 483.0)),
]
out = dict(
  generated_by='research/how-deep-can-fish-live/digitize.py',
  seawater_mosm=1100,
  pre2014=dict(source='Yancey et al. 2014, PNAS 111:4461, Fig. 3 (circles: bathyal and abyssal teleosts)',
               points=pre2014),
  kermadec2014=dict(source='Yancey et al. 2014, abstract: 991 +/- 22 mOsmol/kg, five N. kermadecensis from 7,000 m',
                    point=dict(depth=7000, osm=991)),
  check_replotted_2014=dict(depth=round(X25(check2014[0])), osm=round(Y25(check2014[1]))),
  later=dict(source='Yancey 2025, figshare 10.6084/m9.figshare.28692455.v7, Fig. 1 (data modified from Linley et al. 2016); by eye on a 4x enlargement',
             points=[dict(depth=round(X25(x)), osm=round(Y25(y)), who=w, note=n) for w, n, (x, y) in later]),
  others=dict(source='the same figure; other families, not snailfish',
              points=[dict(depth=round(X25(x)), osm=round(Y25(y)), who=w, note=n) for w, n, (x, y) in others]),
)
xs = np.array([p['depth'] for p in pre2014]); ys = np.array([p['osm'] for p in pre2014])
b, a = np.polyfit(xs, ys, 1)
assert abs(a - 326) < 5 and abs(b - 0.0916) < 0.002, (a, b)
b2, a2 = np.polyfit(np.append(xs, 7000), np.append(ys, 991), 1)
assert abs(a2 - 320) < 5 and abs(b2 - 0.0953) < 0.002, (a2, b2)
assert abs(out['check_replotted_2014']['depth'] - 7000) <= 30 and abs(out['check_replotted_2014']['osm'] - 991) <= 5
json.dump(out, open('digitized.json', 'w'), indent=1)
print(f'2014 refit {a:.1f} + {b:.4f} x (printed 326 + 0.0916 x); with snailfish {a2:.1f} + {b2:.4f} x (printed 320 + 0.0953 x)')
print('replotted 2014 snailfish at', out['check_replotted_2014'])