from metaflow import FlowSpec, step, Parameter, batch


class TestBatchFlow(FlowSpec):
    """
    Simple test flow to validate AWS Batch setup.

    This flow:
    1. Generates a list of numbers (0-9)
    2. Processes each number in parallel on AWS Batch (squares it)
    3. Joins the results
    4. Sums all squared values
    """

    num_items = Parameter(
        'num_items',
        help='Number of items to process',
        default=10,
        type=int
    )

    @step
    def start(self):
        """
        Initialize the flow with a list of numbers.
        """
        print(f"Starting test flow with {self.num_items} items")

        # Create a list of numbers to process
        self.numbers = list(range(self.num_items))

        print(f"Generated numbers: {self.numbers}")

        # Fan out to process each number in parallel
        self.next(self.process_number, foreach='numbers')

    @batch(cpu=1, memory=2048, cpu_architecture="ARM64", log_driver="awslogs")
    @step
    def process_number(self):
        """
        Process a single number on AWS Batch (square it).
        This step runs in parallel for each number.
        """
        import time

        # Get the input number
        number = self.input

        print(f"Processing number: {number}")

        # Simulate some work
        time.sleep(1)

        # Square the number
        self.squared = number ** 2
        self.original = number

        print(f"Number {number} squared = {self.squared}")

        # Explicitly verify artifacts are set
        if not hasattr(self, 'squared'):
            raise RuntimeError(f"squared attribute not set for number {number}")
        if not hasattr(self, 'original'):
            raise RuntimeError(f"original attribute not set for number {number}")

        print(f"Artifacts verified: original={self.original}, squared={self.squared}")

        self.next(self.join)

    @step
    def join(self, inputs):
        """
        Collect all results from parallel processing.
        """
        # Collect all squared values
        self.results = []
        for inp in inputs:
            if hasattr(inp, 'squared') and hasattr(inp, 'original'):
                self.results.append({
                    'original': inp.original,
                    'squared': inp.squared
                })

        print(f"Collected {len(self.results)} results from parallel processing")
        print(f"Results: {self.results}")

        # Merge num_items from inputs
        self.merge_artifacts(inputs, include=['num_items'])

        self.next(self.calculate_sum)

    @step
    def calculate_sum(self):
        """
        Calculate the sum of all squared values.
        """
        self.total_sum = sum(r['squared'] for r in self.results)

        print(f"Sum of squared values: {self.total_sum}")

        # Expected sum formula: sum(i^2) for i in 0..n-1 = (n-1)*n*(2n-1)/6
        if self.num_items > 0:
            n = self.num_items
            expected = (n - 1) * n * (2 * n - 1) // 6
            print(f"Expected sum: {expected}")

            if self.total_sum == expected:
                print("✓ Result matches expected value!")
            else:
                print(f"✗ Result mismatch: got {self.total_sum}, expected {expected}")

        self.next(self.end)

    @step
    def end(self):
        """
        Flow completed successfully.
        """
        print(f"\n{'='*60}")
        print(f"Test Batch Flow Completed!")
        print(f"{'='*60}")
        print(f"Items processed: {len(self.results)}")
        print(f"Total sum: {self.total_sum}")
        print(f"{'='*60}\n")


if __name__ == '__main__':
    TestBatchFlow()

print

[MFLOG|0|2025-11-15T12:54:48.689613Z|runtime|f2d7e201-2489-4c18-82b1-a591260de3ce]    Data store error:
[MFLOG|0|2025-11-15T12:54:48.874403Z|runtime|aea850d3-33ff-4ff8-b4a5-2377ee400daf]    No completed attempts of the task was found for task 'TestBatchFlow/1763211230786273/process_number/4'
[MFLOG|0|2025-11-15T12:54:48.874563Z|runtime|c285532a-2ea7-4c40-96e6-e5c40da9f6bf]
[MFLOG|0|2025-11-15T12:54:48.992257Z|runtime|72cfe503-f9f7-4c00-b83e-24e3ea2240ce]Task failed.

[MFLOG|0|2025-11-15T12:53:56.753069Z|runtime|74597a2b-0e20-4b00-be17-aa7c859dafd9][957c29eb-7eb8-4c7a-9b21-494a6b9174e5] Task is starting (status SUBMITTED)...
[MFLOG|0|2025-11-15T12:53:57.827061Z|runtime|b30cf22c-0614-4bf9-96e5-3f548e52775a][957c29eb-7eb8-4c7a-9b21-494a6b9174e5] Task is starting (status RUNNABLE)... 
[MFLOG|0|2025-11-15T12:54:03.267219Z|runtime|4fdc75a8-fcf4-4708-b8f8-842deeab1342][957c29eb-7eb8-4c7a-9b21-494a6b9174e5] Task is starting (status STARTING)... 
[MFLOG|0|2025-11-15T12:54:24.861614Z|runtime|c3beeebb-c535-4c4b-8477-93dea447137d][957c29eb-7eb8-4c7a-9b21-494a6b9174e5] Task is starting (status RUNNING)... 
[MFLOG|0|2025-11-15T12:54:48.434800Z|runtime|3e8c2fc6-4f67-4fe2-a823-1bf83476b035][957c29eb-7eb8-4c7a-9b21-494a6b9174e5] Task finished with exit code 0.