diff --git a/package.json b/package.json
index d358261..2f76f2b 100644
--- a/package.json
+++ b/package.json
@@ -12,7 +12,7 @@
     "copy:readme": "copyfiles ./README.md ./dist",
     "format:write": "prettier --write 'index.ts'",
     "test": "jest",
-    "start": "ts-node ./src/index.ts -s playground/sampleDir/main.ts -p '@custom'",
+    "start": "ts-node ./src/index.ts -s playground/sampleDir/main.ts -p '@myorg'",
     "start:staged": "ts-node ./src/index.ts --staged",
     "start:bar": "ts-node ./src/index.ts --dryRun -s playground/sampleDir/bar.ts",
     "release": "semantic-release",
diff --git a/src/conductor/format-import-statements.ts b/src/conductor/format-import-statements.ts
index 31901a3..a8daca4 100644
--- a/src/conductor/format-import-statements.ts
+++ b/src/conductor/format-import-statements.ts
@@ -1,23 +1,18 @@
 import { ImportCategories } from '../types';
 
+const categoriesOrder = ['thirdPartyImports', 'userLibraryImports', 'differentModuleImports', 'sameModuleImports'];
+
 export function formatImportStatements(importCategories: ImportCategories) {
-  const { differentModuleImports, sameModuleImports, thirdPartyImports, userLibraryImports } = importCategories;
-  let result = '';
+  const [first, ...otherCategories] = Object.entries(importCategories)
+    .filter(([, imports]) => imports.size > 0)
+    .sort(([a], [b]) => categoriesOrder.indexOf(a) - categoriesOrder.indexOf(b))
+    .map(([, imports]) => imports);
 
-  function updateResult(sortedPot: Map<string, string>, spaceBefore = true) {
-    if (sortedPot.size > 0 && spaceBefore) {
-      result += '\n\n';
-    }
-    [...sortedPot.values()].forEach(
-      (fullImportLiteral: string, index: number) =>
-        (result += index === sortedPot.size - 1 ? `${fullImportLiteral}` : `${fullImportLiteral}\n`)
-    );
-  }
+  let result = first ? [...first.values()].join('\n') : '';
 
-  updateResult(thirdPartyImports, false);
-  updateResult(userLibraryImports, thirdPartyImports.size > 0);
-  updateResult(differentModuleImports, thirdPartyImports.size > 0 || userLibraryImports.size > 0);
-  updateResult(sameModuleImports, thirdPartyImports.size > 0 || userLibraryImports.size > 0 || differentModuleImports.size > 0);
+  for (const imports of otherCategories) {
+    result += '\n\n' + [...imports.values()].join('\n');
+  }
 
   return result;
 }
diff --git a/src/conductor/get-import-statement-map.ts b/src/conductor/get-import-statement-map.ts
index 0610d04..904675a 100644
--- a/src/conductor/get-import-statement-map.ts
+++ b/src/conductor/get-import-statement-map.ts
@@ -10,9 +10,9 @@ export function getImportStatementMap(importNodes: ts.Node[]): Map<string, strin
 
   importNodes.forEach((node: ts.Node) => {
     const importSegments = node.getChildren();
-    let completeImportStatement = node.getFullText();
-    if (completeImportStatement.startsWith('\n')) {
-      completeImportStatement = completeImportStatement.replace('\n', '');
+    let importStatement = node.getFullText();
+    if (importStatement.startsWith('\n')) {
+      importStatement = importStatement.replace(/^\n*/, '');
     }
     const importLiteral = importSegments.find((segment) => segment.kind === ts.SyntaxKind.StringLiteral)?.getText();
 
@@ -21,13 +21,13 @@ export function getImportStatementMap(importNodes: ts.Node[]): Map<string, strin
     }
 
     const existingImport = importStatementMap.get(importLiteral);
-    const canMerge = autoMerge && existingImport && [existingImport, completeImportStatement].every((i) => !i.includes('*'));
+    const canMerge = autoMerge && existingImport && [existingImport, importStatement].every((i) => !i.includes('*'));
 
     if (canMerge) {
-      importStatementMap.set(importLiteral, mergeImportStatements(existingImport, completeImportStatement));
+      importStatementMap.set(importLiteral, mergeImportStatements(existingImport, importStatement));
     } else {
       const key = existingImport ? `${importLiteral}_${Math.random()}` : importLiteral;
-      importStatementMap.set(key, completeImportStatement);
+      importStatementMap.set(key, importStatement);
     }
   });
 
diff --git a/src/conductor/optimize-imports.ts b/src/conductor/optimize-imports.ts
index 198b001..e5c843b 100644
--- a/src/conductor/optimize-imports.ts
+++ b/src/conductor/optimize-imports.ts
@@ -43,15 +43,14 @@ export async function optimizeImports(filePath: string): Promise<string> {
   }
 
   let fileComment = getFileComment(fileContent);
+  // Remove the comment from the file content.
   if (fileComment) {
     fileContent = fileContent.replace(fileComment, '');
-    fileComment += '\n';
   }
 
   const rootNode = ts.createSourceFile(filePath, fileContent, ts.ScriptTarget.Latest, true);
   const importNodes = collectImportNodes(rootNode);
   const importStatementMap = getImportStatementMap(importNodes);
-
   if (importStatementMap.size === 0) {
     return actions.none;
   }
@@ -63,7 +62,14 @@ export async function optimizeImports(filePath: string): Promise<string> {
   const lastImport = importNodes.pop();
   const contentWithoutImportStatements = fileContent.slice(lastImport.end);
 
-  const updatedContent = `${fileComment}${result}${contentWithoutImportStatements}`;
+  let updatedContent = `${result}${contentWithoutImportStatements}`;
+
+  if (fileComment) {
+    // Add the comment back to the file content.
+    fileContent = `${fileComment}${fileContent}`;
+    updatedContent = `${fileComment}\n` + updatedContent;
+  }
+
   const fileHasChanged = updatedContent !== fileContent;
   if (fileHasChanged) {
     !dryRun && writeFileSync(filePath, updatedContent);
diff --git a/src/types.ts b/src/types.ts
index 3830b63..3cfdf86 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,8 +1,5 @@
 export interface ImportCategories {
-  thirdPartyImports: Map<string, string>;
-  userLibraryImports: Map<string, string>;
-  differentModuleImports: Map<string, string>;
-  sameModuleImports: Map<string, string>;
+  [key: string]: Map<string, string>;
 }
 
 export interface Config {