#!/usr/bin/env python3
"""
Table Normalizer - Standardizes cropped table regions into a fixed canonical geometry.
This reduces device-to-device UI variation by forcing all tables into the exact same dimensions.
"""

import cv2
import numpy as np
import json
from pathlib import Path
from typing import Optional


def normalize_table(
    table_image: np.ndarray,
    target_width: int = 1400,
    target_height: int = 700,
    debug: bool = False,
    output_path: Optional[str] = None,
    image_name: str = "normalized_table.jpg"
) -> np.ndarray:
    """
    Normalizes a cropped table image to exactly target_width x target_height.
    
    Args:
        table_image: The cropped table image (numpy array).
        target_width: Target width in pixels (default: 1400).
        target_height: Target height in pixels (default: 700).
        debug: If True, prints logging information and saves output if output_path is provided.
        output_path: Directory path to save the normalized table (used if debug=True).
        image_name: Filename to use when saving the debug image.
        
    Returns:
        The normalized table image (numpy array).
    """
    if table_image is None or table_image.size == 0:
        raise ValueError("Invalid or empty table image provided.")

    orig_height, orig_width = table_image.shape[:2]

    # Determine interpolation method automatically
    # INTER_AREA is best for shrinking (downscaling)
    # INTER_CUBIC is best for enlarging (upscaling)
    if target_width < orig_width or target_height < orig_height:
        interpolation = cv2.INTER_AREA
        inter_name = "INTER_AREA"
    else:
        interpolation = cv2.INTER_CUBIC
        inter_name = "INTER_CUBIC"

    # Resize directly to target dimensions, ignoring aspect ratio
    normalized_table = cv2.resize(
        table_image, 
        (target_width, target_height), 
        interpolation=interpolation
    )

    if debug:
        print(f"Original Table: {orig_width}x{orig_height}")
        print(f"Normalized Table: {target_width}x{target_height}")
        print(f"Interpolation: {inter_name}")
        
        if output_path:
            out_dir = Path(output_path)
            out_dir.mkdir(parents=True, exist_ok=True)
            out_file = out_dir / image_name
            cv2.imwrite(str(out_file), normalized_table)
            print(f"Saved normalized table to: {out_file}")

    return normalized_table


def process_batch(input_dir: str, layouts_json: str, output_dir: str, target_width: int = 1400, target_height: int = 700):
    """
    Process a batch of images by cropping them according to layouts_json and normalizing them.
    """
    import json
    
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)
    
    try:
        with open(layouts_json, 'r') as f:
            layouts = json.load(f)
    except Exception as e:
        print(f"Error loading layouts JSON: {e}")
        return
        
    print(f"Normalizing tables to {target_width}x{target_height}")
    print(f"Output: {output_path.absolute()}")
    print("=" * 70)
    
    success_count = 0
    total = len(layouts)
    
    for i, (img_filename, layout_data) in enumerate(layouts.items(), 1):
        img_path = input_path / img_filename
        
        if not img_path.exists():
            print(f"✗ [{i:2d}/{total}] Missing image: {img_filename}")
            continue
            
        screenshot = cv2.imread(str(img_path))
        if screenshot is None:
            print(f"✗ [{i:2d}/{total}] Error reading: {img_filename}")
            continue
            
        bounds = layout_data.get("table_bounds")
        if not bounds:
            print(f"✗ [{i:2d}/{total}] No bounds found in JSON for: {img_filename}")
            continue
            
        # Crop the table
        table_crop = screenshot[
            bounds["top"]:bounds["bottom"],
            bounds["left"]:bounds["right"]
        ]
        
        # Normalize
        normalized_table = normalize_table(
            table_image=table_crop,
            target_width=target_width,
            target_height=target_height
        )
        
        # Save
        out_file = output_path / f"norm_{img_filename}"
        cv2.imwrite(str(out_file), normalized_table)
        success_count += 1
        print(f"✓ [{i:2d}/{total}] {img_filename:30s} -> {out_file.name}")
        
    print("=" * 70)
    print(f"\n✅ Normalization Complete!")
    print(f"   Tables processed: {success_count}/{total}")
    print(f"   Output: {output_path.absolute()}")

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='Normalize cropped tables to a fixed geometry')
    parser.add_argument('--input', '-i', required=True, help='Directory of clean resized images')
    parser.add_argument('--layouts', '-l', required=True, help='Path to table_layouts.json')
    parser.add_argument('--output', '-o', default='./normalized_tables', help='Output directory')
    parser.add_argument('--width', '-w', type=int, default=1400, help='Target width')
    parser.add_argument('--height', '-H', type=int, default=700, help='Target height')
    
    args = parser.parse_args()
    
    process_batch(args.input, args.layouts, args.output, args.width, args.height)

if __name__ == "__main__":
    main()