SoftArgmaxVisualizationViewController.swift

//
//  ViewController.swift
//  ArgmaxChartApp
//
//  Created by Doyoung Gwak on 2021/06/21.
//

import UIKit

class ViewController: UIViewController {
    
    var numberOfBar = 8 {
        willSet {
            bars.forEach { $0.removeFromSuperview() }
            bars.removeAll()
            valueLabels.forEach { $0.removeFromSuperview() }
            valueLabels.removeAll()
            
            values = (0..<numberOfBar).map { _ in Float.random(in: -1.0...6.0) }
            
            bars = values.map { _ in
                let bar = UIView()
                bar.backgroundColor = .darkGray
                view.addSubview(bar)
                return bar
            }
            
            valueLabels = values.map { _ in
                let label = UILabel()
                label.textAlignment = .center
                label.numberOfLines = 1
                label.font = .boldSystemFont(ofSize: 30)
                view.addSubview(label)
                return label
            }
            
            barTitleLabel1.text = "Argmax"
            barTitleLabel1.font = .systemFont(ofSize: 24, weight: .heavy)
            barTitleLabel1.textAlignment = .right
            view.addSubview(barTitleLabel1)
            barTitleLabel2.text = "Soft-Argmax"
            barTitleLabel2.font = .systemFont(ofSize: 24, weight: .heavy)
            barTitleLabel2.textAlignment = .right
            view.addSubview(barTitleLabel2)
            
            hBar1.backgroundColor = .systemGray
            view.addSubview(hBar1)
            hBar2.backgroundColor = .systemGray
            view.addSubview(hBar2)
            
            vBars1 = values.map { _ in
                let bar = UIView()
                bar.backgroundColor = .systemGray
                view.addSubview(bar)
                return bar
            }
            
            vBars2 = values.map { _ in
                let bar = UIView()
                bar.backgroundColor = .systemGray
                view.addSubview(bar)
                return bar
            }
            
            vBarLabels1 = values.map { _ in
                let label = UILabel()
                label.textAlignment = .center
                label.numberOfLines = 1
                label.font = .systemFont(ofSize: 20, weight: .semibold)
                label.textColor = .lightGray
                view.addSubview(label)
                return label
            }
            
            vBarLabels2 = values.map { _ in
                let label = UILabel()
                label.textAlignment = .center
                label.numberOfLines = 1
                label.font = .systemFont(ofSize: 20, weight: .semibold)
                label.textColor = .lightGray
                view.addSubview(label)
                return label
            }
            
            
            dotView1.layer.cornerRadius = dotViewWidth/2
            dotView2.layer.cornerRadius = dotViewWidth/2
            dotView1.backgroundColor = UIColor.white
            dotView2.backgroundColor = UIColor.white
            view.addSubview(dotView1)
            view.addSubview(dotView2)
            
            updateUIs()
        }
    }
    var values: [Float] = []
    var bars: [UIView] = []
    var valueLabels: [UILabel] = []
    let maxValue: CGFloat = 7.0
    let minValue: CGFloat = -1.0
    
    let barTitleLabel1 = UILabel()
    let barTitleLabel2 = UILabel()
    let hBar1 = UIView()
    let hBar2 = UIView()
    var vBars1: [UIView] = []
    var vBars2: [UIView] = []
    var vBarLabels1: [UILabel] = []
    var vBarLabels2: [UILabel] = []
    
    let dotView1 = UIView()
    let dotView2 = UIView()

    override func viewDidLoad() {
        super.viewDidLoad()
        
        numberOfBar = 8
    }
    
    let topMargin: CGFloat = 44
    let bottomMargin: CGFloat = 340.0
    let horizentalMargin: CGFloat = 204.0
    let horizentalGap: CGFloat = 8.0
    let labelHeight: CGFloat = 50.0
    let barTitleLabelLeftMargin: CGFloat = 8
    let barTitleLabelRightMargin: CGFloat = 4
    let barTitleLabelHeight: CGFloat = 40.0
    
    let dotViewWidth: CGFloat = 16
    
    
    func updateUIs() {
        updateBars()
        updateLabels()
        updateXaxis1()
        updateXaxis2()
        updateDots()
    }
    
    func updateBars() {
        let startingY: CGFloat = topMargin + (view.frame.height - topMargin - bottomMargin) * (maxValue/(maxValue - minValue))
        let barWidth: CGFloat = (view.frame.width - horizentalMargin*2 + horizentalGap) / CGFloat(numberOfBar) - horizentalGap
        
        zip(bars, values).enumerated().forEach { (offset, element) in
            let x: CGFloat = horizentalMargin + (barWidth + horizentalGap) * CGFloat(offset)
            let y: CGFloat = startingY
            let h: CGFloat = ((view.frame.height - topMargin - bottomMargin) * (maxValue/(maxValue - minValue))) * (-CGFloat(element.1)/maxValue)
            element.0.frame = CGRect(x: x, y: y, width: barWidth, height: h)
        }
    }

    func updateLabels() {
        let startingY: CGFloat = topMargin + (view.frame.height - topMargin - bottomMargin) * (maxValue/(maxValue - minValue))
        let barWidth: CGFloat = (view.frame.width - horizentalMargin*2 + horizentalGap) / CGFloat(numberOfBar) - horizentalGap
        
        zip(valueLabels, values).enumerated().forEach { (offset, element) in
            let x: CGFloat = horizentalMargin + (barWidth + horizentalGap) * CGFloat(offset)
            let y: CGFloat = startingY
            element.0.frame = CGRect(x: x, y: y - labelHeight/2.0, width: barWidth, height: labelHeight)
            element.0.text = String(format: "%.3f", element.1)
        }
    }
    
    func updateXaxis1() {
        let barHeight: CGFloat = 2.0
        let x: CGFloat = horizentalMargin - horizentalGap / 2
        let y: CGFloat = view.frame.height - bottomMargin + bottomMargin * (1.0/3.0)
        let w: CGFloat = view.frame.width - horizentalMargin*2 + horizentalGap
        hBar1.frame = CGRect(x: x, y: y, width: w, height: barHeight)
        
        let vBarHeight: CGFloat = 12.0
        let vBarWidth: CGFloat = 2.0
        vBars1.enumerated().forEach {
            $0.element.frame = CGRect(x: x + (w/CGFloat(numberOfBar)) * (CGFloat($0.offset) + 0.5) - vBarWidth/2, y: y + barHeight/2 - vBarHeight/2, width: vBarWidth, height: vBarHeight)
        }
        
        vBarLabels1.enumerated().forEach {
            $0.element.text = "\($0.offset)"
            $0.element.frame = CGRect(x: x + (w/CGFloat(numberOfBar))*CGFloat($0.offset), y: y + barHeight/2 + vBarHeight/2 + 8, width: (w/CGFloat(numberOfBar)), height: 20.0)
        }
        
        barTitleLabel1.frame = CGRect(x: barTitleLabelLeftMargin, y: y + barHeight/2 - barTitleLabelHeight/2, width: x - barTitleLabelLeftMargin - barTitleLabelRightMargin, height: barTitleLabelHeight)
    }
    
    func updateXaxis2() {
        let barHeight: CGFloat = 2.0
        let x: CGFloat = horizentalMargin - horizentalGap / 2
        let y: CGFloat = view.frame.height - bottomMargin + bottomMargin * (2.0/3.0)
        let w: CGFloat = view.frame.width - horizentalMargin*2 + horizentalGap
        hBar2.frame = CGRect(x: x, y: y, width: w, height: barHeight)
        
        let vBarHeight: CGFloat = 12.0
        let vBarWidth: CGFloat = 2.0
        vBars2.enumerated().forEach {
            $0.element.frame = CGRect(x: x + (w/CGFloat(numberOfBar)) * (CGFloat($0.offset) + 0.5) - vBarWidth/2, y: y + barHeight/2 - vBarHeight/2, width: vBarWidth, height: vBarHeight)
        }
        vBarLabels2.enumerated().forEach {
            $0.element.text = "\($0.offset)"
            $0.element.frame = CGRect(x: x + (w/CGFloat(numberOfBar))*CGFloat($0.offset), y: y + barHeight/2 + vBarHeight/2 + 8, width: (w/CGFloat(numberOfBar)), height: 20.0)
        }
        
        barTitleLabel2.frame = CGRect(x: barTitleLabelLeftMargin, y: y + barHeight/2 - barTitleLabelHeight/2, width: x - barTitleLabelLeftMargin - barTitleLabelRightMargin, height: barTitleLabelHeight)
    }
    
    func updateDots() {
        let index1 = argmax(values: values)
        dotView1.frame = CGRect(x: vBars1[Int(index1)].center.x - dotViewWidth/2,
                                y: vBars1[Int(index1)].center.y - dotViewWidth/2,
                                width: dotViewWidth, height: dotViewWidth)
        
        let index2 = softargmax(values: values)
        print(index2)
        let vBarWidth: CGFloat = 2.0
        let x: CGFloat = horizentalMargin - horizentalGap / 2
        let w: CGFloat = view.frame.width - horizentalMargin*2 + horizentalGap
        let datX2 = x + (w/CGFloat(numberOfBar)) * (CGFloat(index2) + 0.5) - vBarWidth/2
        dotView2.frame = CGRect(x: datX2,
                                y: vBars2[Int(index2)].center.y - dotViewWidth/2,
                                width: dotViewWidth, height: dotViewWidth)
    }
    
    override func viewWillLayoutSubviews() {
        super.viewWillLayoutSubviews()
        updateUIs()
    }
    
    var selectedValueIndex: Int? = nil
    var touchedStartedPoint: CGPoint = .zero
    var touchedStartedValue: Float = 0

    override func touchesBegan(_ touches: Set<UITouch>, with event: UIEvent?) {
        super.touchesBegan(touches, with: event)
        
        guard let touch = touches.first else {
            return
        }

        let p = touch.location(in: view)
        bars.enumerated().forEach {
            if $0.element.frame.minX <= p.x && p.x <= $0.element.frame.maxX {
                selectedValueIndex = $0.offset
            }
        }
        
        if let selectedValueIndex = selectedValueIndex {
            touchedStartedPoint = p
            touchedStartedValue = values[selectedValueIndex]
        }
    }
    
    override func touchesMoved(_ touches: Set<UITouch>, with event: UIEvent?) {
        super.touchesMoved(touches, with: event)
        guard let touch = touches.first else {
            return
        }
        guard let selectedValueIndex = selectedValueIndex else { return }
        
        let p = touch.location(in: view)
        let maxHeight = view.frame.height - topMargin - bottomMargin
        let unitHeight = (maxHeight * (maxValue/(maxValue - minValue))) / maxValue
        let vectorY = -(p.y - touchedStartedPoint.y)
        let vectorValue = (vectorY / unitHeight)
        
        values[selectedValueIndex] = touchedStartedValue + Float(vectorValue)
        updateUIs()
    }
    
    override func touchesEnded(_ touches: Set<UITouch>, with event: UIEvent?) {
        super.touchesEnded(touches, with: event)
        selectedValueIndex = nil
    }
    
    override func touchesCancelled(_ touches: Set<UITouch>, with event: UIEvent?) {
        super.touchesCancelled(touches, with: event)
        selectedValueIndex = nil
    }

}

import Accelerate

func argmax(values: [Float]) -> UInt {
    let stride = vDSP_Stride(1)

    let n = vDSP_Length(values.count)

    var c: Float = .nan
    var i: vDSP_Length = 0

    vDSP_maxvi(values,
               stride,
               &c,
               &i,
               n)

    return i
}

func softargmax(values: [Float]) -> Float {
    var x: [Float] = softmax(values)
    x *= (0..<x.count).map { Float($0) }
    return vDSP.sum(x)
}

func softmax(_ x: [Float]) -> [Float] {
    var x = x
    let len = vDSP_Length(x.count)
    
    // Find the maximum value in the input array.
    var max: Float = 0
    vDSP_maxv(x, 1, &max, len)
    
    // Subtract the maximum from all the elements in the array.
    // Now the highest value in the array is 0.
    max = -max
    vDSP_vsadd(x, 1, &max, &x, 1, len)
    
    // Exponentiate all the elements in the array.
    var count = Int32(x.count)
    vvexpf(&x, x, &count)
    
    // Compute the sum of all exponentiated values.
    var sum: Float = 0
    vDSP_sve(x, 1, &sum, len)
    
    // Divide each element by the sum. This normalizes the array contents
    // so that they all add up to 1.
    vDSP_vsdiv(x, 1, &sum, &x, 1, len)
    
    return x
}

extension Array where Element == Float {
    static func *=(lhs: inout Array<Float>, rhs: Array<Float>) {
        let stride = vDSP_Stride(1)
        let n = vDSP_Length(lhs.count)
        vDSP_vmul(rhs, stride, lhs, stride, &lhs, stride, n)
    }
    
    static func /=(lhs: inout Array<Float>, rhs: Float) {
        let newrhs = Array(repeating: 1 / rhs, count: lhs.count)
        let stride = vDSP_Stride(1)
        let n = vDSP_Length(lhs.count)
        vDSP_vmul(newrhs, stride, lhs, stride, &lhs, stride, n)
    }
}