import Foundation
import Accelerate

public final class PCA: Codable, @unchecked Sendable {
    public let inputDimension: Int
    public let outputDimension: Int
    public let mean: [Float]
    public let components: [[Float]]
    public let explainedVariance: [Float]
    public let whiteningEnabled: Bool
    public let sampleCount: Int
    private let trainingData: [[Float]]

    enum CodingKeys: String, CodingKey {
        case inputDimension, outputDimension, mean, components, explainedVariance, whiteningEnabled, sampleCount, trainingData
    }

    public init(inputDimension: Int, outputDimension: Int, mean: [Float], components: [[Float]], explainedVariance: [Float], whiteningEnabled: Bool = false, sampleCount: Int = 0, trainingData: [[Float]] = []) {
        self.inputDimension = inputDimension
        self.outputDimension = outputDimension
        self.mean = mean
        self.components = components
        self.explainedVariance = explainedVariance
        self.whiteningEnabled = whiteningEnabled
        self.sampleCount = sampleCount
        self.trainingData = trainingData
    }

    public func transform(_ input: [Float]) throws -> [Float] {
        guard input.count == inputDimension else {
            throw PCAError.dimensionMismatch
        }
        var centered = [Float](repeating: 0, count: inputDimension)
        for i in 0..<inputDimension {
            centered[i] = input[i] - mean[i]
        }
        var result = [Float](repeating: 0, count: outputDimension)
        for j in 0..<outputDimension {
            var dot: Float = 0
            for i in 0..<inputDimension {
                dot += centered[i] * components[j][i]
            }
            result[j] = dot
        }
        return result
    }

    public func transformWhitened(_ input: [Float]) throws -> [Float] {
        var result = try transform(input)
        for j in 0..<outputDimension {
            let denom = sqrt(max(explainedVariance[j], 1e-10))
            result[j] /= denom
        }
        return result
    }

    public func explainedVarianceRatio() -> [Float] {
        let total = explainedVariance.reduce(0, +)
        guard total > 0 else { return explainedVariance.map { _ in 0 } }
        return explainedVariance.map { $0 / total }
    }

    public func cumulativeExplainedVarianceRatio() -> [Float] {
        let ratios = explainedVarianceRatio()
        var cum: [Float] = []
        var sum: Float = 0
        for r in ratios {
            sum += r
            cum.append(sum)
        }
        return cum
    }

    public func save(to url: URL) throws {
        let encoder = JSONEncoder()
        let data = try encoder.encode(self)
        try data.write(to: url)
    }

    public static func load(from url: URL) throws -> PCA {
        let data = try Data(contentsOf: url)
        let decoder = JSONDecoder()
        return try decoder.decode(PCA.self, from: data)
    }

    public static func train(data: [[Float]], outputDimension: Int, whitening: Bool = false) throws -> PCA {
        guard let first = data.first else { throw PCAError.noData }
        let n = data.count
        let d = first.count
        let k = min(outputDimension, d, n)

        guard k > 0 else { throw PCAError.invalidDimension }

        var mean = [Float](repeating: 0, count: d)
        for i in 0..<n {
            for j in 0..<d {
                mean[j] += data[i][j]
            }
        }
        for j in 0..<d {
            mean[j] /= Float(n)
        }

        var A = [Float](repeating: 0, count: n * d)
        for j in 0..<d {
            for i in 0..<n {
                A[i + j * n] = data[i][j] - mean[j]
            }
        }

        let m = n
        var m32 = Int32(m)
        var n32 = Int32(d)
        var lda = Int32(m)
        var ldu = Int32(1)
        var ldvt = Int32(d)
        var s = [Float](repeating: 0, count: min(m, d))
        var u = [Float](repeating: 0, count: 1)
        var vt = [Float](repeating: 0, count: d * d)
        var lwork = Int32(-1)
        var work: [Float] = [0]
        var info = Int32(0)
        var jobU = Int8(78)
        var jobVT = Int8(65)

        sgesvd_(&jobU, &jobVT, &m32, &n32, &A, &lda, &s, &u, &ldu, &vt, &ldvt, &work, &lwork, &info)
        guard info == 0 else { throw PCAError.svdFailed }

        lwork = Int32(work[0])
        work = [Float](repeating: 0, count: Int(lwork))

        sgesvd_(&jobU, &jobVT, &m32, &n32, &A, &lda, &s, &u, &ldu, &vt, &ldvt, &work, &lwork, &info)
        guard info == 0 else { throw PCAError.svdFailed }

        var components: [[Float]] = []
        var explainedVariance: [Float] = []
        for i in 0..<k {
            var comp = [Float](repeating: 0, count: d)
            for j in 0..<d {
                comp[j] = vt[i + j * d]
            }
            components.append(comp)
            explainedVariance.append(s[i] * s[i] / Float(n - 1))
        }

        return PCA(
            inputDimension: d,
            outputDimension: k,
            mean: mean,
            components: components,
            explainedVariance: explainedVariance,
            whiteningEnabled: whitening,
            sampleCount: n,
            trainingData: data
        )
    }

    public func incrementalUpdate(newSamples: [[Float]]) throws -> PCA {
        let combined = trainingData + newSamples
        return try PCA.train(data: combined, outputDimension: outputDimension, whitening: whiteningEnabled)
    }

    public func partialFit(sample: [Float]) throws -> PCA {
        return try incrementalUpdate(newSamples: [sample])
    }
}

public enum PCAError: Error, LocalizedError {
    case noData
    case invalidDimension
    case dimensionMismatch
    case svdFailed

    public var errorDescription: String? {
        switch self {
        case .noData: return "No data provided for PCA training"
        case .invalidDimension: return "Invalid output dimension"
        case .dimensionMismatch: return "Input dimension does not match model"
        case .svdFailed: return "SVD computation failed"
        }
    }
}