diff --git a/editors/editoria/src/config/configMobile.js b/editors/editoria/src/config/configMobile.js
index a3f5ea244c3d0839adbc3f835869a55e83ee6e49..45eb49d4720104d145838676f563c449efca07e8 100644
--- a/editors/editoria/src/config/configMobile.js
+++ b/editors/editoria/src/config/configMobile.js
@@ -34,6 +34,7 @@ import {
 } from 'wax-prosemirror-services';
 
 import { WaxSelectionPlugin } from 'wax-prosemirror-plugins';
+import { DefaultSchema } from 'wax-prosemirror-utilities';
 
 import invisibles, {
   space,
@@ -63,7 +64,7 @@ export default {
         'Images',
         'CodeBlock',
         'Tables',
-        'TrackChange',
+        // 'TrackChange',
       ],
     },
     {
@@ -72,6 +73,7 @@ export default {
     },
   ],
 
+  SchemaService: DefaultSchema,
   RulesService: [emDash, ellipsis],
   ShortCutsService: {},
   TitleService: { updateTitle },
diff --git a/editors/editoria/src/layout/EditoriaLayout.js b/editors/editoria/src/layout/EditoriaLayout.js
index c01a058061daf7f9e21b0113e33def786f2b6a49..5871d1f6d83dfd915f70fb9bc00fa7237a1db5ee 100644
--- a/editors/editoria/src/layout/EditoriaLayout.js
+++ b/editors/editoria/src/layout/EditoriaLayout.js
@@ -196,38 +196,6 @@ const getNotes = main => {
   return notes;
 };
 
-const getCommentsTracks = main => {
-  const marks = DocumentHelpers.findInlineNodes(main.state.doc);
-  const commentsTracks = [];
-  marks.map(node => {
-    if (node.node.marks.length > 0) {
-      node.node.marks.filter(mark => {
-        if (
-          mark.type.name === 'comment' ||
-          mark.type.name === 'insertion' ||
-          mark.type.name === 'deletion' ||
-          mark.type.name === 'format_change'
-        ) {
-          mark.pos = node.pos;
-          commentsTracks.push(mark);
-        }
-      });
-    }
-  });
-  return commentsTracks;
-};
-
-const getTrackBlockNodes = main => {
-  const allBlockNodes = DocumentHelpers.findBlockNodes(main.state.doc);
-  const trackBlockNodes = [];
-  allBlockNodes.map(node => {
-    if (node.node.attrs.track && node.node.attrs.track.length > 0) {
-      trackBlockNodes.push(node);
-    }
-  });
-  return trackBlockNodes;
-};
-
 const LeftSideBar = ComponentPlugin('leftSideBar');
 const MainMenuToolBar = ComponentPlugin('mainMenuToolBar');
 const NotesArea = ComponentPlugin('notesArea');
@@ -258,8 +226,10 @@ const EditoriaLayout = ({ editor }) => {
     };
   }
   const notes = main && getNotes(main);
-  const commentsTracks = main && getCommentsTracks(main).length;
-  const trackBlockNodes = main && getTrackBlockNodes(main).length;
+  const commentsTracksCount =
+    main && DocumentHelpers.getCommentsTracksCount(main);
+  const trackBlockNodesCount =
+    main && DocumentHelpers.getTrackBlockNodesCount(main);
 
   const areNotes = notes && !!notes.length && notes.length > 0;
 
@@ -302,7 +272,7 @@ const EditoriaLayout = ({ editor }) => {
                 <CommentsContainer>
                   <CommentTrackToolsContainer>
                     <CommentTrackTools>
-                      {commentsTracks + trackBlockNodes} COMMENTS AND
+                      {commentsTracksCount + trackBlockNodesCount} COMMENTS AND
                       SUGGESTIONS
                       <CommentTrackOptions>
                         <CommentTrackToolBar />
diff --git a/wax-prosemirror-components/package.json b/wax-prosemirror-components/package.json
index 60b004bee2f97e85e2a682657bc7d39f34acf13b..458c4a92d9320cccbded83cc587248586acfab5d 100644
--- a/wax-prosemirror-components/package.json
+++ b/wax-prosemirror-components/package.json
@@ -4,7 +4,7 @@
   "version": "0.0.29",
   "description": "Wax prosemirror UI components",
   "license": "MIT",
-  "main": "dist/index.js",
+  "main": "index.js",
   "files": [
     "dist"
   ],
diff --git a/wax-prosemirror-components/src/components/findAndReplace/ExandedFindAndReplaceComponent.js b/wax-prosemirror-components/src/components/findAndReplace/ExandedFindAndReplaceComponent.js
index 37702d5961905ae420945b43dea7c45961eaed35..d9ee659ce34e0f1ab7318e7e974966e6c062510d 100644
--- a/wax-prosemirror-components/src/components/findAndReplace/ExandedFindAndReplaceComponent.js
+++ b/wax-prosemirror-components/src/components/findAndReplace/ExandedFindAndReplaceComponent.js
@@ -139,13 +139,18 @@ const PreviousNextButton = styled.span`
 
 const ExandedFindAndReplaceComponent = ({
   close,
+  findNextMatch,
+  findPreviousMatch,
   matchCaseOption,
   nonExpandedText,
+  setMatchCaseValue,
 }) => {
-  const { app, view } = useContext(WaxContext);
+  const { app, view, activeViewId } = useContext(WaxContext);
   const searchRef = useRef(null);
   const replaceRef = useRef(null);
   const [searchValue, setSearchValue] = useState(nonExpandedText);
+  const [matchCaseSearch, setMatchCaseSearch] = useState(false);
+  const [match, setMatch] = useState([]);
   const [replaceValue, setReplaceValue] = useState('');
   const [counterText, setCounterText] = useState('0 of 0');
   const findAndReplacePlugin = app.PmPlugins.get('findAndReplacePlugin');
@@ -158,7 +163,7 @@ const ExandedFindAndReplaceComponent = ({
 
   const delayedSearch = useCallback(
     debounce(() => searchDocument(), 300),
-    [searchValue],
+    [searchValue, matchCaseSearch],
   );
 
   const onChangeSearchInput = () => {
@@ -167,18 +172,68 @@ const ExandedFindAndReplaceComponent = ({
 
   useEffect(() => {
     delayedSearch();
-  }, [searchValue, delayedSearch, JSON.stringify(allStates)]);
+  }, [searchValue, delayedSearch, matchCaseSearch, JSON.stringify(allStates)]);
+
+  const setCounterSearches = counter => {
+    if (counter === 0) return setCounterText('0 of 0');
+    setCounterText(`0 of ${counter}`);
+
+    const {
+      state: {
+        selection: { from, to },
+      },
+    } = view[activeViewId];
+
+    const results = helpers.getAllResultsByView(
+      view,
+      searchValue,
+      matchCaseSearch,
+    );
+    const resultsFrom = helpers.getResultsFrom(results);
+    let counterMatch = 0;
+    if (activeViewId === 'main') {
+      const matchFound = results.main.filter(result => {
+        return from === result.from && to === result.to;
+      });
+
+      setMatch(matchFound);
+      if (matchFound.length === 1) {
+        setCounterText(`${resultsFrom.main.indexOf(from) + 1} of ${counter}`);
+      }
+    } else {
+      if (resultsFrom.main) counterMatch = resultsFrom.main.length;
+      const notesIds = helpers.getNotesIds(view.main);
+
+      for (let i = 0; i < notesIds.length; i += 1) {
+        if (resultsFrom[notesIds[i]] && activeViewId !== notesIds[i]) {
+          counterMatch += resultsFrom[notesIds[i]].length;
+        }
+        if (resultsFrom[notesIds[i]] && activeViewId === notesIds[i]) {
+          const matchFound = results[notesIds[i]].filter(result => {
+            return from === result.from && to === result.to;
+          });
+          setMatch(matchFound);
+
+          if (matchFound.length === 1) {
+            counterMatch += resultsFrom[notesIds[i]].indexOf(from) + 1;
+            setCounterText(`${counterMatch} of ${counter}`);
+          }
+          break;
+        }
+      }
+    }
+  };
 
   const searchDocument = () => {
-    setCounterText('0 of 0');
     let counter = 0;
     findAndReplacePlugin.props.setSearchText(searchValue);
-    counter = helpers.getMatchesByView(view, searchValue);
+    findAndReplacePlugin.props.setSearchMatchCase(matchCaseSearch);
+    counter = helpers.getMatchesByView(view, searchValue, matchCaseSearch);
 
-    if (counter > 0) setCounterText(`1 of ${counter}`);
+    setCounterSearches(counter);
 
     if (searchRef.current === document.activeElement) {
-      each(view, (singleView, viewId) => {
+      eachRight(view, (singleView, viewId) => {
         singleView.dispatch(singleView.state.tr);
       });
     }
@@ -189,9 +244,15 @@ const ExandedFindAndReplaceComponent = ({
   };
 
   const replace = () => {
-    // const { from, to } = results[0];
-    // dispatch(state.tr.insertText(replace, from, to));
+    if (match.length === 1) {
+      const { from, to } = match[0];
+      view[activeViewId].dispatch(
+        view[activeViewId].state.tr.insertText(replaceValue, from, to),
+      );
+      findNextMatch(searchValue, matchCaseOption);
+    }
   };
+
   const replaceAll = () => {
     each(view, (singleView, viewId) => {
       const results = helpers.findMatches(singleView.state.doc, searchValue);
@@ -214,12 +275,10 @@ const ExandedFindAndReplaceComponent = ({
     close();
   };
 
-  const findNext = () => {
-    // console.log('next');
-  };
-
-  const findPrevious = () => {
-    // console.log('previous');
+  const matchCase = () => {
+    setMatchCaseSearch(!matchCaseOption);
+    setMatchCaseValue(!matchCaseOption);
+    searchRef.current.focus();
   };
 
   return (
@@ -254,16 +313,25 @@ const ExandedFindAndReplaceComponent = ({
           checked={matchCaseOption}
           label="Case Sensitive"
           name="case-sensitive"
+          onChange={matchCase}
         />
       </CheckBoxWrapper>
       <ControlContainer>
         <ButtonReplace onClick={replace}>Replace</ButtonReplace>
         <ButtonReplaceAll onClick={replaceAll}>Replace All</ButtonReplaceAll>
         <PreviousNextContainer>
-          <PreviousNextButton onClick={findPrevious} role="button" tabIndex="0">
+          <PreviousNextButton
+            onClick={() => findPreviousMatch(searchValue, matchCaseOption)}
+            role="button"
+            tabIndex="0"
+          >
             <StyledIcon name="navigatePrevious" />
           </PreviousNextButton>
-          <PreviousNextButton onClick={findNext} role="button" tabIndex="0">
+          <PreviousNextButton
+            onClick={() => findNextMatch(searchValue, matchCaseOption)}
+            role="button"
+            tabIndex="0"
+          >
             <StyledIcon name="navigateNext" />
           </PreviousNextButton>
         </PreviousNextContainer>
diff --git a/wax-prosemirror-components/src/components/findAndReplace/FindAndReplaceComponent.js b/wax-prosemirror-components/src/components/findAndReplace/FindAndReplaceComponent.js
index 5e6e7ebc1e9ef36aaf6b1b45c22a6d5889f14514..8bbe0b533c06850f35b8d88b92001a65847200fe 100644
--- a/wax-prosemirror-components/src/components/findAndReplace/FindAndReplaceComponent.js
+++ b/wax-prosemirror-components/src/components/findAndReplace/FindAndReplaceComponent.js
@@ -1,9 +1,16 @@
 /* eslint react/prop-types: 0 */
-import React, { useState } from 'react';
+import React, { useState, useContext, useEffect } from 'react';
+import { TextSelection } from 'prosemirror-state';
+import { WaxContext } from 'wax-prosemirror-core';
 import FindComponent from './FindComponent';
 import ExandedFindAndReplaceComponent from './ExandedFindAndReplaceComponent';
+import helpers from './helpers';
+
+let lastActiveViewId;
+let lastSelection;
 
 const FindAndReplaceComponent = ({ close }) => {
+  const { view, activeViewId } = useContext(WaxContext);
   const [isExpanded, setExpanded] = useState(false);
   const [nonExpandedText, setNonExpandedText] = useState('');
   const [matchCaseOption, setMatchCaseOption] = useState(false);
@@ -19,16 +26,197 @@ const FindAndReplaceComponent = ({ close }) => {
     setMatchCaseOption(matchCase);
   };
 
+  useEffect(() => {
+    if (view[activeViewId].state.selection.from !== 0) {
+      lastSelection = view[activeViewId].state.selection;
+      lastActiveViewId = activeViewId;
+    }
+  }, []);
+
+  const nextInNotes = (notesIds, results, findViewWithMatches) => {
+    for (let i = 0; i < notesIds.length; i += 1) {
+      if (results[notesIds[i]].length > 0 && notesIds[i] !== lastActiveViewId) {
+        helpers.clearViewSelection(view, lastActiveViewId);
+        helpers.moveToMatch(view, notesIds[i], results, 0);
+        lastActiveViewId = findViewWithMatches;
+        lastSelection = view[lastActiveViewId].state.selection;
+
+        break;
+      }
+    }
+  };
+
+  const findNextMatch = (searchValue, matchCaseSearch) => {
+    const counter = helpers.getMatchesByView(
+      view,
+      searchValue,
+      matchCaseSearch,
+    );
+    if (counter === 0) return;
+
+    lastActiveViewId = activeViewId;
+    lastSelection = view[activeViewId].state.selection;
+    const results = helpers.getAllResultsByView(
+      view,
+      searchValue,
+      matchCaseSearch,
+    );
+    const resultsFrom = helpers.getResultsFrom(results);
+    const notesIds = helpers.getNotesIds(view.main);
+    const findViewWithMatches = helpers.findViewWithMatches(
+      results,
+      view,
+      lastActiveViewId,
+    );
+    /* if no matches are found on focused view */
+    if (!resultsFrom[lastActiveViewId]) {
+      view[findViewWithMatches].dispatch(
+        view[findViewWithMatches].state.tr.setSelection(
+          new TextSelection(view[findViewWithMatches].state.tr.doc.resolve(1)),
+        ),
+      );
+      view[findViewWithMatches].focus();
+      lastActiveViewId = findViewWithMatches;
+      lastSelection = view[lastActiveViewId].state.selection;
+    }
+
+    const found = helpers.getClosestMatch(
+      lastSelection.from,
+      resultsFrom[lastActiveViewId],
+    );
+
+    const position = resultsFrom[lastActiveViewId].indexOf(found);
+    /* User selection lesser than found */
+    if (lastSelection.from < found) {
+      helpers.moveToMatch(view, lastActiveViewId, results, position);
+    }
+    /* User selection greater than found move to next if not already at the end of results for the view */
+    if (
+      lastSelection.from >= found &&
+      position < resultsFrom[lastActiveViewId].length - 1
+    ) {
+      helpers.moveToMatch(view, lastActiveViewId, results, position + 1);
+    }
+
+    /* Last result of the specific view. Move to next view */
+    if (
+      (lastSelection.from === found &&
+        position === resultsFrom[lastActiveViewId].length - 1) ||
+      (lastSelection.from >= found &&
+        position === resultsFrom[lastActiveViewId].length - 1)
+    ) {
+      /* End of results in notes move to main if results exist */
+      if (
+        notesIds.indexOf(lastActiveViewId) === notesIds.length - 1 &&
+        results.main.length > 0
+      ) {
+        setTimeout(() => {
+          lastActiveViewId = findViewWithMatches;
+          lastSelection = view[lastActiveViewId].state.selection;
+        }, 50);
+        helpers.moveToMatch(view, 'main', results, 0);
+        helpers.clearViewSelection(view, lastActiveViewId);
+      } else {
+        nextInNotes(notesIds, results, findViewWithMatches);
+      }
+    }
+  };
+
+  const findPreviousMatch = (searchValue, matchCaseSearch) => {
+    const counter = helpers.getMatchesByView(
+      view,
+      searchValue,
+      matchCaseSearch,
+    );
+    if (counter === 0) return;
+
+    lastActiveViewId = activeViewId;
+    lastSelection = view[activeViewId].state.selection;
+    const results = helpers.getAllResultsByView(
+      view,
+      searchValue,
+      matchCaseSearch,
+    );
+    const resultsFrom = helpers.getResultsFrom(results);
+    const notesIds = helpers.getNotesIds(view.main);
+
+    const found = helpers.getClosestMatch(
+      lastSelection.from,
+      resultsFrom[lastActiveViewId],
+      false,
+    );
+    const position = resultsFrom[lastActiveViewId].indexOf(found);
+
+    /* User selection lesser than found */
+    if (lastSelection.from > found) {
+      helpers.moveToMatch(view, lastActiveViewId, results, position);
+    }
+
+    if (lastSelection.from <= found && position !== 0) {
+      helpers.moveToMatch(view, lastActiveViewId, results, position - 1);
+    }
+
+    if (lastSelection.from === found && position === 0) {
+      if (lastActiveViewId === 'main') {
+        for (let i = notesIds.length - 1; i >= 0; i -= 1) {
+          if (
+            results[notesIds[i]].length > 0 &&
+            notesIds[i] !== lastActiveViewId
+          ) {
+            helpers.moveToMatch(
+              view,
+              notesIds[i],
+              results,
+              results[notesIds[i]].length - 1,
+            );
+            lastSelection = view[activeViewId].state.selection;
+            lastActiveViewId = activeViewId;
+            helpers.clearViewSelection(view, lastActiveViewId);
+            break;
+          }
+        }
+      } else if (
+        notesIds[notesIds.length - 1] === activeViewId &&
+        notesIds.length > 1
+      ) {
+        for (let i = notesIds.length - 1; i >= 0; i -= 1) {
+          if (
+            results[notesIds[i]].length > 0 &&
+            notesIds[i] !== lastActiveViewId
+          ) {
+            helpers.moveToMatch(
+              view,
+              notesIds[i],
+              results,
+              results[notesIds[i]].length - 1,
+            );
+            lastSelection = view[activeViewId].state.selection;
+            lastActiveViewId = activeViewId;
+            helpers.clearViewSelection(view, lastActiveViewId);
+            break;
+          } else {
+            console.log('go to main', lastActiveViewId);
+          }
+        }
+      }
+    }
+  };
+
   return isExpanded ? (
     <ExandedFindAndReplaceComponent
       close={close}
+      findNextMatch={findNextMatch}
+      findPreviousMatch={findPreviousMatch}
       matchCaseOption={matchCaseOption}
       nonExpandedText={nonExpandedText}
+      setMatchCaseValue={getMatchCaseOption}
     />
   ) : (
     <FindComponent
       close={close}
       expand={expand}
+      findNextMatch={findNextMatch}
+      findPreviousMatch={findPreviousMatch}
       setMatchCaseValue={getMatchCaseOption}
       setPreviousSearcValue={getNonExpandedText}
     />
diff --git a/wax-prosemirror-components/src/components/findAndReplace/FindComponent.js b/wax-prosemirror-components/src/components/findAndReplace/FindComponent.js
index fcc59b27aed63f1e3271dd76c1aaf0cc85ca9705..b36e2aa366c717a2e7114cff608009c96e5d8d73 100644
--- a/wax-prosemirror-components/src/components/findAndReplace/FindComponent.js
+++ b/wax-prosemirror-components/src/components/findAndReplace/FindComponent.js
@@ -7,7 +7,6 @@ import React, {
   useCallback,
   useEffect,
 } from 'react';
-import { TextSelection } from 'prosemirror-state';
 import { debounce, each, eachRight } from 'lodash';
 import styled from 'styled-components';
 import { grid } from '@pubsweet/ui-toolkit';
@@ -94,13 +93,13 @@ const Svg = styled.svg.attrs(() => ({
   width: 24px;
 `;
 
-let lastActiveViewId;
-let lastSelection;
 const FindComponent = ({
   close,
   expand,
   setPreviousSearcValue,
   setMatchCaseValue,
+  findNextMatch,
+  findPreviousMatch,
 }) => {
   const { app, view, activeViewId } = useContext(WaxContext);
   const searchRef = useRef(null);
@@ -119,7 +118,7 @@ const FindComponent = ({
 
   const delayedSearch = useCallback(
     debounce(() => searchDocument(), 300),
-    [searchValue, matchCaseSearch],
+    [searchValue, matchCaseSearch, activeViewId],
   );
 
   const onChange = () => {
@@ -127,11 +126,6 @@ const FindComponent = ({
   };
 
   useEffect(() => {
-    if (view[activeViewId].state.selection.from !== 0) {
-      lastSelection = view[activeViewId].state.selection;
-      lastActiveViewId = activeViewId;
-    }
-
     delayedSearch();
     if (isFirstRun) {
       setTimeout(() => {
@@ -141,20 +135,61 @@ const FindComponent = ({
     }
   }, [searchValue, delayedSearch, matchCaseSearch, JSON.stringify(allStates)]);
 
-  const searchDocument = () => {
-    setCounterText('0 of 0');
-    let counter = 0;
-    findAndReplacePlugin.props.setSearchText(searchValue);
-    findAndReplacePlugin.props.setSearchMatchCase(matchCaseSearch);
-    counter = helpers.getMatchesByView(view, searchValue, matchCaseSearch);
+  const setCounterSearches = counter => {
+    if (counter === 0) return setCounterText('0 of 0');
+    setCounterText(`0 of ${counter}`);
+
+    const {
+      state: {
+        selection: { from, to },
+      },
+    } = view[activeViewId];
+
     const results = helpers.getAllResultsByView(
       view,
       searchValue,
       matchCaseSearch,
     );
-    if (results.main) {
+    const resultsFrom = helpers.getResultsFrom(results);
+    let counterMatch = 0;
+    if (activeViewId === 'main') {
+      const match = results.main.filter(result => {
+        return from === result.from && to === result.to;
+      });
+
+      if (match.length === 1) {
+        setCounterText(`${resultsFrom.main.indexOf(from) + 1} of ${counter}`);
+      }
+    } else {
+      if (resultsFrom.main) counterMatch = resultsFrom.main.length;
+      const notesIds = helpers.getNotesIds(view.main);
+
+      for (let i = 0; i < notesIds.length; i += 1) {
+        if (resultsFrom[notesIds[i]] && activeViewId !== notesIds[i]) {
+          counterMatch += resultsFrom[notesIds[i]].length;
+        }
+        if (resultsFrom[notesIds[i]] && activeViewId === notesIds[i]) {
+          const match = results[notesIds[i]].filter(result => {
+            return from === result.from && to === result.to;
+          });
+
+          if (match.length === 1) {
+            counterMatch += resultsFrom[notesIds[i]].indexOf(from) + 1;
+            setCounterText(`${counterMatch} of ${counter}`);
+          }
+          break;
+        }
+      }
     }
-    if (counter > 0) setCounterText(`1 of ${counter}`);
+  };
+
+  const searchDocument = () => {
+    let counter = 0;
+    findAndReplacePlugin.props.setSearchText(searchValue);
+    findAndReplacePlugin.props.setSearchMatchCase(matchCaseSearch);
+    counter = helpers.getMatchesByView(view, searchValue, matchCaseSearch);
+
+    setCounterSearches(counter);
 
     if (searchRef.current === document.activeElement) {
       eachRight(view, (singleView, viewId) => {
@@ -182,157 +217,6 @@ const FindComponent = ({
     searchRef.current.focus();
   };
 
-  const findNext = () => {
-    lastActiveViewId = activeViewId;
-    lastSelection = view[activeViewId].state.selection;
-    const results = helpers.getAllResultsByView(
-      view,
-      searchValue,
-      matchCaseSearch,
-    );
-    const resultsFrom = helpers.getResultsFrom(results);
-    const notesIds = helpers.getNotesIds(view.main);
-    const findViewWithMatches = helpers.findViewWithMatches(
-      results,
-      view,
-      lastActiveViewId,
-    );
-    /* if no matches are found on focused view */
-    if (!resultsFrom[lastActiveViewId]) {
-      view[findViewWithMatches].dispatch(
-        view[findViewWithMatches].state.tr.setSelection(
-          new TextSelection(view[findViewWithMatches].state.tr.doc.resolve(1)),
-        ),
-      );
-      view[findViewWithMatches].focus();
-      lastActiveViewId = findViewWithMatches;
-      lastSelection = view[lastActiveViewId].state.selection;
-    }
-
-    const found = helpers.getClosestMatch(
-      lastSelection.from,
-      resultsFrom[lastActiveViewId],
-    );
-    const position = resultsFrom[lastActiveViewId].indexOf(found);
-    /* User selection lesser than found */
-    if (lastSelection.from < found) {
-      helpers.moveToMatch(view, lastActiveViewId, results, position);
-    }
-    /* User selection greater than found move to next if not already at the end of results for the view */
-    if (
-      lastSelection.from >= found &&
-      position < resultsFrom[lastActiveViewId].length - 1
-    ) {
-      helpers.moveToMatch(view, lastActiveViewId, results, position + 1);
-    }
-
-    /* Last result of the specific view. Move to next view */
-    if (
-      lastSelection.from === found &&
-      position === resultsFrom[lastActiveViewId].length - 1
-    ) {
-      /* End of results in notes move to main if results exist */
-      if (
-        notesIds.indexOf(lastActiveViewId) === notesIds.length - 1 &&
-        results.main.length > 0
-      ) {
-        setTimeout(() => {
-          lastActiveViewId = findViewWithMatches;
-          lastSelection = view[lastActiveViewId].state.selection;
-        }, 50);
-        helpers.moveToMatch(view, 'main', results, 0);
-        helpers.clearViewSelection(view, lastActiveViewId);
-      } else {
-        for (let i = 0; i < notesIds.length; i += 1) {
-          if (
-            results[notesIds[i]].length > 0 &&
-            notesIds[i] !== lastActiveViewId
-          ) {
-            helpers.clearViewSelection(view, lastActiveViewId);
-            helpers.moveToMatch(view, notesIds[i], results, 0);
-            lastActiveViewId = findViewWithMatches;
-            lastSelection = view[lastActiveViewId].state.selection;
-
-            break;
-          }
-        }
-      }
-    }
-  };
-
-  const findPrevious = () => {
-    lastActiveViewId = activeViewId;
-    lastSelection = view[activeViewId].state.selection;
-    const results = helpers.getAllResultsByView(
-      view,
-      searchValue,
-      matchCaseSearch,
-    );
-    const resultsFrom = helpers.getResultsFrom(results);
-    const notesIds = helpers.getNotesIds(view.main);
-
-    const found = helpers.getClosestMatch(
-      lastSelection.from,
-      resultsFrom[lastActiveViewId],
-      false,
-    );
-    const position = resultsFrom[lastActiveViewId].indexOf(found);
-
-    /* User selection lesser than found */
-    if (lastSelection.from > found) {
-      helpers.moveToMatch(view, lastActiveViewId, results, position);
-    }
-
-    if (lastSelection.from <= found && position !== 0) {
-      helpers.moveToMatch(view, lastActiveViewId, results, position - 1);
-    }
-
-    if (lastSelection.from === found && position === 0) {
-      if (lastActiveViewId === 'main') {
-        for (let i = notesIds.length - 1; i >= 0; i -= 1) {
-          if (
-            results[notesIds[i]].length > 0 &&
-            notesIds[i] !== lastActiveViewId
-          ) {
-            helpers.moveToMatch(
-              view,
-              notesIds[i],
-              results,
-              results[notesIds[i]].length - 1,
-            );
-            lastSelection = view[activeViewId].state.selection;
-            lastActiveViewId = activeViewId;
-            helpers.clearViewSelection(view, lastActiveViewId);
-            break;
-          }
-        }
-      } else if (
-        notesIds[notesIds.length - 1] === activeViewId &&
-        notesIds.length > 1
-      ) {
-        for (let i = notesIds.length - 1; i >= 0; i -= 1) {
-          if (
-            results[notesIds[i]].length > 0 &&
-            notesIds[i] !== lastActiveViewId
-          ) {
-            helpers.moveToMatch(
-              view,
-              notesIds[i],
-              results,
-              results[notesIds[i]].length - 1,
-            );
-            lastSelection = view[activeViewId].state.selection;
-            lastActiveViewId = activeViewId;
-            helpers.clearViewSelection(view, lastActiveViewId);
-            break;
-          } else {
-            console.log('go to main', lastActiveViewId);
-          }
-        }
-      }
-    }
-  };
-
   return (
     <Wrapper>
       <SingleRow>
@@ -352,10 +236,18 @@ const FindComponent = ({
             <path d="M2.5,4v3h5v12h3V7h5V4H2.5z M21.5,9h-9v3h3v7h3v-7h3V9z" />
           </Svg>
         </IconWrapper>
-        <IconWrapper onClick={findPrevious} role="button" tabIndex="0">
+        <IconWrapper
+          onClick={() => findPreviousMatch(searchValue, matchCaseSearch)}
+          role="button"
+          tabIndex="0"
+        >
           <StyledIcon name="navigatePrevious" />
         </IconWrapper>
-        <IconWrapper onClick={findNext} role="button" tabIndex="0">
+        <IconWrapper
+          onClick={() => findNextMatch(searchValue, matchCaseSearch)}
+          role="button"
+          tabIndex="0"
+        >
           <StyledIcon name="navigateNext" />
         </IconWrapper>
 
diff --git a/wax-prosemirror-components/src/components/findAndReplace/helpers.js b/wax-prosemirror-components/src/components/findAndReplace/helpers.js
index 3860959bc1553df32b014a4cbc2f0dc34053feef..88c373fec3bed0d535be2d4a171f2d40d5015b5f 100644
--- a/wax-prosemirror-components/src/components/findAndReplace/helpers.js
+++ b/wax-prosemirror-components/src/components/findAndReplace/helpers.js
@@ -37,7 +37,7 @@ const findMatches = (doc, searchValue, matchCase) => {
     }
   });
   mergedTextNodes.forEach(({ text, pos }) => {
-    const search = RegExp(searchValue, matchCase ? 'gu' : 'gui');
+    const search = RegExp(escapeRegExp(searchValue), matchCase ? 'gu' : 'gui');
     let m;
     // eslint-disable-next-line no-cond-assign
     while ((m = search.exec(text))) {
@@ -54,6 +54,10 @@ const findMatches = (doc, searchValue, matchCase) => {
   return results;
 };
 
+const escapeRegExp = string => {
+  return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+};
+
 const getMatchesByView = (views, searchValue, matchCase) => {
   let allResults = 0;
   each(views, (singleView, viewId) => {
diff --git a/wax-prosemirror-components/src/components/trackChanges/TrackChangeOptionsComponent.js b/wax-prosemirror-components/src/components/trackChanges/TrackChangeOptionsComponent.js
index 5e9e65003706ef4bd3591b383b8a18f76e83e33e..ae95ecb265d236b97de8b47033bdc942ab1edab2 100644
--- a/wax-prosemirror-components/src/components/trackChanges/TrackChangeOptionsComponent.js
+++ b/wax-prosemirror-components/src/components/trackChanges/TrackChangeOptionsComponent.js
@@ -163,12 +163,22 @@ const getTrackBlockNodes = main => {
 };
 
 const getComments = main => {
-  const comments = DocumentHelpers.findChildrenByMark(
+  const comments = [];
+  const commentsNodes = DocumentHelpers.findChildrenByMark(
     main.state.doc,
     main.state.schema.marks.comment,
     true,
   );
-  return comments;
+  commentsNodes.map(node => {
+    if (node.node.marks.length > 0) {
+      node.node.marks.filter(mark => {
+        if (mark.type.name === 'comment') {
+          comments.push(mark);
+        }
+      });
+    }
+  });
+  return [...new Set(comments.map(item => item.attrs.id))];
 };
 
 const TrackChangeOptionsComponent = ({
diff --git a/wax-prosemirror-components/src/ui/inputs/CheckBox.js b/wax-prosemirror-components/src/ui/inputs/CheckBox.js
index 13cbf70bed205d4c364a11b35b9c064e99aa08c9..970f02f7c8d1ef0336ff57894ceedd8793d2c0df 100644
--- a/wax-prosemirror-components/src/ui/inputs/CheckBox.js
+++ b/wax-prosemirror-components/src/ui/inputs/CheckBox.js
@@ -96,11 +96,17 @@ const CheckboxTitle = styled.div`
 `;
 
 const CheckBox = props => {
-  const { checked, name, label } = props;
+  const { checked, name, label, onChange } = props;
   return (
     <CheckBoxContainer>
       <CheckBoxLabel>
-        <input defaultChecked={checked} id={name} name={name} type="checkbox" />
+        <input
+          defaultChecked={checked}
+          id={name}
+          name={name}
+          onChange={onChange}
+          type="checkbox"
+        />
         <CheckboxCustom />
         <CheckboxTitle>{label}</CheckboxTitle>
       </CheckBoxLabel>
diff --git a/wax-prosemirror-core/package.json b/wax-prosemirror-core/package.json
index 994092fee0d5d9c394cc87bd091f39287f12bae4..0c08578f75155ea5fa5c04e8c2cf06f375d544a8 100644
--- a/wax-prosemirror-core/package.json
+++ b/wax-prosemirror-core/package.json
@@ -4,7 +4,7 @@
   "version": "0.0.29",
   "description": "Wax prosemirror core",
   "license": "MIT",
-  "main": "dist/index.js",
+  "main": "index.js",
   "files": [
     "dist"
   ],
diff --git a/wax-prosemirror-core/src/Wax.js b/wax-prosemirror-core/src/Wax.js
index d11128ade6545a3272cf685724cac5934d7de3d0..09f15c7de3ff401a91269811555dd098f66949eb 100644
--- a/wax-prosemirror-core/src/Wax.js
+++ b/wax-prosemirror-core/src/Wax.js
@@ -1,10 +1,12 @@
 /* eslint react/prop-types: 0 */
 import React, { useEffect, useState } from 'react';
 import debounce from 'lodash/debounce';
+
 import { DOMSerializer, DOMParser } from 'prosemirror-model';
-import { DefaultSchema } from 'wax-prosemirror-utilities';
+
 import WaxProvider from './WaxContext';
 import Application from './Application';
+
 import WaxView from './WaxView';
 import defaultPlugins from './plugins/defaultPlugins';
 import Placeholder from './plugins/placeholder';
@@ -29,9 +31,10 @@ const serializer = schema => {
   };
 };
 
+let schema;
 const createApplication = props => {
   const application = Application.create(props);
-  application.getSchema();
+  schema = application.getSchema();
   application.bootServices();
   return application;
 };
@@ -65,8 +68,9 @@ const Wax = props => {
   } = props;
 
   if (!application) return null;
-  const WaxOnchange = onChange ? onChange : value => true;
-  const { schema } = application.schema;
+  // const { schema } = application.schema;
+  const WaxOnchange = onChange || (v => true);
+
   const editorContent = value || '';
 
   finalPlugins = defaultPlugins.concat([
@@ -135,7 +139,7 @@ const Wax = props => {
 };
 
 Wax.defaultProps = {
-  config: { SchemaService: DefaultSchema, services: [] },
+  config: { services: [] },
 };
 
 export default Wax;
diff --git a/wax-prosemirror-plugins/package.json b/wax-prosemirror-plugins/package.json
index 1d0c9e56f943ed39f99c803d8db34e5c734e382b..0452847f2321229ea9c58e14a68208c694f15be2 100644
--- a/wax-prosemirror-plugins/package.json
+++ b/wax-prosemirror-plugins/package.json
@@ -4,7 +4,7 @@
   "version": "0.0.29",
   "description": "Wax prosemirror plugins",
   "license": "MIT",
-  "main": "dist/index.js",
+  "main": "index.js",
   "files": [
     "dist"
   ],
diff --git a/wax-prosemirror-plugins/src/findAndReplace/FindAndReplacePlugin.js b/wax-prosemirror-plugins/src/findAndReplace/FindAndReplacePlugin.js
index 3eafc64377729d445f34d4c76e60d6c01d5eee67..403d392932c4e2901aae2b0ef8081354e384e407 100644
--- a/wax-prosemirror-plugins/src/findAndReplace/FindAndReplacePlugin.js
+++ b/wax-prosemirror-plugins/src/findAndReplace/FindAndReplacePlugin.js
@@ -42,7 +42,7 @@ const findMatches = (doc, searchValue) => {
     }
   });
   mergedTextNodes.forEach(({ text, pos }) => {
-    const search = RegExp(searchValue, matchCase ? 'gu' : 'gui');
+    const search = RegExp(escapeRegExp(searchValue), matchCase ? 'gu' : 'gui');
     let m;
     // eslint-disable-next-line no-cond-assign
     while ((m = search.exec(text))) {
@@ -59,6 +59,10 @@ const findMatches = (doc, searchValue) => {
   return results;
 };
 
+const escapeRegExp = string => {
+  return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+};
+
 export default props => {
   return new Plugin({
     key: findAndReplacePlugin,
diff --git a/wax-prosemirror-schema/package.json b/wax-prosemirror-schema/package.json
index 2b15b3322437f9bcc45620875b14e992ef45a7ad..814499940ff3938081b18d3fb1d793456ee09648 100644
--- a/wax-prosemirror-schema/package.json
+++ b/wax-prosemirror-schema/package.json
@@ -4,7 +4,7 @@
   "version": "0.0.29",
   "description": "Wax prosemirror schema",
   "license": "MIT",
-  "main": "dist/index.js",
+  "main": "index.js",
   "files": [
     "dist"
   ],
diff --git a/wax-prosemirror-services/package.json b/wax-prosemirror-services/package.json
index 244e08209b53e4f64d7277d257c44cc96e201080..9a313454edcaa8aeb0d8df35a706505e7e880e6b 100644
--- a/wax-prosemirror-services/package.json
+++ b/wax-prosemirror-services/package.json
@@ -4,7 +4,7 @@
   "version": "0.0.29",
   "description": "Wax prosemirror services",
   "license": "MIT",
-  "main": "dist/index.js",
+  "main": "index.js",
   "files": [
     "dist"
   ],
diff --git a/wax-prosemirror-services/src/NoteService/Editor.js b/wax-prosemirror-services/src/NoteService/Editor.js
index 79d3f2766bd41c470d5e0454813b1c5ee9fb9875..54faedc6aa7fdd18654dcd4dcbac796fbc3769e4 100644
--- a/wax-prosemirror-services/src/NoteService/Editor.js
+++ b/wax-prosemirror-services/src/NoteService/Editor.js
@@ -2,7 +2,7 @@
 import React, { useEffect, useRef, useContext, useMemo } from 'react';
 import { filter } from 'lodash';
 import { EditorView } from 'prosemirror-view';
-import { EditorState } from 'prosemirror-state';
+import { EditorState, TextSelection } from 'prosemirror-state';
 import { StepMap } from 'prosemirror-transform';
 import { baseKeymap } from 'prosemirror-commands';
 import { keymap } from 'prosemirror-keymap';
@@ -44,6 +44,13 @@ export default ({ node, view }) => {
             // the parent editor is focused.
             if (noteView.hasFocus()) noteView.focus();
           },
+          blur: view => {
+            view.dispatch(
+              view.state.tr.setSelection(
+                TextSelection.create(view.state.doc, 0),
+              ),
+            );
+          },
         },
         transformPasted: slice => {
           return transformPasted(slice, noteView);
diff --git a/wax-prosemirror-utilities/package.json b/wax-prosemirror-utilities/package.json
index b38b207df6e04939c852d45e4cb3ed00a00a9dc4..b37dd6985a08eb050a7c8a4ded598804a2830f67 100644
--- a/wax-prosemirror-utilities/package.json
+++ b/wax-prosemirror-utilities/package.json
@@ -4,7 +4,7 @@
   "version": "0.0.29",
   "description": "Wax prosemirror utilities",
   "license": "MIT",
-  "main": "dist/index.js",
+  "main": "index.js",
   "files": [
     "dist"
   ],
diff --git a/wax-prosemirror-utilities/src/document/DocumentHelpers.js b/wax-prosemirror-utilities/src/document/DocumentHelpers.js
index b5b0f7b9eeeeb2302f0c116f53bed69ceba0f39e..cea418a1119ab7ecb6d110b4c69877953b51b29f 100644
--- a/wax-prosemirror-utilities/src/document/DocumentHelpers.js
+++ b/wax-prosemirror-utilities/src/document/DocumentHelpers.js
@@ -1,5 +1,3 @@
-import { isObject } from 'lodash';
-
 const findMark = (state, PMmark, toArr = false) => {
   const {
     selection: { $from, $to },
@@ -27,6 +25,43 @@ const findMark = (state, PMmark, toArr = false) => {
   return markFound;
 };
 
+const getCommentsTracksCount = main => {
+  const marks = findInlineNodes(main.state.doc);
+  const commentsTracksFormat = [];
+  const insertionsDeletions = [];
+  marks.map(node => {
+    if (node.node.marks.length > 0) {
+      node.node.marks.filter(mark => {
+        if (
+          mark.type.name === 'comment' ||
+          mark.type.name === 'format_change'
+        ) {
+          commentsTracksFormat.push(mark);
+        } else if (
+          mark.type.name === 'insertion' ||
+          mark.type.name === 'deletion'
+        ) {
+          insertionsDeletions.push(mark);
+        }
+      });
+    }
+  });
+  const unique = [...new Set(commentsTracksFormat.map(item => item.attrs.id))];
+  const total = unique.length + insertionsDeletions.length;
+  return total;
+};
+
+const getTrackBlockNodesCount = main => {
+  const allBlockNodes = findBlockNodes(main.state.doc);
+  const trackBlockNodes = [];
+  allBlockNodes.map(node => {
+    if (node.node.attrs.track && node.node.attrs.track.length > 0) {
+      trackBlockNodes.push(node);
+    }
+  });
+  return trackBlockNodes.length;
+};
+
 /* TODO */
 /* this is a hacky workaround for now to find marks
   that are pm will break them. Correct way is to be done
@@ -172,4 +207,6 @@ export default {
   findFragmentedMark,
   findAllMarksWithSameId,
   findMarkPosition,
+  getCommentsTracksCount,
+  getTrackBlockNodesCount,
 };