diff --git a/api/src/main/java/org/openmrs/module/emrapi/db/HibernateVisitDAO.java b/api/src/main/java/org/openmrs/module/emrapi/db/HibernateVisitDAO.java
index 7d13c4db..99f9da93 100644
--- a/api/src/main/java/org/openmrs/module/emrapi/db/HibernateVisitDAO.java
+++ b/api/src/main/java/org/openmrs/module/emrapi/db/HibernateVisitDAO.java
@@ -8,7 +8,12 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Repository;
 
-import java.util.*;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 @Repository
@@ -30,7 +35,7 @@ public List<VisitWithDiagnoses> getVisitsByPatientId(Patient patient) {
                 "ORDER BY v.startDatetime DESC";
 
          List<Visit> visits = sessionFactory.getCurrentSession()
-                .createQuery(hqlVisit)  // For Hibernate versions < 5.2
+                .createQuery(hqlVisit)
                 .setParameter("patientId", patient.getId())
                 .setParameter("encounterTypeUuid", visitNoteEncounterTypeUuid).list();
 
@@ -39,12 +44,12 @@ public List<VisitWithDiagnoses> getVisitsByPatientId(Patient patient) {
                 "WHERE e.visit.id IN :visitIds";
 
         List<Integer> visitIds = visits.stream()
-                .map(Visit::getId)  // Extract the IDs from Visit objects
+                .map(Visit::getId)
                 .collect(Collectors.toList());
 
         List<Diagnosis> diagnoses = sessionFactory.getCurrentSession()
                 .createQuery(hqlDiagnosis)
-                .setParameterList("visitIds", visitIds)  // Use the list of IDs
+                .setParameterList("visitIds", visitIds)
                 .list();
 
         Map<Visit, Set<Diagnosis>> visitToDiagnosesMap = new HashMap<>();
diff --git a/api/src/main/java/org/openmrs/module/emrapi/visit/VisitWithDiagnoses.java b/api/src/main/java/org/openmrs/module/emrapi/visit/VisitWithDiagnoses.java
index 4d569d64..03d68058 100644
--- a/api/src/main/java/org/openmrs/module/emrapi/visit/VisitWithDiagnoses.java
+++ b/api/src/main/java/org/openmrs/module/emrapi/visit/VisitWithDiagnoses.java
@@ -2,7 +2,8 @@
 
 import lombok.Getter;
 import lombok.Setter;
-import org.openmrs.*;
+import org.openmrs.Diagnosis;
+import org.openmrs.Visit;
 
 import java.util.Set;