import Foundation

public enum VisionSampleType {
    case syntheticFlat
    case syntheticGradient
    case syntheticEdge
    case syntheticNatural
}

public struct VisionSample {
    public let type: VisionSampleType
    public let name: String
    public let patchEmbeddings: [Float]
    public let numPatches: Int
    public let patchDim: Int
    public let description: String
}

public final class VisionSampleGenerator {
    public let patchDim: Int = 768
    public let numPatches: Int = 256
    
    public init() {}
    
    public func generate(type: VisionSampleType) -> VisionSample {
        switch type {
        case .syntheticFlat:
            return generateFlat()
        case .syntheticGradient:
            return generateGradient()
        case .syntheticEdge:
            return generateEdge()
        case .syntheticNatural:
            return generateNatural()
        }
    }
    
    public func generateAll() -> [VisionSample] {
        return [
            generateFlat(),
            generateGradient(),
            generateEdge(),
            generateNatural()
        ]
    }
    
    private func generateFlat() -> VisionSample {
        let patches: [Float] = Array(repeating: 0.3, count: numPatches * patchDim)
        return VisionSample(
            type: .syntheticFlat,
            name: "synthetic_flat",
            patchEmbeddings: patches,
            numPatches: numPatches,
            patchDim: patchDim,
            description: "Flat uniform patch values (0.3)"
        )
    }
    
    private func generateGradient() -> VisionSample {
        var patches: [Float] = []
        for p in 0..<numPatches {
            let row = p / 16
            let col = p % 16
            let gradientH = Float(row) / 16.0
            let gradientV = Float(col) / 16.0
            for d in 0..<patchDim {
                let spatialWeight = (gradientH + gradientV) / 2.0
                let featureWeight = Float(d) / Float(patchDim)
                let value = spatialWeight * 0.4 + featureWeight * 0.2 - 0.3
                patches.append(value)
            }
        }
        return VisionSample(
            type: .syntheticGradient,
            name: "synthetic_gradient",
            patchEmbeddings: patches,
            numPatches: numPatches,
            patchDim: patchDim,
            description: "Linear gradient (spatial + feature)"
        )
    }
    
    private func generateEdge() -> VisionSample {
        var patches: [Float] = []
        for p in 0..<numPatches {
            let row = p / 16
            let col = p % 16
            let isEdge = (row == 0 || row == 15 || col == 0 || col == 15)
            let isCorner = (row == 0 && col == 0) || (row == 15 && col == 15) ||
                           (row == 0 && col == 15) || (row == 15 && col == 0)
            for d in 0..<patchDim {
                var value: Float = 0.1
                if isEdge {
                    value = 0.5 + Float.random(in: -0.1...0.1)
                }
                if isCorner {
                    value = 0.8 + Float.random(in: -0.1...0.1)
                }
                let featureMod = Float(d % 128) / 128.0 * 0.1
                patches.append(value + featureMod)
            }
        }
        return VisionSample(
            type: .syntheticEdge,
            name: "synthetic_edge",
            patchEmbeddings: patches,
            numPatches: numPatches,
            patchDim: patchDim,
            description: "Edge detection pattern (border/corners)"
        )
    }
    
    private func generateNatural() -> VisionSample {
        var patches: [Float] = []
        let imagenetMean: Float = 0.485
        let imagenetStd: Float = 0.229
        for _ in 0..<numPatches {
            for d in 0..<patchDim {
                let base = imagenetMean + Float.random(in: -imagenetStd...imagenetStd)
                let channelBias = Float(d % 3) * 0.05
                let spatialVar = Float.random(in: -0.05...0.05)
                patches.append(base + channelBias + spatialVar)
            }
        }
        return VisionSample(
            type: .syntheticNatural,
            name: "synthetic_natural",
            patchEmbeddings: patches,
            numPatches: numPatches,
            patchDim: patchDim,
            description: "Natural image statistics (ImageNet mean/std)"
        )
    }
}

public struct VisionTestResult {
    public let modelName: String
    public let sampleName: String
    public let outputShape: (Int, Int)
    public let min: Float
    public let max: Float
    public let mean: Float
    public let std: Float
    public let nanCount: Int
    public let infCount: Int
    public let cosineSimilarity: Float?
    public let forwardTimeMs: Double
    public let passed: Bool
}