import XCTest
@testable import MarkBase

class TwelveBNaNDebugTest: XCTestCase {
    
    func test12BNaNDebug() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  12B NaN Debug Test - Find Exact NaN Positions")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        
        print("Step 1: Run forward pass and find NaN indices...")
        let logits = try model.forwardOptimized(tokenId: 2, position: 0)
        
        var nanIndices: [Int] = []
        var nanValues: [Float] = []
        var extremeIndices: [Int] = []
        var extremeValues: [Float] = []
        
        for i in 0..<logits.count {
            let v = logits[i]
            if v.isNaN {
                nanIndices.append(i)
                nanValues.append(v)
            }
            if abs(v) > 100.0 {
                extremeIndices.append(i)
                extremeValues.append(v)
            }
        }
        
        print("  Total logits: \(logits.count)")
        print("  NaN count: \(nanIndices.count)")
        print("  Extreme values (>100): \(extremeIndices.count)")
        print()
        
        if nanIndices.count > 0 {
            print("Step 2: NaN Details:")
            print("  NaN at indices: \(nanIndices)")
            print()
            
            print("Step 3: Analyze logits around NaN positions:")
            for idx in nanIndices.prefix(3) {
                let start = max(0, idx - 5)
                let end = min(logits.count - 1, idx + 5)
                print("  Around index \(idx):")
                for i in start...end {
                    let marker = i == idx ? " [NaN]" : ""
                    print("    logits[\(i)] = \(logits[i])\(marker)")
                }
                print()
            }
        }
        
        if extremeIndices.count > 0 {
            print("Step 4: Extreme values (>100) that might cause NaN:")
            print("  Extreme indices: \(extremeIndices.prefix(10))")
            print("  Extreme values: \(extremeValues.prefix(10))")
            print()
        }
        
        // Check overall statistics
        var minVal: Float = Float.infinity
        var maxVal: Float = -Float.infinity
        var sumVal: Float = 0
        
        for v in logits {
            if v.isFinite {
                if v < minVal { minVal = v }
                if v > maxVal { maxVal = v }
                sumVal += v
            }
        }
        
        let meanVal = sumVal / Float(logits.count - nanIndices.count)
        
        print("Step 5: Logit Statistics (excluding NaN/Inf):")
        print("  Min: \(minVal)")
        print("  Max: \(maxVal)")
        print("  Mean: \(meanVal)")
        print("  Range: \(maxVal - minVal)")
        print()
        
        // Check if softcapping is applied
        print("Step 6: Config Check:")
        let configPath = modelPath + "/config.json"
        if let configData = FileManager.default.contents(atPath: configPath) {
            let config = try JSONDecoder().decode(ConfigDebug.self, from: configData)
            print("  final_logit_softcapping: \(config.textConfig.finalLogitSoftcapping)")
            print("  rms_norm_eps: \(config.textConfig.rmsNormEps)")
            print()
            
            if config.textConfig.finalLogitSoftcapping > 0 {
                print("  ⚠️ Softcapping ACTIVE: logits /= (1 + |logits| / \(config.textConfig.finalLogitSoftcapping))")
                print("  💡 Large logits could cause division by zero or NaN")
            }
        }
        
        print("═══════════════════════════════════════════════════════════════════\n")
        
        XCTAssertLessThanOrEqual(nanIndices.count, 10, "Should have minimal NaN")
    }
}

struct ConfigDebug: Decodable {
    let textConfig: TextConfigDebug
    
    enum CodingKeys: String, CodingKey {
        case textConfig = "text_config"
    }
}

struct TextConfigDebug: Decodable {
    let finalLogitSoftcapping: Float
    let rmsNormEps: Float
    
    enum CodingKeys: String, CodingKey {
        case finalLogitSoftcapping = "final_logit_softcapping"
        case rmsNormEps = "rms_norm_eps"
    }
}