import 'package:body_detection/models/pose.dart';
import 'package:body_detection/models/pose_landmark.dart';
import 'package:body_detection/models/pose_landmark_type.dart';
import 'package:flutter/widgets.dart';

class PosePainter extends CustomPainter {
  final Pose? pose;
  final Size imageSize;
  PosePainter({required this.pose, required this.imageSize, Key? key});

  final pointPaint = Paint()..color = const Color.fromRGBO(255, 255, 255, 0.8);
  final leftPointPaint = Paint()..color = const Color.fromRGBO(223, 157, 80, 1);
  final rightPointPaint = Paint()..color = const Color.fromRGBO(100, 208, 218, 1);
  final linePaint = Paint()
    ..color = const Color.fromARGB(228, 0, 0, 0)
    ..strokeWidth = 3;

  @override
  void paint(Canvas canvas, Size size) {
    if (pose == null) return;

    final double hRatio = imageSize.width == 0 ? 1 : size.width / imageSize.width;
    final double vRatio = imageSize.height == 0 ? 1 : size.height / imageSize.height;

    offsetForPart(PoseLandmark part) => Offset(part.position.x * hRatio, part.position.y * vRatio);

    // Landmark connections
    final landmarksByType = {for (final it in pose!.landmarks) it.type: it};
    for (final connection in connections) {
      if (landmarksByType[connection[0]] != null &&
          landmarksByType[connection[1]] != null &&
          landmarksByType[connection[0]]!.inFrameLikelihood > 0.5 &&
          landmarksByType[connection[1]]!.inFrameLikelihood > 0.5) {
        final point1 = offsetForPart(landmarksByType[connection[0]]!);
        final point2 = offsetForPart(landmarksByType[connection[1]]!);
        canvas.drawLine(point1, point2, linePaint);
      }
    }
  }

  @override
  bool shouldRepaint(PosePainter oldDelegate) {
    return oldDelegate.pose != pose || oldDelegate.imageSize != imageSize;
  }

  static const List<List<PoseLandmarkType>> connections = [
    [PoseLandmarkType.leftShoulder, PoseLandmarkType.rightShoulder],
    [PoseLandmarkType.leftShoulder, PoseLandmarkType.leftHip],
    [PoseLandmarkType.rightShoulder, PoseLandmarkType.rightHip],
    [PoseLandmarkType.rightShoulder, PoseLandmarkType.rightElbow],
    [PoseLandmarkType.rightWrist, PoseLandmarkType.rightElbow],
    [PoseLandmarkType.leftHip, PoseLandmarkType.rightHip],
    [PoseLandmarkType.leftHip, PoseLandmarkType.leftKnee],
    [PoseLandmarkType.rightHip, PoseLandmarkType.rightKnee],
    [PoseLandmarkType.rightKnee, PoseLandmarkType.rightAnkle],
    [PoseLandmarkType.leftKnee, PoseLandmarkType.leftAnkle],
    [PoseLandmarkType.leftElbow, PoseLandmarkType.leftShoulder],
    [PoseLandmarkType.leftWrist, PoseLandmarkType.leftElbow],
  ];
}