param(
    [Parameter(Mandatory = $true)]
    [string]$Manifest,
    [Parameter(Mandatory = $true)]
    [string]$Model,
    [Parameter(Mandatory = $true)]
    [string]$OutputRoot,
    [string[]]$Seeds = @('0', '1', '2'),
    [int]$Steps = 12000,
    [int]$Images = 2500,
    [int]$Layers = 3,
    [int]$Size = 224,
    [string]$EvaluationSplit = "validation"
)

$ErrorActionPreference = "Stop"
$repo = (Resolve-Path (Join-Path $PSScriptRoot "..")).Path
$manifestPath = (Resolve-Path $Manifest).Path
$modelPath = (Resolve-Path $Model).Path
$manifestHash = (Get-FileHash -LiteralPath $manifestPath -Algorithm SHA256).Hash.ToLowerInvariant()
$modelHash = (Get-FileHash -LiteralPath $modelPath -Algorithm SHA256).Hash.ToLowerInvariant()
$outputPath = [System.IO.Path]::GetFullPath($OutputRoot)
New-Item -ItemType Directory -Force -Path $outputPath | Out-Null
$cargo = (Get-Command cargo).Source
$seedValues = @(
    $Seeds | ForEach-Object { $_ -split ',' } | ForEach-Object {
        $value = 0L
        if (-not [long]::TryParse($_.Trim(), [ref]$value) -or $value -lt 0) {
            throw "invalid seed: $_"
        }
        $value
    }
)
if ($seedValues.Count -eq 0 -or @($seedValues | Select-Object -Unique).Count -ne $seedValues.Count) {
    throw 'provide at least one seed, without duplicates'
}

function Run-Cargo(
    [string[]]$Arguments,
    [string]$Stdout,
    [string]$Stderr,
    [string]$LogLevel
) {
    $env:RUST_LOG = $LogLevel
    Push-Location $repo
    $oldErrorPreference = $ErrorActionPreference
    try {
        # Windows PowerShell may wrap native stderr as ErrorRecord objects;
        # Continue preserves the native exit code while the two streams are
        # retained independently.
        $ErrorActionPreference = "Continue"
        & $cargo @Arguments 1> $Stdout 2> $Stderr
        $exitCode = $LASTEXITCODE
    } finally {
        $ErrorActionPreference = $oldErrorPreference
        Pop-Location
    }
    if ($exitCode -ne 0) {
        throw "cargo exited with code $exitCode; see $Stdout and $Stderr"
    }
}

foreach ($seed in $seedValues) {
    $runDir = Join-Path $outputPath "${Layers}l${Size}-seed${seed}"
    New-Item -ItemType Directory -Force -Path $runDir | Out-Null
    $decoder = Join-Path $runDir "decoder.bin"
    $training = Join-Path $runDir "training.json"
    $quality = Join-Path $runDir "quality-imagenette-$EvaluationSplit.json"

    if (-not ((Test-Path -LiteralPath $training) -and (Test-Path -LiteralPath $decoder))) {
        Write-Output "training seed $seed at $(Get-Date -Format o)"
        Run-Cargo `
            -Arguments @(
                "run", "--release", "--example", "train_decoder", "--",
                $manifestPath, $modelPath, "$Steps", "$Images", "$Layers", "$Size", "$seed", $runDir
            ) `
            -Stdout (Join-Path $runDir "training.stdout.log") `
            -Stderr (Join-Path $runDir "training.stderr.log") `
            -LogLevel "info"
    } else {
        Write-Output "seed $seed training outputs already exist; preserving them"
    }

    $record = Get-Content -Raw -LiteralPath $training | ConvertFrom-Json
    if (
        $record.seed -ne $seed -or
        $record.steps -ne $Steps -or
        $record.requested_images -ne $Images -or
        $record.encoder_layers -ne $Layers -or
        $record.image_size -ne $Size -or
        $record.dataset_manifest_sha256 -ne $manifestHash -or
        $record.model_sha256 -ne $modelHash
    ) {
        throw "existing training record does not match requested cell: $training"
    }

    if (-not (Test-Path -LiteralPath $quality)) {
        Write-Output "evaluating seed $seed at $(Get-Date -Format o)"
        Run-Cargo `
            -Arguments @(
                "run", "--release", "--example", "evaluate_decoder", "--",
                "--manifest", $manifestPath,
                "--model", $modelPath,
                "--decoder", $decoder,
                "--split", $EvaluationSplit,
                "--layers", "$Layers",
                "--size", "$Size",
                "--output", $quality,
                "--samples", "12"
            ) `
            -Stdout (Join-Path $runDir "evaluation.stdout.log") `
            -Stderr (Join-Path $runDir "evaluation.stderr.log") `
            -LogLevel "warn"
    } else {
        Write-Output "seed $seed quality output already exists; preserving it"
    }

    $qualityRecord = Get-Content -Raw -LiteralPath $quality | ConvertFrom-Json
    if (
        $qualityRecord.split -ne $EvaluationSplit -or
        $qualityRecord.encoder_layers -ne $Layers -or
        $qualityRecord.image_size -ne $Size -or
        $qualityRecord.manifest_sha256 -ne $manifestHash -or
        $qualityRecord.model_sha256 -ne $modelHash -or
        $qualityRecord.decoder_sha256 -ne $record.decoder_sha256
    ) {
        throw "quality record does not match its training cell: $quality"
    }

    $safeTensors = Join-Path $runDir "decoder.safetensors"
    if (-not (Test-Path -LiteralPath $safeTensors)) {
        & python (Join-Path $repo "tools\convert_decoder.py") `
            --training $training `
            --input $decoder `
            --output $safeTensors
        if ($LASTEXITCODE -ne 0) {
            throw "decoder SafeTensors conversion failed for seed $seed"
        }
    }
}

Write-Output "decoder matrix complete at $(Get-Date -Format o)"