from annotypes import add_call_types

from malcolm.core import PartRegistrar
from malcolm.modules import ADCore, builtin, scanning


@builtin.util.no_save("numImagesPerSeries")
class TetrAMMDriverPart(ADCore.parts.DetectorDriverPart):
    def setup(self, registrar: PartRegistrar) -> None:
        super().setup(registrar)
        registrar.hook(scanning.hooks.PostRunReadyHook, self.on_post_run_ready)
        registrar.hook(scanning.hooks.PostRunArmedHook, self.on_post_run_armed)

    @add_call_types
    def on_run(self, context: scanning.hooks.AContext) -> None:
        # this override stops waiting on the start_future as we have to press
        # stop in post_run below
        pass

    @add_call_types
    def on_post_run_ready(self, context: scanning.hooks.AContext) -> None:
        # currently the AD support does not know how many frames the detector
        # has taken and never finishes the Acquire. We know that the file
        # writer has completed at post run so stop the AD Acquisition
        child = context.block_view(self.mri)
        child.stop()
        context.wait_all_futures(self.start_future)

    @add_call_types
    def on_post_run_armed(self, context: scanning.hooks.AContext) -> None:
        # Stop the acquisition as per post_run_ready
        # TODO: this should call configure too, will fail for 3D scans at
        # present
        child = context.block_view(self.mri)
        child.stop()
        context.wait_all_futures(self.start_future)

@@ -1,12 +1,35 @@
+from math import ceil%0A%0A
 from annotyp
@@ -78,16 +78,28 @@
 e import
+ NumberMeta,
  PartReg
@@ -104,16 +104,24 @@
 egistrar
+, Widget
 %0Afrom ma
@@ -168,16 +168,140 @@
 anning%0A%0A
+APartName = builtin.parts.APartName%0AAMri = builtin.parts.AMri%0ATETRAMM_BASE_FREQ = 100000.0%0ATETRAMM_MIN_VALUES_PER_READ = 5%0A%0A
 %0A@builti
@@ -335,16 +335,33 @@
 rSeries%22
+, %22valuesPerRead%22
 )%0Aclass 
@@ -424,245 +424,1504 @@
 def 
-setup(self, registrar: PartRegistrar) -%3E None:%0A        super().setup(registrar)%0A        registrar.hook(scanning.hooks.PostRunReadyHook, self.on_post_run_ready)%0A        registrar.hook(scanning.hooks.PostRunArmedHook, self.on_post_run_arme
+__init__(self, name: APartName, mri: AMri) -%3E None:%0A        super().__init__(name, mri)%0A        self.targetSamplesPerFrame = NumberMeta(%0A            %22uint32%22,%0A            %22Target samples that each frame contains, %22%0A            %220 for maximum, -1 for not controlling it.%22,%0A            tags=%5BWidget.TEXTINPUT.tag()%5D,%0A        ).create_attribute_model()%0A%0A    def setup(self, registrar: PartRegistrar) -%3E None:%0A        super().setup(registrar)%0A        registrar.add_attribute_model(%22targetSamplesPerFrame%22,%0A                                      self.targetSamplesPerFrame,%0A                                      self.targetSamplesPerFrame.set_value)%0A        registrar.hook(scanning.hooks.PostRunReadyHook, self.on_post_run_ready)%0A        registrar.hook(scanning.hooks.PostRunArmedHook, self.on_post_run_armed)%0A        registrar.hook(scanning.hooks.PostConfigureHook,%0A                       self.on_post_configure)%0A%0A    @add_call_types%0A    def on_post_configure(self, context: scanning.hooks.AContext):%0A        child = context.block_view(self.mri)%0A        if self.targetSamplesPerFrame.value == 0:%0A            child.valuesPerRead.put_value(TETRAMM_MIN_VALUES_PER_READ)%0A        elif self.targetSamplesPerFrame.value %3E 0:%0A            values_per_read = ceil(TETRAMM_BASE_FREQ * child.exposure.value%0A                                   / self.targetSamplesPerFrame.value)%0A            values_per_read = max(values_per_read, TETRAMM_MIN_VALUES_PER_READ)%0A            child.valuesPerRead.put_value(values_per_rea
 d)%0A%0A

add attribute targetSamplesPerFrame