Retain The State Of On Off Button In A UITableViewCell - Swift - swift

Question

I have been searching through many posts on SO but couldn't find an answer to this one.
I have a Table view listing various items. Each cell in the table view has a button that swaps an image around when clicked on, effectively working as an "on" or "off" button to show a user which items in the list they have selected. I have a variable inside my custom Cell Prototype class which stores a value of true or false which is updated every time the button is clicked on.
There is a "Done" button in the Table View that when tapped on calls an unwind Segue to go back to the first View Controller.
When the user taps on the Enter Table View button (on the first View Controller) to display the Table the buttons all go back to their default state, am guessing because each time the segue to the Table View Screen happens it creates a new instance of the Table to be displayed.
What I'm trying to achieve is that the state of the button (either on or off) is retained when going back into the Table screen. I've tried for a while sending an Integer value back from the TableCell Class (using a delegate) to the first View controller and then passing that value back into the Table View controller when the forward segue is called in order to have a "retained from the previous state value" that can be compared against when the cells are created to indicate if a button had been clicked or not. Couldn't get it to work though to save the state of the buttons. The list of Items in the Table will also change depending on what a user adds.
Some Screen shots and the code are below. The code hasn't got the delegate i was trying included (as it didn't work) but if its needed I will edit the post to include it. Any help would be hugely appreciated. Many Thanks!!
Code I have so far is:
//Main View Controller//
class ViewController: UIViewController {
override func viewDidLoad() {
super.viewDidLoad()
}
//segue to the Table View Screen
override func prepare(for segue: UIStoryboardSegue, sender: Any?) {
if segue.identifier == "EnterTableView" {
let newTableView = segue.destination as! TableViewController
}
}
// Unwind Segue Called on Exit From Table View
#IBAction func unwindToMainViewController (_segue:UIStoryboardSegue) {
}
}
// Table View Controller //
class TableViewController: UITableViewController {
var dataArray = ["A", "B", "C", "D"]
override func viewDidLoad() {
super.viewDidLoad()
}
// MARK: - Table view data source
override func numberOfSections(in tableView: UITableView) -> Int {
//create one section for table
return 1
}
override func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
//create number of rows based on the number of items in the dataArray set above
return dataArray.count
}
// an array that will contain all the cells
var cellArray = [UITableViewCell] ()
override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
//create a new cell based on the cell class "TableViewCell"
let cell = tableView.dequeueReusableCell(withIdentifier: "TableViewCell", for: indexPath) as? TableViewCell
// variable to contain a single element from the dataArray - updates every time a new cell is created
let dataArrayForCells = dataArray [indexPath.row]
//set cell label text to show the value from the dataArrayForCells
cell!.label.text = dataArrayForCells
cellArray.append(cell!)
return cell!
}
// Table Cell Class//
class TableViewCell: UITableViewCell {
#IBOutlet weak var label: UILabel! // label to hold display some text
#IBOutlet weak var buttonImage: UIImageView! //image for button
var buttonClicked = true //variable to contain weather the has been clicked
//if button is tapped on run the below
#IBAction func aButton(_ sender: Any) {
//if button is clicked is true, swap image to red ("on") button, set buttonClicked value to false
if buttonClicked {
buttonImage.image = #imageLiteral(resourceName: "Rec Button Red")
buttonClicked = false
}
// if buttonClicked value is false swap image to grey ("off") button set buttonClicked value back to true
else {
buttonImage.image = #imageLiteral(resourceName: "Rec Button Grey")
buttonClicked = true
}
}