From 72c4974a4393f8e74b08f2906a79646c370b9879 Mon Sep 17 00:00:00 2001
From: Ilya <i.tykhonkov@gmail.com>
Date: Fri, 1 Sep 2017 17:33:02 +0300
Subject: [PATCH] changes after review

---
 Segmentio/Source/Cells/SegmentioCell.swift | 2 +-
 Segmentio/Source/Segmentio.swift           | 8 ++++++--
 Segmentio/Source/SegmentioOptions.swift    | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/Segmentio/Source/Cells/SegmentioCell.swift b/Segmentio/Source/Cells/SegmentioCell.swift
index 187e6fd..3b83d1f 100644
--- a/Segmentio/Source/Cells/SegmentioCell.swift
+++ b/Segmentio/Source/Cells/SegmentioCell.swift
@@ -130,7 +130,7 @@ class SegmentioCell: UICollectionViewCell {
         configurateBadgeWithCount(content.badgeCount, color: content.badgeColor)
     }
     
-    func configure(selected: Bool, selectedImage:UIImage? = nil, image:UIImage? = nil) {
+    func configure(selected: Bool, selectedImage: UIImage? = nil, image: UIImage? = nil) {
         cellSelected = selected
         
         let selectedState = options.states.selectedState
diff --git a/Segmentio/Source/Segmentio.swift b/Segmentio/Source/Segmentio.swift
index f4da738..74c5638 100644
--- a/Segmentio/Source/Segmentio.swift
+++ b/Segmentio/Source/Segmentio.swift
@@ -96,7 +96,7 @@ open class Segmentio: UIView {
         collectionView.bounces = true
         collectionView.isScrollEnabled = segmentioOptions.scrollEnabled
         collectionView.backgroundColor = .clear
-			collectionView.accessibilityIdentifier = "segmentio_collection_view"
+        collectionView.accessibilityIdentifier = "segmentio_collection_view"
         
         segmentioCollectionView = collectionView
         
@@ -546,7 +546,11 @@ extension Segmentio: UICollectionViewDataSource {
             isLastCell: indexPath.row == segmentioItems.count - 1
         )
         
-        cell.configure(selected: (indexPath.row == selectedSegmentioIndex), selectedImage:content.selectedImage, image:content.image)
+        cell.configure(
+            selected: (indexPath.row == selectedSegmentioIndex),
+            selectedImage: content.selectedImage,
+            image: content.image
+        )
         
         return cell
     }
diff --git a/Segmentio/Source/SegmentioOptions.swift b/Segmentio/Source/SegmentioOptions.swift
index 3addf9c..2ca66ca 100644
--- a/Segmentio/Source/SegmentioOptions.swift
+++ b/Segmentio/Source/SegmentioOptions.swift
@@ -21,7 +21,7 @@ public struct SegmentioItem {
     public init(title: String?, image: UIImage?, selectedImage: UIImage? = nil) {
         self.title = title
         self.image = image
-        self.selectedImage = selectedImage
+        self.selectedImage = selectedImage ?? image
     }
     
     public mutating func addBadge(_ count: Int, color: UIColor) {