"""Standalone inference API and CLI for StrandSeg-Lite."""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np
import torch

from model import load_model


IMAGE_MEAN = np.asarray([0.485, 0.456, 0.406], dtype=np.float32)
IMAGE_STD = np.asarray([0.229, 0.224, 0.225], dtype=np.float32)


@dataclass(frozen=True)
class Prediction:
    probability: np.ndarray
    mask: np.ndarray
    threshold: float


def _letterbox(image: np.ndarray, width: int, height: int) -> tuple[np.ndarray, tuple[int, int, int, int]]:
    source_height, source_width = image.shape[:2]
    scale = min(width / source_width, height / source_height)
    resized_width = max(1, int(round(source_width * scale)))
    resized_height = max(1, int(round(source_height * scale)))
    pad_x = (width - resized_width) // 2
    pad_y = (height - resized_height) // 2
    resized = cv2.resize(image, (resized_width, resized_height), interpolation=cv2.INTER_AREA)
    canvas = np.zeros((height, width, 3), dtype=np.uint8)
    canvas[pad_y : pad_y + resized_height, pad_x : pad_x + resized_width] = resized
    return canvas, (pad_x, pad_y, resized_width, resized_height)


class StrandSegLite:
    def __init__(self, checkpoint: str | Path = "model.pt", device: str | torch.device = "cpu"):
        self.device = torch.device(device)
        self.model, self.metadata = load_model(checkpoint, self.device)
        config = dict(self.metadata.get("model_config") or {})
        self.input_width = int(config.get("input_width", 960))
        self.input_height = int(config.get("input_height", 544))
        self.threshold = float(self.metadata.get("threshold", 0.9204))

    @torch.inference_mode()
    def predict(self, image: np.ndarray, *, color: str = "rgb", threshold: float | None = None) -> Prediction:
        if image.ndim != 3 or image.shape[2] != 3:
            raise ValueError("image must be an HxWx3 uint8 array")
        if image.dtype != np.uint8:
            raise ValueError("image must have dtype uint8")
        color = color.lower()
        if color == "rgb":
            bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        elif color == "bgr":
            bgr = image
        else:
            raise ValueError("color must be 'rgb' or 'bgr'")
        original_height, original_width = bgr.shape[:2]
        prepared, (pad_x, pad_y, resized_width, resized_height) = _letterbox(
            bgr, self.input_width, self.input_height
        )
        rgb = cv2.cvtColor(prepared, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        rgb = (rgb - IMAGE_MEAN) / IMAGE_STD
        tensor = torch.from_numpy(np.transpose(rgb, (2, 0, 1))).unsqueeze(0).to(self.device)
        probability = torch.sigmoid(self.model(tensor))[0, 0].cpu().numpy()
        probability = probability[pad_y : pad_y + resized_height, pad_x : pad_x + resized_width]
        probability = cv2.resize(probability, (original_width, original_height), interpolation=cv2.INTER_LINEAR)
        selected_threshold = self.threshold if threshold is None else float(threshold)
        mask = (probability >= selected_threshold).astype(np.uint8)
        return Prediction(probability=probability.astype(np.float32), mask=mask, threshold=selected_threshold)

    def predict_file(self, path: str | Path, *, threshold: float | None = None) -> Prediction:
        encoded = np.fromfile(str(path), dtype=np.uint8)
        image = cv2.imdecode(encoded, cv2.IMREAD_COLOR) if encoded.size else None
        if image is None:
            raise ValueError(f"Could not read image: {path}")
        return self.predict(image, color="bgr", threshold=threshold)


def main() -> None:
    parser = argparse.ArgumentParser(description="Run StrandSeg-Lite on one image")
    parser.add_argument("image")
    parser.add_argument("--checkpoint", default="model.pt")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--threshold", type=float)
    parser.add_argument("--output-mask", default="mask.png")
    parser.add_argument("--output-probability", default="probability.png")
    args = parser.parse_args()
    result = StrandSegLite(args.checkpoint, args.device).predict_file(args.image, threshold=args.threshold)
    cv2.imwrite(args.output_mask, result.mask * 255)
    cv2.imwrite(args.output_probability, np.clip(result.probability * 255.0, 0, 255).astype(np.uint8))
    print(f"threshold={result.threshold:.4f} foreground_pixels={int(result.mask.sum())}")


if __name__ == "__main__":
    main()