# =============================================================================
#  CUBIE TEP - LLM FAULT-DETECTION BENCHMARK HARNESS
#  Calls the Anthropic API directly. Measures per-IDV FDR and FAR for an
#  LLM monitoring sliding windows of TEP sensor data.
#
#  Output is a row you can drop straight into Cubie_TEP_Demo.ps1's leaderboard.
# =============================================================================
#
#  Cost notes (rough, 2026-05-28 pricing):
#    Default run (~25 calls, Haiku 4.5):     under $0.05
#    --Full (~192 calls, Haiku 4.5):         ~$0.40
#    --Full --Model claude-sonnet-4-6:       ~$2.00
#    --Full --Model claude-opus-4-7:         ~$15
#
#  Auth:
#    1. Set environment variable: $env:ANTHROPIC_API_KEY = "sk-ant-..."
#       (or save it in C:\Users\NickV\OneDrive\Desktop\RAKKIT STACK\.env.portal
#        as ANTHROPIC_API_KEY=sk-ant-... and the script will pick it up)
#
#  Usage:
#    Default smoke run (5 windows per file, Haiku, zero-shot):
#       pwsh -File Cubie_TEP_LLM_Benchmark.ps1
#
#    Full sweep (all 48 windows per file, Haiku, zero-shot):
#       pwsh -File Cubie_TEP_LLM_Benchmark.ps1 -Full
#
#    Few-shot mode (includes 2 example windows per class in the prompt):
#       pwsh -File Cubie_TEP_LLM_Benchmark.ps1 -FewShot
#
#    Different model:
#       pwsh -File Cubie_TEP_LLM_Benchmark.ps1 -Model claude-sonnet-4-6
#
#    Dry run (prints prompts, never hits API, $0):
#       pwsh -File Cubie_TEP_LLM_Benchmark.ps1 -DryRun

param(
    [string]   $RepoDir   = "C:\Users\NickV\cubie-tf-completion",
    [string]   $Model     = "claude-haiku-4-5-20251001",
    [string[]] $Models    = $null,
    [int]      $WindowLen = 20,
    [int]      $InjectionSample = 161,
    [int]      $WindowsPerFile  = 5,
    [switch]   $Full,
    [switch]   $FewShot,
    [switch]   $AllClaude,
    [switch]   $DryRun
)

# Resolve which models to run.
# Priority: -Models > -AllClaude > -Model (single)
if ($AllClaude -and -not $Models) {
    $Models = @("claude-haiku-4-5-20251001", "claude-sonnet-4-6", "claude-opus-4-7")
}
if (-not $Models) {
    $Models = @($Model)
}

$ErrorActionPreference = "Continue"
$startTime = Get-Date

# --- presentation helpers ---------------------------------------------------
function Banner($t) {
    Write-Host ""
    Write-Host ("=" * 78) -ForegroundColor Cyan
    Write-Host ("  " + $t)  -ForegroundColor Cyan
    Write-Host ("=" * 78) -ForegroundColor Cyan
    Write-Host ""
}
function Note ($m) { Write-Host "  . $m" -ForegroundColor Gray }
function Ok   ($m) { Write-Host "  + $m" -ForegroundColor Green }
function Fail ($m) { Write-Host "  ! $m" -ForegroundColor Red }

Banner "CUBIE TEP - LLM FAULT-DETECTION BENCHMARK"
Note "Model:           $Model"
Note "Window length:   $WindowLen samples (= $($WindowLen * 3) minutes)"
Note "Injection point: sample $InjectionSample (pre-injection = NORMAL, post = FAULT)"
if ($Full) { $WindowsPerFile = 0 }
Note "Windows / file:  $(if ($WindowsPerFile -eq 0) {'ALL'} else {$WindowsPerFile})  $(if($FewShot){'(few-shot mode)'} else {'(zero-shot mode)'})"
if ($DryRun) { Note "DRY RUN: prompts printed, NO API CALLS" }
Write-Host ""

# --- resolve API key --------------------------------------------------------
$anthropicCredential = $env:ANTHROPIC_API_KEY
if (-not $anthropicCredential) {
    $envFile = "C:\Users\NickV\OneDrive\Desktop\RAKKIT STACK\.env.portal"
    if (Test-Path $envFile) {
        $match = Get-Content $envFile | Select-String -Pattern "^ANTHROPIC_API_KEY=(.*)$"
        if ($match) { $anthropicCredential = $match.Matches[0].Groups[1].Value.Trim() }
    }
}
if (-not $anthropicCredential -and -not $DryRun) {
    Fail "ANTHROPIC_API_KEY not set."
    Fail "Set with:    `$env:ANTHROPIC_API_KEY = 'sk-ant-...'"
    Fail "Or use -DryRun to preview prompts without calling the API."
    exit 1
}
if ($anthropicCredential) {
    $masked = ($anthropicCredential[0..6] -join "") + "..." + ($anthropicCredential[-4..-1] -join "")
    Ok "API key loaded ($masked)"
}

# --- load CSV files ---------------------------------------------------------
$dataDir = Join-Path $RepoDir "data\tep"
$files = @(
    @{ Name = "d00_test"; Path = "$dataDir\d00_test.csv"; ExpectedFault = $null  },
    @{ Name = "d03_test"; Path = "$dataDir\d03_test.csv"; ExpectedFault = "IDV-3" },
    @{ Name = "d09_test"; Path = "$dataDir\d09_test.csv"; ExpectedFault = "IDV-9" },
    @{ Name = "d15_test"; Path = "$dataDir\d15_test.csv"; ExpectedFault = "IDV-15"}
)
$loadedFiles = @()
foreach ($f in $files) {
    if (-not (Test-Path $f.Path)) {
        Fail "Missing $($f.Path) - run Cubie_TEP_Demo.ps1 first to auto-stage."
        exit 1
    }
    # CSV format: header row + faultNumber, simulationRun, sample, xmeas_1..xmeas_52
    # CRITICAL: skip first 3 cols - column 0 is the fault label, which would leak the answer.
    $rows = @()
    $first = $true
    foreach ($line in (Get-Content $f.Path)) {
        if ($first) { $first = $false; continue }  # header
        $parts = $line.Trim() -split ','
        if ($parts.Count -lt 25) { continue }
        try {
            $vals = @()
            for ($i = 3; $i -lt 25; $i++) { $vals += [double]$parts[$i] }
            $rows += ,$vals
        } catch { continue }
    }
    $f.Rows = $rows
    $loadedFiles += $f
    Note ("Loaded $($f.Name): {0} samples (xmeas_1..xmeas_22, no metadata leak)" -f $rows.Count)
}
Write-Host ""

# --- window construction ---------------------------------------------------
function Get-Window($rows, $startIdx, $len) {
    $rows[$startIdx..($startIdx + $len - 1)]
}

# Variable labels (canonical TEP, first 22 of 52)
$varLabels = @(
    "XMEAS(1) A feed (stream 1)",   "XMEAS(2) D feed (stream 2)",
    "XMEAS(3) E feed (stream 3)",   "XMEAS(4) A+C feed (stream 4)",
    "XMEAS(5) recycle flow",        "XMEAS(6) reactor feed rate",
    "XMEAS(7) reactor pressure",    "XMEAS(8) reactor level",
    "XMEAS(9) reactor temp",        "XMEAS(10) purge rate (stream 9)",
    "XMEAS(11) sep temp",           "XMEAS(12) sep level",
    "XMEAS(13) sep pressure",       "XMEAS(14) sep underflow",
    "XMEAS(15) stripper level",     "XMEAS(16) stripper pressure",
    "XMEAS(17) stripper underflow", "XMEAS(18) stripper temp",
    "XMEAS(19) stripper steam flow","XMEAS(20) compressor work",
    "XMEAS(21) reactor CW out temp","XMEAS(22) condenser CW out temp"
)

function Format-Window($win) {
    $sb = New-Object System.Text.StringBuilder
    for ($v = 0; $v -lt 22; $v++) {
        $vals = @()
        for ($t = 0; $t -lt $win.Count; $t++) {
            $vals += "{0,8:F3}" -f $win[$t][$v]
        }
        [void]$sb.AppendLine(("  {0,-32}  {1}" -f $varLabels[$v], ($vals -join "")))
    }
    $sb.ToString()
}

# --- prompt builder --------------------------------------------------------
$systemPrompt = @"
You are a chemical process control engineer monitoring the Tennessee Eastman Process
(TEP), a benchmark simulation of a fluid catalytic cracking unit published by
Downs and Vogel (1993, Eastman Chemical Co.). The TEP has 52 sensor variables
and 20 documented fault modes (IDV-1 through IDV-20). The process runs with
closed-loop control that compensates quickly when faults occur, so fault
signatures can be subtle.

You will be shown a sliding window of the 22 primary measurement variables
(XMEAS(1) through XMEAS(22)), sampled every 3 minutes. Decide whether a
process fault is currently active in this window.

Reply with EXACTLY one of these single tokens, no other text, no explanation:
  FAULT     if you believe a fault is present anywhere in this window
  NORMAL    if you believe the process is operating normally throughout
"@

function Build-UserPrompt($win, $exampleNormal, $exampleFault) {
    $sb = New-Object System.Text.StringBuilder
    if ($exampleNormal) {
        [void]$sb.AppendLine("EXAMPLE 1 - process operating NORMALLY (no fault):")
        [void]$sb.AppendLine((Format-Window $exampleNormal))
        [void]$sb.AppendLine("Answer: NORMAL")
        [void]$sb.AppendLine("")
    }
    if ($exampleFault) {
        [void]$sb.AppendLine("EXAMPLE 2 - a fault is ACTIVE in the window:")
        [void]$sb.AppendLine((Format-Window $exampleFault))
        [void]$sb.AppendLine("Answer: FAULT")
        [void]$sb.AppendLine("")
    }
    [void]$sb.AppendLine("Now classify THIS window:")
    [void]$sb.AppendLine((Format-Window $win))
    [void]$sb.AppendLine("Answer (FAULT or NORMAL):")
    $sb.ToString()
}

# --- API caller ------------------------------------------------------------
function Call-Anthropic($userPrompt, $modelId) {
    if ($DryRun) {
        Write-Host "--- DRY RUN [$modelId]: would send prompt of $($userPrompt.Length) chars ---" -ForegroundColor DarkGray
        return @{ ok = $true; text = "NORMAL"; usage = @{} }
    }
    $body = @{
        model      = $modelId
        max_tokens = 16
        system     = $systemPrompt
        messages   = @(@{ role = "user"; content = $userPrompt })
    } | ConvertTo-Json -Depth 8 -Compress

    $headers = @{
        "x-api-key"         = $anthropicCredential
        "anthropic-version" = "2023-06-01"
        "content-type"      = "application/json"
    }
    try {
        $resp = Invoke-RestMethod -Method Post `
            -Uri "https://api.anthropic.com/v1/messages" `
            -Headers $headers -Body $body -ErrorAction Stop
        $text = ""
        if ($resp.content) { $text = ($resp.content | ForEach-Object { $_.text } | Out-String).Trim() }
        return @{ ok = $true; text = $text; usage = $resp.usage }
    } catch {
        return @{ ok = $false; text = "ERROR: $_"; usage = @{} }
    }
}

# --- few-shot examples (sampled from d00 and d03) ---------------------------
$exNormal = $null; $exFault = $null
if ($FewShot) {
    $exNormal = Get-Window ($loadedFiles | Where-Object Name -eq "d00_test").Rows 0 $WindowLen
    $exFault  = Get-Window ($loadedFiles | Where-Object Name -eq "d03_test").Rows ($InjectionSample + 50) $WindowLen
    Note "Few-shot: 1 NORMAL example from d00, 1 FAULT example from d03 (post-injection)"
}

# --- per-model benchmark ---------------------------------------------------
function Pct($num, $den) {
    if ($den -eq 0) { return "n/a" }
    "{0:N1}%" -f (100.0 * $num / $den)
}

function Run-OneModel($modelId) {
    $modelStart = Get-Date
    Banner "MODEL: $modelId"

    $results = @{
        "d00_test" = @{ predFault = 0; predNormal = 0; totalWindows = 0; preInject = @(); postInject = @() }
        "d03_test" = @{ predFault = 0; predNormal = 0; totalWindows = 0; preInject = @(); postInject = @() }
        "d09_test" = @{ predFault = 0; predNormal = 0; totalWindows = 0; preInject = @(); postInject = @() }
        "d15_test" = @{ predFault = 0; predNormal = 0; totalWindows = 0; preInject = @(); postInject = @() }
    }
    $totalCalls = 0
    $totalInputTokens = 0
    $totalOutputTokens = 0

    foreach ($f in $loadedFiles) {
        $expected = if ($f.ExpectedFault) { $f.ExpectedFault } else { "none" }
        Write-Host ""
        Write-Host ("  --- $($f.Name)  (expected fault: $expected) ---") -ForegroundColor Cyan

        $maxStart = $f.Rows.Count - $WindowLen
        $stride = if ($WindowsPerFile -eq 0) { $WindowLen } else { [int]($maxStart / $WindowsPerFile) }
        if ($stride -lt 1) { $stride = 1 }
        $starts = @()
        for ($s = 0; $s -le $maxStart; $s += $stride) { $starts += $s }
        if ($WindowsPerFile -gt 0 -and $starts.Count -gt $WindowsPerFile) {
            $starts = $starts[0..($WindowsPerFile - 1)]
        }

        foreach ($s in $starts) {
            $win = Get-Window $f.Rows $s $WindowLen
            $prompt = Build-UserPrompt $win $exNormal $exFault
            $r = Call-Anthropic $prompt $modelId
            $totalCalls++
            if ($r.usage.input_tokens)  { $totalInputTokens  += $r.usage.input_tokens }
            if ($r.usage.output_tokens) { $totalOutputTokens += $r.usage.output_tokens }
            $endIdx = $s + $WindowLen - 1
            $pred = if ($r.text -match "FAULT") { "FAULT" }
                    elseif ($r.text -match "NORMAL") { "NORMAL" }
                    else { "?($($r.text))" }
            $actual = if ($endIdx -lt $InjectionSample -or $null -eq $f.ExpectedFault) { "NORMAL" } else { "FAULT" }
            $correct = ($pred -eq $actual)
            $color = if ($correct) { "Green" } elseif ($pred -like "?*") { "Yellow" } else { "Red" }
            Write-Host ("    s={0,-4} [{1,3}..{2,-3}]   pred={3,-9} actual={4,-8} {5}" -f $s, $s, $endIdx, $pred, $actual, $(if ($correct) {"YES"} else {"NO"})) -ForegroundColor $color

            $results[$f.Name].totalWindows++
            if ($pred -eq "FAULT")  { $results[$f.Name].predFault++ }
            if ($pred -eq "NORMAL") { $results[$f.Name].predNormal++ }
            if ($endIdx -lt $InjectionSample) {
                $results[$f.Name].preInject += @{ pred = $pred; correct = ($pred -eq "NORMAL") }
            } else {
                $results[$f.Name].postInject += @{ pred = $pred; correct = ($pred -eq $actual) }
            }
        }
    }

    # Compute per-IDV FDR + FAR
    $farD00 = $results["d00_test"]
    $fdr = @{}
    foreach ($key in @("d03_test", "d09_test", "d15_test")) {
        $post = $results[$key].postInject
        if ($post.Count -eq 0) {
            $fdr[$key] = @{ tp = 0; n = 0 }
        } else {
            $tp = ($post | Where-Object { $_.pred -eq "FAULT" }).Count
            $fdr[$key] = @{ tp = $tp; n = $post.Count }
        }
    }

    $elapsedModel = (Get-Date) - $modelStart
    Write-Host ""
    Write-Host ("  d00 FAR:  $($farD00.predFault) / $($farD00.totalWindows) = $(Pct $farD00.predFault $farD00.totalWindows)") -ForegroundColor Yellow
    Write-Host ("  IDV-3 :   $($fdr['d03_test'].tp) / $($fdr['d03_test'].n)  = $(Pct $fdr['d03_test'].tp $fdr['d03_test'].n)") -ForegroundColor Yellow
    Write-Host ("  IDV-9 :   $($fdr['d09_test'].tp) / $($fdr['d09_test'].n)  = $(Pct $fdr['d09_test'].tp $fdr['d09_test'].n)") -ForegroundColor Yellow
    Write-Host ("  IDV-15:   $($fdr['d15_test'].tp) / $($fdr['d15_test'].n) = $(Pct $fdr['d15_test'].tp $fdr['d15_test'].n)") -ForegroundColor Yellow
    Write-Host ("  Calls: $totalCalls   tokens in/out: $totalInputTokens / $totalOutputTokens   walltime: {0:N1}s" -f $elapsedModel.TotalSeconds)

    return @{
        Model = $modelId
        Mode = $(if ($FewShot) { "few-shot" } else { "zero-shot" })
        FarD00 = @{ tp = $farD00.predFault; n = $farD00.totalWindows; pct = Pct $farD00.predFault $farD00.totalWindows }
        Fdr3 = @{ tp = $fdr['d03_test'].tp; n = $fdr['d03_test'].n; pct = Pct $fdr['d03_test'].tp $fdr['d03_test'].n }
        Fdr9 = @{ tp = $fdr['d09_test'].tp; n = $fdr['d09_test'].n; pct = Pct $fdr['d09_test'].tp $fdr['d09_test'].n }
        Fdr15 = @{ tp = $fdr['d15_test'].tp; n = $fdr['d15_test'].n; pct = Pct $fdr['d15_test'].tp $fdr['d15_test'].n }
        Calls = $totalCalls
        TokensIn = $totalInputTokens
        TokensOut = $totalOutputTokens
        WalltimeSec = [math]::Round($elapsedModel.TotalSeconds, 2)
    }
}

# --- multi-model loop ------------------------------------------------------
$allResults = @()
foreach ($modelId in $Models) {
    $allResults += Run-OneModel $modelId
}

# --- combined leaderboard --------------------------------------------------
Banner "COMBINED LEADERBOARD"
$mode = if ($FewShot) { "few-shot" } else { "zero-shot" }
Write-Host ("  Drop these rows into Cubie_TEP_Demo.ps1's leaderboard:") -ForegroundColor Cyan
Write-Host ""
Write-Host ("  Model                                     IDV-3    IDV-9    IDV-15   d00 FAR") -ForegroundColor DarkGray
Write-Host ("  ----                                      -----    -----    ------   -------")
foreach ($r in $allResults) {
    Write-Host ("  LLM $mode  ({0,-35}) {1,6}  {2,6}  {3,6}   {4}" -f $r.Model, $r.Fdr3.pct, $r.Fdr9.pct, $r.Fdr15.pct, $r.FarD00.pct)
}
Write-Host ""
$totalAllCalls = ($allResults | Measure-Object -Property Calls -Sum).Sum
$totalAllIn   = ($allResults | Measure-Object -Property TokensIn -Sum).Sum
$totalAllOut  = ($allResults | Measure-Object -Property TokensOut -Sum).Sum
$elapsed = (Get-Date) - $startTime
Write-Host ("  Total API calls across all models: $totalAllCalls")
Write-Host ("  Total tokens in/out: $totalAllIn / $totalAllOut")
Write-Host ("  Total walltime: {0:N1}s" -f $elapsed.TotalSeconds)
Write-Host ""

# --- JSON output -----------------------------------------------------------
$outJson = Join-Path $RepoDir "data\tep\layouts\llm_benchmark_$(Get-Date -f 'yyyyMMdd_HHmmss').json"
$jsonOut = @{
    timestamp        = (Get-Date -Format "o")
    mode             = $mode
    window_len       = $WindowLen
    windows_per_file = $(if ($WindowsPerFile -eq 0) { "all" } else { $WindowsPerFile })
    injection_sample = $InjectionSample
    total_walltime_sec = [math]::Round($elapsed.TotalSeconds, 2)
    total_calls      = $totalAllCalls
    total_tokens_in  = $totalAllIn
    total_tokens_out = $totalAllOut
    models           = $allResults
} | ConvertTo-Json -Depth 8
Set-Content -Path $outJson -Value $jsonOut -Encoding UTF8
Ok "JSON results written to $outJson"
