#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
NYU Depth V2 Eigen test split (654 张) 相对深度精度评估。

输入：
    --mat     nyu_depth_v2_labeled.mat（GT 直接从这里读，不落盘）
    --split   export_nyu654_test.py 生成的 split.json
    --pred    一个或多个 <名称>=<NPZ 目录>，NPZ 由 infer 脚本 --dump-npz 生成

评估协议：
    1. 有效像素：rawDepths > 0 且 depths 在 (0, max_depth]；
    2. 预测上采样到 GT 分辨率（bilinear）；
    3. 相对深度对齐：
       - disparity（默认，MiDaS 协议）：拟合 a*pred + b ≈ 1/gt，再取倒数得米制深度；
       - depth：拟合 a*pred + b ≈ gt；
       最小二乘在有效像素上求解；
    4. 逐图计算 AbsRel / SqRel / RMSE / RMSE_log / δ<1.25 / 1.25² / 1.25³，
       对 654 张取平均；
    5. 第一个 --pred 作为基准，其余输出相对掉点 (%)。

注意：未做 Eigen crop，绝对数值不能直接与论文表格对比；同一批图上不同后端的
相对掉点不受影响。

用法：
    ../../venv/bin/python src/tool/evaluate_nyu654_depth.py \
        --mat ../../datasets/nyu_depth_v2_labeled.mat \
        --split datasets/nyu654_test/split.json \
        --pred fp32=results/nyu654_fp32/npz axonnx=results/nyu654_axonnx/npz \
        --out results/nyu654_eval
"""

import argparse
import json
from pathlib import Path

import cv2
import h5py
import numpy as np
from tqdm import tqdm

METRICS = ('AbsRel', 'SqRel', 'RMSE', 'RMSE_log', 'd1', 'd2', 'd3')
SIGMA = (1.25, 1.25 ** 2, 1.25 ** 3)


def solve_scale_shift(pred, target, mask):
    """在 mask 上最小二乘求解 a, b 使 a*pred + b ≈ target。"""
    p = pred[mask]
    t = target[mask]
    A = np.stack([p, np.ones_like(p)], axis=1)
    (a, b), *_ = np.linalg.lstsq(A, t, rcond=None)
    return float(a), float(b)


def compute_metrics(pred_depth, gt_depth, mask, align, max_depth):
    clipped = 0.0
    if align == 'disparity':
        gt_target = 1.0 / gt_depth
        a, b = solve_scale_shift(pred_depth, gt_target, mask)
        if a <= 0:
            raise ValueError(f'disparity 对齐斜率为负 (a={a:.4f})，'
                             f'预测不是"越大越近"的视差量')
        aligned = a * pred_depth + b
        # 极少数图（低纹理/异常预测）会拟出负视差，裁剪到最大深度的下限
        neg = aligned[mask] <= 0
        clipped = float(neg.mean())
        aligned = np.maximum(aligned, 1.0 / max_depth)
        pred = 1.0 / aligned
    else:
        a, b = solve_scale_shift(pred_depth, gt_depth, mask)
        pred = a * pred_depth + b

    pred = pred[mask]
    gt = gt_depth[mask]
    if np.any(pred <= 0):
        raise ValueError('对齐后预测深度出现非正值')

    diff = pred - gt
    out = {
        'AbsRel': float(np.mean(np.abs(diff) / gt)),
        'SqRel': float(np.mean(diff ** 2 / gt)),
        'RMSE': float(np.sqrt(np.mean(diff ** 2))),
        'RMSE_log': float(np.sqrt(np.mean((np.log(pred) - np.log(gt)) ** 2))),
    }
    ratio = np.maximum(pred / gt, gt / pred)
    for name, s in zip(('d1', 'd2', 'd3'), SIGMA):
        out[name] = float(np.mean(ratio < s))
    return out, a, b, clipped


def load_predictions(path):
    """读取预测：汇总 NPZ（names/depths）、分块汇总目录，或旧的逐图 NPZ 目录。"""
    path = Path(path)
    if path.is_file():
        files = [path]
    elif path.is_dir():
        files = sorted(path.glob('*.npz'))
        if not files:
            raise ValueError(f'目录中没有 NPZ: {path}')
    else:
        raise FileNotFoundError(f'预测路径不存在: {path}')

    preds = {}
    for fp in files:
        with np.load(fp, allow_pickle=False) as data:
            if 'depths' in data and 'names' in data:
                for name, depth in zip(data['names'], data['depths']):
                    preds[str(name)] = depth
            elif 'depth' in data:
                preds[fp.name[:-len('.npz')]] = data['depth']
            else:
                raise ValueError(f'无法识别的 NPZ 格式: {fp}')
    return preds


def main():
    parser = argparse.ArgumentParser(description='NYU 654 相对深度评测')
    parser.add_argument('--mat', required=True)
    parser.add_argument('--split', required=True, help='split.json')
    parser.add_argument('--pred', nargs='+', required=True, help='名称=NPZ目录，第一个为基准')
    parser.add_argument('--out', required=True, help='结果目录')
    parser.add_argument('--align', choices=['disparity', 'depth'], default='disparity')
    parser.add_argument('--max-depth', type=float, default=10.0)
    args = parser.parse_args()

    pred_sets = []
    for item in args.pred:
        name, _, path = item.partition('=')
        if not path:
            raise ValueError(f'--pred 需要 名称=路径 格式: {item}')
        pred_sets.append((name, load_predictions(path)))

    split = json.load(open(args.split))
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    per_image = {name: [] for name, _ in pred_sets}
    failed = {name: [] for name, _ in pred_sets}
    with h5py.File(args.mat, 'r') as f:
        for record in tqdm(split, desc=f'eval({args.align})', ncols=80):
            i = record['mat_index']
            gt = np.asarray(f['depths'][i], dtype=np.float32)
            raw = np.asarray(f['rawDepths'][i], dtype=np.float32)
            mask = (raw > 0) & (gt > 0) & (gt <= args.max_depth)

            for name, preds in pred_sets:
                if record['file'] not in preds:
                    raise KeyError(f'{name} 缺少预测: {record["file"]}')
                pred = np.asarray(preds[record['file']], dtype=np.float32)
                if pred.ndim == 3:
                    pred = pred[0]
                pred = cv2.resize(pred, (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_LINEAR)
                try:
                    m, a, b, clipped = compute_metrics(pred, gt, mask, args.align, args.max_depth)
                except ValueError as error:
                    failed[name].append({'file': record['file'], 'reason': str(error)})
                    continue
                m.update({'file': record['file'], 'scale': a, 'shift': b,
                          'clipped_frac': clipped, 'valid_pixels': int(mask.sum())})
                per_image[name].append(m)

    results = {}
    for name, _ in pred_sets:
        rows = per_image[name]
        results[name] = {k: float(np.mean([r[k] for r in rows])) for k in METRICS}
        results[name]['p90_absrel'] = float(np.percentile([r['AbsRel'] for r in rows], 90))
        results[name]['clipped_images'] = int(sum(r['clipped_frac'] > 0 for r in rows))
        results[name]['clipped_pixel_frac'] = float(np.mean([r['clipped_frac'] for r in rows]))
        results[name]['evaluated_images'] = len(rows)
        results[name]['failed_images'] = len(failed[name])

    base = pred_sets[0][0]
    drops = {}
    for name, _ in pred_sets[1:]:
        drops[name] = {}
        for k in METRICS:
            b, v = results[base][k], results[name][k]
            drops[name][k] = float((v - b) / b * 100)
        b, v = results[base]['p90_absrel'], results[name]['p90_absrel']
        drops[name]['p90_absrel'] = float((v - b) / b * 100)

    summary = {
        'dataset': 'nyu_depth_v2_labeled / eigen test split',
        'num_images': len(split),
        'align': args.align,
        'max_depth': args.max_depth,
        'eigen_crop': False,
        'results': results,
        'baseline': base,
        'drop_percent': drops,
        'failed': failed,
        'per_image': {name: per_image[name] for name, _ in pred_sets},
    }
    with open(out_dir / 'metrics.json', 'w') as fp:
        json.dump(summary, fp, indent=2)

    print(f'\n{"指标":<10}' + ''.join(f'{n:>14}' for n, _ in pred_sets))
    for k in METRICS + ('p90_absrel',):
        print(f'{k:<10}' + ''.join(f'{results[n][k]:>14.4f}' for n, _ in pred_sets))
    for name, d in drops.items():
        print(f'\n{name} 相对 {base} 掉点 (%): ' +
              ', '.join(f'{k}={v:+.2f}' for k, v in d.items()))
    for name, _ in pred_sets:
        if failed[name]:
            print(f'\n{name}: {len(failed[name])} 张对齐失败（方向反了或输出异常），未计入均值')
    print(f'\n结果已保存: {out_dir / "metrics.json"}')


if __name__ == '__main__':
    main()