Genetics.swift

//
//  Genetics.swift
//  Genetics
//
//  Created by Julius Parishy on 12/1/14.
//  Copyright (c) 2014 Julius Parishy. All rights reserved.
//

import Cocoa
import CoreGraphics

struct Color {
    let r, g, b, a: CGFloat
}

func ==(lhs: Color, rhs: Color) -> Bool {
    return (lhs.r == rhs.r && lhs.g == rhs.g && lhs.b == rhs.b && lhs.a == rhs.a)
}

func closeness(original: CGFloat, next: CGFloat) -> CGFloat {
    let vals = (original, next)
    switch vals {
        case (0, 0):
            return 1.0
        
        case (0, _):
            return 1.0 - (abs(next - original) / next)
        
        default:
            return 1.0 - (min(original, abs(original - next)) / original)
    }
}

func closeness(original: Color, next: Color) -> CGFloat {
    let r = closeness(original.r, next.r)
    let g = closeness(original.g, next.g)
    let b = closeness(original.b, next.b)
    
    return (r + g + b) / 3.0
}

func bitmapContext(size: NSSize) -> CGContextRef {

    let width = UInt(size.width)
    let height = UInt(size.height)
    let bytesPerPixel: UInt = 4
    let dataLength = Int(width * bytesPerPixel * height)
    
    let colorSpace = CGColorSpaceCreateWithName(kCGColorSpaceGenericRGB)
    
    let data = UnsafeMutablePointer<UInt8>.null()
    
    let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.PremultipliedFirst.rawValue)
    let context = CGBitmapContextCreate(data, width, height, 8, width * bytesPerPixel, colorSpace, bitmapInfo)
    
    return context
}

func pixels(nsImage: NSImage) -> [Color] {
    let image = nsImage.CGImageForProposedRect(nil, context: nil, hints: nil)?.takeRetainedValue()
    
    let imageRep = nsImage.representations.first as NSImageRep
    let size = imageRep.size
    
    let width = Int(size.width)
    let height = Int(size.height)
    
    let context = bitmapContext(NSSize(width: CGFloat(width), height: CGFloat(height)))
    
    let rect = CGRect(x: 0.0, y: 0.0, width: CGFloat(width), height: CGFloat(height))
    CGContextDrawImage(context, rect, image)
    
    let renderedData = UnsafePointer<UInt8>(CGBitmapContextGetData(context))
    
    let numberOfPixels = Int(width * height)
    var pixels = [Color]()
    
    for i in 0..<numberOfPixels {
        let index = i * 4
        
        let r = CGFloat(renderedData[index + 1]) / 255.0
        let g = CGFloat(renderedData[index + 2]) / 255.0
        let b = CGFloat(renderedData[index + 3]) / 255.0
        let a = CGFloat(renderedData[index + 0]) / 255.0
        
        let pixel = Color(r: r, g: g, b: b, a: a)
        pixels.append(pixel)
    }
    
    return pixels
}

func zip<T>(a1: [T], a2: [T]) -> [(T, T)] {
    let smaller = a1.count < a2.count ? a1 : a2
    var zipped = [(T, T)]()
    for i in 0..<smaller.count {
        zipped.append((a1[i], a2[i]))
    }
    return zipped
}

func mapReduce<T, R>(values: [T], numberOfThreads: Int, map: (T) -> (R), initialValue: R, reduce: (R, R) -> (R), completion: (R) -> ()) {
    
    let segments = segment(values, numberOfThreads)
    
    var outputs = [[R]]()
    outputs.reserveCapacity(segments.count)
    
    let group = dispatch_group_create()
    
    var i = 0
    for s in segments {
        
        let queue = dispatch_queue_create("SegmentQueue_\(i)", DISPATCH_QUEUE_CONCURRENT)
        
        dispatch_group_async(group, queue) {
            
            var output = [R]()
            output.reserveCapacity(s.count)
            
            for v in s {
                let mv = map(v)
                output.append(mv)
            }
            
            outputs.append(output)
        }
        
        i += 1
    }
    
    let queue = dispatch_get_main_queue()
    
    dispatch_group_notify(group, queue) {
        
        var sum: R = initialValue
        
        for output in outputs {
            for v in output {
                sum = reduce(sum, v)
            }
        }
        
        completion(sum)
        
        return ()
    }
}

func _closeness(original: NSImage, next: NSImage, completion: (CGFloat) -> ()) {
    let originalPixels = pixels(original)
    let nextPixels = pixels(next)
    
    var sum: CGFloat = 0.0

    for i in (0..<originalPixels.count) {
        let op = originalPixels[i]
        let np = nextPixels[i]
        
        sum += closeness(op, np)
    }
    
    let c = sum / CGFloat(originalPixels.count)
    completion(c)
}

#if !DEBUG
let closeness: (NSImage, NSImage, (CGFloat) -> ()) -> () = {
    (original, next, completion) in
    
    let begin = NSDate()
    _closeness(original, next) {
        value in
        let end = NSDate()
        
        let duration = end.timeIntervalSinceDate(begin)
        println("Calculated closeness=\(value) in \(duration) seconds")
    
        completion(value)
    }
}
#else
let closeness = _closeness
#endif

func segment<T>(all: [T], segmentCount: Int) -> [[T]] {
    
    let perSegment = all.count / segmentCount
    
    var segments: [[T]] = [[T]]()
    var currentSegment = [T]()
    
    for i in (0..<all.count) {
        let a = all[i]
        
        if currentSegment.count == perSegment {
            segments.append(currentSegment)
            currentSegment = [T]()
        }
        
        currentSegment.append(a)
    }
    
    return segments
}

func mostSimilarImage(original: NSImage, current: NSImage, next: NSImage) -> NSImage {

    var currentCloseness: CGFloat = 0.0
    var nextCloseness: CGFloat = 0.0
    
    closeness(original, current) {
        c in
        currentCloseness = c
    }
    
    closeness(original, next) {
        c in
        nextCloseness = c
    }
    
    return nextCloseness > currentCloseness ? next : current
}

extension NSColor {
    class func randomComponent() -> CGFloat {
        return CGFloat(arc4random()) / CGFloat(UINT32_MAX)
    }
    
    class func random() -> NSColor {
        let r = randomComponent()
        let g = randomComponent()
        let b = randomComponent()
        let a = randomComponent()
        
        return NSColor(deviceRed: r, green: g, blue: b, alpha: a)
    }
    
    func mutate() -> NSColor {
        let r = evolution(self.redComponent) {
            return NSColor.randomComponent()
        }
        
        let g = evolution(self.greenComponent) {
            return NSColor.randomComponent()
        }
        
        let b = evolution(self.blueComponent) {
            return NSColor.randomComponent()
        }
        
        let a = evolution(self.alphaComponent) {
            return NSColor.randomComponent()
        }
        
        return NSColor(calibratedRed: r, green: g, blue: b, alpha: a)
    }
}

struct PolygonInfo {
    
    let size: NSSize
    let color: NSColor
    let points: [CGPoint]
    
    func mutate() -> PolygonInfo {
        
        return evolution(self) {
            let color = evolution(self.color) {
                return self.color.mutate()
            }
            
            let points: [CGPoint] = self.points.map {
                p in
                return evolution(p) {
                    let x = evolution(p.x) {
                        return CGFloat(arc4random_uniform(UInt32(self.size.width)))
                    }
                    
                    let y = evolution(p.x) {
                        return CGFloat(arc4random_uniform(UInt32(self.size.width)))
                    }
                    
                    return CGPoint(x: x, y: y)
                }
            }
            
            return PolygonInfo(size: self.size, color: color, points: points)
        }
    }
    
    static func random(size: NSSize) -> PolygonInfo {
        
        let numberOfPoints = 3 + arc4random_uniform(3)
        let points: [CGPoint] = (0..<numberOfPoints).map() { _ in
            let x = CGFloat(arc4random_uniform(UInt32(size.width)))
            let y = CGFloat(arc4random_uniform(UInt32(size.height)))
            return CGPoint(x: x, y: y)
        }
        
        let sorted = points.sorted {
            p1, p2 in
            
            let distanceToOrigin: (CGPoint) -> CGFloat = {
                p in
                return sqrt(p.x * p.x + p.y * p.y)
            }
            
            return distanceToOrigin(p1) < distanceToOrigin(p2)
        }

        return PolygonInfo(size: size, color: NSColor.random(), points: sorted)
    }
}

typealias DNA = [PolygonInfo]

let DefaultLevelOfEvolutionRandomness: CGFloat = 0.5

func evolution<T>(original: T, task: () -> (T)) -> T {
    let randomNumber = CGFloat(arc4random()) / CGFloat(UINT32_MAX)
    if randomNumber <= DefaultLevelOfEvolutionRandomness {
        return task()
    } else {
        return original
    }
}

func mutate(currentDNA: DNA, size: NSSize) -> DNA {
    
    let evolvedDNA: DNA = currentDNA.map {
        (polygonInfo) in
        return evolution(polygonInfo) {
            return polygonInfo.mutate()
        }
    }
    
    return evolution(evolvedDNA) {
        return evolvedDNA + [ PolygonInfo.random(size) ]
    }
}

func render(dna: DNA, size: NSSize) -> NSImage {
    
    let context = bitmapContext(size)
    
    CGContextSetFillColorWithColor(context, NSColor.whiteColor().CGColor)
    
    let rect = CGRect(x: 0, y: 0, width: size.width, height: size.height)
    CGContextFillRect(context, rect)
    
    for polygonInfo in dna {
        
        CGContextSetFillColorWithColor(context, polygonInfo.color.CGColor)
    
        CGContextBeginPath(context)
            
        let first = polygonInfo.points[0]
        CGContextMoveToPoint(context, first.x, first.y)
        
        let count = polygonInfo.points.count
        let points = polygonInfo.points[1..<count]
        
        for p in points {
            CGContextAddLineToPoint(context, p.x, p.y)
        }
        
        CGContextClosePath(context)
        CGContextFillPath(context)
    }
    
    let image = CGBitmapContextCreateImage(context)
    return NSImage(CGImage: image, size: size)
}

func step(originalImage: NSImage, currentImage: NSImage?, dna: DNA) -> (NSImage, DNA) {
    
    let imageRep = originalImage.representations.first as NSImageRep
    let size = imageRep.size
    let mutatedDNA = mutate(dna, size)
    
    let rendered = render(mutatedDNA, size)
    
    if let current = currentImage {
        let image = mostSimilarImage(originalImage, current, rendered)
        if image == current {
            return (current, dna)
        } else {
            return (rendered, mutatedDNA)
        }
    } else {
        return (rendered, mutatedDNA)
    }
}