$ErrorActionPreference = "Stop"

Add-Type -AssemblyName System.Drawing

$repo = $PSScriptRoot
$outputDir = Join-Path $repo "samples\roi_preview"

if (Test-Path $outputDir) {
    Remove-Item $outputDir -Recurse -Force
}
New-Item -ItemType Directory -Path $outputDir | Out-Null

$imageExts = @(".jpg", ".jpeg", ".png", ".bmp", ".webp")

function New-Box {
    param(
        [double]$X1,
        [double]$Y1,
        [double]$X2,
        [double]$Y2
    )
    [PSCustomObject]@{
        X1 = $X1
        Y1 = $Y1
        X2 = $X2
        Y2 = $Y2
        W = [Math]::Max(0.0, $X2 - $X1)
        H = [Math]::Max(0.0, $Y2 - $Y1)
        Cx = ($X1 + $X2) / 2.0
        Cy = ($Y1 + $Y2) / 2.0
        Area = [Math]::Max(0.0, $X2 - $X1) * [Math]::Max(0.0, $Y2 - $Y1)
    }
}

function Clip-Box {
    param(
        $Box,
        [double]$Width,
        [double]$Height
    )
    $x1 = [Math]::Min([Math]::Max($Box.X1, 0.0), $Width)
    $y1 = [Math]::Min([Math]::Max($Box.Y1, 0.0), $Height)
    $x2 = [Math]::Min([Math]::Max($Box.X2, 0.0), $Width)
    $y2 = [Math]::Min([Math]::Max($Box.Y2, 0.0), $Height)
    if ($x2 -le $x1 -or $y2 -le $y1) {
        return $null
    }
    return New-Box -X1 $x1 -Y1 $y1 -X2 $x2 -Y2 $y2
}

function Find-Image {
    param(
        [string]$ImageDir,
        [string]$Stem
    )
    foreach ($ext in $imageExts) {
        $path = Join-Path $ImageDir ($Stem + $ext)
        if (Test-Path $path) {
            return $path
        }
    }
    return $null
}

function Load-Boxes {
    param(
        [string]$LabelPath,
        [int]$ImageWidth,
        [int]$ImageHeight
    )
    $boxes = @()
    foreach ($line in Get-Content $LabelPath) {
        $text = $line.Trim()
        if (-not $text) {
            continue
        }
        $parts = $text -split "\s+"
        if ($parts.Length -lt 5) {
            continue
        }
        $xc = [double]$parts[1] * $ImageWidth
        $yc = [double]$parts[2] * $ImageHeight
        $w = [double]$parts[3] * $ImageWidth
        $h = [double]$parts[4] * $ImageHeight
        $box = New-Box -X1 ($xc - $w / 2.0) -Y1 ($yc - $h / 2.0) -X2 ($xc + $w / 2.0) -Y2 ($yc + $h / 2.0)
        $clipped = Clip-Box -Box $box -Width $ImageWidth -Height $ImageHeight
        if ($null -ne $clipped -and $clipped.Area -gt 1.0) {
            $boxes += $clipped
        }
    }
    return $boxes | Sort-Object Cx, Cy
}

function Should-Pair {
    param($Left, $Right)
    $widthRef = [Math]::Max($Left.W, $Right.W)
    $heightRef = [Math]::Max($Left.H, $Right.H)
    if ($widthRef -le 0 -or $heightRef -le 0) {
        return $false
    }
    $dx = [Math]::Abs($Left.Cx - $Right.Cx)
    $dy = [Math]::Abs($Left.Cy - $Right.Cy)
    $areaRatio = if ($Right.Area -gt 0) { $Left.Area / $Right.Area } else { [double]::PositiveInfinity }
    return (
        $dx -le ($widthRef * 3.2) -and
        $dy -le ($heightRef * 1.2) -and
        $areaRatio -ge 0.4 -and
        $areaRatio -le 2.5
    )
}

function Get-Groups {
    param($Boxes)
    $groups = @()
    $used = @{}
    $candidates = @()
    for ($i = 0; $i -lt $Boxes.Count; $i++) {
        for ($j = $i + 1; $j -lt $Boxes.Count; $j++) {
            if (Should-Pair $Boxes[$i] $Boxes[$j]) {
                $score = [Math]::Abs($Boxes[$i].Cx - $Boxes[$j].Cx) + 0.5 * [Math]::Abs($Boxes[$i].Cy - $Boxes[$j].Cy)
                $candidates += [PSCustomObject]@{ Score = $score; I = $i; J = $j }
            }
        }
    }
    foreach ($candidate in ($candidates | Sort-Object Score)) {
        if ($used.ContainsKey($candidate.I) -or $used.ContainsKey($candidate.J)) {
            continue
        }
        $used[$candidate.I] = $true
        $used[$candidate.J] = $true
        $groups += ,@($candidate.I, $candidate.J)
    }
    for ($i = 0; $i -lt $Boxes.Count; $i++) {
        if (-not $used.ContainsKey($i)) {
            $groups += ,@($i)
        }
    }
    return $groups
}

function Estimate-PersonFromSingle {
    param($Box)
    $personW = [Math]::Max(4.2 * $Box.W, 2.8 * $Box.H)
    $personH = [Math]::Max(7.6 * $Box.H, 3.4 * $Box.W)
    $personCx = $Box.Cx
    $personY2 = $Box.Y2 + 0.08 * $personH
    $personX1 = $personCx - $personW / 2.0
    $personY1 = $personY2 - $personH
    return New-Box -X1 $personX1 -Y1 $personY1 -X2 ($personX1 + $personW) -Y2 $personY2
}

function Estimate-PersonFromPair {
    param($A, $B)
    $ux1 = [Math]::Min($A.X1, $B.X1)
    $uy1 = [Math]::Min($A.Y1, $B.Y1)
    $ux2 = [Math]::Max($A.X2, $B.X2)
    $uy2 = [Math]::Max($A.Y2, $B.Y2)
    $uw = $ux2 - $ux1
    $uh = $uy2 - $uy1
    $personW = [Math]::Max(1.95 * $uw, 2.6 * $uh)
    $personH = [Math]::Max(7.8 * $uh, 2.9 * $uw)
    $personCx = ($ux1 + $ux2) / 2.0
    $personY2 = $uy2 + 0.08 * $personH
    $personX1 = $personCx - $personW / 2.0
    $personY1 = $personY2 - $personH
    return New-Box -X1 $personX1 -Y1 $personY1 -X2 ($personX1 + $personW) -Y2 $personY2
}

function Get-RoiFromPerson {
    param($Person)
    $roiX = $Person.X1 - 0.24 * $Person.W
    $roiY = $Person.Y1 + 0.64 * $Person.H
    $roiW = 1.48 * $Person.W
    $roiH = 0.58 * $Person.H
    $roiX = $roiX - 0.08 * $roiW
    $roiY = $roiY - 0.08 * $roiH
    $roiW = $roiW * 1.16
    $roiH = $roiH * 1.18
    return New-Box -X1 $roiX -Y1 $roiY -X2 ($roiX + $roiW) -Y2 ($roiY + $roiH)
}

function Clamp-Roi {
    param(
        $Roi,
        [int]$ImageWidth,
        [int]$ImageHeight
    )
    $clipped = Clip-Box -Box $Roi -Width $ImageWidth -Height $ImageHeight
    if ($null -eq $clipped) {
        return $null
    }
    $x1 = [Math]::Max(0, [Math]::Min([int][Math]::Floor($clipped.X1), $ImageWidth - 1))
    $y1 = [Math]::Max(0, [Math]::Min([int][Math]::Floor($clipped.Y1), $ImageHeight - 1))
    $x2 = [Math]::Max($x1 + 1, [Math]::Min([int][Math]::Ceiling($clipped.X2), $ImageWidth))
    $y2 = [Math]::Max($y1 + 1, [Math]::Min([int][Math]::Ceiling($clipped.Y2), $ImageHeight))
    return New-Box -X1 $x1 -Y1 $y1 -X2 $x2 -Y2 $y2
}

function Resize-RoiToTarget {
    param(
        $Roi,
        [int]$ImageWidth,
        [int]$ImageHeight,
        [double]$ObjectArea,
        [double]$MinRatio,
        [double]$MaxRatio
    )
    $adjusted = $Roi
    $targetRatio = ($MinRatio + $MaxRatio) / 2.0
    for ($iter = 0; $iter -lt 3; $iter++) {
        $ratio = if ($adjusted.Area -gt 0) { $ObjectArea / $adjusted.Area } else { 0.0 }
        if ($ratio -ge $MinRatio -and $ratio -le $MaxRatio) {
            return $adjusted
        }
        $scale = [Math]::Sqrt($ratio / $targetRatio)
        if ($ratio -lt $MinRatio) {
            $scale = [Math]::Max(0.6, [Math]::Min(0.95, $scale))
        }
        else {
            $scale = [Math]::Min(1.8, [Math]::Max(1.05, $scale))
        }
        $newW = $adjusted.W * $scale
        $newH = $adjusted.H * $scale
        $cx = $adjusted.Cx
        $cy = $adjusted.Cy
        $adjusted = Clamp-Roi (New-Box -X1 ($cx - $newW / 2.0) -Y1 ($cy - $newH / 2.0) -X2 ($cx + $newW / 2.0) -Y2 ($cy + $newH / 2.0)) $ImageWidth $ImageHeight
        if ($null -eq $adjusted) {
            return $null
        }
    }
    return $adjusted
}

function Draw-Rect {
    param(
        [System.Drawing.Graphics]$Graphics,
        [System.Drawing.Pen]$Pen,
        $Box
    )
    $Graphics.DrawRectangle($Pen, [float]$Box.X1, [float]$Box.Y1, [float]$Box.W, [float]$Box.H)
}

$samples = @(
    @{ Dataset = "datasets\openimages-shoes-yolo"; Split = "train"; Stem = "00015a7cf95ec19d"; Label = "openimages_single" },
    @{ Dataset = "datasets\openimages-shoes-yolo"; Split = "train"; Stem = "0036655159bdef7f"; Label = "openimages_pair" },
    @{ Dataset = "datasets\openimages-shoes-yolo"; Split = "train"; Stem = "00003223e04e2e66"; Label = "openimages_mixed" },
    @{ Dataset = "datasets\ppe-shoes"; Split = "train"; Stem = "image1001"; Label = "ppe_pair_a" },
    @{ Dataset = "datasets\ppe-shoes"; Split = "train"; Stem = "image1011"; Label = "ppe_pair_b" }
)

$summary = @()

foreach ($sample in $samples) {
    $imageDir = Join-Path $repo (Join-Path $sample.Dataset ("images\" + $sample.Split))
    $labelPath = Join-Path $repo (Join-Path $sample.Dataset ("labels\" + $sample.Split + "\" + $sample.Stem + ".txt"))
    $imagePath = Find-Image -ImageDir $imageDir -Stem $sample.Stem
    if (-not $imagePath -or -not (Test-Path $labelPath)) {
        continue
    }

    $bitmap = [System.Drawing.Bitmap]::new($imagePath)
    try {
        $boxes = Load-Boxes -LabelPath $labelPath -ImageWidth $bitmap.Width -ImageHeight $bitmap.Height
        if ($boxes.Count -eq 0) {
            continue
        }

        $groups = Get-Groups -Boxes $boxes
        $canvas = [System.Drawing.Bitmap]::new($bitmap)
        $graphics = [System.Drawing.Graphics]::FromImage($canvas)
        $graphics.SmoothingMode = [System.Drawing.Drawing2D.SmoothingMode]::HighQuality
        $shoePen = [System.Drawing.Pen]::new([System.Drawing.Color]::Lime, 3)
        $roiPen = [System.Drawing.Pen]::new([System.Drawing.Color]::Red, 4)
        $font = [System.Drawing.Font]::new("Arial", 18, [System.Drawing.FontStyle]::Bold)
        $brush = [System.Drawing.SolidBrush]::new([System.Drawing.Color]::Yellow)

        try {
            for ($i = 0; $i -lt $boxes.Count; $i++) {
                Draw-Rect -Graphics $graphics -Pen $shoePen -Box $boxes[$i]
            }

            $roiIndex = 0
            foreach ($group in $groups) {
                if ($group.Count -eq 2) {
                    $person = Estimate-PersonFromPair $boxes[$group[0]] $boxes[$group[1]]
                    $roi = Get-RoiFromPerson $person
                    $minRatio = 0.18
                    $maxRatio = 0.40
                    $mode = "pair"
                }
                else {
                    $person = Estimate-PersonFromSingle $boxes[$group[0]]
                    $roi = Get-RoiFromPerson $person
                    $minRatio = 0.10
                    $maxRatio = 0.26
                    $mode = "single"
                }
                $roi = Clamp-Roi $roi $bitmap.Width $bitmap.Height
                if ($null -eq $roi) {
                    continue
                }
                $objectArea = 0.0
                foreach ($idx in $group) {
                    $objectArea += $boxes[$idx].Area
                }
                $roi = Resize-RoiToTarget $roi $bitmap.Width $bitmap.Height $objectArea $minRatio $maxRatio
                if ($null -eq $roi) {
                    continue
                }

                Draw-Rect -Graphics $graphics -Pen $roiPen -Box $roi
                $graphics.DrawString("ROI$roiIndex", $font, $brush, [float]($roi.X1 + 4), [float]([Math]::Max(0, $roi.Y1 - 28)))

                $cropRect = [System.Drawing.Rectangle]::new([int]$roi.X1, [int]$roi.Y1, [int]$roi.W, [int]$roi.H)
                $crop = $bitmap.Clone($cropRect, $bitmap.PixelFormat)
                $cropGraphics = [System.Drawing.Graphics]::FromImage($crop)
                $cropPen = [System.Drawing.Pen]::new([System.Drawing.Color]::Lime, 3)
                try {
                    foreach ($box in $boxes) {
                        if ($box.Cx -lt $roi.X1 -or $box.Cx -gt $roi.X2 -or $box.Cy -lt $roi.Y1 -or $box.Cy -gt $roi.Y2) {
                            continue
                        }
                        $local = New-Box -X1 ($box.X1 - $roi.X1) -Y1 ($box.Y1 - $roi.Y1) -X2 ($box.X2 - $roi.X1) -Y2 ($box.Y2 - $roi.Y1)
                        $local = Clip-Box -Box $local -Width $roi.W -Height $roi.H
                        if ($null -ne $local) {
                            Draw-Rect -Graphics $cropGraphics -Pen $cropPen -Box $local
                        }
                    }
                }
                finally {
                    $cropGraphics.Dispose()
                    $cropPen.Dispose()
                }

                $cropPath = Join-Path $outputDir ($sample.Label + "_roi" + $roiIndex + ".jpg")
                $crop.Save($cropPath, [System.Drawing.Imaging.ImageFormat]::Jpeg)
                $crop.Dispose()

                $areaRatio = [Math]::Round($objectArea / $roi.Area, 3)
                $summary += [PSCustomObject]@{
                    Sample = $sample.Label
                    Roi = "roi$roiIndex"
                    Mode = $mode
                    Boxes = $group.Count
                    AreaRatio = $areaRatio
                    Crop = $cropPath
                }
                $roiIndex++
            }

            $overlayPath = Join-Path $outputDir ($sample.Label + "_overlay.jpg")
            $canvas.Save($overlayPath, [System.Drawing.Imaging.ImageFormat]::Jpeg)
        }
        finally {
            $graphics.Dispose()
            $shoePen.Dispose()
            $roiPen.Dispose()
            $font.Dispose()
            $brush.Dispose()
            $canvas.Dispose()
        }
    }
    finally {
        $bitmap.Dispose()
    }
}

$summaryPath = Join-Path $outputDir "summary.txt"
$summary | Sort-Object Sample, Roi | Format-Table -AutoSize | Out-String | Set-Content -Path $summaryPath -Encoding UTF8
Write-Output "Preview images written to: $outputDir"
Write-Output "Summary written to: $summaryPath"