PhysiGo/app/lib/exercises/exercises_validation/widgets/pose_detector.dart

import 'dart:async';

import 'package:body_detection/models/pose_landmark.dart';
import 'package:body_detection/models/pose_landmark_type.dart';
import 'package:physigo/exercises/exercises_validation/models/exercise.dart';
import 'package:rxdart/rxdart.dart';
import 'package:body_detection/body_detection.dart';
import 'package:body_detection/models/image_result.dart';
import 'package:body_detection/models/pose.dart';
import 'package:flutter/material.dart';

import 'pose_painter.dart';

typedef MeanFilteredData = Iterable<List<double>>;
typedef LandmarkVariations = List<List<double>>;

enum StepExercise {
  notInPlace,
  ready,
  start,
  end,
}

// TODO: authorized joint in exercise
// TODO: correct axis depending on phone orientation. Determine axis based on skeleton(shoulder line, ...)

class PoseDetector extends StatefulWidget {
  final Exercise exercise;
  const PoseDetector({required this.exercise, Key? key}) : super(key: key);

  @override
  State<PoseDetector> createState() => _PoseDetectorState();
}

class _PoseDetectorState extends State<PoseDetector> {
  static const buffer = 5;
  final List<StreamSubscription> _streamSubscriptions = [];
  final List<StreamController> _streamControllers = [];
  final StreamController<Pose> _streamController = StreamController.broadcast();
  final StreamController<StepExercise> _stepExerciseStream = StreamController.broadcast();
  Image? _cameraImage;
  Pose? _detectedPose;
  Size _imageSize = Size.zero;
  late Future<void> _startCamera;
  late Stream<MeanFilteredData> _meanFilterStream;
  late Stream<int> _repCounter;

  @override
  initState() {
    super.initState();
    _streamControllers.add(_streamController);
    _streamControllers.add(_stepExerciseStream);
    _startCamera = _startCameraStream();
    _meanFilterStream = _getMeanFilterStream(_streamController.stream);
    _streamSubscriptions.add(
        CombineLatestStream.combine2(_stepExerciseStream.stream, _meanFilterStream, (a, b) => [a, b]).listen((value) {
      StepExercise stepExercise = value.first as StepExercise;
      MeanFilteredData meanFilteredData = value.last as MeanFilteredData;
      final isStartOfExerciseMovement = widget.exercise.isAtStartMovement(meanFilteredData);
      final isEndOfExerciseMovement = widget.exercise.isAtEndMovement(meanFilteredData);
      if (stepExercise == StepExercise.notInPlace && isStartOfExerciseMovement) {
        _stepExerciseStream.add(StepExercise.ready);
      }
      if ((stepExercise == StepExercise.ready || stepExercise == StepExercise.start) && isEndOfExerciseMovement) {
        _stepExerciseStream.add(StepExercise.end);
      }
      if (stepExercise == StepExercise.end && isStartOfExerciseMovement) {
        _stepExerciseStream.add(StepExercise.start);
      }
    }));
    _stepExerciseStream.add(StepExercise.notInPlace);
    _repCounter = _stepExerciseStream.stream
        .where((event) => event == StepExercise.start)
        .scan((int accumulated, value, index) => accumulated + 1, 0);
  }

  Stream<MeanFilteredData> _getMeanFilterStream(Stream<Pose> stream) {
    return stream
        .where((pose) => pose.landmarks.isNotEmpty)
        .map((pose) => pose.landmarks.where((landmark) => authorizedType.contains(landmark.type)).toList())
        // Get last [buffer] poses
        .bufferCount(buffer, 1)
        // Swap matrix [buffer] * [authorizedType.length]
        .map(_swapMatrixDimensions)
        // For every landmarks, get meanFilter of size [buffer]
        .map((filteredLandmarks) => filteredLandmarks.map(_meanFilter));
  }

  List<double> _meanFilter(List<PoseLandmark> landmarks) {
    return landmarks
        .map((landmark) => landmark.position)
        .map((position) => [
              position.x / buffer,
              position.y / buffer,
              position.z / buffer,
            ])
        .reduce((value, element) => [
              value[0] + element[0],
              value[1] + element[1],
              value[2] + element[2],
            ]);
  }

  List<List<T>> _swapMatrixDimensions<T>(List<List<T>> matrix) {
    final height = matrix.length;
    final width = matrix[0].length;
    List<List<T>> newMatrix = [];
    for (int col = 0; col < width; col++) {
      List<T> newRow = [];
      for (int row = 0; row < height; row++) {
        newRow.add(matrix[row][col]);
      }
      newMatrix.add(newRow);
    }
    return newMatrix;
  }

  Future<void> _startCameraStream() async {
    await BodyDetection.startCameraStream(onFrameAvailable: _handleCameraImage, onPoseAvailable: _handlePose);
    await BodyDetection.enablePoseDetection();
  }

  Future<void> _stopCameraStream() async {
    await BodyDetection.disablePoseDetection();
    await BodyDetection.stopCameraStream();
  }

  void _handleCameraImage(ImageResult result) {
    if (!mounted) return;

    // To avoid a memory leak issue.
    // https://github.com/flutter/flutter/issues/60160
    PaintingBinding.instance?.imageCache?.clear();
    PaintingBinding.instance?.imageCache?.clearLiveImages();

    final image = Image.memory(
      result.bytes,
      gaplessPlayback: true,
      fit: BoxFit.contain,
    );

    setState(() {
      _cameraImage = image;
      _imageSize = result.size;
    });
  }

  void _handlePose(Pose? pose) {
    if (!mounted) return;
    if (pose != null) _streamController.add(pose);
    setState(() {
      _detectedPose = pose;
    });
  }

  @override
  void dispose() {
    _stopCameraStream();
    for (final ss in _streamSubscriptions) {
      ss.cancel();
    }
    for (final sc in _streamControllers) {
      sc.close();
    }
    super.dispose();
  }

  @override
  Widget build(BuildContext context) {
    return FutureBuilder<void>(
      future: _startCamera,
      builder: (context, snapshot) {
        if (snapshot.connectionState == ConnectionState.waiting) {
          return const Center(child: CircularProgressIndicator());
        }
        return Column(
          children: [
            Center(
              child: CustomPaint(
                child: _cameraImage,
                foregroundPainter: PosePainter(
                  pose: _detectedPose,
                  imageSize: _imageSize,
                ),
              ),
            ),
            StreamBuilder<StepExercise>(
              stream: _stepExerciseStream.stream,
              builder: (context, snapshot) {
                Color color;
                if (!snapshot.hasData) {
                  color = Colors.black;
                } else {
                  switch (snapshot.data!) {
                    case StepExercise.notInPlace:
                      color = Colors.black;
                      break;
                    case StepExercise.ready:
                      color = Colors.green;
                      break;
                    case StepExercise.start:
                      color = Colors.blue;
                      break;
                    case StepExercise.end:
                      color = Colors.red;
                      break;
                  }
                }
                return Container(
                  height: 100,
                  width: 100,
                  color: color,
                );
              },
            ),
            StreamBuilder<int>(
              stream: _repCounter,
              builder: (context, snapshot) {
                var repCounter = 0;
                if (snapshot.hasData) {
                  repCounter = snapshot.data!;
                }
                return Column(
                  children: [
                    Text(
                      "${repCounter % widget.exercise.reps}/${widget.exercise.reps}",
                      style: TextStyle(fontSize: 40),
                    ),
                    Text(
                      "${repCounter ~/ widget.exercise.reps}/${widget.exercise.series}",
                      style: TextStyle(fontSize: 40),
                    ),
                  ],
                );
              },
            )
          ],
        );
      },
    );
  }

  static const authorizedType = [
    PoseLandmarkType.nose,
    PoseLandmarkType.leftShoulder,
    PoseLandmarkType.rightShoulder,
    PoseLandmarkType.leftElbow,
    PoseLandmarkType.rightElbow,
    PoseLandmarkType.leftWrist,
    PoseLandmarkType.rightWrist,
    PoseLandmarkType.leftHip,
    PoseLandmarkType.rightHip,
    PoseLandmarkType.leftKnee,
    PoseLandmarkType.rightKnee,
    PoseLandmarkType.leftAnkle,
    PoseLandmarkType.rightAnkle,
  ];
}


/*

GETTING IN POSITION:
  CHECK IF EVERY NECESSARY JOINT ARE ON SCREEN (reliability > 0.8)
  CHECK IF START POSITION IS OKAY (for squat, if knee, hip, shoulder aligned)

COUTING REPETITION:
  FROM BEGINNING TO END:
    - BEGINNING: defined by start position, get position of interesting joint
    - END: defined by positions/distance interesting joints (knee and hip same level for squat,
          elbow and should same level for push up)
 */