from matplotlib.figure import Figure
import numpy as np

from karabo_data.geometry2 import AGIPD_1MGeometry

def test_snap_assemble_data():
    geom = AGIPD_1MGeometry.from_quad_positions(quad_pos=[
        (-525, 625),
        (-550, -10),
        (520, -160),
        (542.5, 475),
    ])

    stacked_data = np.zeros((16, 512, 128))
    img, centre = geom.position_modules_fast(stacked_data)
    assert img.shape == (1256, 1092)
    assert tuple(centre) == (631, 550)
    assert np.isnan(img[0, 0])
    assert img[50, 50] == 0

def test_write_read_crystfel_file(tmpdir):
    geom = AGIPD_1MGeometry.from_quad_positions(quad_pos=[
        (-525, 625),
        (-550, -10),
        (520, -160),
        (542.5, 475),
    ])
    path = str(tmpdir / 'test.geom')
    geom.write_crystfel_geom(path)

    # We need to add some experiment details before cfelpyutils will read the
    # file
    with open(path, 'r') as f:
        contents = f.read()
    with open(path, 'w') as f:
        f.write('clen = 0.119\n')
        f.write('adu_per_eV = 0.0075\n')
        f.write(contents)

    loaded = AGIPD_1MGeometry.from_crystfel_geom(path)
    np.testing.assert_allclose(loaded.modules[0][0].corner_pos,
                               geom.modules[0][0].corner_pos)
    np.testing.assert_allclose(loaded.modules[0][0].fs_vec,
                               geom.modules[0][0].fs_vec)

def test_inspect():
    geom = AGIPD_1MGeometry.from_quad_positions(quad_pos=[
        (-525, 625),
        (-550, -10),
        (520, -160),
        (542.5, 475),
    ])
    # Smoketest
    fig = geom.inspect()
    assert isinstance(fig, Figure)

def test_compare():
    geom1 = AGIPD_1MGeometry.from_quad_positions(quad_pos=[
        (-525, 625),
        (-550, -10),
        (520, -160),
        (542.5, 475),
    ])
    geom2 = AGIPD_1MGeometry.from_quad_positions(quad_pos=[
        (-527, 625),
        (-548, -10),
        (520, -162),
        (542.5, 473),
    ])
    # Smoketest
    fig = geom1.compare(geom2)
    assert isinstance(fig, Figure)

def test_to_distortion_array():
    geom = AGIPD_1MGeometry.from_quad_positions(quad_pos=[
        (-525, 625),
        (-550, -10),
        (520, -160),
        (542.5, 475),
    ])
    # Smoketest
    distortion = geom.to_distortion_array()
    assert isinstance(distortion, np.ndarray)
    assert distortion.shape == (8192, 128, 4, 3)

    # Coordinates in m; max x & y should be ~ 12 cm from centre
    assert .08 < distortion[..., 1].max() < .15
    assert .08 < distortion[..., 2].max() < .15

@@ -2388,16 +2388,34 @@
 tes in m
+, origin at corner
 ; max x 
@@ -2434,25 +2434,108 @@
 e ~ 
-12 cm from centre
+25cm%0A    assert .20 %3C distortion%5B..., 1%5D.max() %3C .30%0A    assert .20 %3C distortion%5B..., 2%5D.max() %3C .30
 %0A   
@@ -2534,35 +2534,36 @@
  .30%0A    assert 
+-
 .0
-8
+1
  %3C distortion%5B..
@@ -2569,26 +2569,26 @@
 .., 1%5D.m
-ax
+in
 () %3C .
+0
 1
-5
 %0A    ass
@@ -2595,11 +2595,12 @@
 ert 
+-
 .0
-8
+1
  %3C d
@@ -2622,15 +2622,15 @@
 2%5D.m
-ax
+in
 () %3C .
+0
 1
-5
 %0A

Update test for distortion array