<commit_before>//
//  RowSortedMatrix.swift
//  SwiftyMath
//
//  Created by <NAME> on 2017/10/16.
//  Copyright © 2017年 <NAME>. All rights reserved.
//

import Foundation

internal typealias ComputationSpecializedRing = 𝐙

internal final class MatrixImpl<R: Ring>: Hashable, CustomStringConvertible {
    typealias Component = MatrixComponent<R>
    typealias Table =  [Int : [(Int, R)]]
    
    enum Alignment: String, Codable {
        case Rows, Cols
    }
    
    var rows: Int
    var cols: Int
    
    var align: Alignment
    var table: Table
    
    private init(_ rows: Int, _ cols: Int, _ align: Alignment, _ table: Table) {
        if _isDebugAssertConfiguration() {
            assert(table.values.forAll{ !$0.isEmpty })
            assert(table.values.forAll{ $0.forAll{ $0.1 != .zero }})
            
            switch align {
            case .Rows:
                assert(table.keys.forAll{ (0 ..< rows).contains($0) })
                assert(table.values.forAll{ $0.forAll{ (0 ..< cols).contains($0.0) } })
                assert(table.values.forAll{ $0.map{ $0.0 } == $0.map{ $0.0 }.unique().sorted() })
            case .Cols:
                assert(table.keys.forAll{ (0 ..< cols).contains($0) })
                assert(table.values.forAll{ $0.forAll{ (0 ..< rows).contains($0.0) } })
                assert(table.values.forAll{ $0.map{ $0.0 } == $0.map{ $0.0 }.unique().sorted() })
            }
        }
        
        self.rows = rows
        self.cols = cols
        self.align = align
        self.table = table
    }
    
    convenience init(rows: Int, cols: Int, align: Alignment = .Rows, grid: [R]) {
        let components = grid.enumerated().compactMap{ (k, a) -> Component? in
            (a != .zero) ? MatrixComponent(k / cols, k % cols, a) : nil
        }
        self.init(rows: rows, cols: cols, align: align, components: components)
    }
    
    convenience init(rows: Int, cols: Int, align: Alignment = .Rows, generator g: (Int, Int) -> R) {
        let components = (0 ..< rows).flatMap { i -> [Component] in
            (0 ..< cols).compactMap { j -> Component? in
                let a = g(i, j)
                return (a != .zero) ? Component(i, j, a) : nil
            }
        }
        self.init(rows: rows, cols: cols, align: align, components: components)
    }
    
    convenience init<S: Sequence>(rows: Int, cols: Int, align: Alignment = .Rows, components: S) where S.Element == Component {
        self.init(rows, cols, align, MatrixImpl.generateTable(align, components))
    }
    
    static func generateTable<S: Sequence>(_ align: Alignment, _ components: S) -> Table where S.Element == Component {
        let filtered = components.filter{ $0.value != .zero }
        switch align {
        case .Rows:
            return filtered.group{ c in c.row }.mapValues{ l in l.sorted{ c in c.col }.map{ c in (c.col, c.value) } }
        case .Cols:
            return filtered.group{ c in c.col }.mapValues{ l in l.sorted{ c in c.row }.map{ c in (c.row, c.value) } }
        }
    }
    
    subscript(i: Int, j: Int) -> R {
        get {
            switch align {
            case .Rows: return table[i]?.binarySearch(j, { $0.0 } )?.element.1 ?? .zero
            case .Cols: return table[j]?.binarySearch(i, { $0.0 } )?.element.1 ?? .zero
            }
        } set {
            switch align {
            case .Rows:
                let row = MatrixImpl.addRows(table[i] ?? [], [(j, newValue - self[i, j])])
                table[i] = row
            case .Cols:
                let col = MatrixImpl.addRows(table[j] ?? [], [(i, newValue - self[i, j])])
                table[j] = col
            }
        }
    }
    
    func copy() -> MatrixImpl<R> {
        return MatrixImpl(rows, cols, align, table)
    }
    
    func switchAlignment(_ align: Alignment) {
        if self.align == align {
            return
        }
        
        let comps = components
        
        self.align = align
        self.table = MatrixImpl.generateTable(align, comps)
    }
    
    var density: Double {
        let c = table.sum{ (_, list) in list.count }
        return Double(c) / Double(rows * cols)
    }
    
    var sparsity: Double {
        return 1.0 - density
    }
    
    var isZero: Bool {
        return table.isEmpty
    }
    
    var isDiagonal: Bool {
        return table.forAll { (i, list) in (list.count == 1) && list.first!.0 == i }
    }
    
    static func identity(size n: Int, align: Alignment) -> MatrixImpl<R> {
        let components = (0 ..< n).map{ i in MatrixComponent(i, i, R.identity)}
        return MatrixImpl(rows: n, cols: n, align: align, components: components)
    }
    
    var components: [Component] {
        switch align {
        case .Rows:
            return table.flatMap{ (i, list) in list.map{(j, a) in MatrixComponent(i, j, a)} }
        case .Cols:
            return table.flatMap{ (j, list) in list.map{(i, a) in MatrixComponent(i, j, a)} }
        }
    }
    
    func components(ofRow i: Int) -> [Component] {
        switch align {
        case .Rows:
            return table[i].map { list in list.map{ (j, r) in MatrixComponent(i, j, r) } } ?? []
        case .Cols:
            return table.compactMap{ (j, list) in list.binarySearch(i, { $0.0 } ).map{ MatrixComponent(i, j, $0.1.1) } }
        }
    }
    
    func components(ofCol j: Int) -> [Component] {
        switch align {
        case .Cols:
            return table[j].map { list in list.map{ (i, r) in MatrixComponent(i, j, r) } } ?? []
        case .Rows:
            return table.compactMap{ (i, list) in list.binarySearch(j, { $0.0 } ).map{ MatrixComponent(i, j, $0.1.1) } }
        }
    }
    
    var grid: [R] {
        var grid = Array(repeating: R.zero, count: rows * cols)
        for c in components {
            grid[c.row * cols + c.col] = c.value
        }
        return grid
    }
    
    @discardableResult
    func transpose() -> MatrixImpl<R> {
        (rows, cols) = (cols, rows)
        align = (align == .Rows) ? .Cols : .Rows
        return self
    }
    
    var transposed: MatrixImpl<R> {
        return self.copy().transpose()
    }
    
    func submatrix(rowRange: CountableRange<Int>) -> MatrixImpl<R> {
        return submatrix(rowRange, 0 ..< cols)
    }
    
    func submatrix(colRange: CountableRange<Int>) -> MatrixImpl<R> {
        return submatrix(0 ..< rows, colRange)
    }
    
    @_specialize(where R == ComputationSpecializedRing)
    func submatrix(_ rowRange: CountableRange<Int>, _ colRange: CountableRange<Int>) -> MatrixImpl<R> {
        assert(0 <= rowRange.lowerBound && rowRange.upperBound <= rows)
        assert(0 <= colRange.lowerBound && colRange.upperBound <= cols)

        return submatrix({i in rowRange.contains(i)}, {j in colRange.contains(j)})
    }
    
    func submatrix(_ rowCond: (Int) -> Bool, _ colCond: (Int) -> Bool) -> MatrixImpl<R> {
        let (sRows, sCols, iList, jList): (Int, Int, [Int], [Int])
        
        switch align {
        case .Rows:
            (iList, jList) = ((0 ..< rows).filter(rowCond), (0 ..< cols).filter(colCond))
            (sRows, sCols) = (iList.count, jList.count)
        case .Cols:
            (iList, jList) = ((0 ..< cols).filter(colCond), (0 ..< rows).filter(rowCond))
            (sRows, sCols) = (jList.count, iList.count)
        }
        
        let subTable = table.compactMap{ (i, list) -> (Int, [(Int, R)])? in
            guard let i1 = iList.binarySearch(i) else {
                return nil
            }
            let subList = list.compactMap{ (j, a) -> (Int, R)? in
                guard let j1 = jList.binarySearch(j) else {
                    return nil
                }
                return (j1, a)
            }
            return !subList.isEmpty ? (i1, subList) : nil
        }
        
        return MatrixImpl(sRows, sCols, align, Dictionary(pairs: subTable))
    }
    
    func mapValues<R2>(_ f: (R) -> R2) -> MatrixImpl<R2> {
        typealias M = MatrixImpl<R2>
        let mapped = table.mapValues { $0.map{ ($0, f($1)) } }
        let a = (align == .Rows) ? M.Alignment.Rows : M.Alignment.Cols
        return M(rows, cols, a, mapped)
    }
    
    @_specialize(where R == ComputationSpecializedRing)
    static func ==(a: MatrixImpl, b: MatrixImpl) -> Bool {
        if (a.rows, a.cols) != (b.rows, b.cols) {
            return false
        }
        return Set(a.components) == Set(b.components)
    }
    
    @_specialize(where R == ComputationSpecializedRing)
    static func +(a: MatrixImpl, b: MatrixImpl) -> MatrixImpl<R> {
        assert( (a.rows, a.cols) == (b.rows, b.cols) )
        
        b.switchAlignment(a.align)
        
        let iList = Set(a.table.keys).union(b.table.keys)
        let table = iList.compactMap { i in
            MatrixImpl.addRows(a.table[i] ?? [], b.table[i] ?? []).map{ (i, $0) }
        }<commit_msg>fix MatrixImpl.mapValue, accepts mapping to zero.<commit_after>//
//  RowSortedMatrix.swift
//  SwiftyMath
//
//  Created by <NAME> on 2017/10/16.
//  Copyright © 2017年 <NAME>. All rights reserved.
//

import Foundation

internal typealias ComputationSpecializedRing = 𝐙

internal final class MatrixImpl<R: Ring>: Hashable, CustomStringConvertible {
    typealias Component = MatrixComponent<R>
    typealias Table =  [Int : [(Int, R)]]
    
    enum Alignment: String, Codable {
        case Rows, Cols
    }
    
    var rows: Int
    var cols: Int
    
    var align: Alignment
    var table: Table
    
    private init(_ rows: Int, _ cols: Int, _ align: Alignment, _ table: Table) {
        if _isDebugAssertConfiguration() {
            assert(table.values.forAll{ !$0.isEmpty })
            assert(table.values.forAll{ $0.forAll{ $0.1 != .zero }})
            
            switch align {
            case .Rows:
                assert(table.keys.forAll{ (0 ..< rows).contains($0) })
                assert(table.values.forAll{ $0.forAll{ (0 ..< cols).contains($0.0) } })
                assert(table.values.forAll{ $0.map{ $0.0 } == $0.map{ $0.0 }.unique().sorted() })
            case .Cols:
                assert(table.keys.forAll{ (0 ..< cols).contains($0) })
                assert(table.values.forAll{ $0.forAll{ (0 ..< rows).contains($0.0) } })
                assert(table.values.forAll{ $0.map{ $0.0 } == $0.map{ $0.0 }.unique().sorted() })
            }
        }
        
        self.rows = rows
        self.cols = cols
        self.align = align
        self.table = table
    }
    
    convenience init(rows: Int, cols: Int, align: Alignment = .Rows, grid: [R]) {
        let components = grid.enumerated().compactMap{ (k, a) -> Component? in
            (a != .zero) ? MatrixComponent(k / cols, k % cols, a) : nil
        }
        self.init(rows: rows, cols: cols, align: align, components: components)
    }
    
    convenience init(rows: Int, cols: Int, align: Alignment = .Rows, generator g: (Int, Int) -> R) {
        let components = (0 ..< rows).flatMap { i -> [Component] in
            (0 ..< cols).compactMap { j -> Component? in
                let a = g(i, j)
                return (a != .zero) ? Component(i, j, a) : nil
            }
        }
        self.init(rows: rows, cols: cols, align: align, components: components)
    }
    
    convenience init<S: Sequence>(rows: Int, cols: Int, align: Alignment = .Rows, components: S) where S.Element == Component {
        self.init(rows, cols, align, MatrixImpl.generateTable(align, components))
    }
    
    static func generateTable<S: Sequence>(_ align: Alignment, _ components: S) -> Table where S.Element == Component {
        let filtered = components.filter{ $0.value != .zero }
        switch align {
        case .Rows:
            return filtered.group{ c in c.row }.mapValues{ l in l.sorted{ c in c.col }.map{ c in (c.col, c.value) } }
        case .Cols:
            return filtered.group{ c in c.col }.mapValues{ l in l.sorted{ c in c.row }.map{ c in (c.row, c.value) } }
        }
    }
    
    subscript(i: Int, j: Int) -> R {
        get {
            switch align {
            case .Rows: return table[i]?.binarySearch(j, { $0.0 } )?.element.1 ?? .zero
            case .Cols: return table[j]?.binarySearch(i, { $0.0 } )?.element.1 ?? .zero
            }
        } set {
            switch align {
            case .Rows:
                let row = MatrixImpl.addRows(table[i] ?? [], [(j, newValue - self[i, j])])
                table[i] = row
            case .Cols:
                let col = MatrixImpl.addRows(table[j] ?? [], [(i, newValue - self[i, j])])
                table[j] = col
            }
        }
    }
    
    func copy() -> MatrixImpl<R> {
        return MatrixImpl(rows, cols, align, table)
    }
    
    func switchAlignment(_ align: Alignment) {
        if self.align == align {
            return
        }
        
        let comps = components
        
        self.align = align
        self.table = MatrixImpl.generateTable(align, comps)
    }
    
    var density: Double {
        let c = table.sum{ (_, list) in list.count }
        return Double(c) / Double(rows * cols)
    }
    
    var sparsity: Double {
        return 1.0 - density
    }
    
    var isZero: Bool {
        return table.isEmpty
    }
    
    var isDiagonal: Bool {
        return table.forAll { (i, list) in (list.count == 1) && list.first!.0 == i }
    }
    
    static func identity(size n: Int, align: Alignment) -> MatrixImpl<R> {
        let components = (0 ..< n).map{ i in MatrixComponent(i, i, R.identity)}
        return MatrixImpl(rows: n, cols: n, align: align, components: components)
    }
    
    var components: [Component] {
        switch align {
        case .Rows:
            return table.flatMap{ (i, list) in list.map{(j, a) in MatrixComponent(i, j, a)} }
        case .Cols:
            return table.flatMap{ (j, list) in list.map{(i, a) in MatrixComponent(i, j, a)} }
        }
    }
    
    func components(ofRow i: Int) -> [Component] {
        switch align {
        case .Rows:
            return table[i].map { list in list.map{ (j, r) in MatrixComponent(i, j, r) } } ?? []
        case .Cols:
            return table.compactMap{ (j, list) in list.binarySearch(i, { $0.0 } ).map{ MatrixComponent(i, j, $0.1.1) } }
        }
    }
    
    func components(ofCol j: Int) -> [Component] {
        switch align {
        case .Cols:
            return table[j].map { list in list.map{ (i, r) in MatrixComponent(i, j, r) } } ?? []
        case .Rows:
            return table.compactMap{ (i, list) in list.binarySearch(j, { $0.0 } ).map{ MatrixComponent(i, j, $0.1.1) } }
        }
    }
    
    var grid: [R] {
        var grid = Array(repeating: R.zero, count: rows * cols)
        for c in components {
            grid[c.row * cols + c.col] = c.value
        }
        return grid
    }
    
    @discardableResult
    func transpose() -> MatrixImpl<R> {
        (rows, cols) = (cols, rows)
        align = (align == .Rows) ? .Cols : .Rows
        return self
    }
    
    var transposed: MatrixImpl<R> {
        return self.copy().transpose()
    }
    
    func submatrix(rowRange: CountableRange<Int>) -> MatrixImpl<R> {
        return submatrix(rowRange, 0 ..< cols)
    }
    
    func submatrix(colRange: CountableRange<Int>) -> MatrixImpl<R> {
        return submatrix(0 ..< rows, colRange)
    }
    
    @_specialize(where R == ComputationSpecializedRing)
    func submatrix(_ rowRange: CountableRange<Int>, _ colRange: CountableRange<Int>) -> MatrixImpl<R> {
        assert(0 <= rowRange.lowerBound && rowRange.upperBound <= rows)
        assert(0 <= colRange.lowerBound && colRange.upperBound <= cols)

        return submatrix({i in rowRange.contains(i)}, {j in colRange.contains(j)})
    }
    
    func submatrix(_ rowCond: (Int) -> Bool, _ colCond: (Int) -> Bool) -> MatrixImpl<R> {
        let (sRows, sCols, iList, jList): (Int, Int, [Int], [Int])
        
        switch align {
        case .Rows:
            (iList, jList) = ((0 ..< rows).filter(rowCond), (0 ..< cols).filter(colCond))
            (sRows, sCols) = (iList.count, jList.count)
        case .Cols:
            (iList, jList) = ((0 ..< cols).filter(colCond), (0 ..< rows).filter(rowCond))
            (sRows, sCols) = (jList.count, iList.count)
        }
        
        let subTable = table.compactMap{ (i, list) -> (Int, [(Int, R)])? in
            guard let i1 = iList.binarySearch(i) else {
                return nil
            }
            let subList = list.compactMap{ (j, a) -> (Int, R)? in
                guard let j1 = jList.binarySearch(j) else {
                    return nil
                }
                return (j1, a)
            }
            return !subList.isEmpty ? (i1, subList) : nil
        }
        
        return MatrixImpl(sRows, sCols, align, Dictionary(pairs: subTable))
    }
    
    func mapValues<R2>(_ f: (R) -> R2) -> MatrixImpl<R2> {
        typealias M = MatrixImpl<R2>
        let mapped = table.mapValues { list in
                list.map { (i, r) in (i, f(r)) }
                    .exclude { (_, r) in r == .zero }
            }.exclude{ (_, list) in list.isEmpty }
        let a = (align == .Rows) ? M.Alignment.Rows : M.Alignment.Cols
        return M(rows, cols, a, mapped)
    }
    
    @_specialize(where R == ComputationSpecializedRing)
    static func ==(a: MatrixImpl, b: MatrixImpl) -> Bool {
        if (a.rows, a.cols) != (b.rows, b.cols) {
            return false
        }
        return Set(a.components) == Set(b.components)
    }
    
    @_specialize(where R == ComputationSpecializedRing)
    static func +(a: MatrixImpl, b: MatrixImpl) -> MatrixImpl<R> {
        assert( (a.rows, a.cols) == (b.rows, b.cols) )
        
        b.switchAlignment(a.align)
        
        let iList = Set(a.table.keys).union(b.table.keys)
        let table = iList.compactMap { i in
            MatrixImpl.addRows(a.table[i] ?? [], b.table[i] ?? []).map{ (i, $0) }
        }
{'fraction_non_alphanumeric': 0.14537861915367484, 'fraction_numerical': 0.009187082405345211, 'mean_word_length': 1.8037777083983766, 'pattern_counts': {'":': 0, '<': 81, '<?xml version=': 0, '>': 93, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5383962', 'n_tokens_mistral': 5957, 'n_tokens_neox': 5846, 'n_words': 1501}