name: Tests

on:
  push:
    branches:
      - main
  pull_request:
    branches:
      - main

# Cancels all previous workflow runs for pull requests that have not completed.
concurrency:
  # The concurrency group contains the workflow name and the branch name for pull requests
  # or the commit hash for any other events.
  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
  cancel-in-progress: true

jobs:
  changes:
    name: "Check for changes"
    runs-on: ubuntu-latest
    outputs:
      changes: ${{ steps.changes.outputs.src }}
    steps:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
      - uses: dorny/paths-filter@v3
        id: changes
        with:
          filters: |
            python: &python
            - 'pytensor/**/*.py'
            - 'tests/**/*.py'
            - 'pytensor/**/*.pyx'
            - 'tests/**/*.pyx'
            - '*.py'
            src:
            - *python
            - 'pytensor/**/*.c'
            - 'tests/**/*.c'
            - 'pytensor/**/*.h'
            - 'tests/**/*.h'
            - '.github/workflows/*.yml'
            - 'setup.cfg'
            - 'requirements.txt'
            - '.pre-commit-config.yaml'

  style:
    name: Check code style
    needs: changes
    runs-on: ubuntu-latest
    if: ${{ needs.changes.outputs.changes == 'true' }}
    strategy:
      matrix:
        python-version: ["3.10", "3.12"]
    steps:
    - uses: actions/checkout@v4
    - uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
    - uses: pre-commit/action@v3.0.1

  test_ubuntu:
    name: "Test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}"
    needs:
      - changes
      - style
    runs-on: ubuntu-latest
    if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }}
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.10", "3.12"]
        fast-compile: [0,1]
        float32: [0,1]
        install-numba: [0]
        install-jax: [0]
        part:
          - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse"
          - "tests/scan"
          - "tests/sparse"
          - "tests/tensor --ignore=tests/tensor/conv --ignore=tests/tensor/rewriting --ignore=tests/tensor/test_math.py --ignore=tests/tensor/test_basic.py --ignore=tests/tensor/test_blas.py --ignore=tests/tensor/test_math_scipy.py --ignore=tests/tensor/test_inplace.py --ignore=tests/tensor/test_elemwise.py"
          - "tests/tensor/conv"
          - "tests/tensor/rewriting"
          - "tests/tensor/test_math.py"
          - "tests/tensor/test_basic.py tests/tensor/test_inplace.py"          
          - "tests/tensor/test_blas.py tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py"
        exclude:
          - python-version: "3.10"
            fast-compile: 1
          - python-version: "3.10"
            float32: 1
          - python-version: "3.10"
            part: "tests/tensor/test_math.py"
          - fast-compile: 1
            float32: 1
        include:
          - install-numba: 1
            python-version: "3.10"
            fast-compile: 0
            float32: 0
            part: "tests/link/numba"
          - install-numba: 1
            python-version: "3.12"
            fast-compile: 0
            float32: 0
            part: "tests/link/numba"
          - install-jax: 1
            python-version: "3.10"
            fast-compile: 0
            float32: 0
            part: "tests/link/jax"
          - install-jax: 1
            python-version: "3.12"
            fast-compile: 0
            float32: 0
            part: "tests/link/jax"
    steps:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
      - name: Set up Python ${{ matrix.python-version }}
        uses: conda-incubator/setup-miniconda@v3
        with:
          miniforge-variant: Mambaforge
          miniforge-version: latest
          python-version: ${{ matrix.python-version }}

      - name: Create matrix id
        id: matrix-id
        env:
          MATRIX_CONTEXT: ${{ toJson(matrix) }}
        run: |
          echo $MATRIX_CONTEXT
          export MATRIX_ID=`echo $MATRIX_CONTEXT | md5sum | cut -c 1-32`
          echo $MATRIX_ID
          echo "id=$MATRIX_ID" >> $GITHUB_OUTPUT

      - name: Install dependencies
        shell: bash -l {0}
        run: |
          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
          if [[ $INSTALL_NUMBA == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
          if [[ $INSTALL_JAX == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
          pip install -e ./
          mamba list && pip freeze
          python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
          python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
          INSTALL_NUMBA: ${{ matrix.install-numba }}
          INSTALL_JAX: ${{ matrix.install-jax }}

      - name: Run tests
        shell: bash -l {0}
        run: |
          if [[ $FAST_COMPILE == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi
          if [[ $FLOAT32 == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,floatX=float32; fi
          export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
          python -m pytest -r A --verbose --runslow --durations=50 --cov=pytensor/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART --benchmark-skip
        env:
          MATRIX_ID: ${{ steps.matrix-id.outputs.id }}
          MKL_THREADING_LAYER: GNU
          MKL_NUM_THREADS: 1
          OMP_NUM_THREADS: 1
          PART: ${{ matrix.part }}
          FAST_COMPILE: ${{ matrix.fast-compile }}
          FLOAT32: ${{ matrix.float32 }}

      - name: Upload coverage file
        uses: actions/upload-artifact@v3
        with:
          name: coverage
          path: coverage/coverage-${{ steps.matrix-id.outputs.id }}.xml

  benchmarks:
      name: "Benchmarks"
      needs:
        - changes
        - style
      runs-on: ubuntu-latest
      if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }}
      strategy:
        fail-fast: false
      steps:
        - uses: actions/checkout@v4
          with:
            fetch-depth: 0
        - name: Set up Python 3.10
          uses: conda-incubator/setup-miniconda@v3
          with:
            miniforge-variant: Mambaforge
            miniforge-version: latest
        - name: Install dependencies
          shell: bash -l {0}
          run: |
            mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark
            pip install -e ./
            mamba list && pip freeze
            python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
            python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'
          env:
            PYTHON_VERSION: 3.10
        - name: Download previous benchmark data
          uses: actions/cache@v4
          with:
            path: ./cache
            key: ${{ runner.os }}-benchmark
        - name: Run benchmarks
          shell: bash -l {0}
          run: |
            export PYTENSOR_FLAGS=mode=FAST_COMPILE,warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
            python -m pytest --runslow --benchmark-only --benchmark-json output.json
        - name: Store benchmark result
          uses: benchmark-action/github-action-benchmark@v1
          with:
            name: Python Benchmark with pytest-benchmark
            tool: 'pytest'
            output-file-path: output.json
            external-data-json-path: ./cache/benchmark-data.json
            alert-threshold: '200%'
            github-token: ${{ secrets.GITHUB_TOKEN }}
            comment-on-alert: false
            fail-on-alert: false
            auto-push: false

  all-checks:
    if: ${{ always() }}
    runs-on: ubuntu-latest
    name: "All tests"
    needs: [changes, style, test_ubuntu]
    steps:
      - name: Check build matrix status
        if: ${{ needs.changes.outputs.changes == 'true' && (needs.style.result != 'success' || needs.test_ubuntu.result != 'success') }}
        run: exit 1

  upload-coverage:
    runs-on: ubuntu-latest
    name: "Upload coverage"
    needs: [changes, all-checks]
    if: ${{ needs.changes.outputs.changes == 'true' && needs.all-checks.result == 'success' }}
    steps:
      - uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.12"

      - name: Install dependencies
        run: |
          python -m pip install -U coverage>=5.1 coveralls

      - name: Download coverage file
        uses: actions/download-artifact@v3
        with:
          name: coverage
          path: coverage

      - name: Upload coverage to Codecov
        uses: codecov/codecov-action@v4
        with:
          directory: ./coverage/
          fail_ci_if_error: true
          token: ${{ secrets.CODECOV_TOKEN }}
