#!/usr/bin/env python3

import argparse
import json
import os
import hashlib
import time
from typing import Dict, Any

import cv2
import numpy as np


# ------------------------------------------------------------
# Utilities
# ------------------------------------------------------------

def load_intrinsics_npz(npz_path: str):
    data = np.load(npz_path)

    for k in ('camera_matrix', 'mtx', 'K'):
        if k in data:
            K = data[k].astype(np.float32)
            break
    else:
        raise KeyError('Camera matrix not found in npz')

    for k in ('dist_coeffs', 'dist', 'D'):
        if k in data:
            D = data[k].astype(np.float32).reshape(-1, 1)
            break
    else:
        D = np.zeros((5,1), dtype=np.float32)

    return K, D


# ------------------------------------------------------------

def get_aruco_detector(dict_name: str):
    mapping = {
        'DICT_4X4_250': cv2.aruco.DICT_4X4_250,
        'DICT_5X5_100': cv2.aruco.DICT_5X5_100,
        'DICT_6X6_250': cv2.aruco.DICT_6X6_250,
        'DICT_ARUCO_ORIGINAL': cv2.aruco.DICT_ARUCO_ORIGINAL,
    }

    dict_id = mapping.get(dict_name, cv2.aruco.DICT_4X4_250)

    dictionary = cv2.aruco.getPredefinedDictionary(dict_id)

    try:
        params = cv2.aruco.DetectorParameters()
    except Exception:
        params = cv2.aruco.DetectorParameters_create()

    try:
        detector = cv2.aruco.ArucoDetector(dictionary, params)
        return detector, None
    except Exception:
        return None, (dictionary, params)


# ------------------------------------------------------------

def detect_markers(image, detector_tuple):
    detector, fallback = detector_tuple

    if detector is not None:
        corners, ids, rejected = detector.detectMarkers(image)
    else:
        dictionary, params = fallback
        corners, ids, rejected = cv2.aruco.detectMarkers(
            image,
            dictionary,
            parameters=params
        )

    return corners, ids, rejected


# ------------------------------------------------------------

def compute_sharpness(gray_roi):
    if gray_roi.size == 0:
        return 0.0

    return float(cv2.Laplacian(gray_roi, cv2.CV_64F).var())


# ------------------------------------------------------------

def compute_local_stats(gray_roi):
    if gray_roi.size == 0:
        return 0.0, 0.0

    mean = float(np.mean(gray_roi))
    std = float(np.std(gray_roi))

    return mean, std


# ------------------------------------------------------------

def hash_file(path):
    sha = hashlib.sha256()

    with open(path, 'rb') as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            sha.update(chunk)

    return sha.hexdigest()


# ------------------------------------------------------------

def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('-i', '--image', required=True)
    parser.add_argument('-npz', '--intrinsics', required=True)
    parser.add_argument('-cameraId', '--cameraId', required=True)

    parser.add_argument(
        '--dict',
        default='DICT_4X4_250'
    )

    parser.add_argument(
        '-o',
        '--output',
        default=None
    )

    args = parser.parse_args()

    image = cv2.imread(args.image)

    if image is None:
        raise RuntimeError(f'Cannot read image: {args.image}')

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    h, w = gray.shape[:2]

    K, D = load_intrinsics_npz(args.intrinsics)

    detector_tuple = get_aruco_detector(args.dict)

    corners_list, ids, rejected = detect_markers(gray, detector_tuple)

    observations = []

    if ids is not None:

        ids = ids.flatten().tolist()

        for i, marker_id in enumerate(ids):

            corners = corners_list[i].reshape((4,2)).astype(np.float32)

            center = corners.mean(axis=0)

            area_px = float(cv2.contourArea(corners))

            perimeter_px = float(cv2.arcLength(corners, True))

            x_min = float(np.min(corners[:,0]))
            x_max = float(np.max(corners[:,0]))
            y_min = float(np.min(corners[:,1]))
            y_max = float(np.max(corners[:,1]))

            bbox_x = int(max(0, np.floor(x_min)))
            bbox_y = int(max(0, np.floor(y_min)))
            bbox_w = int(min(w - bbox_x, np.ceil(x_max - x_min)))
            bbox_h = int(min(h - bbox_y, np.ceil(y_max - y_min)))

            roi = gray[
                bbox_y:bbox_y+bbox_h,
                bbox_x:bbox_x+bbox_w
            ]

            sharpness = compute_sharpness(roi)

            mean_gray, std_gray = compute_local_stats(roi)

            edge_lengths = []

            for k in range(4):
                p1 = corners[k]
                p2 = corners[(k+1)%4]
                edge_lengths.append(float(np.linalg.norm(p1 - p2)))

            edge_ratio = (
                max(edge_lengths) / max(1e-6, min(edge_lengths))
            )

            obs = {
                'marker_id': int(marker_id),

                'corners_px': corners.tolist(),

                'center_px': center.tolist(),

                'bbox_px': {
                    'x': x_min,
                    'y': y_min,
                    'w': x_max - x_min,
                    'h': y_max - y_min
                },

                'quality': {
                    'area_px': area_px,
                    'perimeter_px': perimeter_px,
                    'sharpness_laplacian': sharpness,
                    'mean_gray': mean_gray,
                    'std_gray': std_gray,
                    'edge_ratio': edge_ratio,
                    'edge_lengths_px': edge_lengths
                }
            }

            observations.append(obs)

    output = {
        'schema_version': '1.0',

        'created_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),

        'camera': {
            'camera_id': args.cameraId,
            'intrinsics_file': os.path.abspath(args.intrinsics),
            'camera_matrix': K.tolist(),
            'distortion_coefficients': D.reshape(-1).tolist()
        },

        'image': {
            'image_file': os.path.abspath(args.image),
            'image_sha256': hash_file(args.image),
            'width_px': int(w),
            'height_px': int(h)
        },

        'aruco': {
            'dictionary': args.dict,
            'num_detected_markers': len(observations),
            'num_rejected_candidates': int(len(rejected))
        },

        'observations': observations
    }

    if args.output is None:
        base = os.path.splitext(args.image)[0]
        out_json = f'{base}_aruco_detection.json'
    else:
        out_json = args.output

    with open(out_json, 'w', encoding='utf-8') as f:
        json.dump(output, f, indent=2)

    print(f'Saved: {out_json}')


# ------------------------------------------------------------

if __name__ == '__main__':
    main()