# -*- coding: utf-8 -*-
"""
Copyright (C) 2024–2025 Amlogic, Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import cv2
import glob
import argparse
import numpy as np
from pathlib import Path
from amlnnlite.api import AMLNNLite


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))


def nms_xyxy(boxes, scores, iou_thres=0.5):
    if len(boxes) == 0:
        return []

    boxes = np.asarray(boxes, dtype=np.float32)
    scores = np.asarray(scores, dtype=np.float32)

    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h
        union = areas[i] + areas[order[1:]] - inter
        iou = inter / np.maximum(union, 1e-6)

        inds = np.where(iou <= iou_thres)[0]
        order = order[inds + 1]

    return keep


def preprocess(img_path, input_size=(320, 320)):
    img = cv2.imread(img_path)
    if img is None:
        return None, None, None

    orig = img.copy()
    img = cv2.resize(img, input_size)
    img = img.astype(np.float32) / 255.0
    img = np.expand_dims(img, axis=0)  
    return img, orig, orig.shape[:2]


def postprocess_qrcode(outputs, orig_shape, conf_thres=0.8, nms_thres=0.5, pad=40):
    h0, w0 = orig_shape
    sx = w0 / 320.0
    sy = h0 / 320.0

    out = outputs[0] if isinstance(outputs, (list, tuple)) else outputs
    out = np.asarray(out)

    if out.ndim == 4 and out.shape[0] == 1 and out.shape[1] == 1:
        out = out[0, 0]                 
        pred = out.transpose(1, 0)      
    elif out.ndim == 3 and out.shape[0] == 1:
        out = out[0]                    
        pred = out.transpose(1, 0)      
    else:
        raise ValueError(f"Unexpected output shape: {out.shape}")

    boxes_xyxy_320 = []
    scores = []

    for i in range(pred.shape[0]):
        x, y, w, h, raw_score = pred[i]

        score = sigmoid(raw_score)

        if score < conf_thres:
            continue

        x1 = x - w / 2.0
        y1 = y - h / 2.0
        x2 = x + w / 2.0
        y2 = y + h / 2.0

        boxes_xyxy_320.append([float(x1), float(y1), float(x2), float(y2)])
        scores.append(float(score))

    keep = nms_xyxy(boxes_xyxy_320, scores, iou_thres=nms_thres)

    results = []
    for idx in keep:
        x1, y1, x2, y2 = boxes_xyxy_320[idx]
        score = scores[idx]

        x1 = x1 * sx
        y1 = y1 * sy
        x2 = x2 * sx
        y2 = y2 * sy

        x1p = int(max(0, x1 - pad-10))
        y1p = int(max(0, y1 - pad-15))
        x2p = int(min(w0 - 1, x2 + pad-10))
        y2p = int(min(h0 - 1, y2 + pad))

        results.append({
            "box": [x1p, y1p, x2p, y2p],
            "score": score,
        })

    return results


def decode_qrcodes(orig, dets):
    detector = cv2.QRCodeDetector()
    decoded = []

    h, w = orig.shape[:2]

    for det in dets:
        x1, y1, x2, y2 = map(int, det["box"])
        score = det["score"]

        x1 = max(0, min(x1, w - 1))
        y1 = max(0, min(y1, h - 1))
        x2 = max(0, min(x2, w - 1))
        y2 = max(0, min(y2, h - 1))

        if x2 <= x1 or y2 <= y1:
            print(f"skip invalid box: {[x1, y1, x2, y2]}, score={score:.3f}")
            continue

        crop = orig[y1:y2, x1:x2].copy()
        if crop.size == 0:
            print(f"skip empty crop: {[x1, y1, x2, y2]}, score={score:.3f}")
            continue

        ch, cw = crop.shape[:2]
        if cw < 20 or ch < 20:
            print(f"skip too small crop: {[x1, y1, x2, y2]}, size=({cw},{ch}), score={score:.3f}")
            continue

        try:
            text, points, _ = detector.detectAndDecode(crop)
        except cv2.error as e:
            print(f"skip cv2 decode error: {[x1, y1, x2, y2]}, score={score:.3f}, err={e}")
            continue

        decoded.append({
            "box": [x1, y1, x2, y2],
            "score": score,
            "text": text,
            "points": points,
        })

    return decoded


def draw_results(img, results):
    vis = img.copy()
    for r in results:
        x1, y1, x2, y2 = r["box"]
        score = r["score"]
        text = r["text"]

        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            vis, f"{score:.3f}", (x1, max(0, y1 - 8)),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2
        )

        if text:
            cv2.putText(
                vis, text[:60], (x1, min(img.shape[0] - 10, y2 + 25)),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2
            )
    return vis


def main():
    parser = argparse.ArgumentParser(description="QRCode AMLNNLite Demo")
    parser.add_argument('--board-work-path', type=str, default='/data/local/tmp')
    parser.add_argument('--model-path', required=True, help='Path to .adla model')
    parser.add_argument('--image-dir', required=True, help='Directory of test images')
    parser.add_argument('--run-cycles', type=int, default=1, help='Inference cycles')
    parser.add_argument('--loglevel', type=str, default='WARNING',
                        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'])
    parser.add_argument('--conf-thres', type=float, default=0.8)
    parser.add_argument('--nms-thres', type=float, default=0.5)
    parser.add_argument('--pad', type=int, default=40)
    args = parser.parse_args()

    amlnn = AMLNNLite()
    amlnn.config(
        board_work_path=args.board_work_path,
        model_path=args.model_path,
        run_cycles=args.run_cycles,
        loglevel=args.loglevel
    )
    amlnn.init()

    image_files = sorted(glob.glob(os.path.join(args.image_dir, "*.[jp][pn][g]")))
    if not image_files:
        print(f"No images found in {args.image_dir}")
        amlnn.uninit()
        return

    res_dir = "qrcode_result"
    os.makedirs(res_dir, exist_ok=True)

    for idx, img_path in enumerate(image_files, start=1):
        print("=" * 60)
        print(f"Processing image {idx}/{len(image_files)}: {Path(img_path).name}")
        print("=" * 60)

        inp, orig, orig_shape = preprocess(img_path)
        if inp is None:
            print(f"Failed to read: {img_path}")
            continue

        outputs = amlnn.inference(inp, inputs_data_format='NHWC')
        dets = postprocess_qrcode(
            outputs,
            orig_shape,
            conf_thres=args.conf_thres,
            nms_thres=args.nms_thres,
            pad=args.pad
        )
        results = decode_qrcodes(orig, dets)

        if len(results) == 0:
            print("    No objects detected")
        else:
            print(f"    Detected {len(results)} objects:")
            for i, r in enumerate(results, 1):
                print(f"      {i}. score={r['score']:.3f}")
                print(f"         box={r['box']}")
                print(f"         text={r['text']}")

        vis = draw_results(orig, results)
        save_path = os.path.join(res_dir, Path(img_path).name)
        cv2.imwrite(save_path, vis)
        print(f"    Result saved to: {save_path}")

    if args.loglevel == 'INFO':
        print("\nPerformance analysis visualization starting...")

    amlnn.visualize()
    amlnn.uninit()


if __name__ == "__main__":
    main()