From 42f55d33d20c986f6dcdbb0cf67d1fb94c19fd2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joa=CC=83o=20Pereira?= <knight4@gmail.com>
Date: Thu, 17 Aug 2017 15:11:43 +0100
Subject: [PATCH] Add selected image

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

diff --git a/Segmentio/Source/Cells/SegmentioCell.swift b/Segmentio/Source/Cells/SegmentioCell.swift
index af70d89..187e6fd 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) {
+    func configure(selected: Bool, selectedImage:UIImage? = nil, image:UIImage? = nil) {
         cellSelected = selected
         
         let selectedState = options.states.selectedState
@@ -140,6 +140,12 @@ class SegmentioCell: UICollectionViewCell {
             segmentTitleLabel?.textColor = selected ? selectedState.titleTextColor : defaultState.titleTextColor
             segmentTitleLabel?.font = selected ? selectedState.titleFont : defaultState.titleFont
         }
+        
+        backgroundColor = selected ? selectedState.backgroundColor : defaultState.backgroundColor
+        
+        if (style != .onlyLabel) {
+            segmentImageView?.image = selected ? selectedImage : image
+        }
     }
     
     func configurateBadgeWithCount(_ badgeCount: Int?, color: UIColor?) {
diff --git a/Segmentio/Source/Segmentio.swift b/Segmentio/Source/Segmentio.swift
index b4797b9..783eafc 100644
--- a/Segmentio/Source/Segmentio.swift
+++ b/Segmentio/Source/Segmentio.swift
@@ -536,14 +536,16 @@ extension Segmentio: UICollectionViewDataSource {
             withReuseIdentifier: segmentioStyle.rawValue,
             for: indexPath) as! SegmentioCell
         
+        let content = segmentioItems[indexPath.row]
+        
         cell.configure(
-            content: segmentioItems[indexPath.row],
+            content: content,
             style: segmentioStyle,
             options: segmentioOptions,
             isLastCell: indexPath.row == segmentioItems.count - 1
         )
         
-        cell.configure(selected: (indexPath.row == selectedSegmentioIndex))
+        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 9c6f641..3addf9c 100644
--- a/Segmentio/Source/SegmentioOptions.swift
+++ b/Segmentio/Source/SegmentioOptions.swift
@@ -14,12 +14,14 @@ public struct SegmentioItem {
     
     public var title: String?
     public var image: UIImage?
+    public var selectedImage: UIImage?
     public var badgeCount: Int?
     public var badgeColor: UIColor?
 
-    public init(title: String?, image: UIImage?) {
+    public init(title: String?, image: UIImage?, selectedImage: UIImage? = nil) {
         self.title = title
         self.image = image
+        self.selectedImage = selectedImage
     }
     
     public mutating func addBadge(_ count: Int, color: UIColor) {