#!/usr/bin/env python3
"""
KYC Document Cropping & Rotation - Complete Inference Pipeline

Models:
- Segmentation (corner detection): https://huggingface.co/Jwalit/kyc-document-corner-detector
- Rotation classifier: https://huggingface.co/Jwalit/kyc-document-rotation-classifier

Usage:
    python inference_pipeline.py --image path/to/image.jpg --output out.jpg
"""

import argparse, warnings
from pathlib import Path
import numpy as np
import cv2
from PIL import Image
import torch, torch.nn as nn
from torchvision import transforms

warnings.filterwarnings("ignore")

# ── Model Classes ──
class SegModel(nn.Module):
    def __init__(self):
        super().__init__()
        from torchvision import models
        self.enc = models.mobilenet_v3_small(weights=None).features
        self.dec = nn.Sequential(
            nn.Conv2d(576,256,3,padding=1), nn.ReLU(),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(256,128,3,padding=1), nn.ReLU(),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(128,64,3,padding=1), nn.ReLU(),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(64,32,3,padding=1), nn.ReLU(),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(32,16,3,padding=1), nn.ReLU(),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(16,1,3,padding=1),
        )
    def forward(self, x):
        return self.dec(self.enc(x))

class RotModel(nn.Module):
    def __init__(self):
        super().__init__()
        from torchvision import models
        self.m = models.mobilenet_v3_small(weights=None)
        self.m.classifier[3] = nn.Linear(self.m.classifier[3].in_features, 4)
    def forward(self, x):
        return self.m(x)

# ── Preprocessing ──
seg_tf = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
])
rot_tf = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
])

# ── Helpers ──
def mask_to_corners(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None
    cnt = max(contours, key=cv2.contourArea)
    peri = cv2.arcLength(cnt, True)
    for eps in [0.02, 0.05, 0.1]:
        approx = cv2.approxPolyDP(cnt, eps*peri, True)
        if len(approx) == 4:
            pts = approx.reshape(4,2).astype(np.float32)
            s, d = pts.sum(axis=1), np.diff(pts, axis=1).flatten()
            rect = np.zeros((4,2), dtype=np.float32)
            rect[0] = pts[np.argmin(s)]
            rect[2] = pts[np.argmax(s)]
            rect[1] = pts[np.argmin(d)]
            rect[3] = pts[np.argmax(d)]
            return rect
    rect = cv2.minAreaRect(cnt)
    return cv2.boxPoints(rect).astype(np.float32)

def predict_corners(seg_model, img_path, device='cpu'):
    img = Image.open(img_path).convert('RGB')
    orig_w, orig_h = img.size
    t = seg_tf(img).unsqueeze(0).to(device)
    seg_model.eval()
    with torch.no_grad():
        pred = torch.sigmoid(seg_model(t))[0,0].cpu().numpy()
    mask = (pred > 0.5).astype(np.uint8) * 255
    mask = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
    return mask_to_corners(mask)

def predict_rotation(rot_model, img_path, device='cpu'):
    img = Image.open(img_path).convert('RGB')
    t = rot_tf(img).unsqueeze(0).to(device)
    rot_model.eval()
    with torch.no_grad():
        pred = torch.argmax(rot_model(t), 1).item()
    return [0, 90, 180, 270][pred]

def crop_and_rotate(img_path, corners, angle):
    img = cv2.imread(str(img_path))
    if img is None:
        return None
    if corners is not None:
        pts = corners.astype(np.float32)
        w = int(max(np.linalg.norm(pts[1]-pts[0]), np.linalg.norm(pts[2]-pts[3])))
        h = int(max(np.linalg.norm(pts[3]-pts[0]), np.linalg.norm(pts[2]-pts[1])))
        dst = np.array([[0,0],[w-1,0],[w-1,h-1],[0,h-1]], dtype=np.float32)
        M = cv2.getPerspectiveTransform(pts, dst)
        cropped = cv2.warpPerspective(img, M, (w, h))
    else:
        cropped = img.copy()
    if angle == 90:
        cropped = cv2.rotate(cropped, cv2.ROTATE_90_COUNTERCLOCKWISE)
    elif angle == 180:
        cropped = cv2.rotate(cropped, cv2.ROTATE_180)
    elif angle == 270:
        cropped = cv2.rotate(cropped, cv2.ROTATE_90_CLOCKWISE)
    return cropped

def process_image(img_path, seg_weights, rot_weights, output_path=None, device='cpu'):
    seg_model = SegModel()
    seg_model.load_state_dict(torch.load(seg_weights, map_location=device))
    seg_model.to(device)
    rot_model = RotModel()
    rot_model.load_state_dict(torch.load(rot_weights, map_location=device))
    rot_model.to(device)
    print(f"Processing: {img_path}")
    corners = predict_corners(seg_model, img_path, device)
    angle = predict_rotation(rot_model, img_path, device)
    print(f"  Rotation: {angle}°, Corners: {corners is not None}")
    result = crop_and_rotate(img_path, corners, angle)
    if output_path:
        cv2.imwrite(str(output_path), result)
        print(f"  Saved: {output_path}")
    return result

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--seg-weights", default="pytorch_model.bin")
    parser.add_argument("--rot-weights", required=True)
    parser.add_argument("--output", default="output.jpg")
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()
    process_image(args.image, args.seg_weights, args.rot_weights, args.output, args.device)

if __name__ == "__main__":
    main()